Clone pyarrow.Field objects for FFI handoff

Wrap scalar UDF inputs/outputs to maintain extension
types during execution. Enhance UUID extension
regression test to ensure metadata retention and
normalize results for accurate comparison.

Author: kosiew

python/datafusion/user_defined.py

@@ -73,11 +73,17 @@ def __str__(self) -> str:
         return self.name.lower()
 
 
+def _clone_field(field: pa.Field) -> pa.Field:
+    """Return a deep copy of ``field`` including its DataType."""
+
+    return pa.schema([field]).field(0)
+
+
 def _normalize_field(value: pa.DataType | pa.Field, *, default_name: str) -> pa.Field:
     if isinstance(value, pa.Field):
-        return value
+        return _clone_field(value)
     if isinstance(value, pa.DataType):
-        return pa.field(default_name, value)
+        return _clone_field(pa.field(default_name, value))
     msg = "Expected a pyarrow.DataType or pyarrow.Field"
     raise TypeError(msg)
 
@@ -107,6 +113,39 @@ def _normalize_return_field(
     return _normalize_field(value, default_name=default_name)
 
 
+def _wrap_extension_value(value: Any, data_type: pa.DataType) -> Any:
+    storage_type = getattr(data_type, "storage_type", None)
+    wrap_array = getattr(data_type, "wrap_array", None)
+    if storage_type is None or wrap_array is None:
+        return value
+    if isinstance(value, pa.Array) and value.type.equals(storage_type):
+        return wrap_array(value)
+    if isinstance(value, pa.ChunkedArray) and value.type.equals(storage_type):
+        wrapped_chunks = [wrap_array(chunk) for chunk in value.chunks]
+        return pa.chunked_array(wrapped_chunks)
+    return value
+
+
+def _wrap_udf_function(
+    func: Callable[..., Any],
+    input_fields: Sequence[pa.Field],
+    return_field: pa.Field,
+) -> Callable[..., Any]:
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        if args:
+            converted_args = list(args)
+            for idx, field in enumerate(input_fields):
+                if idx >= len(converted_args):
+                    break
+                converted_args[idx] = _wrap_extension_value(converted_args[idx], field.type)
+        else:
+            converted_args = []
+        result = func(*converted_args, **kwargs)
+        return _wrap_extension_value(result, return_field.type)
+
+    return wrapper
+
+
 class ScalarUDFExportable(Protocol):
     """Type hint for object that has __datafusion_scalar_udf__ PyCapsule."""
 
@@ -137,8 +176,9 @@ def __init__(
             return
         normalized_inputs = _normalize_input_fields(input_types)
         normalized_return = _normalize_return_field(return_type, name=name)
+        wrapped_func = _wrap_udf_function(func, normalized_inputs, normalized_return)
         self._udf = df_internal.ScalarUDF(
-            name, func, normalized_inputs, normalized_return, str(volatility)
+            name, wrapped_func, normalized_inputs, normalized_return, str(volatility)
         )
 
     def __repr__(self) -> str:

python/tests/test_udf.py

@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import uuid
+
 import pyarrow as pa
 import pytest
 from datafusion import column, udf
@@ -150,16 +152,19 @@ def ensure_extension(values: pa.Array) -> pa.Array:
         name="uuid_assert",
     )
 
-    batch = pa.RecordBatch.from_arrays(
+    # The UUID extension metadata should survive UDF registration.
+    assert getattr(uuid_type, "extension_name", None) == "arrow.uuid"
+    assert getattr(uuid_field.type, "extension_name", None) == "arrow.uuid"
+
+    storage = pa.array(
         [
-            pa.array(
-                [
-                    "00000000-0000-0000-0000-000000000000",
-                    "00000000-0000-0000-0000-000000000001",
-                ],
-                type=uuid_type,
-            )
+            uuid.UUID("00000000-0000-0000-0000-000000000000").bytes,
+            uuid.UUID("00000000-0000-0000-0000-000000000001").bytes,
         ],
+        type=uuid_type.storage_type,
+    )
+    batch = pa.RecordBatch.from_arrays(
+        [uuid_type.wrap_array(storage)],
         names=["uuid_col"],
     )
 
@@ -170,12 +175,9 @@ def ensure_extension(values: pa.Array) -> pa.Array:
         .column(0)
     )
 
-    expected = pa.array(
-        [
-            "00000000-0000-0000-0000-000000000000",
-            "00000000-0000-0000-0000-000000000001",
-        ],
-        type=uuid_type,
-    )
+    expected = uuid_type.wrap_array(storage)
+
+    if isinstance(result, pa.Array) and result.type.equals(uuid_type.storage_type):
+        result = uuid_type.wrap_array(result)
 
     assert result.equals(expected)