Revert "Implement metadata-aware PySimpleScalarUDF"

This reverts commit 24549bd.

Author: kosiew

python/datafusion/user_defined.py

@@ -22,13 +22,17 @@
 import functools
 from abc import ABCMeta, abstractmethod
 from enum import Enum
-from typing import Any, Callable, Optional, Protocol, Sequence, overload
+from typing import TYPE_CHECKING, Any, Callable, Optional, Protocol, TypeVar, 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.
 
@@ -73,40 +77,6 @@ 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."""
 
@@ -123,9 +93,9 @@ class ScalarUDF:
     def __init__(
         self,
         name: str,
-        func: Callable[..., Any],
-        input_types: pa.DataType | pa.Field | Sequence[pa.DataType | pa.Field],
-        return_type: pa.DataType | pa.Field,
+        func: Callable[..., _R],
+        input_types: pa.DataType | list[pa.DataType],
+        return_type: _R,
         volatility: Volatility | str,
     ) -> None:
         """Instantiate a scalar user-defined function (UDF).
@@ -135,10 +105,10 @@ def __init__(
         if hasattr(func, "__datafusion_scalar_udf__"):
             self._udf = df_internal.ScalarUDF.from_pycapsule(func)
             return
-        normalized_inputs = _normalize_input_fields(input_types)
-        normalized_return = _normalize_return_field(return_type, name=name)
+        if isinstance(input_types, pa.DataType):
+            input_types = [input_types]
         self._udf = df_internal.ScalarUDF(
-            name, func, normalized_inputs, normalized_return, str(volatility)
+            name, func, input_types, return_type, str(volatility)
         )
 
     def __repr__(self) -> str:
@@ -157,18 +127,18 @@ def __call__(self, *args: Expr) -> Expr:
     @overload
     @staticmethod
     def udf(
-        input_types: list[pa.DataType | pa.Field],
-        return_type: pa.DataType | pa.Field,
+        input_types: list[pa.DataType],
+        return_type: _R,
         volatility: Volatility | str,
         name: Optional[str] = None,
     ) -> Callable[..., ScalarUDF]: ...
 
     @overload
     @staticmethod
     def udf(
-        func: Callable[..., Any],
-        input_types: list[pa.DataType | pa.Field],
-        return_type: pa.DataType | pa.Field,
+        func: Callable[..., _R],
+        input_types: list[pa.DataType],
+        return_type: _R,
         volatility: Volatility | str,
         name: Optional[str] = None,
     ) -> ScalarUDF: ...
@@ -194,11 +164,10 @@ 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 | 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.
+            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.
             volatility (Volatility | str): See `Volatility` for allowed values.
             name (Optional[str]): A descriptive name for the function.
 

python/tests/test_udf.py

@@ -124,58 +124,3 @@ 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,16 +22,13 @@ use pyo3::types::PyCapsule;
 use pyo3::{prelude::*, types::PyTuple};
 
 use datafusion::arrow::array::{make_array, Array, ArrayData, ArrayRef};
-use datafusion::arrow::datatypes::{DataType, Field};
+use datafusion::arrow::datatypes::DataType;
 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::ptr_eq::PtrEq;
-use datafusion::logical_expr::{
-    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
-    Volatility,
-};
+use datafusion::logical_expr::ScalarUDF;
+use datafusion::logical_expr::{create_udf, ColumnarValue};
 
 use crate::errors::to_datafusion_err;
 use crate::errors::{py_datafusion_err, PyDataFusionResult};
@@ -83,73 +80,6 @@ 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)]
@@ -164,19 +94,17 @@ impl PyScalarUDF {
     fn new(
         name: &str,
         func: PyObject,
-        input_types: PyArrowType<Vec<Field>>,
-        return_type: PyArrowType<Field>,
+        input_types: PyArrowType<Vec<DataType>>,
+        return_type: PyArrowType<DataType>,
         volatility: &str,
     ) -> PyResult<Self> {
-        let volatility = parse_volatility(volatility)?;
-        let scalar_impl = PySimpleScalarUDF::new(
+        let function = create_udf(
             name,
             input_types.0,
             return_type.0,
-            volatility,
+            parse_volatility(volatility)?,
             to_scalar_function_impl(func),
         );
-        let function = ScalarUDF::new_from_impl(scalar_impl);
         Ok(Self { function })
     }