pyarrow: Wrap capsule extraction errors into an explicit error

arrow-pyarrow-testing/Cargo.toml

@@ -47,5 +47,6 @@ publish = false
 [dependencies]
 # Note no dependency on arrow, to ensure arrow-pyarrow can be used by itself
 arrow-array = { path = "../arrow-array" }
+arrow-schema = { path = "../arrow-schema" }
 arrow-pyarrow = { path = "../arrow-pyarrow" }
 pyo3 = { version = "0.29.0", default-features = false }

arrow-pyarrow-testing/tests/pyarrow.rs

@@ -41,6 +41,7 @@ use arrow_array::builder::{BinaryViewBuilder, StringViewBuilder};
 use arrow_array::{
     Array, ArrayRef, BinaryViewArray, Int32Array, RecordBatch, StringArray, StringViewArray,
 };
+use arrow_schema::Schema;
 use arrow_pyarrow::{FromPyArrow, ToPyArrow};
 use pyo3::exceptions::PyTypeError;
 use pyo3::types::{PyAnyMethods, PyModule};
@@ -84,7 +85,6 @@ fn test_to_pyarrow_byte_view() {
         ])
         .unwrap();
 
-        println!("input: {input:?}");
         let res = Python::attach(|py| {
             let py_input = input.to_pyarrow(py)?;
             let records = RecordBatch::from_pyarrow_bound(&py_input)?;
@@ -125,6 +125,34 @@ value = NotATuple()
     });
 }
 
+#[test]
+fn test_from_pyarrow_non_capsule() {
+    Python::initialize();
+
+    Python::attach(|py| {
+        let code = CString::new(
+            r#"
+class NotACapsule:
+    def __arrow_c_schema__(self):
+        return 1
+
+value = NotACapsule()
+"#,
+        )
+            .unwrap();
+
+        let module = PyModule::from_code(py, code.as_c_str(), c"test.py", c"test_module").unwrap();
+        let value = module.getattr("value").unwrap();
+
+        let err = Schema::from_pyarrow_bound(&value).unwrap_err();
+        assert!(err.is_instance_of::<PyTypeError>(py));
+        assert_eq!(
+            err.to_string(),
+            "TypeError: Expected __arrow_c_schema__ to return a capsule."
+        );
+    });
+}
+
 fn binary_view_column(num_variadic_buffers: usize) -> BinaryViewArray {
     let long_scalar = b"but soft what light through yonder window breaks".as_slice();
     let mut builder = BinaryViewBuilder::new().with_fixed_block_size(long_scalar.len() as u32);

arrow-pyarrow/src/lib.rs

@@ -61,6 +61,7 @@
 
 use std::convert::{From, TryFrom};
 use std::ffi::CStr;
+use std::ptr::NonNull;
 use std::sync::Arc;
 
 use arrow_array::ffi;
@@ -74,19 +75,15 @@ use arrow_data::ArrayData;
 use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef};
 use pyo3::exceptions::{PyTypeError, PyValueError};
 use pyo3::ffi::Py_uintptr_t;
-use pyo3::import_exception;
 use pyo3::prelude::*;
 use pyo3::sync::PyOnceLock;
-use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple, PyType};
+use pyo3::types::{PyCapsule, PyDict, PyList, PyType};
+use pyo3::{CastError, import_exception};
 
 import_exception!(pyarrow, ArrowException);
 /// Represents an exception raised by PyArrow.
 pub type PyArrowException = ArrowException;
 
-const ARROW_ARRAY_STREAM_CAPSULE_NAME: &CStr = c"arrow_array_stream";
-const ARROW_SCHEMA_CAPSULE_NAME: &CStr = c"arrow_schema";
-const ARROW_ARRAY_CAPSULE_NAME: &CStr = c"arrow_array";
-
 fn to_py_err(err: ArrowError) -> PyErr {
     PyArrowException::new_err(err.to_string())
 }
@@ -131,56 +128,17 @@ fn validate_class(expected: &Bound<PyType>, value: &Bound<PyAny>) -> PyResult<()
     Ok(())
 }
 
-fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyResult<()> {
-    let capsule_name = capsule.name()?;
-    if capsule_name.is_none() {
-        return Err(PyValueError::new_err(
-            "Expected schema PyCapsule to have name set.",
-        ));
-    }
-
-    let capsule_name = unsafe { capsule_name.unwrap().as_cstr() }
-        .to_str()
-        .map_err(|e| PyValueError::new_err(e.to_string()))?;
-    if capsule_name != name {
-        return Err(PyValueError::new_err(format!(
-            "Expected name '{name}' in PyCapsule, instead got '{capsule_name}'",
-        )));
-    }
-
-    Ok(())
-}
-
-fn extract_arrow_c_array_capsules<'py>(
-    value: &Bound<'py, PyAny>,
-) -> PyResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)> {
-    let tuple = value.call_method0("__arrow_c_array__")?;
-
-    if !tuple.is_instance_of::<PyTuple>() {
-        return Err(PyTypeError::new_err(
-            "Expected __arrow_c_array__ to return a tuple of (schema, array) capsules.",
-        ));
-    }
-
-    tuple.extract().map_err(|_| {
-        PyTypeError::new_err(
-            "Expected __arrow_c_array__ to return a tuple of (schema, array) capsules.",
-        )
-    })
-}
-
 impl FromPyArrow for DataType {
     fn from_pyarrow_bound(value: &Bound<PyAny>) -> PyResult<Self> {
         // Newer versions of PyArrow as well as other libraries with Arrow data implement this
         // method, so prefer it over _export_to_c.
         // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
-        if value.hasattr("__arrow_c_schema__")? {
-            let capsule = value.call_method0("__arrow_c_schema__")?.extract()?;
-            validate_pycapsule(&capsule, "arrow_schema")?;
-
-            let schema_ptr = capsule
-                .pointer_checked(Some(ARROW_SCHEMA_CAPSULE_NAME))?
-                .cast::<FFI_ArrowSchema>();
+        if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_schema__")? {
+            let schema_ptr = extract_capsule::<FFI_ArrowSchema>(
+                &capsule,
+                c"arrow_schema",
+                "__arrow_c_schema__",
+            )?;
             return unsafe { DataType::try_from(schema_ptr.as_ref()) }.map_err(to_py_err);
         }
 
@@ -204,13 +162,12 @@ impl FromPyArrow for Field {
         // Newer versions of PyArrow as well as other libraries with Arrow data implement this
         // method, so prefer it over _export_to_c.
         // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
-        if value.hasattr("__arrow_c_schema__")? {
-            let capsule = value.call_method0("__arrow_c_schema__")?.extract()?;
-            validate_pycapsule(&capsule, "arrow_schema")?;
-
-            let schema_ptr = capsule
-                .pointer_checked(Some(ARROW_SCHEMA_CAPSULE_NAME))?
-                .cast::<FFI_ArrowSchema>();
+        if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_schema__")? {
+            let schema_ptr = extract_capsule::<FFI_ArrowSchema>(
+                &capsule,
+                c"arrow_schema",
+                "__arrow_c_schema__",
+            )?;
             return unsafe { Field::try_from(schema_ptr.as_ref()) }.map_err(to_py_err);
         }
 
@@ -234,13 +191,12 @@ impl FromPyArrow for Schema {
         // Newer versions of PyArrow as well as other libraries with Arrow data implement this
         // method, so prefer it over _export_to_c.
         // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
-        if value.hasattr("__arrow_c_schema__")? {
-            let capsule = value.call_method0("__arrow_c_schema__")?.extract()?;
-            validate_pycapsule(&capsule, "arrow_schema")?;
-
-            let schema_ptr = capsule
-                .pointer_checked(Some(ARROW_SCHEMA_CAPSULE_NAME))?
-                .cast::<FFI_ArrowSchema>();
+        if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_schema__")? {
+            let schema_ptr = extract_capsule::<FFI_ArrowSchema>(
+                &capsule,
+                c"arrow_schema",
+                "__arrow_c_schema__",
+            )?;
             return unsafe { Schema::try_from(schema_ptr.as_ref()) }.map_err(to_py_err);
         }
 
@@ -264,23 +220,11 @@ impl FromPyArrow for ArrayData {
         // Newer versions of PyArrow as well as other libraries with Arrow data implement this
         // method, so prefer it over _export_to_c.
         // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
-        if value.hasattr("__arrow_c_array__")? {
-            let (schema_capsule, array_capsule) = extract_arrow_c_array_capsules(value)?;
-
-            validate_pycapsule(&schema_capsule, "arrow_schema")?;
-            validate_pycapsule(&array_capsule, "arrow_array")?;
-
-            let schema_ptr = schema_capsule
-                .pointer_checked(Some(ARROW_SCHEMA_CAPSULE_NAME))?
-                .cast::<FFI_ArrowSchema>();
-            let array = unsafe {
-                FFI_ArrowArray::from_raw(
-                    array_capsule
-                        .pointer_checked(Some(ARROW_ARRAY_CAPSULE_NAME))?
-                        .cast::<FFI_ArrowArray>()
-                        .as_ptr(),
-                )
-            };
+        if let Some((schema_capsule, array_capsule)) = call_arrow_c_array_method_if_exists(value)? {
+            let schema_ptr =
+                extract_capsule(&schema_capsule, c"arrow_schema", "__arrow_c_array__")?;
+            let array_ptr = extract_capsule(&array_capsule, c"arrow_array", "__arrow_c_array__")?;
+            let array = unsafe { FFI_ArrowArray::from_raw(array_ptr.as_ptr()) };
             return unsafe { ffi::from_ffi(array, schema_ptr.as_ref()) }.map_err(to_py_err);
         }
 
@@ -342,25 +286,18 @@ impl FromPyArrow for RecordBatch {
         // method, so prefer it over _export_to_c.
         // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
 
-        if value.hasattr("__arrow_c_array__")? {
-            let (schema_capsule, array_capsule) = extract_arrow_c_array_capsules(value)?;
-
-            validate_pycapsule(&schema_capsule, "arrow_schema")?;
-            validate_pycapsule(&array_capsule, "arrow_array")?;
-
-            let schema_ptr = schema_capsule
-                .pointer_checked(Some(ARROW_SCHEMA_CAPSULE_NAME))?
-                .cast::<FFI_ArrowSchema>();
-            let array_ptr = array_capsule
-                .pointer_checked(Some(ARROW_ARRAY_CAPSULE_NAME))?
-                .cast::<FFI_ArrowArray>();
+        if let Some((schema_capsule, array_capsule)) = call_arrow_c_array_method_if_exists(value)? {
+            let schema_ptr =
+                extract_capsule(&schema_capsule, c"arrow_schema", "__arrow_c_array__")?;
+            let array_ptr = extract_capsule(&array_capsule, c"arrow_array", "__arrow_c_array__")?;
             let ffi_array = unsafe { FFI_ArrowArray::from_raw(array_ptr.as_ptr()) };
             let array_data =
                 unsafe { ffi::from_ffi(ffi_array, schema_ptr.as_ref()) }.map_err(to_py_err)?;
             if !matches!(array_data.data_type(), DataType::Struct(_)) {
-                return Err(PyTypeError::new_err(
-                    "Expected Struct type from __arrow_c_array.",
-                ));
+                return Err(PyTypeError::new_err(format!(
+                    "Expected Struct type from __arrow_c_array__, found {}.",
+                    array_data.data_type()
+                )));
             }
             let options = RecordBatchOptions::default().with_row_count(Some(array_data.len()));
             let array = StructArray::from(array_data);
@@ -417,19 +354,10 @@ impl FromPyArrow for ArrowArrayStreamReader {
         // Newer versions of PyArrow as well as other libraries with Arrow data implement this
         // method, so prefer it over _export_to_c.
         // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
-        if value.hasattr("__arrow_c_stream__")? {
-            let capsule = value.call_method0("__arrow_c_stream__")?.extract()?;
-
-            validate_pycapsule(&capsule, "arrow_array_stream")?;
-
-            let stream = unsafe {
-                FFI_ArrowArrayStream::from_raw(
-                    capsule
-                        .pointer_checked(Some(ARROW_ARRAY_STREAM_CAPSULE_NAME))?
-                        .cast::<FFI_ArrowArrayStream>()
-                        .as_ptr(),
-                )
-            };
+        if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_stream__")? {
+            let stream_ptr =
+                extract_capsule(&capsule, c"arrow_array_stream", "__arrow_c_stream__")?;
+            let stream = unsafe { FFI_ArrowArrayStream::from_raw(stream_ptr.as_ptr()) };
 
             let stream_reader = ArrowArrayStreamReader::try_new(stream)
                 .map_err(|err| PyValueError::new_err(err.to_string()))?;
@@ -445,8 +373,7 @@ impl FromPyArrow for ArrowArrayStreamReader {
         // make the conversion through PyArrow's private API
         // this changes the pointer's memory and is thus unsafe.
         // In particular, `_export_to_c` can go out of bounds
-        let args = PyTuple::new(value.py(), [&raw mut stream as Py_uintptr_t])?;
-        value.call_method1("_export_to_c", args)?;
+        value.call_method1("_export_to_c", (&raw mut stream as Py_uintptr_t,))?;
 
         ArrowArrayStreamReader::try_new(stream)
             .map_err(|err| PyValueError::new_err(err.to_string()))
@@ -628,3 +555,62 @@ impl<T> From<T> for PyArrowType<T> {
         Self(s)
     }
 }
+
+fn call_capsule_method_if_exists<'py>(
+    object: &Bound<'py, PyAny>,
+    method_name: &'static str,
+) -> PyResult<Option<Bound<'py, PyCapsule>>> {
+    let Some(method) = object.getattr_opt(method_name)? else {
+        return Ok(None);
+    };
+    Ok(Some(method.call0()?.extract().map_err(
+        |e: CastError| {
+            wrapping_type_error(
+                object.py(),
+                e.into(),
+                format!("Expected {method_name} to return a capsule."),
+            )
+        },
+    )?))
+}
+
+fn call_arrow_c_array_method_if_exists<'py>(
+    object: &Bound<'py, PyAny>,
+) -> PyResult<Option<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)>> {
+    let Some(method) = object.getattr_opt("__arrow_c_array__")? else {
+        return Ok(None);
+    };
+    Ok(Some(method.call0()?.extract().map_err(|e| {
+        wrapping_type_error(
+            object.py(),
+            e,
+            "Expected __arrow_c_array__ to return a tuple of (schema, array) capsules.".into(),
+        )
+    })?))
+}
+
+fn extract_capsule<T>(
+    capsule: &Bound<PyCapsule>,
+    capsule_name: &CStr,
+    method_name: &'static str,
+) -> PyResult<NonNull<T>> {
+    Ok(capsule
+        .pointer_checked(Some(capsule_name))
+        .map_err(|e| {
+            wrapping_type_error(
+                capsule.py(),
+                e,
+                format!(
+                    "Expected {method_name} to return a {} capsule.",
+                    capsule_name.to_str().unwrap(),
+                ),
+            )
+        })?
+        .cast::<T>())
+}
+
+fn wrapping_type_error(py: Python<'_>, error: PyErr, message: String) -> PyErr {
+    let e = PyTypeError::new_err(message);
+    e.set_cause(py, Some(error));
+    e
+}