Implement Arrow PyCapsule Interface

.github/workflows/integration.yml

@@ -128,7 +128,7 @@ jobs:
  key: ${{ runner.os }}-${{ matrix.arch }}-target-maturin-cache-${{ matrix.rust }}-
  - uses: actions/setup-python@v4
  with:
- python-version: '3.7'
+ python-version: '3.8'
  - name: Upgrade pip and setuptools
  run: pip install --upgrade pip setuptools wheel virtualenv
  - name: Create virtualenv and install dependencies

arrow-pyarrow-integration-testing/README.md

@@ -25,6 +25,7 @@ Note that this crate uses two languages and an external ABI:
 * `Rust`
 * `Python`
 * C ABI privately exposed by `Pyarrow`.
+* PyCapsule ABI publicly exposed by `pyarrow`
 
 ## Basic idea
 
@@ -36,6 +37,7 @@ we can use pyarrow's interface to move pointers from and to Rust.
 ## Relevant literature
 
 * [Arrow's CDataInterface](https://arrow.apache.org/docs/format/CDataInterface.html)
+* [Arrow PyCapsule Interface](https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html)
 * [Rust's FFI](https://doc.rust-lang.org/nomicon/ffi.html)
 * [Pyarrow private binds](https://github.com/apache/arrow/blob/ae1d24efcc3f1ac2a876d8d9f544a34eb04ae874/python/pyarrow/array.pxi#L1226)
 * [PyO3](https://docs.rs/pyo3/0.12.1/pyo3/index.html)

arrow-pyarrow-integration-testing/tests/test_sql.py

@@ -27,6 +27,8 @@
 
 import arrow_pyarrow_integration_testing as rust
 
+PYARROW_PRE_14 = int(pa.__version__.split('.')[0]) < 14
+
 
 @contextlib.contextmanager
 def no_pyarrow_leak():
@@ -113,13 +115,49 @@ def assert_pyarrow_leak():
 _unsupported_pyarrow_types = [
 ]
 
+# As of pyarrow 14, pyarrow implements the Arrow PyCapsule interface
+# (https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html).
+# This defines that Arrow consumers should allow any object that has specific "dunder"
+# methods, `__arrow_c_*_`. These wrapper classes ensure that arrow-rs is able to handle
+# _any_ class, without pyarrow-specific handling.
+class SchemaWrapper:
+ def __init__(self, schema):
+ self.schema = schema
+
+ def __arrow_c_schema__(self):
+ return self.schema.__arrow_c_schema__()
+
+
+class ArrayWrapper:
+ def __init__(self, array):
+ self.array = array
+
+ def __arrow_c_array__(self):
+ return self.array.__arrow_c_array__()
+
+
+class StreamWrapper:
+ def __init__(self, stream):
+ self.stream = stream
+
+ def __arrow_c_stream__(self):
+ return self.stream.__arrow_c_stream__()
+
 
 @pytest.mark.parametrize("pyarrow_type", _supported_pyarrow_types, ids=str)
 def test_type_roundtrip(pyarrow_type):
  restored = rust.round_trip_type(pyarrow_type)
  assert restored == pyarrow_type
  assert restored is not pyarrow_type
 
+@pytest.mark.skipif(PYARROW_PRE_14, reason="requires pyarrow 14")
+@pytest.mark.parametrize("pyarrow_type", _supported_pyarrow_types, ids=str)
+def test_type_roundtrip_pycapsule(pyarrow_type):
+ wrapped = SchemaWrapper(pyarrow_type)
+ restored = rust.round_trip_type(wrapped)
+ assert restored == pyarrow_type
+ assert restored is not pyarrow_type
+
 
 @pytest.mark.parametrize("pyarrow_type", _unsupported_pyarrow_types, ids=str)
 def test_type_roundtrip_raises(pyarrow_type):
@@ -138,6 +176,20 @@ def test_field_roundtrip(pyarrow_type):
  field = rust.round_trip_field(pyarrow_field)
  assert field == pyarrow_field
 
+@pytest.mark.skipif(PYARROW_PRE_14, reason="requires pyarrow 14")
+@pytest.mark.parametrize('pyarrow_type', _supported_pyarrow_types, ids=str)
+def test_field_roundtrip_pycapsule(pyarrow_type):
+ pyarrow_field = pa.field("test", pyarrow_type, nullable=True)
+ wrapped = SchemaWrapper(pyarrow_field)
+ field = rust.round_trip_field(wrapped)
+ assert field == wrapped.schema
+
+ if pyarrow_type != pa.null():
+ # A null type field may not be non-nullable
+ pyarrow_field = pa.field("test", pyarrow_type, nullable=False)
+ field = rust.round_trip_field(wrapped)
+ assert field == wrapped.schema
+
 def test_field_metadata_roundtrip():
  metadata = {"hello": "World! 😊", "x": "2"}
  pyarrow_field = pa.field("test", pa.int32(), metadata=metadata)
@@ -163,6 +215,17 @@ def test_primitive_python():
  del b
 
 
+@pytest.mark.skipif(PYARROW_PRE_14, reason="requires pyarrow 14")
+def test_primitive_python_pycapsule():
+ """
+ Python -> Rust -> Python
+ """
+ a = pa.array([1, 2, 3])
+ wrapped = ArrayWrapper(a)
+ b = rust.double(wrapped)
+ assert b == pa.array([2, 4, 6])
+
+
 def test_primitive_rust():
  """
  Rust -> Python -> Rust
@@ -433,6 +496,33 @@ def test_record_batch_reader():
  got_batches = list(b)
  assert got_batches == batches
 
+@pytest.mark.skipif(PYARROW_PRE_14, reason="requires pyarrow 14")
+def test_record_batch_reader_pycapsule():
+ """
+ Python -> Rust -> Python
+ """
+ schema = pa.schema([('ints', pa.list_(pa.int32()))], metadata={b'key1': b'value1'})
+ batches = [
+ pa.record_batch([[[1], [2, 42]]], schema),
+ pa.record_batch([[None, [], [5, 6]]], schema),
+ ]
+ a = pa.RecordBatchReader.from_batches(schema, batches)
+ wrapped = StreamWrapper(a)
+ b = rust.round_trip_record_batch_reader(wrapped)
+
+ assert b.schema == schema
+ got_batches = list(b)
+ assert got_batches == batches
+
+ # Also try the boxed reader variant
+ a = pa.RecordBatchReader.from_batches(schema, batches)
+ wrapped = StreamWrapper(a)
+ b = rust.boxed_reader_roundtrip(wrapped)
+ assert b.schema == schema
+ got_batches = list(b)
+ assert got_batches == batches
+
+
 def test_record_batch_reader_error():
  schema = pa.schema([('ints', pa.list_(pa.int32()))])
 
@@ -453,24 +543,64 @@ def iter_batches():
  with pytest.raises(ValueError, match="invalid utf-8"):
  rust.round_trip_record_batch_reader(reader)
 
+
+@pytest.mark.skipif(PYARROW_PRE_14, reason="requires pyarrow 14")
+def test_record_batch_pycapsule():
+ """
+ Python -> Rust -> Python
+ """
+ schema = pa.schema([('ints', pa.list_(pa.int32()))], metadata={b'key1': b'value1'})
+ batch = pa.record_batch([[[1], [2, 42]]], schema)
+ wrapped = StreamWrapper(batch)
+ b = rust.round_trip_record_batch_reader(wrapped)
+ new_table = b.read_all()
+ new_batches = new_table.to_batches()
+
+ assert len(new_batches) == 1
+ new_batch = new_batches[0]
+
+ assert batch == new_batch
+ assert batch.schema == new_batch.schema
+
+
+@pytest.mark.skipif(PYARROW_PRE_14, reason="requires pyarrow 14")
+def test_table_pycapsule():
+ """
+ Python -> Rust -> Python
+ """
+ schema = pa.schema([('ints', pa.list_(pa.int32()))], metadata={b'key1': b'value1'})
+ batches = [
+ pa.record_batch([[[1], [2, 42]]], schema),
+ pa.record_batch([[None, [], [5, 6]]], schema),
+ ]
+ table = pa.Table.from_batches(batches)
+ wrapped = StreamWrapper(table)
+ b = rust.round_trip_record_batch_reader(wrapped)
+ new_table = b.read_all()
+
+ assert table.schema == new_table.schema
+ assert table == new_table
+ assert len(table.to_batches()) == len(new_table.to_batches())
+
+
 def test_reject_other_classes():
  # Arbitrary type that is not a PyArrow type
  not_pyarrow = ["hello"]
 
  with pytest.raises(TypeError, match="Expected instance of pyarrow.lib.Array, got builtins.list"):
  rust.round_trip_array(not_pyarrow)
- 
+
  with pytest.raises(TypeError, match="Expected instance of pyarrow.lib.Schema, got builtins.list"):
  rust.round_trip_schema(not_pyarrow)
- 
+
  with pytest.raises(TypeError, match="Expected instance of pyarrow.lib.Field, got builtins.list"):
  rust.round_trip_field(not_pyarrow)
- 
+
  with pytest.raises(TypeError, match="Expected instance of pyarrow.lib.DataType, got builtins.list"):
  rust.round_trip_type(not_pyarrow)
 
  with pytest.raises(TypeError, match="Expected instance of pyarrow.lib.RecordBatch, got builtins.list"):
  rust.round_trip_record_batch(not_pyarrow)
- 
+
  with pytest.raises(TypeError, match="Expected instance of pyarrow.lib.RecordBatchReader, got builtins.list"):
  rust.round_trip_record_batch_reader(not_pyarrow)

arrow-schema/src/ffi.rs

@@ -351,6 +351,8 @@ impl Drop for FFI_ArrowSchema {
  }
 }
 
+unsafe impl Send for FFI_ArrowSchema {}
+
 impl TryFrom<&FFI_ArrowSchema> for DataType {
  type Error = ArrowError;