This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new be299a1efd4e [SPARK-57361][PYTHON] Refactor SQL_ARROW_UDTF
be299a1efd4e is described below
commit be299a1efd4e30ddbee15e874bb6cfd2ff95f179
Author: Yicong Huang <[email protected]>
AuthorDate: Thu Jun 11 18:05:20 2026 +0000
[SPARK-57361][PYTHON] Refactor SQL_ARROW_UDTF
### What changes were proposed in this pull request?
This PR refactors `SQL_ARROW_UDTF` so that the worker uses the plain
`ArrowStreamSerializer` for pure Arrow stream I/O, moving the per-batch
transformation logic from `ArrowStreamArrowUDTFSerializer` into `read_udtf()`
in `worker.py`:
- The input-side flattening (struct columns at `table_arg_offsets` are
flattened into `pa.RecordBatch`, other columns are passed as `pa.Array`) moves
from `ArrowStreamArrowUDTFSerializer.load_stream` into the `func` execution
block (same style as the refactored eval types in `read_udfs()`).
- The output-side type coercion (`ArrowBatchTransformer.enforce_schema`
against the Arrow return type) and struct wrapping
(`ArrowBatchTransformer.wrap_struct`) move from the serializer `dump_stream`
chain into the `evaluate` wrapper, which now yields ready-to-write record
batches instead of `(batch, arrow_return_type)` tuples.
`ArrowStreamArrowUDTFSerializer` itself is left in place and will be
removed in a follow-up once it has no remaining usages.
### Why are the changes needed?
Part of [SPARK-55388](https://issues.apache.org/jira/browse/SPARK-55388).
Keeping serializers as pure Arrow stream I/O and concentrating
eval-type-specific logic in `worker.py` makes the per-eval-type data flow
explicit and removes serializer subclasses that exist only to carry
per-eval-type transforms.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests (`pyspark.sql.tests.arrow.test_arrow_udtf`,
`pyspark.sql.tests.test_udtf`). No behavior change: replaying identical worker
input through the old and new code paths produces byte-identical worker output
(modulo the timing section) across 9 scenario/UDTF combinations.
ASV benchmark comparison (`bench_eval_type.ArrowUDTFTimeBench`, `-a
repeat=3`, 3 runs per side, averaged). before = `upstream/master`, after = this
PR.
```text
scenario udtf before after diff
------------------ -------------- ------- ------- -------
sm_batch_few_col identity_udtf 1.37ms 1.35ms -1.5%
sm_batch_few_col filter_udtf 1.81ms 1.84ms +1.7%
sm_batch_few_col count_udtf 1.07ms 1.05ms -2.2%
sm_batch_many_col identity_udtf 2.42ms 2.33ms -3.9%
sm_batch_many_col filter_udtf 3.50ms 3.43ms -2.2%
sm_batch_many_col count_udtf 1.14ms 1.11ms -2.6%
lg_batch_few_col identity_udtf 2.00ms 2.21ms +10.7%
lg_batch_few_col filter_udtf 3.33ms 3.40ms +2.2%
lg_batch_few_col count_udtf 1.20ms 1.20ms +0.3%
lg_batch_many_col identity_udtf 9.72ms 9.85ms +1.4%
lg_batch_many_col filter_udtf 15.57ms 15.70ms +0.9%
lg_batch_many_col count_udtf 3.54ms 3.49ms -1.5%
pure_ints identity_udtf 3.20ms 3.19ms -0.2%
pure_ints filter_udtf 4.77ms 4.38ms -8.1%
pure_ints count_udtf 1.90ms 1.85ms -2.6%
pure_strings identity_udtf 5.16ms 4.79ms -7.2%
pure_strings filter_udtf 9.15ms 9.02ms -1.4%
pure_strings count_udtf 2.28ms 2.25ms -1.2%
```
The `lg_batch_few_col / identity_udtf` cell is a noise artifact: its
per-run confidence intervals overlap (both sides spike to ~2.3ms intermittently
on this machine), and re-running it in isolation with min-of-30 direct timing
shows the refactored path is faster (min 2.34ms vs 2.59ms, median 2.36ms vs
2.76ms).
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #56440 from Yicong-Huang/refactor/arrow-udtf.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/worker.py | 174 +++++++++++++++++++++++------------------------
1 file changed, 86 insertions(+), 88 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 47a80173f86e..74319b08927a 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -87,7 +87,6 @@ from pyspark.sql.pandas.serializers import (
TransformWithStateInPySparkRowInitStateSerializer,
ArrowStreamAggPandasUDFSerializer,
ArrowStreamUDTFSerializer,
- ArrowStreamArrowUDTFSerializer,
)
from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type
from pyspark.sql.types import (
@@ -1050,9 +1049,9 @@ def read_udtf(pickleSer, udtf_info, eval_type,
runner_conf, eval_conf):
else:
ser = ArrowStreamUDTFSerializer()
elif eval_type == PythonEvalType.SQL_ARROW_UDTF:
- # Read the table argument offsets
- # Use PyArrow-native serializer for Arrow UDTFs with potential UDT
support
- ser =
ArrowStreamArrowUDTFSerializer(table_arg_offsets=eval_conf.table_arg_offsets)
+ # Pure Arrow stream I/O; table-arg flattening and output coercion
+ # are handled in the mapper below.
+ ser = ArrowStreamSerializer(write_start_stream=True)
else:
# Each row is a group so do not batch but send one by one.
ser = BatchedSerializer(CPickleSerializer(), 1)
@@ -1971,113 +1970,112 @@ def read_udtf(pickleSer, udtf_info, eval_type,
runner_conf, eval_conf):
return mapper, None, ser, ser
elif eval_type == PythonEvalType.SQL_ARROW_UDTF:
+ import pyarrow as pa
- def wrap_pyarrow_udtf(f, return_type):
- import pyarrow as pa
-
- arrow_return_type = to_arrow_type(
- return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
- )
- return_type_size = len(return_type)
+ arrow_return_type = to_arrow_type(
+ return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
+ )
+ return_type_size = len(return_type)
+ target_schema = pa.schema(list(arrow_return_type))
- def verify_result(result):
- # Validate the output schema when the result has columns
- if result.num_columns != return_type_size:
- raise PySparkRuntimeError(
- errorClass="UDTF_RETURN_SCHEMA_MISMATCH",
- messageParameters={
- "expected": str(return_type_size),
- "actual": str(result.num_columns),
- "func": f.__name__,
- },
- )
+ def verify_result(result: pa.RecordBatch, method_name: str) ->
pa.RecordBatch:
+ # Validate the output schema when the result has columns
+ if result.num_columns != return_type_size:
+ raise PySparkRuntimeError(
+ errorClass="UDTF_RETURN_SCHEMA_MISMATCH",
+ messageParameters={
+ "expected": str(return_type_size),
+ "actual": str(result.num_columns),
+ "func": method_name,
+ },
+ )
+ return result
- # We verify the type of the result and do type corerion
- # in the serializer
- return result
+ def convert_to_arrow(res: Any, method_name: str) ->
Iterator[pa.RecordBatch]:
+ # Check whether the result of a PyArrow UDTF is iterable before
processing
+ if res is None:
+ res = iter([])
+ elif not isinstance(res, Iterable):
+ raise PySparkRuntimeError(
+ errorClass="UDTF_RETURN_NOT_ITERABLE",
+ messageParameters={
+ "type": type(res).__name__,
+ "func": method_name,
+ },
+ )
- # Wrap the exception thrown from the UDTF in a PySparkRuntimeError.
- def func(*a: Any) -> Any:
- try:
- return f(*a)
- except SkipRestOfInputTableException:
- raise
- except Exception as e:
+ # Handle PyArrow Tables/RecordBatches directly
+ is_empty = True
+ for item in res:
+ is_empty = False
+ if isinstance(item, pa.Table):
+ yield from item.to_batches()
+ elif isinstance(item, pa.RecordBatch):
+ yield item
+ else:
+ # Arrow UDTF should only return Arrow types
(RecordBatch/Table)
raise PySparkRuntimeError(
- errorClass="UDTF_EXEC_ERROR",
- messageParameters={"method_name": f.__name__, "error":
str(e)},
+ errorClass="UDTF_ARROW_TYPE_CONVERSION_ERROR",
+ messageParameters={},
)
- def check_return_value(res):
- # Check whether the result of a PyArrow UDTF is iterable
before processing
- if res is not None:
- if not isinstance(res, Iterable):
- raise PySparkRuntimeError(
- errorClass="UDTF_RETURN_NOT_ITERABLE",
- messageParameters={
- "type": type(res).__name__,
- "func": f.__name__,
- },
- )
- return res
- else:
- return iter([])
-
- def convert_to_arrow(data: Iterable):
- data_iter = check_return_value(data)
-
- # Handle PyArrow Tables/RecordBatches directly
- is_empty = True
- for item in data_iter:
- is_empty = False
- if isinstance(item, pa.Table):
- yield from item.to_batches()
- elif isinstance(item, pa.RecordBatch):
- yield item
- else:
- # Arrow UDTF should only return Arrow types
(RecordBatch/Table)
- raise PySparkRuntimeError(
- errorClass="UDTF_ARROW_TYPE_CONVERSION_ERROR",
- messageParameters={},
- )
+ if is_empty:
+ yield pa.RecordBatch.from_pylist([], schema=target_schema)
- if is_empty:
- yield pa.RecordBatch.from_pylist([],
schema=pa.schema(list(arrow_return_type)))
-
- def evaluate(*args: pa.RecordBatch):
- # For Arrow UDTFs, unpack the RecordBatches and pass them to
the function
- for batch in convert_to_arrow(func(*args)):
- yield verify_result(batch), arrow_return_type
+ def evaluate(method: Callable, *args: pa.RecordBatch) ->
Iterator[pa.RecordBatch]:
+ # Wrap the exception thrown from the UDTF in a PySparkRuntimeError.
+ try:
+ res = method(*args)
+ except SkipRestOfInputTableException:
+ raise
+ except Exception as e:
+ raise PySparkRuntimeError(
+ errorClass="UDTF_EXEC_ERROR",
+ messageParameters={"method_name": method.__name__,
"error": str(e)},
+ )
- return evaluate
+ for batch in convert_to_arrow(res, method.__name__):
+ coerced = ArrowBatchTransformer.enforce_schema(
+ verify_result(batch, method.__name__), target_schema,
safecheck=True
+ )
+ yield ArrowBatchTransformer.wrap_struct(coerced)
- eval_func_kwargs_support, args_kwargs_offsets = wrap_kwargs_support(
+ eval_method, args_kwargs_offsets = wrap_kwargs_support(
getattr(udtf, "eval"), udtf_info.args, udtf_info.kwargs
)
- eval = wrap_pyarrow_udtf(eval_func_kwargs_support, return_type)
-
- if hasattr(udtf, "terminate"):
- terminate = wrap_pyarrow_udtf(getattr(udtf, "terminate"),
return_type)
- else:
- terminate = None
+ terminate = getattr(udtf, "terminate", None)
+ cleanup = getattr(udtf, "cleanup", None)
- cleanup = getattr(udtf, "cleanup") if hasattr(udtf, "cleanup") else
None
+ table_arg_offsets = (
+ set(eval_conf.table_arg_offsets) if eval_conf.table_arg_offsets
else set()
+ )
- def mapper(_, it):
+ def func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
+ """Apply Arrow UDTF"""
try:
- for a in it:
+ for batch in data:
+ # Pre-processing: for each column, flatten struct columns
at
+ # table_arg_offsets into RecordBatch, keep other columns
as Array.
+ columns = [
+ (
+ ArrowBatchTransformer.flatten_struct(batch,
column_index=i)
+ if i in table_arg_offsets
+ else batch.column(i)
+ )
+ for i in range(batch.num_columns)
+ ]
# For PyArrow UDTFs, pass RecordBatches directly (no row
conversion needed)
- yield from eval(*[a[o] for o in args_kwargs_offsets])
+ yield from evaluate(eval_method, *[columns[o] for o in
args_kwargs_offsets])
if terminate is not None:
- yield from terminate()
+ yield from evaluate(terminate)
except SkipRestOfInputTableException:
if terminate is not None:
- yield from terminate()
+ yield from evaluate(terminate)
finally:
if cleanup is not None:
cleanup()
- return mapper, None, ser, ser
+ return func, None, ser, ser
else:
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
