refactor: simplify PyCapsule interface using pyo3-arrow (#364)

* Simplify pycapsule interface using pyo3-arrow

* docs: improve docstring

* chore: remove useless clones

---------

Co-authored-by: Eric Jolibois <em.jolibois@gmail.com>

Author: kylebarron

Cargo.lock

@@ -32,6 +32,7 @@ calamine = { version = "^0.30.0", features = ["dates"] }
 chrono = { version = "^0.4.41", default-features = false }
 log = "0.4.27"
 pyo3 = { version = "^0.25", features = ["abi3-py39"] }
+pyo3-arrow = { version = "0.11", default-features = false }
 pyo3-log = "^0.12.4"
 
 [dev-dependencies]

Cargo.toml

@@ -142,18 +142,24 @@ def to_polars(self) -> "pl.DataFrame":
         return pl.DataFrame(self)
 
     def __arrow_c_schema__(self) -> object:
-        """Arrow PyCapsule Interface: Export schema as a PyCapsule.
+        """Export the schema as an `ArrowSchema` `PyCapsule`.
 
-        This method allows zero-copy data exchange with Arrow-compatible libraries
-        like Polars without requiring PyArrow as a dependency.
+        https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowschema-export
+
+        The Arrow PyCapsule Interface enables zero-copy data exchange with
+        Arrow-compatible libraries without requiring PyArrow as a dependency.
         """
         return self._sheet.__arrow_c_schema__()
 
     def __arrow_c_array__(self, requested_schema: object | None = None) -> tuple[object, object]:
-        """Arrow PyCapsule Interface: Export array and schema as PyCapsules.
+        """Export the schema and data as a pair of `ArrowSchema` and `ArrowArray` `PyCapsules`.
+
+        The optional `requested_schema` parameter allows for potential schema conversion.
+
+        https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowarray-export
 
-        Returns data as a tuple of (schema_capsule, array_capsule) for zero-copy
-        data exchange with Arrow-compatible libraries.
+        The Arrow PyCapsule Interface enables zero-copy data exchange with
+        Arrow-compatible libraries without requiring PyArrow as a dependency.
         """
         return self._sheet.__arrow_c_array__(requested_schema)
 
@@ -240,18 +246,24 @@ def to_polars(self) -> "pl.DataFrame":
         return pl.DataFrame(self)
 
     def __arrow_c_schema__(self) -> object:
-        """Arrow PyCapsule Interface: Export schema as a PyCapsule.
+        """Export the schema as an `ArrowSchema` `PyCapsule`.
 
-        This method allows zero-copy data exchange with Arrow-compatible libraries
-        like Polars without requiring PyArrow as a dependency.
+        https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowschema-export
+
+        The Arrow PyCapsule Interface enables zero-copy data exchange with
+        Arrow-compatible libraries without requiring PyArrow as a dependency.
         """
         return self._table.__arrow_c_schema__()
 
     def __arrow_c_array__(self, requested_schema: object | None = None) -> tuple[object, object]:
-        """Arrow PyCapsule Interface: Export array and schema as PyCapsules.
+        """Export the schema and data as a pair of `ArrowSchema` and `ArrowArray` `PyCapsules`.
+
+        The optional `requested_schema` parameter allows for potential schema conversion.
+
+        https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowarray-export
 
-        Returns data as a tuple of (schema_capsule, array_capsule) for zero-copy
-        data exchange with Arrow-compatible libraries.
+        The Arrow PyCapsule Interface enables zero-copy data exchange with
+        Arrow-compatible libraries without requiring PyArrow as a dependency.
         """
         return self._table.__arrow_c_array__(requested_schema)
 

python/fastexcel/__init__.py

@@ -92,9 +92,23 @@ class _ExcelSheet:
         therefore converted to None.
         """
     def __arrow_c_schema__(self) -> object:
-        """Arrow PyCapsule Interface: Export schema as a PyCapsule"""
+        """Export the schema as an `ArrowSchema` `PyCapsule`.
+
+        https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowschema-export
+
+        The Arrow PyCapsule Interface enables zero-copy data exchange with
+        Arrow-compatible libraries without requiring PyArrow as a dependency.
+        """
     def __arrow_c_array__(self, requested_schema: object = None) -> tuple[object, object]:
-        """Arrow PyCapsule Interface: Export array and schema as PyCapsules"""
+        """Export the schema and data as a pair of `ArrowSchema` and `ArrowArray` `PyCapsules`.
+
+        The optional `requested_schema` parameter allows for potential schema conversion.
+
+        https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowarray-export
+
+        The Arrow PyCapsule Interface enables zero-copy data exchange with
+        Arrow-compatible libraries without requiring PyArrow as a dependency.
+        """
 
 class _ExcelTable:
     @property
@@ -126,9 +140,24 @@ class _ExcelTable:
     def to_arrow(self) -> "pa.RecordBatch":
         """Converts the table to a pyarrow `RecordBatch`"""
     def __arrow_c_schema__(self) -> object:
-        """Arrow PyCapsule Interface: Export schema as a PyCapsule"""
+        """Export the schema as an `ArrowSchema` `PyCapsule`.
+
+        https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowschema-export
+
+        The Arrow PyCapsule Interface enables zero-copy data exchange with
+        Arrow-compatible libraries without requiring PyArrow as a dependency.
+        """
+
     def __arrow_c_array__(self, requested_schema: object = None) -> tuple[object, object]:
-        """Arrow PyCapsule Interface: Export array and schema as PyCapsules"""
+        """Export the schema and data as a pair of `ArrowSchema` and `ArrowArray` `PyCapsules`.
+
+        The optional `requested_schema` parameter allows for potential schema conversion.
+
+        https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowarray-export
+
+        The Arrow PyCapsule Interface enables zero-copy data exchange with
+        Arrow-compatible libraries without requiring PyArrow as a dependency.
+        """
 
 class _ExcelReader:
     """A class representing an open Excel file and allowing to read its sheets"""

python/fastexcel/_fastexcel.pyi

@@ -1,4 +1,3 @@
-mod arrow_capsule;
 mod data;
 mod error;
 mod types;

src/arrow_capsule.rs

@@ -1,11 +1,15 @@
 pub(crate) mod column_info;
 pub(crate) mod table;
 
+use arrow_schema::Field;
 use calamine::{CellType, Range, Sheet as CalamineSheet, SheetVisible as CalamineSheetVisible};
 use column_info::{AvailableColumns, ColumnInfoNoDtype};
+use pyo3::types::{PyCapsule, PyTuple};
+use pyo3_arrow::ffi::{to_array_pycapsules, to_schema_pycapsule};
+use std::sync::Arc;
 use std::{cmp, collections::HashSet, fmt::Debug, str::FromStr};
 
-use arrow_array::RecordBatch;
+use arrow_array::{RecordBatch, StructArray};
 #[cfg(feature = "pyarrow")]
 use arrow_pyarrow::ToPyArrow;
 
@@ -15,7 +19,7 @@ use pyo3::{
     types::{PyList, PyString},
 };
 
-use crate::arrow_capsule;
+use crate::data::selected_columns_to_schema;
 use crate::{
     data::{
         ExcelSheetData, record_batch_from_data_and_columns,
@@ -647,29 +651,31 @@ impl ExcelSheet {
         (rb, errors).into_bound_py_any(py)
     }
 
-    /// Arrow PyCapsule Interface: __arrow_c_schema__
-    pub fn __arrow_c_schema__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        let record_batch = RecordBatch::try_from(self)
-            .with_context(|| {
-                format!(
-                    "could not create RecordBatch from sheet \"{}\"",
-                    self.name()
-                )
-            })
-            .into_pyresult()?;
-
-        arrow_capsule::schema_to_pycapsule(py, record_batch.schema().as_ref())
-            .map(|capsule| capsule.into_any())
-    }
-
-    /// Arrow PyCapsule Interface: __arrow_c_array__
+    /// Export the schema as an [`ArrowSchema`] [`PyCapsule`].
+    ///
+    /// <https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowschema-export>
+    ///
+    /// [`ArrowSchema`]: arrow_array::ffi::FFI_ArrowSchema
+    /// [`PyCapsule`]: pyo3::types::PyCapsule
+    pub fn __arrow_c_schema__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyCapsule>> {
+        let schema = selected_columns_to_schema(&self.selected_columns);
+        Ok(to_schema_pycapsule(py, &schema)?)
+    }
+
+    /// Export the schema and data as a pair of [`ArrowSchema`] and [`ArrowArray`] [`PyCapsules`]
+    ///
+    /// The optional `requested_schema` parameter allows for potential schema conversion.
+    ///
+    /// <https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowarray-export>
+    ///
+    /// [`ArrowSchema`]: arrow_array::ffi::FFI_ArrowSchema
+    /// [`ArrowArray`]: arrow_array::ffi::FFI_ArrowArray
+    /// [`PyCapsules`]: pyo3::types::PyCapsule
     pub fn __arrow_c_array__<'py>(
         &self,
         py: Python<'py>,
-        requested_schema: Option<&Bound<'py, PyAny>>,
-    ) -> PyResult<Bound<'py, PyAny>> {
-        let _ = requested_schema; // TODO: Support schema conversion if needed
-
+        requested_schema: Option<Bound<'py, PyCapsule>>,
+    ) -> PyResult<Bound<'py, PyTuple>> {
         let record_batch = RecordBatch::try_from(self)
             .with_context(|| {
                 format!(
@@ -679,13 +685,14 @@ impl ExcelSheet {
             })
             .into_pyresult()?;
 
-        let (schema_capsule, array_capsule) =
-            arrow_capsule::record_batch_to_pycapsules(py, &record_batch)?;
-
-        Ok(
-            pyo3::types::PyTuple::new(py, [schema_capsule.into_any(), array_capsule.into_any()])?
-                .into_any(),
-        )
+        let field = Field::new_struct("", record_batch.schema_ref().fields().clone(), false);
+        let array = Arc::new(StructArray::from(record_batch));
+        Ok(to_array_pycapsules(
+            py,
+            field.into(),
+            array.as_ref(),
+            requested_schema,
+        )?)
     }
 
     pub fn __repr__(&self) -> String {