docs: update Arrow C Array export documentation to clarify return values

Author: kosiew

docs/source/user-guide/io/arrow.rst

@@ -35,7 +35,12 @@ To import an Arrow table, use :py:func:`datafusion.context.SessionContext.from_a
 This will accept any Python object that implements
 `__arrow_c_stream__ <https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowstream-export>`_
 or `__arrow_c_array__ <https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowarray-export>`_
-and returns a ``StructArray``. Common pyarrow sources you can use are:
+and returns a ``StructArray``.
+
+.. note::
+``__arrow_c_array__`` must return ``(schema_capsule, array_capsule)``.
+
+Common pyarrow sources you can use are:
 
 - `Array <https://arrow.apache.org/docs/python/generated/pyarrow.Array.html>`_ (but it must return a Struct Array)
 - `Record Batch <https://arrow.apache.org/docs/python/generated/pyarrow.RecordBatch.html>`_

python/datafusion/__init__.py

@@ -30,7 +30,12 @@
 except ImportError:
     import importlib_metadata
 
-from . import functions, object_store, substrait, unparser
+from . import functions, object_store, unparser
+
+try:  # pragma: no cover - optional dependency
+    from . import substrait
+except ImportError:  # pragma: no cover - substrait not built
+    substrait = None
 
 # The following imports are okay to remain as opaque to the user.
 from ._internal import Config

python/datafusion/context.py

@@ -64,6 +64,8 @@ class ArrowArrayExportable(Protocol):
     """Type hint for object exporting Arrow C Array via Arrow PyCapsule Interface.
 
     https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
+
+    The method should return ``(schema_capsule, array_capsule)``.
     """
 
     def __arrow_c_array__(  # noqa: D105

python/datafusion/record_batch.py

@@ -57,8 +57,15 @@ def __arrow__(self, *args: object, **kwargs: object) -> pa.RecordBatch:
     def __arrow_c_array__(
         self, requested_schema: object | None = None
     ) -> tuple[object, object]:
-        """Arrow C Data Interface export."""
-        return self.record_batch.__arrow_c_array__(requested_schema)
+        """Arrow C Data Interface export.
+
+        Returns:
+            tuple[object, object]: ``(schema_capsule, array_capsule)``
+        """
+        schema_capsule, array_capsule = self.record_batch.__arrow_c_array__(
+            requested_schema
+        )
+        return schema_capsule, array_capsule
 
 
 class RecordBatchStream:

src/record_batch.rs

@@ -98,7 +98,7 @@ impl PyRecordBatch {
         let data = struct_array.to_data();
         let array = FFI_ArrowArray::new(&data);
         let schema =
-            FFI_ArrowSchema::try_from(data.data_type()).map_err(PyDataFusionError::from)?;
+            FFI_ArrowSchema::try_from(self.batch.schema().as_ref()).map_err(PyDataFusionError::from)?;
 
         let array_ptr = Box::into_raw(Box::new(array));
         let schema_ptr = Box::into_raw(Box::new(schema));