From 9497c5c3d25599dc2608ff62612741134f4e28da Mon Sep 17 00:00:00 2001 From: Thomas Tanon Date: Fri, 3 Apr 2026 17:42:43 +0200 Subject: [PATCH 1/8] pyarrow: Wrap capsule extraction errors into an explicit error Remove validate_pycapsule: pointer_checked already raises an exception if the capsule type is wrong, and we wrap it into a nicer error --- arrow-pyarrow-testing/tests/pyarrow.rs | 3 +- arrow-pyarrow/src/lib.rs | 205 +++++++++++++------------ 2 files changed, 104 insertions(+), 104 deletions(-) diff --git a/arrow-pyarrow-testing/tests/pyarrow.rs b/arrow-pyarrow-testing/tests/pyarrow.rs index 6f3606478c72..09a4f9756822 100644 --- a/arrow-pyarrow-testing/tests/pyarrow.rs +++ b/arrow-pyarrow-testing/tests/pyarrow.rs @@ -84,7 +84,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)?; @@ -120,7 +119,7 @@ value = NotATuple() assert!(err.is_instance_of::(py)); assert_eq!( err.to_string(), - "TypeError: Expected __arrow_c_array__ to return a tuple of (schema, array) capsules." + "TypeError: Expected __arrow_c_array__ to return a tuple of (arrow_schema, arrow_array) capsules." ); }); } diff --git a/arrow-pyarrow/src/lib.rs b/arrow-pyarrow/src/lib.rs index c0d91d081113..414f2794d2ea 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/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; @@ -77,15 +78,12 @@ 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}; 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 +129,16 @@ fn validate_class(expected: &Bound, value: &Bound) -> PyResult<() Ok(()) } -fn validate_pycapsule(capsule: &Bound, 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::() { - 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) -> PyResult { // 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::(); + let schema_ptr = extract_capsule_from_method::( + value, + "__arrow_c_schema__", + )?; return unsafe { DataType::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -205,12 +163,10 @@ impl FromPyArrow for Field { // 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::(); + let schema_ptr = extract_capsule_from_method::( + value, + "__arrow_c_schema__", + )?; return unsafe { Field::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -235,12 +191,10 @@ impl FromPyArrow for Schema { // 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::(); + let schema_ptr = extract_capsule_from_method::( + value, + "__arrow_c_schema__", + )?; return unsafe { Schema::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -265,22 +219,12 @@ impl FromPyArrow for ArrayData { // 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::(); - let array = unsafe { - FFI_ArrowArray::from_raw( - array_capsule - .pointer_checked(Some(ARROW_ARRAY_CAPSULE_NAME))? - .cast::() - .as_ptr(), - ) - }; + let (schema_ptr, array_ptr) = + extract_capsule_pair_from_method::( + value, + "__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); } @@ -343,23 +287,17 @@ impl FromPyArrow for RecordBatch { // 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::(); - let array_ptr = array_capsule - .pointer_checked(Some(ARROW_ARRAY_CAPSULE_NAME))? - .cast::(); + let (schema_ptr, array_ptr) = + extract_capsule_pair_from_method::( + value, + "__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.", + format!("Expected Struct type from __arrow_c_array__, found {}.", array_data.data_type()), )); } let options = RecordBatchOptions::default().with_row_count(Some(array_data.len())); @@ -418,18 +356,11 @@ impl FromPyArrow for ArrowArrayStreamReader { // 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::() - .as_ptr(), - ) - }; + let stream_ptr = extract_capsule_from_method::( + value, + "__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 +376,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 +558,74 @@ impl From for PyArrowType { Self(s) } } + +trait PyCapsuleType { + const NAME: &CStr; +} + +impl PyCapsuleType for FFI_ArrowSchema { + const NAME: &CStr = c"arrow_schema"; +} + +impl PyCapsuleType for FFI_ArrowArray { + const NAME: &CStr = c"arrow_array"; +} + +impl PyCapsuleType for FFI_ArrowArrayStream { + const NAME: &CStr = c"arrow_array_stream"; +} + +fn extract_capsule_from_method( + object: &Bound<'_, PyAny>, + method_name: &'static str, +) -> PyResult> { + (|| { + Ok(object + .call_method0(method_name)? + .extract::>()? + .pointer_checked(Some(T::NAME))? + .cast::()) + })() + .map_err(|e| { + wrapping_type_error( + object.py(), + e, + format!( + "Expected {method_name} to return a {} capsule.", + T::NAME.to_str().unwrap(), + ), + ) + }) +} + +fn extract_capsule_pair_from_method( + object: &Bound<'_, PyAny>, + method_name: &'static str, +) -> PyResult<(NonNull, NonNull)> { + (|| { + let (c1, c2) = object + .call_method0(method_name)? + .extract::<(Bound<'_, PyCapsule>, Bound<'_, PyCapsule>)>()?; + Ok(( + c1.pointer_checked(Some(T1::NAME))?.cast::(), + c2.pointer_checked(Some(T2::NAME))?.cast::(), + )) + })() + .map_err(|e| { + wrapping_type_error( + object.py(), + e, + format!( + "Expected {method_name} to return a tuple of ({}, {}) capsules.", + T1::NAME.to_str().unwrap(), + T2::NAME.to_str().unwrap() + ), + ) + }) +} + +fn wrapping_type_error(py: Python<'_>, error: PyErr, message: String) -> PyErr { + let e = PyTypeError::new_err(message); + e.set_cause(py, Some(error)); + e +} From c1a28607549e299559c0e118cef3e27339439cdd Mon Sep 17 00:00:00 2001 From: Thomas Tanon Date: Thu, 16 Apr 2026 21:25:55 +0200 Subject: [PATCH 2/8] Fix use after free error --- arrow-pyarrow/src/lib.rs | 141 ++++++++++++++++++--------------------- 1 file changed, 65 insertions(+), 76 deletions(-) diff --git a/arrow-pyarrow/src/lib.rs b/arrow-pyarrow/src/lib.rs index 414f2794d2ea..0f8275755d0e 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/arrow-pyarrow/src/lib.rs @@ -75,16 +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, PyType}; +use pyo3::{CastError, import_exception}; import_exception!(pyarrow, ArrowException); /// Represents an exception raised by PyArrow. pub type PyArrowException = ArrowException; - fn to_py_err(err: ArrowError) -> PyErr { PyArrowException::new_err(err.to_string()) } @@ -135,10 +134,8 @@ impl FromPyArrow for DataType { // 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 schema_ptr = extract_capsule_from_method::( - value, - "__arrow_c_schema__", - )?; + let capsule = call_capsule_method(value, "__arrow_c_schema__")?; + let schema_ptr = extract_capsule::(&capsule, "__arrow_c_schema__")?; return unsafe { DataType::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -163,10 +160,8 @@ impl FromPyArrow for Field { // 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 schema_ptr = extract_capsule_from_method::( - value, - "__arrow_c_schema__", - )?; + let capsule = call_capsule_method(value, "__arrow_c_schema__")?; + let schema_ptr = extract_capsule::(&capsule, "__arrow_c_schema__")?; return unsafe { Field::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -191,10 +186,8 @@ impl FromPyArrow for Schema { // 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 schema_ptr = extract_capsule_from_method::( - value, - "__arrow_c_schema__", - )?; + let capsule = call_capsule_method(value, "__arrow_c_schema__")?; + let schema_ptr = extract_capsule::(&capsule, "__arrow_c_schema__")?; return unsafe { Schema::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -219,11 +212,10 @@ impl FromPyArrow for ArrayData { // 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_ptr, array_ptr) = - extract_capsule_pair_from_method::( - value, - "__arrow_c_array__", - )?; + let (schema_capsule, array_capsule) = + call_capsule_pair_method(value, "__arrow_c_array__")?; + let schema_ptr = extract_capsule(&schema_capsule, "__arrow_c_array__")?; + let array_ptr = extract_capsule(&array_capsule, "__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); } @@ -287,18 +279,18 @@ impl FromPyArrow for RecordBatch { // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html if value.hasattr("__arrow_c_array__")? { - let (schema_ptr, array_ptr) = - extract_capsule_pair_from_method::( - value, - "__arrow_c_array__", - )?; + let (schema_capsule, array_capsule) = + call_capsule_pair_method(value, "__arrow_c_array__")?; + let schema_ptr = extract_capsule(&schema_capsule, "__arrow_c_array__")?; + let array_ptr = extract_capsule(&array_capsule, "__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( - format!("Expected Struct type from __arrow_c_array__, found {}.", array_data.data_type()), - )); + 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); @@ -356,10 +348,8 @@ impl FromPyArrow for ArrowArrayStreamReader { // 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 stream_ptr = extract_capsule_from_method::( - value, - "__arrow_c_stream__", - )?; + let capsule = call_capsule_method(value, "__arrow_c_stream__")?; + let stream_ptr = extract_capsule(&capsule, "__arrow_c_stream__")?; let stream = unsafe { FFI_ArrowArrayStream::from_raw(stream_ptr.as_ptr()) }; let stream_reader = ArrowArrayStreamReader::try_new(stream) @@ -559,6 +549,35 @@ impl From for PyArrowType { } } +fn call_capsule_method<'py>( + object: &Bound<'py, PyAny>, + method_name: &'static str, +) -> PyResult> { + object + .call_method0(method_name)? + .extract() + .map_err(|e: CastError| { + wrapping_type_error( + object.py(), + e.into(), + format!("Expected {method_name} to return a capsule."), + ) + }) +} + +fn call_capsule_pair_method<'py>( + object: &Bound<'py, PyAny>, + method_name: &'static str, +) -> PyResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)> { + object.call_method0(method_name)?.extract().map_err(|e| { + wrapping_type_error( + object.py(), + e, + format!("Expected {method_name} to return a pair of capsule."), + ) + }) +} + trait PyCapsuleType { const NAME: &CStr; } @@ -575,53 +594,23 @@ impl PyCapsuleType for FFI_ArrowArrayStream { const NAME: &CStr = c"arrow_array_stream"; } -fn extract_capsule_from_method( - object: &Bound<'_, PyAny>, +fn extract_capsule( + capsule: &Bound, method_name: &'static str, ) -> PyResult> { - (|| { - Ok(object - .call_method0(method_name)? - .extract::>()? - .pointer_checked(Some(T::NAME))? - .cast::()) - })() - .map_err(|e| { - wrapping_type_error( - object.py(), - e, - format!( - "Expected {method_name} to return a {} capsule.", - T::NAME.to_str().unwrap(), - ), - ) - }) -} - -fn extract_capsule_pair_from_method( - object: &Bound<'_, PyAny>, - method_name: &'static str, -) -> PyResult<(NonNull, NonNull)> { - (|| { - let (c1, c2) = object - .call_method0(method_name)? - .extract::<(Bound<'_, PyCapsule>, Bound<'_, PyCapsule>)>()?; - Ok(( - c1.pointer_checked(Some(T1::NAME))?.cast::(), - c2.pointer_checked(Some(T2::NAME))?.cast::(), - )) - })() - .map_err(|e| { - wrapping_type_error( - object.py(), - e, - format!( - "Expected {method_name} to return a tuple of ({}, {}) capsules.", - T1::NAME.to_str().unwrap(), - T2::NAME.to_str().unwrap() - ), - ) - }) + Ok(capsule + .pointer_checked(Some(T::NAME)) + .map_err(|e| { + wrapping_type_error( + capsule.py(), + e, + format!( + "Expected {method_name} to return a {} capsule.", + T::NAME.to_str().unwrap(), + ), + ) + })? + .cast::()) } fn wrapping_type_error(py: Python<'_>, error: PyErr, message: String) -> PyErr { From 6c51b4382285d9d29523f0b788c40f6bced07b0d Mon Sep 17 00:00:00 2001 From: Thomas Tanon Date: Thu, 16 Apr 2026 21:52:19 +0200 Subject: [PATCH 3/8] Fix error message for RecordBatch from_pyarrow_bound --- arrow-pyarrow-testing/tests/pyarrow.rs | 2 +- arrow-pyarrow/src/lib.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/arrow-pyarrow-testing/tests/pyarrow.rs b/arrow-pyarrow-testing/tests/pyarrow.rs index 09a4f9756822..de1950f2ecbc 100644 --- a/arrow-pyarrow-testing/tests/pyarrow.rs +++ b/arrow-pyarrow-testing/tests/pyarrow.rs @@ -119,7 +119,7 @@ value = NotATuple() assert!(err.is_instance_of::(py)); assert_eq!( err.to_string(), - "TypeError: Expected __arrow_c_array__ to return a tuple of (arrow_schema, arrow_array) capsules." + "TypeError: Expected __arrow_c_array__ to return a pair of capsules." ); }); } diff --git a/arrow-pyarrow/src/lib.rs b/arrow-pyarrow/src/lib.rs index 0f8275755d0e..10ae296105e3 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/arrow-pyarrow/src/lib.rs @@ -573,7 +573,7 @@ fn call_capsule_pair_method<'py>( wrapping_type_error( object.py(), e, - format!("Expected {method_name} to return a pair of capsule."), + format!("Expected {method_name} to return a pair of capsules."), ) }) } From e730f3c24d69413ea4f3980cab5fe216eb70d8c8 Mon Sep 17 00:00:00 2001 From: Thomas Tanon Date: Wed, 17 Jun 2026 07:30:24 +0200 Subject: [PATCH 4/8] Factorize checking if capsule method exists --- arrow-pyarrow/src/lib.rs | 51 ++++++++++++++++++++-------------------- 1 file changed, 26 insertions(+), 25 deletions(-) diff --git a/arrow-pyarrow/src/lib.rs b/arrow-pyarrow/src/lib.rs index 10ae296105e3..bd19cfa5c9dd 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/arrow-pyarrow/src/lib.rs @@ -133,8 +133,7 @@ impl FromPyArrow for DataType { // 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 = call_capsule_method(value, "__arrow_c_schema__")?; + if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_schema__")? { let schema_ptr = extract_capsule::(&capsule, "__arrow_c_schema__")?; return unsafe { DataType::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -159,8 +158,7 @@ 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 = call_capsule_method(value, "__arrow_c_schema__")?; + if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_schema__")? { let schema_ptr = extract_capsule::(&capsule, "__arrow_c_schema__")?; return unsafe { Field::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -185,8 +183,7 @@ 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 = call_capsule_method(value, "__arrow_c_schema__")?; + if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_schema__")? { let schema_ptr = extract_capsule::(&capsule, "__arrow_c_schema__")?; return unsafe { Schema::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -211,9 +208,9 @@ 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) = - call_capsule_pair_method(value, "__arrow_c_array__")?; + if let Some((schema_capsule, array_capsule)) = + call_capsule_pair_method_if_exists(value, "__arrow_c_array__")? + { let schema_ptr = extract_capsule(&schema_capsule, "__arrow_c_array__")?; let array_ptr = extract_capsule(&array_capsule, "__arrow_c_array__")?; let array = unsafe { FFI_ArrowArray::from_raw(array_ptr.as_ptr()) }; @@ -278,9 +275,9 @@ 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) = - call_capsule_pair_method(value, "__arrow_c_array__")?; + if let Some((schema_capsule, array_capsule)) = + call_capsule_pair_method_if_exists(value, "__arrow_c_array__")? + { let schema_ptr = extract_capsule(&schema_capsule, "__arrow_c_array__")?; let array_ptr = extract_capsule(&array_capsule, "__arrow_c_array__")?; let ffi_array = unsafe { FFI_ArrowArray::from_raw(array_ptr.as_ptr()) }; @@ -347,8 +344,7 @@ 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 = call_capsule_method(value, "__arrow_c_stream__")?; + if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_stream__")? { let stream_ptr = extract_capsule(&capsule, "__arrow_c_stream__")?; let stream = unsafe { FFI_ArrowArrayStream::from_raw(stream_ptr.as_ptr()) }; @@ -549,33 +545,38 @@ impl From for PyArrowType { } } -fn call_capsule_method<'py>( +fn call_capsule_method_if_exists<'py>( object: &Bound<'py, PyAny>, method_name: &'static str, -) -> PyResult> { - object - .call_method0(method_name)? - .extract() - .map_err(|e: CastError| { +) -> PyResult>> { + 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_capsule_pair_method<'py>( +fn call_capsule_pair_method_if_exists<'py>( object: &Bound<'py, PyAny>, method_name: &'static str, -) -> PyResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)> { - object.call_method0(method_name)?.extract().map_err(|e| { +) -> PyResult, Bound<'py, PyCapsule>)>> { + let Some(method) = object.getattr_opt(method_name)? else { + return Ok(None); + }; + Ok(Some(method.call0()?.extract().map_err(|e| { wrapping_type_error( object.py(), e, format!("Expected {method_name} to return a pair of capsules."), ) - }) + })?)) } trait PyCapsuleType { From 2ab7223b416075696fb8290b4a7818c6c7bf6277 Mon Sep 17 00:00:00 2001 From: Thomas Tanon Date: Wed, 17 Jun 2026 17:08:23 +0200 Subject: [PATCH 5/8] Remove PyCapsuleType trait --- arrow-pyarrow/src/lib.rs | 54 ++++++++++++++++++++-------------------- 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/arrow-pyarrow/src/lib.rs b/arrow-pyarrow/src/lib.rs index bd19cfa5c9dd..edce0ec9d15f 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/arrow-pyarrow/src/lib.rs @@ -134,7 +134,11 @@ impl FromPyArrow for DataType { // method, so prefer it over _export_to_c. // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_schema__")? { - let schema_ptr = extract_capsule::(&capsule, "__arrow_c_schema__")?; + let schema_ptr = extract_capsule::( + &capsule, + c"arrow_schema", + "__arrow_c_schema__", + )?; return unsafe { DataType::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -159,7 +163,11 @@ impl FromPyArrow for Field { // method, so prefer it over _export_to_c. // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_schema__")? { - let schema_ptr = extract_capsule::(&capsule, "__arrow_c_schema__")?; + let schema_ptr = extract_capsule::( + &capsule, + c"arrow_schema", + "__arrow_c_schema__", + )?; return unsafe { Field::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -184,7 +192,11 @@ impl FromPyArrow for Schema { // method, so prefer it over _export_to_c. // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_schema__")? { - let schema_ptr = extract_capsule::(&capsule, "__arrow_c_schema__")?; + let schema_ptr = extract_capsule::( + &capsule, + c"arrow_schema", + "__arrow_c_schema__", + )?; return unsafe { Schema::try_from(schema_ptr.as_ref()) }.map_err(to_py_err); } @@ -211,8 +223,9 @@ impl FromPyArrow for ArrayData { if let Some((schema_capsule, array_capsule)) = call_capsule_pair_method_if_exists(value, "__arrow_c_array__")? { - let schema_ptr = extract_capsule(&schema_capsule, "__arrow_c_array__")?; - let array_ptr = extract_capsule(&array_capsule, "__arrow_c_array__")?; + 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); } @@ -278,8 +291,9 @@ impl FromPyArrow for RecordBatch { if let Some((schema_capsule, array_capsule)) = call_capsule_pair_method_if_exists(value, "__arrow_c_array__")? { - let schema_ptr = extract_capsule(&schema_capsule, "__arrow_c_array__")?; - let array_ptr = extract_capsule(&array_capsule, "__arrow_c_array__")?; + 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)?; @@ -345,7 +359,8 @@ impl FromPyArrow for ArrowArrayStreamReader { // method, so prefer it over _export_to_c. // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html if let Some(capsule) = call_capsule_method_if_exists(value, "__arrow_c_stream__")? { - let stream_ptr = extract_capsule(&capsule, "__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) @@ -579,35 +594,20 @@ fn call_capsule_pair_method_if_exists<'py>( })?)) } -trait PyCapsuleType { - const NAME: &CStr; -} - -impl PyCapsuleType for FFI_ArrowSchema { - const NAME: &CStr = c"arrow_schema"; -} - -impl PyCapsuleType for FFI_ArrowArray { - const NAME: &CStr = c"arrow_array"; -} - -impl PyCapsuleType for FFI_ArrowArrayStream { - const NAME: &CStr = c"arrow_array_stream"; -} - -fn extract_capsule( +fn extract_capsule( capsule: &Bound, + capsule_name: &CStr, method_name: &'static str, ) -> PyResult> { Ok(capsule - .pointer_checked(Some(T::NAME)) + .pointer_checked(Some(capsule_name)) .map_err(|e| { wrapping_type_error( capsule.py(), e, format!( "Expected {method_name} to return a {} capsule.", - T::NAME.to_str().unwrap(), + capsule_name.to_str().unwrap(), ), ) })? From 1f72b242413d9ff3612c787805bc57b20f9acb02 Mon Sep 17 00:00:00 2001 From: Thomas Tanon Date: Wed, 17 Jun 2026 20:59:14 +0200 Subject: [PATCH 6/8] Keep tuple error message --- arrow-pyarrow-testing/tests/pyarrow.rs | 2 +- arrow-pyarrow/src/lib.rs | 11 +++++------ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/arrow-pyarrow-testing/tests/pyarrow.rs b/arrow-pyarrow-testing/tests/pyarrow.rs index de1950f2ecbc..08800f0bf3fd 100644 --- a/arrow-pyarrow-testing/tests/pyarrow.rs +++ b/arrow-pyarrow-testing/tests/pyarrow.rs @@ -119,7 +119,7 @@ value = NotATuple() assert!(err.is_instance_of::(py)); assert_eq!( err.to_string(), - "TypeError: Expected __arrow_c_array__ to return a pair of capsules." + "TypeError: Expected __arrow_c_array__ to return a tuple of (schema, array) capsules." ); }); } diff --git a/arrow-pyarrow/src/lib.rs b/arrow-pyarrow/src/lib.rs index edce0ec9d15f..94fd444d0ebb 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/arrow-pyarrow/src/lib.rs @@ -221,7 +221,7 @@ impl FromPyArrow for ArrayData { // method, so prefer it over _export_to_c. // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html if let Some((schema_capsule, array_capsule)) = - call_capsule_pair_method_if_exists(value, "__arrow_c_array__")? + call_arrow_c_array_method_if_exists(value)? { let schema_ptr = extract_capsule(&schema_capsule, c"arrow_schema", "__arrow_c_array__")?; @@ -289,7 +289,7 @@ impl FromPyArrow for RecordBatch { // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html if let Some((schema_capsule, array_capsule)) = - call_capsule_pair_method_if_exists(value, "__arrow_c_array__")? + call_arrow_c_array_method_if_exists(value)? { let schema_ptr = extract_capsule(&schema_capsule, c"arrow_schema", "__arrow_c_array__")?; @@ -578,18 +578,17 @@ fn call_capsule_method_if_exists<'py>( )?)) } -fn call_capsule_pair_method_if_exists<'py>( +fn call_arrow_c_array_method_if_exists<'py>( object: &Bound<'py, PyAny>, - method_name: &'static str, ) -> PyResult, Bound<'py, PyCapsule>)>> { - let Some(method) = object.getattr_opt(method_name)? else { + 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, - format!("Expected {method_name} to return a pair of capsules."), + "Expected __arrow_c_array__ to return a tuple of (schema, array) capsules.".into(), ) })?)) } From 56218cd86634582683091ffc41707de88eec630c Mon Sep 17 00:00:00 2001 From: Thomas Tanon Date: Wed, 17 Jun 2026 21:06:43 +0200 Subject: [PATCH 7/8] Add an error test --- arrow-pyarrow-testing/Cargo.toml | 1 + arrow-pyarrow-testing/tests/pyarrow.rs | 29 ++++++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/arrow-pyarrow-testing/Cargo.toml b/arrow-pyarrow-testing/Cargo.toml index d5542e25d65d..c5a15b5cc923 100644 --- a/arrow-pyarrow-testing/Cargo.toml +++ b/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 } diff --git a/arrow-pyarrow-testing/tests/pyarrow.rs b/arrow-pyarrow-testing/tests/pyarrow.rs index 08800f0bf3fd..6aa533add7ec 100644 --- a/arrow-pyarrow-testing/tests/pyarrow.rs +++ b/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}; @@ -124,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::(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); From b4231e0ed712667bde0537c0c2e39376e1b66112 Mon Sep 17 00:00:00 2001 From: Thomas Tanon Date: Wed, 17 Jun 2026 21:32:10 +0200 Subject: [PATCH 8/8] rustfmt --- arrow-pyarrow/src/lib.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/arrow-pyarrow/src/lib.rs b/arrow-pyarrow/src/lib.rs index 94fd444d0ebb..1582f9582151 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/arrow-pyarrow/src/lib.rs @@ -220,9 +220,7 @@ 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 let Some((schema_capsule, array_capsule)) = - call_arrow_c_array_method_if_exists(value)? - { + 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__")?; @@ -288,9 +286,7 @@ impl FromPyArrow for RecordBatch { // method, so prefer it over _export_to_c. // See https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html - if let Some((schema_capsule, array_capsule)) = - call_arrow_c_array_method_if_exists(value)? - { + 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__")?;