feat: add Arrow FFI support for Python object conversions and remove pyarrow_util module

Author: kosiew

src/arrow_ffi.rs

@@ -0,0 +1,124 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Conversions between Python objects exposing Arrow's C Data Interface
+//! and DataFusion types.
+
+use std::ptr::addr_of;
+
+use arrow::array::{make_array, Array, ArrayData};
+use arrow::ffi::{from_ffi, to_ffi, FFI_ArrowArray, FFI_ArrowSchema};
+use arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
+use datafusion::scalar::ScalarValue;
+use pyo3::ffi::Py_uintptr_t;
+use pyo3::types::{PyAnyMethods, PyCapsule, PyCapsuleMethods, PyTuple};
+use pyo3::{Bound, FromPyObject, PyAny, PyObject, PyResult, Python};
+
+use crate::common::data_type::PyScalarValue;
+use crate::errors::PyDataFusionError;
+use crate::utils::validate_pycapsule;
+
+/// Convert a Python object implementing ``__arrow_c_array__`` into ``ArrayData``.
+pub fn array_data_from_py(value: &Bound<'_, PyAny>) -> PyResult<ArrayData> {
+    if value.hasattr("__arrow_c_array__")? {
+        let obj = value.getattr("__arrow_c_array__")?.call0()?;
+        let tuple = obj.downcast::<PyTuple>()?;
+
+        let schema_item = tuple.get_item(0)?;
+        let schema_capsule = schema_item.downcast::<PyCapsule>()?;
+        let array_item = tuple.get_item(1)?;
+        let array_capsule = array_item.downcast::<PyCapsule>()?;
+
+        validate_pycapsule(schema_capsule, "arrow_schema")?;
+        validate_pycapsule(array_capsule, "arrow_array")?;
+
+        let schema_ptr = unsafe { schema_capsule.reference::<FFI_ArrowSchema>() };
+        let array = unsafe { FFI_ArrowArray::from_raw(array_capsule.pointer().cast()) };
+        unsafe { from_ffi(array, schema_ptr) }
+            .map_err(|e| PyDataFusionError::Common(e.to_string()).into())
+    } else {
+        Err(pyo3::exceptions::PyTypeError::new_err(
+            "Object does not implement __arrow_c_array__",
+        ))
+    }
+}
+
+/// Convert a Python object implementing ``__arrow_c_stream__`` into an
+/// ``ArrowArrayStreamReader``.
+pub fn stream_reader_from_py(value: &Bound<'_, PyAny>) -> PyResult<ArrowArrayStreamReader> {
+    if value.hasattr("__arrow_c_stream__")? {
+        let capsule = value.getattr("__arrow_c_stream__")?.call0()?;
+        let capsule = capsule.downcast::<PyCapsule>()?;
+        validate_pycapsule(capsule, "arrow_array_stream")?;
+        let stream = unsafe { FFI_ArrowArrayStream::from_raw(capsule.pointer() as _) };
+        ArrowArrayStreamReader::try_new(stream)
+            .map_err(|err| pyo3::exceptions::PyValueError::new_err(err.to_string()))
+    } else {
+        Err(pyo3::exceptions::PyTypeError::new_err(
+            "Object does not implement __arrow_c_stream__",
+        ))
+    }
+}
+
+impl PyScalarValue {
+    /// Construct a ``PyScalarValue`` from any Python object implementing
+    /// ``__arrow_c_array__``.
+    pub fn from_arrow_ffi(value: &Bound<'_, PyAny>) -> PyResult<Self> {
+        let array_data = array_data_from_py(value)?;
+        let array = make_array(array_data);
+        let scalar = ScalarValue::try_from_array(&array, 0).map_err(PyDataFusionError::from)?;
+        Ok(PyScalarValue(scalar))
+    }
+}
+
+impl<'source> FromPyObject<'source> for PyScalarValue {
+    fn extract_bound(value: &Bound<'source, PyAny>) -> PyResult<Self> {
+        Self::from_arrow_ffi(value)
+    }
+}
+
+/// Convert a ``ScalarValue`` into a Python ``pyarrow.Scalar`` using the
+/// Arrow FFI.
+pub fn scalar_to_pyarrow(scalar: &ScalarValue, py: Python) -> PyResult<PyObject> {
+    let array = scalar.to_array().map_err(PyDataFusionError::from)?;
+    let data = array.to_data();
+    let (ffi_array, ffi_schema) =
+        to_ffi(&data).map_err(|e| PyDataFusionError::Common(e.to_string()))?;
+
+    let pa = py.import("pyarrow")?;
+    let class = pa.getattr("Array")?;
+    let pyarray = class.call_method1(
+        "_import_from_c",
+        (
+            addr_of!(ffi_array) as Py_uintptr_t,
+            addr_of!(ffi_schema) as Py_uintptr_t,
+        ),
+    )?;
+    let pyscalar = pyarray.call_method1("__getitem__", (0,))?;
+    Ok(pyscalar.into())
+}
+
+/// Helper newtype mirroring Arrow's ``PyArrowType`` for convenience when
+/// building Python objects from Rust values.
+#[derive(Debug)]
+pub struct PyArrowType<T>(pub T);
+
+impl<T> From<T> for PyArrowType<T> {
+    fn from(v: T) -> Self {
+        Self(v)
+    }
+}

src/expr.rs

@@ -36,14 +36,14 @@ use datafusion::logical_expr::{
     lit, Between, BinaryExpr, Case, Cast, Expr, Like, Operator, TryCast,
 };
 
+use crate::arrow_ffi::scalar_to_pyarrow;
 use crate::common::data_type::{DataTypeMap, NullTreatment, PyScalarValue, RexType};
 use crate::errors::{py_runtime_err, py_type_err, py_unsupported_variant_err, PyDataFusionResult};
 use crate::expr::aggregate_expr::PyAggregateFunction;
 use crate::expr::binary_expr::PyBinaryExpr;
 use crate::expr::column::PyColumn;
 use crate::expr::literal::PyLiteral;
 use crate::functions::add_builder_fns_to_window;
-use crate::pyarrow_util::scalar_to_pyarrow;
 use crate::sql::logical::PyLogicalPlan;
 
 use self::alias::PyAlias;

src/lib.rs

@@ -33,6 +33,7 @@ pub use datafusion_substrait;
 pub mod catalog;
 pub mod common;
 
+pub mod arrow_ffi;
 #[allow(clippy::borrow_deref_ref)]
 mod config;
 #[allow(clippy::borrow_deref_ref)]
@@ -48,7 +49,6 @@ pub mod expr;
 mod functions;
 pub mod physical_plan;
 mod pyarrow_filter_expression;
-pub mod pyarrow_util;
 mod record_batch;
 pub mod sql;
 pub mod store;

src/pyarrow_filter_expression.rs

@@ -24,8 +24,8 @@ use std::result::Result;
 use datafusion::common::{Column, ScalarValue};
 use datafusion::logical_expr::{expr::InList, Between, BinaryExpr, Expr, Operator};
 
+use crate::arrow_ffi::scalar_to_pyarrow;
 use crate::errors::{PyDataFusionError, PyDataFusionResult};
-use crate::pyarrow_util::scalar_to_pyarrow;
 
 #[derive(Debug)]
 #[repr(transparent)]