| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an | |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
| # KIND, either express or implied. See the License for the | |
| # specific language governing permissions and limitations | |
| # under the License. | |
| import decimal | |
| import gc | |
| import os | |
| import sys | |
| import unittest | |
| import xml.etree.ElementTree as ET | |
| import jpype | |
| import pyarrow as pa | |
| import pyarrow.ipc as ipc | |
| from pyarrow.cffi import ffi | |
| def setup_jvm(): | |
| # 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" | |
| # ] | |
| jpype.startJVM(jpype.getDefaultJVMPath(), "-Djava.class.path=" + jar_path, **kwargs) | |
| class Bridge: | |
| 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): | |
| self.java_allocator.close() | |
| class TestPythonIntegration(unittest.TestCase): | |
| @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. | |
| Empty list of int32 produces a vector with empty child data buffer, however with non-zero capacity. | |
| Using streaming forces the c-data array which represent the child data buffer to be NULL (pointer is 0). | |
| On Java side, an attempt to import such array triggered an exception described in GH-37056. | |
| """ | |
| 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 | |
| self.round_trip_array(lambda: pa.RunEndEncodedArray.from_arrays([3, 5, 10, 12, 19], ["1", "2", "1", None, "3"])) | |
| if __name__ == '__main__': | |
| unittest.main(verbosity=2) |