Teach _wrap_extension_value to handle empty arrays

Return a wrapped empty extension array for chunked storage
arrays with no chunks, preserving extension metadata.
Expand UUID UDF regression to support chunked inputs,
test empty chunked returns, and ensure UUID extension
type remains intact through UDF chaining.

Author: kosiew

python/datafusion/user_defined.py

@@ -126,7 +126,10 @@ def _wrap_extension_value(value: PyArrowArrayT, data_type: pa.DataType) -> PyArr
         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)
+        if not wrapped_chunks:
+            empty_storage = pa.array([], type=storage_type)
+            return wrap_array(empty_storage)
+        return pa.chunked_array(wrapped_chunks, type=data_type)
     return value
 
 

python/tests/test_udf.py

@@ -147,8 +147,12 @@ def test_uuid_extension_chain(ctx) -> None:
         name="uuid_identity",
     )
 
-    def ensure_extension(values: pa.Array) -> pa.Array:
+    def ensure_extension(values: pa.Array | pa.ChunkedArray) -> pa.Array:
+        if isinstance(values, pa.ChunkedArray):
+            assert values.type.equals(uuid_type)
+            return values.combine_chunks()
         assert isinstance(values, pa.ExtensionArray)
+        assert values.type.equals(uuid_type)
         return values
 
     second = udf(
@@ -191,3 +195,34 @@ def ensure_extension(values: pa.Array) -> pa.Array:
         assert result.type.equals(uuid_type)
 
     assert result.equals(expected)
+
+    empty_storage = pa.array([], type=uuid_type.storage_type)
+    empty_batch = pa.RecordBatch.from_arrays(
+        [uuid_type.wrap_array(empty_storage)],
+        names=["uuid_col"],
+    )
+
+    empty_first = udf(
+        lambda values: pa.chunked_array([], type=uuid_type.storage_type),
+        [uuid_field],
+        uuid_field,
+        volatility="immutable",
+        name="uuid_empty_chunk",
+    )
+
+    empty_df = ctx.create_dataframe([[empty_batch]])
+    empty_result = (
+        empty_df.select(second(empty_first(column("uuid_col"))))
+        .collect()[0]
+        .column(0)
+    )
+
+    expected_empty = uuid_type.wrap_array(empty_storage)
+
+    if isinstance(empty_result, pa.ChunkedArray):
+        assert empty_result.type.equals(uuid_type)
+        assert empty_result.combine_chunks().equals(expected_empty)
+    else:
+        assert isinstance(empty_result, pa.ExtensionArray)
+        assert empty_result.type.equals(uuid_type)
+        assert empty_result.equals(expected_empty)