# -*- coding: utf-8 -*- # 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 contextlib import ctypes import gc import pyarrow as pa try: from pyarrow.cffi import ffi except ImportError: ffi = None import pytest try: import pandas as pd import pandas.testing as tm except ImportError: pd = tm = None needs_cffi = pytest.mark.skipif(ffi is None, reason="test needs cffi package installed") assert_schema_released = pytest.raises( ValueError, match="Cannot import released ArrowSchema") assert_array_released = pytest.raises( ValueError, match="Cannot import released ArrowArray") assert_stream_released = pytest.raises( ValueError, match="Cannot import released Arrow Stream") def PyCapsule_IsValid(capsule, name): return ctypes.pythonapi.PyCapsule_IsValid(ctypes.py_object(capsule), name) == 1 @contextlib.contextmanager def registered_extension_type(ext_type): pa.register_extension_type(ext_type) try: yield finally: pa.unregister_extension_type(ext_type.extension_name) class ParamExtType(pa.ExtensionType): def __init__(self, width): self._width = width super().__init__(pa.binary(width), "pyarrow.tests.test_cffi.ParamExtType") @property def width(self): return self._width def __arrow_ext_serialize__(self): return str(self.width).encode() @classmethod def __arrow_ext_deserialize__(cls, storage_type, serialized): width = int(serialized.decode()) return cls(width) def make_schema(): return pa.schema([('ints', pa.list_(pa.int32()))], metadata={b'key1': b'value1'}) def make_extension_schema(): return pa.schema([('ext', ParamExtType(3))], metadata={b'key1': b'value1'}) def make_extension_storage_schema(): # Should be kept in sync with make_extension_schema return pa.schema([('ext', ParamExtType(3).storage_type)], metadata={b'key1': b'value1'}) def make_batch(): return pa.record_batch([[[1], [2, 42]]], make_schema()) def make_extension_batch(): schema = make_extension_schema() ext_col = schema[0].type.wrap_array(pa.array([b"foo", b"bar"], type=pa.binary(3))) return pa.record_batch([ext_col], schema) def make_batches(): schema = make_schema() return [ pa.record_batch([[[1], [2, 42]]], schema), pa.record_batch([[None, [], [5, 6]]], schema), ] def make_serialized(schema, batches): with pa.BufferOutputStream() as sink: with pa.ipc.new_stream(sink, schema) as out: for batch in batches: out.write(batch) return sink.getvalue() @needs_cffi def test_export_import_type(): c_schema = ffi.new("struct ArrowSchema*") ptr_schema = int(ffi.cast("uintptr_t", c_schema)) gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() typ = pa.list_(pa.int32()) typ._export_to_c(ptr_schema) assert pa.total_allocated_bytes() > old_allocated # Delete and recreate C++ object from exported pointer del typ assert pa.total_allocated_bytes() > old_allocated typ_new = pa.DataType._import_from_c(ptr_schema) assert typ_new == pa.list_(pa.int32()) assert pa.total_allocated_bytes() == old_allocated # Now released with assert_schema_released: pa.DataType._import_from_c(ptr_schema) # Invalid format string pa.int32()._export_to_c(ptr_schema) bad_format = ffi.new("char[]", b"zzz") c_schema.format = bad_format with pytest.raises(ValueError, match="Invalid or unsupported format string"): pa.DataType._import_from_c(ptr_schema) # Now released with assert_schema_released: pa.DataType._import_from_c(ptr_schema) @needs_cffi def test_export_import_field(): c_schema = ffi.new("struct ArrowSchema*") ptr_schema = int(ffi.cast("uintptr_t", c_schema)) gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() field = pa.field("test", pa.list_(pa.int32()), nullable=True) field._export_to_c(ptr_schema) assert pa.total_allocated_bytes() > old_allocated # Delete and recreate C++ object from exported pointer del field assert pa.total_allocated_bytes() > old_allocated field_new = pa.Field._import_from_c(ptr_schema) assert field_new == pa.field("test", pa.list_(pa.int32()), nullable=True) assert pa.total_allocated_bytes() == old_allocated # Now released with assert_schema_released: pa.Field._import_from_c(ptr_schema) def check_export_import_array(array_type, exporter, importer): c_schema = ffi.new("struct ArrowSchema*") ptr_schema = int(ffi.cast("uintptr_t", c_schema)) c_array = ffi.new(f"struct {array_type}*") ptr_array = int(ffi.cast("uintptr_t", c_array)) gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() # Type is known up front typ = pa.list_(pa.int32()) arr = pa.array([[1], [2, 42]], type=typ) py_value = arr.to_pylist() exporter(arr, ptr_array) assert pa.total_allocated_bytes() > old_allocated # Delete recreate C++ object from exported pointer del arr arr_new = importer(ptr_array, typ) assert arr_new.to_pylist() == py_value assert arr_new.type == pa.list_(pa.int32()) assert pa.total_allocated_bytes() > old_allocated del arr_new, typ assert pa.total_allocated_bytes() == old_allocated # Now released with assert_array_released: importer(ptr_array, pa.list_(pa.int32())) # Type is exported and imported at the same time arr = pa.array([[1], [2, 42]], type=pa.list_(pa.int32())) py_value = arr.to_pylist() exporter(arr, ptr_array, ptr_schema) # Delete and recreate C++ objects from exported pointers del arr arr_new = importer(ptr_array, ptr_schema) assert arr_new.to_pylist() == py_value assert arr_new.type == pa.list_(pa.int32()) assert pa.total_allocated_bytes() > old_allocated del arr_new assert pa.total_allocated_bytes() == old_allocated # Now released with assert_schema_released: importer(ptr_array, ptr_schema) @needs_cffi def test_export_import_array(): check_export_import_array( "ArrowArray", pa.Array._export_to_c, pa.Array._import_from_c, ) @needs_cffi def test_export_import_device_array(): check_export_import_array( "ArrowDeviceArray", pa.Array._export_to_c_device, pa.Array._import_from_c_device, ) # verify exported struct c_array = ffi.new("struct ArrowDeviceArray*") ptr_array = int(ffi.cast("uintptr_t", c_array)) arr = pa.array([[1], [2, 42]], type=pa.list_(pa.int32())) arr._export_to_c_device(ptr_array) assert c_array.device_type == 1 # ARROW_DEVICE_CPU 1 assert c_array.device_id == -1 assert c_array.array.length == 2 def check_export_import_schema(schema_factory, expected_schema_factory=None): if expected_schema_factory is None: expected_schema_factory = schema_factory c_schema = ffi.new("struct ArrowSchema*") ptr_schema = int(ffi.cast("uintptr_t", c_schema)) gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() schema_factory()._export_to_c(ptr_schema) assert pa.total_allocated_bytes() > old_allocated # Delete and recreate C++ object from exported pointer schema_new = pa.Schema._import_from_c(ptr_schema) assert schema_new == expected_schema_factory() assert pa.total_allocated_bytes() == old_allocated del schema_new assert pa.total_allocated_bytes() == old_allocated # Now released with assert_schema_released: pa.Schema._import_from_c(ptr_schema) # Not a struct type pa.int32()._export_to_c(ptr_schema) with pytest.raises(ValueError, match="ArrowSchema describes non-struct type"): pa.Schema._import_from_c(ptr_schema) # Now released with assert_schema_released: pa.Schema._import_from_c(ptr_schema) @needs_cffi def test_export_import_schema(): check_export_import_schema(make_schema) @needs_cffi def test_export_import_schema_with_extension(): # Extension type is unregistered => the storage type is imported check_export_import_schema(make_extension_schema, make_extension_storage_schema) # Extension type is registered => the extension type is imported with registered_extension_type(ParamExtType(1)): check_export_import_schema(make_extension_schema) @needs_cffi def test_export_import_schema_float_pointer(): # Previous versions of the R Arrow library used to pass pointer # values as a double. c_schema = ffi.new("struct ArrowSchema*") ptr_schema = int(ffi.cast("uintptr_t", c_schema)) match = "Passing a pointer value as a float is unsafe" with pytest.warns(UserWarning, match=match): make_schema()._export_to_c(float(ptr_schema)) with pytest.warns(UserWarning, match=match): schema_new = pa.Schema._import_from_c(float(ptr_schema)) assert schema_new == make_schema() def check_export_import_batch(array_type, exporter, importer, batch_factory): c_schema = ffi.new("struct ArrowSchema*") ptr_schema = int(ffi.cast("uintptr_t", c_schema)) c_array = ffi.new(f"struct {array_type}*") ptr_array = int(ffi.cast("uintptr_t", c_array)) gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() # Schema is known up front batch = batch_factory() schema = batch.schema py_value = batch.to_pydict() exporter(batch, ptr_array) assert pa.total_allocated_bytes() > old_allocated # Delete and recreate C++ object from exported pointer del batch batch_new = importer(ptr_array, schema) assert batch_new.to_pydict() == py_value assert batch_new.schema == schema assert pa.total_allocated_bytes() > old_allocated del batch_new, schema assert pa.total_allocated_bytes() == old_allocated # Now released with assert_array_released: importer(ptr_array, make_schema()) # Type is exported and imported at the same time batch = batch_factory() py_value = batch.to_pydict() batch._export_to_c(ptr_array, ptr_schema) # Delete and recreate C++ objects from exported pointers del batch batch_new = importer(ptr_array, ptr_schema) assert batch_new.to_pydict() == py_value assert batch_new.schema == batch_factory().schema assert pa.total_allocated_bytes() > old_allocated del batch_new assert pa.total_allocated_bytes() == old_allocated # Now released with assert_schema_released: importer(ptr_array, ptr_schema) # Not a struct type pa.int32()._export_to_c(ptr_schema) batch_factory()._export_to_c(ptr_array) with pytest.raises(ValueError, match="ArrowSchema describes non-struct type"): importer(ptr_array, ptr_schema) # Now released with assert_schema_released: importer(ptr_array, ptr_schema) @needs_cffi def test_export_import_batch(): check_export_import_batch( "ArrowArray", pa.RecordBatch._export_to_c, pa.RecordBatch._import_from_c, make_batch, ) @needs_cffi def test_export_import_batch_with_extension(): with registered_extension_type(ParamExtType(1)): check_export_import_batch( "ArrowArray", pa.RecordBatch._export_to_c, pa.RecordBatch._import_from_c, make_extension_batch, ) @needs_cffi def test_export_import_device_batch(): check_export_import_batch( "ArrowDeviceArray", pa.RecordBatch._export_to_c_device, pa.RecordBatch._import_from_c_device, make_batch, ) # verify exported struct c_array = ffi.new("struct ArrowDeviceArray*") ptr_array = int(ffi.cast("uintptr_t", c_array)) batch = make_batch() batch._export_to_c_device(ptr_array) assert c_array.device_type == 1 # ARROW_DEVICE_CPU 1 assert c_array.device_id == -1 assert c_array.array.length == 2 def _export_import_batch_reader(ptr_stream, reader_factory): # Prepare input batches = make_batches() schema = batches[0].schema reader = reader_factory(schema, batches) reader._export_to_c(ptr_stream) # Delete and recreate C++ object from exported pointer del reader, batches reader_new = pa.RecordBatchReader._import_from_c(ptr_stream) assert reader_new.schema == schema got_batches = list(reader_new) del reader_new assert got_batches == make_batches() # Test read_pandas() if pd is not None: batches = make_batches() schema = batches[0].schema expected_df = pa.Table.from_batches(batches).to_pandas() reader = reader_factory(schema, batches) reader._export_to_c(ptr_stream) del reader, batches reader_new = pa.RecordBatchReader._import_from_c(ptr_stream) got_df = reader_new.read_pandas() del reader_new tm.assert_frame_equal(expected_df, got_df) def make_ipc_stream_reader(schema, batches): return pa.ipc.open_stream(make_serialized(schema, batches)) def make_py_record_batch_reader(schema, batches): return pa.RecordBatchReader.from_batches(schema, batches) @needs_cffi @pytest.mark.parametrize('reader_factory', [make_ipc_stream_reader, make_py_record_batch_reader]) def test_export_import_batch_reader(reader_factory): c_stream = ffi.new("struct ArrowArrayStream*") ptr_stream = int(ffi.cast("uintptr_t", c_stream)) gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() _export_import_batch_reader(ptr_stream, reader_factory) assert pa.total_allocated_bytes() == old_allocated # Now released with assert_stream_released: pa.RecordBatchReader._import_from_c(ptr_stream) @needs_cffi def test_export_import_exception_reader(): # See: https://github.com/apache/arrow/issues/37164 c_stream = ffi.new("struct ArrowArrayStream*") ptr_stream = int(ffi.cast("uintptr_t", c_stream)) gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() def gen(): if True: try: raise ValueError('foo') except ValueError as e: raise NotImplementedError('bar') from e else: yield from make_batches() original = pa.RecordBatchReader.from_batches(make_schema(), gen()) original._export_to_c(ptr_stream) reader = pa.RecordBatchReader._import_from_c(ptr_stream) with pytest.raises(NotImplementedError) as exc_info: reader.read_next_batch() # inner *and* outer exception should be present assert 'ValueError: foo' in str(exc_info.value) assert 'NotImplementedError: bar' in str(exc_info.value) # Stacktrace containing line of the raise statement assert 'raise ValueError(\'foo\')' in str(exc_info.value) assert pa.total_allocated_bytes() == old_allocated @needs_cffi def test_imported_batch_reader_error(): c_stream = ffi.new("struct ArrowArrayStream*") ptr_stream = int(ffi.cast("uintptr_t", c_stream)) schema = pa.schema([('foo', pa.int32())]) batches = [pa.record_batch([[1, 2, 3]], schema=schema), pa.record_batch([[4, 5, 6]], schema=schema)] buf = make_serialized(schema, batches) # Open a corrupt/incomplete stream and export it reader = pa.ipc.open_stream(buf[:-16]) reader._export_to_c(ptr_stream) del reader reader_new = pa.RecordBatchReader._import_from_c(ptr_stream) batch = reader_new.read_next_batch() assert batch == batches[0] with pytest.raises(OSError, match="Expected to be able to read 16 bytes " "for message body, got 8"): reader_new.read_next_batch() # Again, but call read_all() reader = pa.ipc.open_stream(buf[:-16]) reader._export_to_c(ptr_stream) del reader reader_new = pa.RecordBatchReader._import_from_c(ptr_stream) with pytest.raises(OSError, match="Expected to be able to read 16 bytes " "for message body, got 8"): reader_new.read_all() @pytest.mark.parametrize('obj', [pa.int32(), pa.field('foo', pa.int32()), pa.schema({'foo': pa.int32()})], ids=['type', 'field', 'schema']) def test_roundtrip_schema_capsule(obj): gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() capsule = obj.__arrow_c_schema__() assert PyCapsule_IsValid(capsule, b"arrow_schema") == 1 assert pa.total_allocated_bytes() > old_allocated obj_out = type(obj)._import_from_c_capsule(capsule) assert obj_out == obj assert pa.total_allocated_bytes() == old_allocated capsule = obj.__arrow_c_schema__() assert pa.total_allocated_bytes() > old_allocated del capsule assert pa.total_allocated_bytes() == old_allocated @pytest.mark.parametrize('arr,schema_accessor,bad_type,good_type', [ (pa.array(['a', 'b', 'c']), lambda x: x.type, pa.int32(), pa.string()), ( pa.record_batch([pa.array(['a', 'b', 'c'])], names=['x']), lambda x: x.schema, pa.schema({'x': pa.int32()}), pa.schema({'x': pa.string()}) ), ], ids=['array', 'record_batch']) def test_roundtrip_array_capsule(arr, schema_accessor, bad_type, good_type): gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() import_array = type(arr)._import_from_c_capsule schema_capsule, capsule = arr.__arrow_c_array__() assert PyCapsule_IsValid(schema_capsule, b"arrow_schema") == 1 assert PyCapsule_IsValid(capsule, b"arrow_array") == 1 arr_out = import_array(schema_capsule, capsule) assert arr_out.equals(arr) assert pa.total_allocated_bytes() > old_allocated del arr_out assert pa.total_allocated_bytes() == old_allocated capsule = arr.__arrow_c_array__() assert pa.total_allocated_bytes() > old_allocated del capsule assert pa.total_allocated_bytes() == old_allocated with pytest.raises(ValueError, match=r"Could not cast.* string to requested .* int32"): arr.__arrow_c_array__(bad_type.__arrow_c_schema__()) schema_capsule, array_capsule = arr.__arrow_c_array__( good_type.__arrow_c_schema__()) arr_out = import_array(schema_capsule, array_capsule) assert schema_accessor(arr_out) == good_type @pytest.mark.parametrize('arr,schema_accessor,bad_type,good_type', [ (pa.array(['a', 'b', 'c']), lambda x: x.type, pa.int32(), pa.string()), ( pa.record_batch([pa.array(['a', 'b', 'c'])], names=['x']), lambda x: x.schema, pa.schema({'x': pa.int32()}), pa.schema({'x': pa.string()}) ), ], ids=['array', 'record_batch']) def test_roundtrip_device_array_capsule(arr, schema_accessor, bad_type, good_type): gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() import_array = type(arr)._import_from_c_device_capsule schema_capsule, capsule = arr.__arrow_c_device_array__() assert PyCapsule_IsValid(schema_capsule, b"arrow_schema") == 1 assert PyCapsule_IsValid(capsule, b"arrow_device_array") == 1 arr_out = import_array(schema_capsule, capsule) assert arr_out.equals(arr) assert pa.total_allocated_bytes() > old_allocated del arr_out assert pa.total_allocated_bytes() == old_allocated capsule = arr.__arrow_c_array__() assert pa.total_allocated_bytes() > old_allocated del capsule assert pa.total_allocated_bytes() == old_allocated with pytest.raises(ValueError, match=r"Could not cast.* string to requested .* int32"): arr.__arrow_c_device_array__(bad_type.__arrow_c_schema__()) schema_capsule, array_capsule = arr.__arrow_c_device_array__( good_type.__arrow_c_schema__()) arr_out = import_array(schema_capsule, array_capsule) assert schema_accessor(arr_out) == good_type # TODO: implement requested_schema for stream @pytest.mark.parametrize('constructor', [ pa.RecordBatchReader.from_batches, # Use a lambda because we need to re-order the parameters lambda schema, batches: pa.Table.from_batches(batches, schema), ], ids=['recordbatchreader', 'table']) def test_roundtrip_reader_capsule(constructor): batches = make_batches() schema = batches[0].schema gc.collect() # Make sure no Arrow data dangles in a ref cycle old_allocated = pa.total_allocated_bytes() obj = constructor(schema, batches) capsule = obj.__arrow_c_stream__() assert PyCapsule_IsValid(capsule, b"arrow_array_stream") == 1 imported_reader = pa.RecordBatchReader._import_from_c_capsule(capsule) assert imported_reader.schema == schema imported_batches = list(imported_reader) assert len(imported_batches) == len(batches) for batch, expected in zip(imported_batches, batches): assert batch.equals(expected) del obj, imported_reader, batch, expected, imported_batches assert pa.total_allocated_bytes() == old_allocated obj = constructor(schema, batches) bad_schema = pa.schema({'ints': pa.int32()}) with pytest.raises(pa.lib.ArrowTypeError, match="Field 0 cannot be cast"): obj.__arrow_c_stream__(bad_schema.__arrow_c_schema__()) # Can work with matching schema matching_schema = pa.schema({'ints': pa.list_(pa.int32())}) capsule = obj.__arrow_c_stream__(matching_schema.__arrow_c_schema__()) imported_reader = pa.RecordBatchReader._import_from_c_capsule(capsule) assert imported_reader.schema == matching_schema for batch, expected in zip(imported_reader, batches): assert batch.equals(expected) def test_roundtrip_batch_reader_capsule_requested_schema(): batch = make_batch() requested_schema = pa.schema([('ints', pa.list_(pa.int64()))]) requested_capsule = requested_schema.__arrow_c_schema__() batch_as_requested = batch.cast(requested_schema) capsule = batch.__arrow_c_stream__(requested_capsule) assert PyCapsule_IsValid(capsule, b"arrow_array_stream") == 1 imported_reader = pa.RecordBatchReader._import_from_c_capsule(capsule) assert imported_reader.schema == requested_schema assert imported_reader.read_next_batch().equals(batch_as_requested) with pytest.raises(StopIteration): imported_reader.read_next_batch() def test_roundtrip_batch_reader_capsule(): batch = make_batch() capsule = batch.__arrow_c_stream__() assert PyCapsule_IsValid(capsule, b"arrow_array_stream") == 1 imported_reader = pa.RecordBatchReader._import_from_c_capsule(capsule) assert imported_reader.schema == batch.schema assert imported_reader.read_next_batch().equals(batch) with pytest.raises(StopIteration): imported_reader.read_next_batch() def test_roundtrip_chunked_array_capsule(): chunked = pa.chunked_array([pa.array(["a", "b", "c"])]) capsule = chunked.__arrow_c_stream__() assert PyCapsule_IsValid(capsule, b"arrow_array_stream") == 1 imported_chunked = pa.ChunkedArray._import_from_c_capsule(capsule) assert imported_chunked.type == chunked.type assert imported_chunked == chunked def test_roundtrip_chunked_array_capsule_requested_schema(): chunked = pa.chunked_array([pa.array(["a", "b", "c"])]) # Requesting the same type should work requested_capsule = chunked.type.__arrow_c_schema__() capsule = chunked.__arrow_c_stream__(requested_capsule) imported_chunked = pa.ChunkedArray._import_from_c_capsule(capsule) assert imported_chunked == chunked # Casting to something else should error if not possible requested_type = pa.binary() requested_capsule = requested_type.__arrow_c_schema__() capsule = chunked.__arrow_c_stream__(requested_capsule) imported_chunked = pa.ChunkedArray._import_from_c_capsule(capsule) assert imported_chunked == chunked.cast(pa.binary()) requested_type = pa.int64() requested_capsule = requested_type.__arrow_c_schema__() with pytest.raises( ValueError, match="Could not cast string to requested type int64" ): chunked.__arrow_c_stream__(requested_capsule) @needs_cffi def test_import_device_no_cuda(): try: import pyarrow.cuda # noqa except ImportError: pass else: pytest.skip("pyarrow.cuda is available") c_array = ffi.new("struct ArrowDeviceArray*") ptr_array = int(ffi.cast("uintptr_t", c_array)) arr = pa.array([1, 2, 3], type=pa.int64()) arr._export_to_c_device(ptr_array) # patch the device type of the struct, this results in an invalid ArrowDeviceArray # but this is just to test we raise am error before actually importing buffers c_array.device_type = 2 # ARROW_DEVICE_CUDA with pytest.raises(ImportError, match="Trying to import data on a CUDA device"): pa.Array._import_from_c_device(ptr_array, arr.type)