feat: enhance Arrow interoperability in RecordBatch and improve capsule naming

Author: kosiew

python/datafusion/record_batch.py

@@ -46,13 +46,8 @@ def to_pyarrow(self) -> pa.RecordBatch:
         """Convert to :py:class:`pa.RecordBatch`."""
         return self.record_batch.to_pyarrow()
 
-    def __arrow__(self, *args: object, **kwargs: object) -> pa.RecordBatch:
-        """Return a :py:class:`pa.RecordBatch` for Arrow interoperability.
-
-        This enables ``datafusion.record_batch.RecordBatch`` instances to be
-        automatically recognized by PyArrow when passed to its APIs.
-        """
-        return self.to_pyarrow()
+    def __arrow__(self, type: object | None = None) -> pa.RecordBatch:  # noqa: D105
+        return self.record_batch.to_pyarrow()
 
     def __arrow_c_array__(
         self, requested_schema: object | None = None
@@ -65,6 +60,29 @@ def __arrow_c_array__(
         schema_capsule, array_capsule = self.record_batch.__arrow_c_array__(
             requested_schema
         )
+
+        # Ensure the returned capsules are named as expected by PyArrow. The
+        # Rust implementation already produces properly named capsules, but some
+        # Python consumers (including PyArrow itself) are strict about the
+        # capsule names matching ``"arrow_schema"`` and ``"arrow_array"``.
+        import ctypes
+
+        def _ensure_name(capsule: object, name: bytes) -> None:
+            pythonapi = ctypes.pythonapi
+            get_name = pythonapi.PyCapsule_GetName
+            get_name.restype = ctypes.c_char_p
+            get_name.argtypes = [ctypes.py_object]
+
+            set_name = pythonapi.PyCapsule_SetName
+            set_name.restype = ctypes.c_int
+            set_name.argtypes = [ctypes.py_object, ctypes.c_char_p]
+
+            if get_name(capsule) != name:
+                set_name(capsule, name)
+
+        _ensure_name(schema_capsule, b"arrow_schema")
+        _ensure_name(array_capsule, b"arrow_array")
+
         return schema_capsule, array_capsule
 
 

python/tests/test_arrow_interop.py

@@ -14,3 +14,19 @@ def test_table_from_batches_with_dataframe(ctx):
 
     expected = pa.Table.from_batches([batch1, batch2])
     assert table.equals(expected)
+
+
+def test_table_from_batches_with_record_batch(ctx):
+    batch = pa.record_batch({"a": pa.array([1, 2]), "b": pa.array(["x", "y"])})
+    df = ctx.create_dataframe([[batch]])
+    rb = df.collect()[0]
+
+    try:
+        table = pa.Table.from_batches([rb])
+    except TypeError as err:  # pragma: no cover - failure path
+        pytest.fail(
+            f"TypeError raised when converting RecordBatch to Arrow Table: {err}"
+        )
+
+    expected = pa.Table.from_batches([batch])
+    assert table.equals(expected)