tustvold commented on code in PR #5070:
URL: https://github.com/apache/arrow-rs/pull/5070#discussion_r1394087460
##########
arrow/src/pyarrow.rs:
##########
@@ -247,6 +329,40 @@ impl<T: ToPyArrow> ToPyArrow for Vec<T> {
impl FromPyArrow for RecordBatch {
fn from_pyarrow(value: &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_array__")? {
+ let tuple = value.getattr("__arrow_c_array__")?.call0()?;
+
+ if !tuple.is_instance_of::<PyTuple>() {
+ return Err(PyTypeError::new_err(
+ "Expected __arrow_c_array__ to return a tuple.",
+ ));
+ }
+
+ let schema_capsule: &PyCapsule =
PyTryInto::try_into(tuple.get_item(0)?)?;
+ let array_capsule: &PyCapsule =
PyTryInto::try_into(tuple.get_item(1)?)?;
+
+ validate_pycapsule(schema_capsule, "arrow_schema")?;
+ validate_pycapsule(array_capsule, "arrow_array")?;
+
+ let schema_ptr = unsafe {
schema_capsule.reference::<FFI_ArrowSchema>() };
+ let array_ptr = array_capsule.pointer() as *mut FFI_ArrowArray;
+ let ffi_array = unsafe { std::ptr::replace(array_ptr,
FFI_ArrowArray::empty()) };
+ let array_data = ffi::from_ffi(ffi_array,
schema_ptr).map_err(to_py_err)?;
+ let array_ref = make_array(array_data);
+
+ if !matches!(array_ref.data_type(), DataType::Struct(_)) {
+ return Err(PyTypeError::new_err(
+ "Expected Struct type from __arrow_c_array.",
+ ));
+ }
+
+ let array =
array_ref.as_any().downcast_ref::<StructArray>().unwrap();
Review Comment:
```suggestion
if !matches!(array_data.data_type(), DataType::Struct(_)) {
return Err(PyTypeError::new_err(
"Expected Struct type from __arrow_c_array.",
));
}
let array = StructArray::from(array_data);
```
##########
arrow/src/pyarrow.rs:
##########
@@ -247,6 +329,40 @@ impl<T: ToPyArrow> ToPyArrow for Vec<T> {
impl FromPyArrow for RecordBatch {
fn from_pyarrow(value: &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_array__")? {
Review Comment:
I think this will allow passing a `StructArray` to a method accepting
`PyArrowType<RecordBatch>` provided the `StructArray` isn't nullable. I don't
think this is a problem, RecordBatch is a bit of an oddity, but just an
observation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
