Add ArrowArrayExportable class and use it to create pyarrow arrays for python UDFs

Author: timsaucer

pyproject.toml

@@ -141,6 +141,7 @@ dev = [
     "maturin>=1.8.1",
     "numpy>1.25.0;python_version<'3.14'",
     "numpy>=2.3.2;python_version>='3.14'",
+    "pyarrow>=19.0.0",
     "pre-commit>=4.3.0",
     "pyyaml>=6.0.3",
     "pytest>=7.4.4",

python/tests/test_udf.py

@@ -18,6 +18,7 @@
 import pyarrow as pa
 import pytest
 from datafusion import column, udf
+from datafusion import functions as f
 
 
 @pytest.fixture
@@ -124,3 +125,26 @@ def udf_with_param(values: pa.Array) -> pa.Array:
     result = df2.collect()[0].column(0)
 
     assert result == pa.array([False, True, True])
+
+
+def test_udf_with_metadata(ctx) -> None:
+    from uuid import UUID
+
+    @udf([pa.string()], pa.uuid(), "stable")
+    def uuid_from_string(uuid_string):
+        return pa.array((UUID(s).bytes for s in uuid_string.to_pylist()), pa.uuid())
+
+    @udf([pa.uuid()], pa.int64(), "stable")
+    def uuid_version(uuid):
+        return pa.array(s.version for s in uuid.to_pylist())
+
+    batch = pa.record_batch({"idx": pa.array(range(5))})
+    results = (
+        ctx.create_dataframe([[batch]])
+        .with_column("uuid_string", f.uuid())
+        .with_column("uuid", uuid_from_string(column("uuid_string")))
+        .select(uuid_version(column("uuid").alias("uuid_version")))
+        .collect()
+    )
+
+    assert results[0][0].to_pylist() == [4, 4, 4, 4, 4]

src/array.rs

@@ -0,0 +1,80 @@
+// 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.
+
+use crate::errors::PyDataFusionResult;
+use crate::utils::validate_pycapsule;
+use arrow::array::{Array, ArrayRef};
+use arrow::datatypes::{Field, FieldRef};
+use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema};
+use arrow::pyarrow::ToPyArrow;
+use pyo3::prelude::{PyAnyMethods, PyCapsuleMethods};
+use pyo3::types::PyCapsule;
+use pyo3::{pyclass, pymethods, Bound, PyObject, PyResult, Python};
+use std::sync::Arc;
+
+/// A Python object which implements the Arrow PyCapsule for importing
+/// into other libraries.
+#[pyclass(name = "ArrowArrayExportable", module = "datafusion", frozen)]
+#[derive(Clone)]
+pub struct PyArrowArrayExportable {
+    array: ArrayRef,
+    field: FieldRef,
+}
+
+#[pymethods]
+impl PyArrowArrayExportable {
+    #[pyo3(signature = (requested_schema=None))]
+    fn __arrow_c_array__<'py>(
+        &'py self,
+        py: Python<'py>,
+        requested_schema: Option<Bound<'py, PyCapsule>>,
+    ) -> PyDataFusionResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)> {
+        let field = if let Some(schema_capsule) = requested_schema {
+            validate_pycapsule(&schema_capsule, "arrow_schema")?;
+
+            let schema_ptr = unsafe { schema_capsule.reference::<FFI_ArrowSchema>() };
+            let desired_field = Field::try_from(schema_ptr)?;
+
+            Arc::new(desired_field)
+        } else {
+            Arc::clone(&self.field)
+        };
+
+        let ffi_schema = FFI_ArrowSchema::try_from(&field)?;
+        let schema_capsule = PyCapsule::new(py, ffi_schema, Some(cr"arrow_schema".into()))?;
+
+        let ffi_array = FFI_ArrowArray::new(&self.array.to_data());
+        let array_capsule = PyCapsule::new(py, ffi_array, Some(cr"arrow_array".into()))?;
+
+        Ok((schema_capsule, array_capsule))
+    }
+}
+
+impl ToPyArrow for PyArrowArrayExportable {
+    fn to_pyarrow(&self, py: Python) -> PyResult<PyObject> {
+        let module = py.import("pyarrow")?;
+        let method = module.getattr("array")?;
+        let array = method.call((self.clone(),), None)?;
+        Ok(array.unbind())
+    }
+}
+
+impl PyArrowArrayExportable {
+    pub fn new(array: ArrayRef, field: FieldRef) -> Self {
+        Self { array, field }
+    }
+}

src/lib.rs

@@ -55,6 +55,7 @@ pub mod store;
 pub mod table;
 pub mod unparser;
 
+mod array;
 #[cfg(feature = "substrait")]
 pub mod substrait;
 #[allow(clippy::borrow_deref_ref)]

src/udf.rs

@@ -15,13 +15,14 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::array::PyArrowArrayExportable;
 use crate::errors::to_datafusion_err;
 use crate::errors::{py_datafusion_err, PyDataFusionResult};
 use crate::expr::PyExpr;
 use crate::utils::{parse_volatility, validate_pycapsule};
 use arrow::datatypes::{Field, FieldRef};
-use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema};
-use datafusion::arrow::array::{make_array, Array, ArrayData, ArrayRef};
+use arrow::pyarrow::ToPyArrow;
+use datafusion::arrow::array::{make_array, ArrayData};
 use datafusion::arrow::datatypes::DataType;
 use datafusion::arrow::pyarrow::FromPyArrow;
 use datafusion::arrow::pyarrow::PyArrowType;
@@ -31,12 +32,10 @@ use datafusion::logical_expr::{
     ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
 };
 use datafusion_ffi::udf::{FFI_ScalarUDF, ForeignScalarUDF};
-use pyo3::ffi::Py_uintptr_t;
 use pyo3::types::PyCapsule;
 use pyo3::{prelude::*, types::PyTuple};
 use std::any::Any;
 use std::hash::{Hash, Hasher};
-use std::ptr::addr_of;
 use std::sync::Arc;
 
 /// This struct holds the Python written function that is a
@@ -92,26 +91,6 @@ impl Hash for PythonFunctionScalarUDF {
     }
 }
 
-fn array_to_pyarrow_with_field(
-    py: Python,
-    array: ArrayRef,
-    field: &FieldRef,
-) -> PyResult<PyObject> {
-    let array = FFI_ArrowArray::new(&array.to_data());
-    let schema = FFI_ArrowSchema::try_from(field).map_err(py_datafusion_err)?;
-
-    let module = py.import("pyarrow")?;
-    let class = module.getattr("Array")?;
-    let array = class.call_method1(
-        "_import_from_c",
-        (
-            addr_of!(array) as Py_uintptr_t,
-            addr_of!(schema) as Py_uintptr_t,
-        ),
-    )?;
-    Ok(array.unbind())
-}
-
 impl ScalarUDFImpl for PythonFunctionScalarUDF {
     fn as_any(&self) -> &dyn Any {
         self
@@ -150,7 +129,9 @@ impl ScalarUDFImpl for PythonFunctionScalarUDF {
                 .zip(args.arg_fields)
                 .map(|(arg, field)| {
                     let array = arg.to_array(num_rows)?;
-                    array_to_pyarrow_with_field(py, array, &field).map_err(to_datafusion_err)
+                    PyArrowArrayExportable::new(array, field)
+                        .to_pyarrow(py)
+                        .map_err(to_datafusion_err)
                 })
                 .collect::<Result<Vec<_>, _>>()?;
             let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?;