Implement metadata-aware PySimpleScalarUDF

kosiew · kosiew · commit d0b02a6cea7c · 2025-10-22T18:34:54.000+08:00
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.
diff --git a/python/datafusion/user_defined.py b/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.
 
diff --git a/python/tests/test_udf.py b/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)
diff --git a/src/udf.rs b/src/udf.rs
@@ -15,20 +15,24 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::fmt;
 use std::sync::Arc;
 
 use datafusion_ffi::udf::{FFI_ScalarUDF, ForeignScalarUDF};
 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 +84,83 @@ fn to_scalar_function_impl(func: PyObject) -> ScalarFunctionImplementation {
     })
 }
 
+#[derive(PartialEq, Eq, Hash)]
+struct PySimpleScalarUDF {
+    name: String,
+    signature: Signature,
+    return_field: Arc<Field>,
+    fun: PtrEq<ScalarFunctionImplementation>,
+}
+
+impl fmt::Debug for PySimpleScalarUDF {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("PySimpleScalarUDF")
+            .field("name", &self.name)
+            .field("signature", &self.signature)
+            .field("return_field", &self.return_field)
+            .finish()
+    }
+}
+
+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 +175,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 })
     }
 
