Implement metadata-aware PySimpleScalarUDF

Enhance scalar UDF definitions to retain Arrow Field
information, including extension metadata, in DataFusion.
Normalize Python UDF signatures to accept pyarrow.Field
objects, ensuring metadata survives the Rust bindings roundtrip.
Add a regression test for UUID-backed UDFs to verify
that the second UDF correctly receives a
pyarrow.ExtensionArray, preventing past metadata loss.

Author: kosiew

python/datafusion/user_defined.py

@@ -22,17 +22,13 @@
 import functools
 from abc import ABCMeta, abstractmethod
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Callable, Optional, Protocol, TypeVar, overload
+from typing import Any, Callable, Optional, Protocol, Sequence, overload
 
 import pyarrow as pa
 
 import datafusion._internal as df_internal
 from datafusion.expr import Expr
 
-if TYPE_CHECKING:
-    _R = TypeVar("_R", bound=pa.DataType)
-
-
 class Volatility(Enum):
     """Defines how stable or volatile a function is.
 
@@ -77,6 +73,40 @@ def __str__(self) -> str:
         return self.name.lower()
 
 
+def _normalize_field(value: pa.DataType | pa.Field, *, default_name: str) -> pa.Field:
+    if isinstance(value, pa.Field):
+        return value
+    if isinstance(value, pa.DataType):
+        return pa.field(default_name, value)
+    msg = "Expected a pyarrow.DataType or pyarrow.Field"
+    raise TypeError(msg)
+
+
+def _normalize_input_fields(
+    values: pa.DataType | pa.Field | Sequence[pa.DataType | pa.Field],
+) -> list[pa.Field]:
+    if isinstance(values, (pa.DataType, pa.Field)):
+        sequence: Sequence[pa.DataType | pa.Field] = [values]
+    elif isinstance(values, Sequence) and not isinstance(values, (str, bytes)):
+        sequence = values
+    else:
+        msg = "input_types must be a DataType, Field, or a sequence of them"
+        raise TypeError(msg)
+
+    return [
+        _normalize_field(value, default_name=f"arg_{idx}") for idx, value in enumerate(sequence)
+    ]
+
+
+def _normalize_return_field(
+    value: pa.DataType | pa.Field,
+    *,
+    name: str,
+) -> pa.Field:
+    default_name = f"{name}_result" if name else "result"
+    return _normalize_field(value, default_name=default_name)
+
+
 class ScalarUDFExportable(Protocol):
     """Type hint for object that has __datafusion_scalar_udf__ PyCapsule."""
 
@@ -93,9 +123,9 @@ class ScalarUDF:
     def __init__(
         self,
         name: str,
-        func: Callable[..., _R],
-        input_types: pa.DataType | list[pa.DataType],
-        return_type: _R,
+        func: Callable[..., Any],
+        input_types: pa.DataType | pa.Field | Sequence[pa.DataType | pa.Field],
+        return_type: pa.DataType | pa.Field,
         volatility: Volatility | str,
     ) -> None:
         """Instantiate a scalar user-defined function (UDF).
@@ -105,10 +135,10 @@ def __init__(
         if hasattr(func, "__datafusion_scalar_udf__"):
             self._udf = df_internal.ScalarUDF.from_pycapsule(func)
             return
-        if isinstance(input_types, pa.DataType):
-            input_types = [input_types]
+        normalized_inputs = _normalize_input_fields(input_types)
+        normalized_return = _normalize_return_field(return_type, name=name)
         self._udf = df_internal.ScalarUDF(
-            name, func, input_types, return_type, str(volatility)
+            name, func, normalized_inputs, normalized_return, str(volatility)
         )
 
     def __repr__(self) -> str:
@@ -127,18 +157,18 @@ def __call__(self, *args: Expr) -> Expr:
     @overload
     @staticmethod
     def udf(
-        input_types: list[pa.DataType],
-        return_type: _R,
+        input_types: list[pa.DataType | pa.Field],
+        return_type: pa.DataType | pa.Field,
         volatility: Volatility | str,
         name: Optional[str] = None,
     ) -> Callable[..., ScalarUDF]: ...
 
     @overload
     @staticmethod
     def udf(
-        func: Callable[..., _R],
-        input_types: list[pa.DataType],
-        return_type: _R,
+        func: Callable[..., Any],
+        input_types: list[pa.DataType | pa.Field],
+        return_type: pa.DataType | pa.Field,
         volatility: Volatility | str,
         name: Optional[str] = None,
     ) -> ScalarUDF: ...
@@ -164,10 +194,11 @@ def udf(*args: Any, **kwargs: Any):  # noqa: D417
                 backed ScalarUDF within a PyCapsule, you can pass this parameter
                 and ignore the rest. They will be determined directly from the
                 underlying function. See the online documentation for more information.
-            input_types (list[pa.DataType]): The data types of the arguments
-                to ``func``. This list must be of the same length as the number of
-                arguments.
-            return_type (_R): The data type of the return value from the function.
+            input_types (list[pa.DataType | pa.Field]): The argument types for ``func``.
+                This list must be of the same length as the number of arguments. Pass
+                :class:`pyarrow.Field` instances to preserve extension metadata.
+            return_type (pa.DataType | pa.Field): The return type of the function. Use a
+                :class:`pyarrow.Field` to preserve metadata on extension arrays.
             volatility (Volatility | str): See `Volatility` for allowed values.
             name (Optional[str]): A descriptive name for the function.
 

python/tests/test_udf.py

@@ -124,3 +124,58 @@ def udf_with_param(values: pa.Array) -> pa.Array:
     result = df2.collect()[0].column(0)
 
     assert result == pa.array([False, True, True])
+
+
+def test_uuid_extension_chain(ctx) -> None:
+    uuid_type = pa.uuid()
+    uuid_field = pa.field("uuid_col", uuid_type)
+
+    first = udf(
+        lambda values: values,
+        [uuid_field],
+        uuid_field,
+        volatility="immutable",
+        name="uuid_identity",
+    )
+
+    def ensure_extension(values: pa.Array) -> pa.Array:
+        assert isinstance(values, pa.ExtensionArray)
+        return values
+
+    second = udf(
+        ensure_extension,
+        [uuid_field],
+        uuid_field,
+        volatility="immutable",
+        name="uuid_assert",
+    )
+
+    batch = pa.RecordBatch.from_arrays(
+        [
+            pa.array(
+                [
+                    "00000000-0000-0000-0000-000000000000",
+                    "00000000-0000-0000-0000-000000000001",
+                ],
+                type=uuid_type,
+            )
+        ],
+        names=["uuid_col"],
+    )
+
+    df = ctx.create_dataframe([[batch]])
+    result = (
+        df.select(second(first(column("uuid_col"))))
+        .collect()[0]
+        .column(0)
+    )
+
+    expected = pa.array(
+        [
+            "00000000-0000-0000-0000-000000000000",
+            "00000000-0000-0000-0000-000000000001",
+        ],
+        type=uuid_type,
+    )
+
+    assert result.equals(expected)

src/udf.rs

@@ -22,13 +22,16 @@ use pyo3::types::PyCapsule;
 use pyo3::{prelude::*, types::PyTuple};
 
 use datafusion::arrow::array::{make_array, Array, ArrayData, ArrayRef};
-use datafusion::arrow::datatypes::DataType;
+use datafusion::arrow::datatypes::{DataType, Field};
 use datafusion::arrow::pyarrow::FromPyArrow;
 use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow};
 use datafusion::error::DataFusionError;
 use datafusion::logical_expr::function::ScalarFunctionImplementation;
-use datafusion::logical_expr::ScalarUDF;
-use datafusion::logical_expr::{create_udf, ColumnarValue};
+use datafusion::logical_expr::ptr_eq::PtrEq;
+use datafusion::logical_expr::{
+    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
+    Volatility,
+};
 
 use crate::errors::to_datafusion_err;
 use crate::errors::{py_datafusion_err, PyDataFusionResult};
@@ -80,6 +83,73 @@ fn to_scalar_function_impl(func: PyObject) -> ScalarFunctionImplementation {
     })
 }
 
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct PySimpleScalarUDF {
+    name: String,
+    signature: Signature,
+    return_field: Arc<Field>,
+    fun: PtrEq<ScalarFunctionImplementation>,
+}
+
+impl PySimpleScalarUDF {
+    fn new(
+        name: impl Into<String>,
+        input_fields: Vec<Field>,
+        return_field: Field,
+        volatility: Volatility,
+        fun: ScalarFunctionImplementation,
+    ) -> Self {
+        let signature_types = input_fields
+            .into_iter()
+            .map(|field| field.data_type().clone())
+            .collect();
+        let signature = Signature::exact(signature_types, volatility);
+        Self {
+            name: name.into(),
+            signature,
+            return_field: Arc::new(return_field),
+            fun: fun.into(),
+        }
+    }
+}
+
+impl ScalarUDFImpl for PySimpleScalarUDF {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        &self.name
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> datafusion::error::Result<DataType> {
+        Ok(self.return_field.data_type().clone())
+    }
+
+    fn return_field_from_args(
+        &self,
+        _args: ReturnFieldArgs,
+    ) -> datafusion::error::Result<Arc<Field>> {
+        Ok(Arc::new(
+            self.return_field
+                .as_ref()
+                .clone()
+                .with_name(self.name.clone()),
+        ))
+    }
+
+    fn invoke_with_args(
+        &self,
+        args: ScalarFunctionArgs,
+    ) -> datafusion::error::Result<ColumnarValue> {
+        (self.fun)(&args.args)
+    }
+}
+
 /// Represents a PyScalarUDF
 #[pyclass(frozen, name = "ScalarUDF", module = "datafusion", subclass)]
 #[derive(Debug, Clone)]
@@ -94,17 +164,19 @@ impl PyScalarUDF {
     fn new(
         name: &str,
         func: PyObject,
-        input_types: PyArrowType<Vec<DataType>>,
-        return_type: PyArrowType<DataType>,
+        input_types: PyArrowType<Vec<Field>>,
+        return_type: PyArrowType<Field>,
         volatility: &str,
     ) -> PyResult<Self> {
-        let function = create_udf(
+        let volatility = parse_volatility(volatility)?;
+        let scalar_impl = PySimpleScalarUDF::new(
             name,
             input_types.0,
             return_type.0,
-            parse_volatility(volatility)?,
+            volatility,
             to_scalar_function_impl(func),
         );
+        let function = ScalarUDF::new_from_impl(scalar_impl);
         Ok(Self { function })
     }