# This test requires Arrow Java to be built in the same source tree

try:

arrow_dir = os.environ["ARROW_SOURCE_DIR"]

except KeyError:

arrow_dir = os.path.join(os.path.dirname(

__file__), '..', '..', '..', '..', '..')

pom_path = os.path.join(arrow_dir, 'java', 'pom.xml')

tree = ET.parse(pom_path)

version = tree.getroot().find(

'POM:version',

namespaces={

'POM': 'http://maven.apache.org/POM/4.0.0'

}).text

jar_path = os.path.join(

arrow_dir, 'java', 'tools', 'target',

'arrow-tools-{}-jar-with-dependencies.jar'.format(version))

jar_path = os.getenv("ARROW_TOOLS_JAR", jar_path)

jar_path += ":{}".format(os.path.join(arrow_dir,

"java", "c/target/arrow-c-data-{}.jar".format(version)))

kwargs = {}

# This will be the default behaviour in jpype 0.8+

kwargs['convertStrings'] = False

# For debugging purpose please uncomment the following, and include *jvm_args, before **kwargs

# in startJVM function call

# jvm_args = [

# "-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=*:5005"

# ]

def __init__(self):

self.java_allocator = jpype.JPackage(

"org").apache.arrow.memory.RootAllocator(sys.maxsize)

self.java_c = jpype.JPackage("org").apache.arrow.c

def java_to_python_field(self, jfield):

c_schema = ffi.new("struct ArrowSchema*")

ptr_schema = int(ffi.cast("uintptr_t", c_schema))

self.java_c.Data.exportField(self.java_allocator, jfield, None,

self.java_c.ArrowSchema.wrap(ptr_schema))

return pa.Field._import_from_c(ptr_schema)

def java_to_python_array(self, vector, dictionary_provider=None):

c_schema = ffi.new("struct ArrowSchema*")

ptr_schema = int(ffi.cast("uintptr_t", c_schema))

c_array = ffi.new("struct ArrowArray*")

ptr_array = int(ffi.cast("uintptr_t", c_array))

self.java_c.Data.exportVector(self.java_allocator, vector, dictionary_provider, self.java_c.ArrowArray.wrap(

ptr_array), self.java_c.ArrowSchema.wrap(ptr_schema))

return pa.Array._import_from_c(ptr_array, ptr_schema)

def java_to_python_record_batch(self, root):

c_schema = ffi.new("struct ArrowSchema*")

ptr_schema = int(ffi.cast("uintptr_t", c_schema))

c_array = ffi.new("struct ArrowArray*")

ptr_array = int(ffi.cast("uintptr_t", c_array))

self.java_c.Data.exportVectorSchemaRoot(self.java_allocator, root, None, self.java_c.ArrowArray.wrap(

ptr_array), self.java_c.ArrowSchema.wrap(ptr_schema))

return pa.RecordBatch._import_from_c(ptr_array, ptr_schema)

def java_to_python_reader(self, reader):

c_stream = ffi.new("struct ArrowArrayStream*")

ptr_stream = int(ffi.cast("uintptr_t", c_stream))

self.java_c.Data.exportArrayStream(self.java_allocator, reader,

self.java_c.ArrowArrayStream.wrap(ptr_stream))

return pa.RecordBatchReader._import_from_c(ptr_stream)

def python_to_java_field(self, field):

c_schema = self.java_c.ArrowSchema.allocateNew(self.java_allocator)

field._export_to_c(c_schema.memoryAddress())

return self.java_c.Data.importField(self.java_allocator, c_schema, None)

def python_to_java_array(self, array, dictionary_provider=None):

c_schema = self.java_c.ArrowSchema.allocateNew(self.java_allocator)

c_array = self.java_c.ArrowArray.allocateNew(self.java_allocator)

array._export_to_c(c_array.memoryAddress(), c_schema.memoryAddress())

return self.java_c.Data.importVector(self.java_allocator, c_array, c_schema, dictionary_provider)

def python_to_java_record_batch(self, record_batch):

c_schema = self.java_c.ArrowSchema.allocateNew(self.java_allocator)

c_array = self.java_c.ArrowArray.allocateNew(self.java_allocator)

record_batch._export_to_c(

c_array.memoryAddress(), c_schema.memoryAddress())

return self.java_c.Data.importVectorSchemaRoot(self.java_allocator, c_array, c_schema, None)

def python_to_java_reader(self, reader):

c_stream = self.java_c.ArrowArrayStream.allocateNew(self.java_allocator)

reader._export_to_c(c_stream.memoryAddress())

return self.java_c.Data.importArrayStream(self.java_allocator, c_stream)

def close(self):

@classmethod

def setUpClass(cls) -> None:

setup_jvm()

def setUp(self):

gc.collect()

self.old_allocated_python = pa.total_allocated_bytes()

self.bridge = Bridge()

def tearDown(self):

self.bridge.close()

gc.collect()

diff_python = pa.total_allocated_bytes() - self.old_allocated_python

self.assertEqual(

pa.total_allocated_bytes(), self.old_allocated_python,

f"PyArrow memory was not adequately released: {diff_python} bytes lost")

def round_trip_field(self, field_generator):

original_field = field_generator()

java_field = self.bridge.python_to_java_field(original_field)

del original_field

new_field = self.bridge.java_to_python_field(java_field)

del java_field

expected = field_generator()

self.assertEqual(expected, new_field)

def round_trip_array(self, array_generator, check_metadata=True):

original_arr = array_generator()

with self.bridge.java_c.CDataDictionaryProvider() as dictionary_provider, \

self.bridge.python_to_java_array(original_arr, dictionary_provider) as vector:

del original_arr

new_array = self.bridge.java_to_python_array(vector, dictionary_provider)

expected = array_generator()

self.assertEqual(expected, new_array)

if check_metadata:

self.assertTrue(new_array.type.equals(expected.type, check_metadata=True))

def round_trip_record_batch(self, rb_generator):

original_rb = rb_generator()

with self.bridge.python_to_java_record_batch(original_rb) as root:

del original_rb

new_rb = self.bridge.java_to_python_record_batch(root)

expected = rb_generator()

self.assertEqual(expected, new_rb)

def round_trip_reader(self, schema, batches):

reader = pa.RecordBatchReader.from_batches(schema, batches)

java_reader = self.bridge.python_to_java_reader(reader)

del reader

py_reader = self.bridge.java_to_python_reader(java_reader)

del java_reader

actual = list(py_reader)

self.assertEqual(batches, actual)

def test_string_array(self):

self.round_trip_array(lambda: pa.array([None, "a", "bb", "ccc"]))

def test_stringview_array(self):

# with nulls short strings

self.round_trip_array(lambda: pa.array([None, "a", "bb", "c"], type=pa.string_view()))

# with nulls long and strings

self.round_trip_array(lambda: pa.array([None, "a", "bb"*10, "c"*13], type=pa.string_view()))

# without nulls short strings

self.round_trip_array(lambda: pa.array(["a", "bb", "c"], type=pa.string_view()))

# without nulls long and strings

self.round_trip_array(lambda: pa.array(["a", "bb"*10, "c"*13], type=pa.string_view()))

# with multiple data buffers

arr1 = pa.array(["a", "bb", "c"], type=pa.string_view())

arr2 = pa.array(["b", "ee" * 10, "f" * 20], type=pa.string_view())

arr3 = pa.array(["c", "abc" * 20, "efg" * 30], type=pa.string_view())

arr4 = pa.array(["d", "abcd" * 100, "efgh" * 200], type=pa.string_view())

self.round_trip_array(lambda: pa.concat_arrays([arr1, arr2, arr3, arr4]))

# empty strings

self.round_trip_array(lambda: pa.array(["", "bb" * 10, "c", "", "d", ""], type=pa.string_view()))

# null value variations

self.round_trip_array(lambda: pa.array(["bb" * 10, None, "", "d", None], type=pa.string_view()))

# empty array

self.round_trip_array(lambda: pa.array([], type=pa.string_view()))

# all null array

self.round_trip_array(lambda: pa.array([None, None, None], type=pa.string_view()))

def test_binaryview_array(self):

# with nulls short binary values

self.round_trip_array(lambda: pa.array([None, bytes([97]), bytes([98, 98]), bytes([99])], type=pa.binary_view()))

# with nulls long binary values

self.round_trip_array(lambda: pa.array([None, bytes([97]), bytes([98, 98] * 10), bytes([99] * 13)], type=pa.binary_view()))

# without nulls short binary values

self.round_trip_array(lambda: pa.array([bytes([97]), bytes([98, 98]), bytes([99])], type=pa.binary_view()))

# without nulls long binary values

self.round_trip_array(lambda: pa.array([bytes([97]), bytes([98, 98] * 10), bytes([99] * 13)], type=pa.binary_view()))

# with multiple data buffers

arr1 = pa.array([bytes([97]), bytes([98, 98]), bytes([99])], type=pa.binary_view())

arr2 = pa.array([bytes([98]), bytes([98, 98] * 10), bytes([99] * 13)], type=pa.binary_view())

arr3 = pa.array([bytes([99]), bytes([98, 100] * 100), bytes([99, 100]) * 30], type=pa.binary_view())

arr4 = pa.array([bytes([100]), bytes([98, 100, 101] * 200), bytes([98, 99]) * 300], type=pa.binary_view())

self.round_trip_array(lambda: pa.concat_arrays([arr1, arr2, arr3, arr4]))

# empty binary values

self.round_trip_array(lambda: pa.array([bytes([]), bytes([97, 97]) * 10, bytes([98]), bytes([]), bytes([97]), bytes([])],

type=pa.binary_view()))

# null value variations

self.round_trip_array(lambda: pa.array([bytes([97, 97]) * 10, None, bytes([]), bytes([99]), None], type=pa.binary_view()))

# empty array

self.round_trip_array(lambda: pa.array([], type=pa.binary_view()))

# all null array

self.round_trip_array(lambda: pa.array([None, None, None], type=pa.binary_view()))

def test_decimal_array(self):

data = [

round(decimal.Decimal(722.82), 2),

round(decimal.Decimal(-934.11), 2),

None,

]

self.round_trip_array(lambda: pa.array(data, pa.decimal128(5, 2)))

def test_int_array(self):

self.round_trip_array(lambda: pa.array([1, 2, 3], type=pa.int32()))

def test_list_array(self):

self.round_trip_array(lambda: pa.array(

[[], [0], [1, 2], [4, 5, 6]], pa.list_(pa.int64())

# disabled check_metadata since the list internal field name ("item")

# is not preserved during round trips (it becomes "$data$").

), check_metadata=False)

def test_empty_list_array(self):

"""Validates GH-37056 fix.

with pa.BufferOutputStream() as bos:

schema = pa.schema([pa.field("f0", pa.list_(pa.int32()), True)])

with ipc.new_stream(bos, schema) as writer:

src = pa.RecordBatch.from_arrays([pa.array([[]])], schema=schema)

writer.write(src)

data_bytes = bos.getvalue()

def recreate_batch():

with pa.input_stream(data_bytes) as ios:

with ipc.open_stream(ios) as reader:

return reader.read_next_batch()

self.round_trip_record_batch(recreate_batch)

def test_struct_array(self):

fields = [

("f1", pa.int32()),

("f2", pa.string()),

]

data = [

{"f1": 1, "f2": "a"},

None,

{"f1": 3, "f2": None},

{"f1": None, "f2": "d"},

{"f1": None, "f2": None},

]

self.round_trip_array(lambda: pa.array(data, type=pa.struct(fields)))

def test_dict(self):

self.round_trip_array(

lambda: pa.array(["a", "b", None, "d"], pa.dictionary(pa.int64(), pa.utf8())))

def test_map(self):

offsets = [0, None, 2, 6]

pykeys = [b"a", b"b", b"c", b"d", b"e", b"f"]

pyitems = [1, 2, 3, None, 4, 5]

keys = pa.array(pykeys, type="binary")

items = pa.array(pyitems, type="i4")

self.round_trip_array(

lambda: pa.MapArray.from_arrays(offsets, keys, items))

def test_field(self):

self.round_trip_field(lambda: pa.field("aa", pa.bool_()))

def test_field_nested(self):

self.round_trip_field(lambda: pa.field(

"test", pa.list_(pa.int32()), nullable=True))

def test_field_metadata(self):

self.round_trip_field(lambda: pa.field("aa", pa.bool_(), {"a": "b"}))

def test_record_batch_with_list(self):

data = [

pa.array([[1], [2], [3], [4, 5, 6]]),

pa.array([1, 2, 3, 4]),

pa.array(['foo', 'bar', 'baz', None]),

pa.array([True, None, False, True])

]

self.round_trip_record_batch(

lambda: pa.RecordBatch.from_arrays(data, ['f0', 'f1', 'f2', 'f3']))

def test_reader_roundtrip(self):

schema = pa.schema([("ints", pa.int64()), ("strs", pa.string())])

data = [

pa.record_batch([[1, 2, 3, None],

["a", "bc", None, ""]],

schema=schema),

pa.record_batch([[None, 4, 5, 6],

[None, "", "def", "g"]],

schema=schema),

]

self.round_trip_reader(schema, data)

def test_reader_complex_roundtrip(self):

schema = pa.schema([

("str_dict", pa.dictionary(pa.int8(), pa.string())),

("int_list", pa.list_(pa.int64())),

])

dictionary = pa.array(["a", "bc", None])

data = [

pa.record_batch([pa.DictionaryArray.from_arrays([0, 2], dictionary),

[[1, 2, 3], None]],

schema=schema),

pa.record_batch([pa.DictionaryArray.from_arrays([None, 1], dictionary),

[[], [4]]],

schema=schema),

]

self.round_trip_reader(schema, data)

def test_listview_array(self):

self.round_trip_array(lambda: pa.array(

[[], [0], [1, 2], [4, 5, 6]], pa.list_view(pa.int64())

# disabled check_metadata since in Java API the listview

# internal field name ("item") is not preserved

# during round trips (it becomes "$data$").

), check_metadata=False)

def test_empty_listview_array(self):

with pa.BufferOutputStream() as bos:

schema = pa.schema([pa.field("f0", pa.list_view(pa.int32()), True)])

with ipc.new_stream(bos, schema) as writer:

src = pa.RecordBatch.from_arrays(

[pa.array([[]], pa.list_view(pa.int32()))], schema=schema)

writer.write(src)

data_bytes = bos.getvalue()

def recreate_batch():

with pa.input_stream(data_bytes) as ios:

with ipc.open_stream(ios) as reader:

return reader.read_next_batch()

self.round_trip_record_batch(recreate_batch)

def test_largelistview_array(self):

self.round_trip_array(lambda: pa.array(

[[], [0], [1, 2], [4, 5, 6]], pa.large_list_view(pa.int64())

# disabled check_metadata since in Java API the listview

# internal field name ("item") is not preserved

# during round trips (it becomes "$data$").

), check_metadata=False)

def test_empty_largelistview_array(self):

with pa.BufferOutputStream() as bos:

schema = pa.schema([pa.field("f0", pa.large_list_view(pa.int32()), True)])

with ipc.new_stream(bos, schema) as writer:

src = pa.RecordBatch.from_arrays(

[pa.array([[]], pa.large_list_view(pa.int32()))], schema=schema)

writer.write(src)

data_bytes = bos.getvalue()

def recreate_batch():

with pa.input_stream(data_bytes) as ios:

with ipc.open_stream(ios) as reader:

return reader.read_next_batch()

self.round_trip_record_batch(recreate_batch)

def test_runendencoded_array(self):

# empty vector

self.round_trip_array(lambda: pa.RunEndEncodedArray.from_arrays([], [], pa.run_end_encoded(pa.int64(), pa.int64())))

# constant null vector

self.round_trip_array(lambda: pa.RunEndEncodedArray.from_arrays([10], [None]))

# constant int vector

self.round_trip_array(lambda: pa.RunEndEncodedArray.from_arrays([10], [10]))

# run end int vector

self.round_trip_array(lambda: pa.RunEndEncodedArray.from_arrays([3, 5, 10, 12, 19], [1, 2, 1, None, 3]))

# run end string vector