This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 37194b426423 [SPARK-57394][PYTHON] Refactor SQL_ARROW_TABLE_UDF
37194b426423 is described below
commit 37194b426423be5874bc733f7d8501d435e4d1eb
Author: Yicong Huang <[email protected]>
AuthorDate: Mon Jul 6 23:50:20 2026 +0000
[SPARK-57394][PYTHON] Refactor SQL_ARROW_TABLE_UDF
### What changes were proposed in this pull request?
This PR refactors `SQL_ARROW_TABLE_UDF` (arrow-optimized Python UDTF,
`udtf(useArrow=True)`) so that the worker uses the plain
`ArrowStreamSerializer` for pure Arrow stream I/O, moving the remaining
per-batch transformation logic from `ArrowStreamUDTFSerializer` into
`read_udtf()` in `worker.py`:
- The input side already received raw Arrow record batches
(`ArrowStreamUDTFSerializer.load_stream` delegates to
`ArrowStreamSerializer.load_stream`), so loading is unchanged.
- The output-side struct wrapping (`ArrowBatchTransformer.wrap_struct`)
moves 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.
The legacy pandas conversion path (`use_legacy_pandas_udtf_conversion`,
using `ArrowStreamPandasUDTFSerializer`) is unchanged.
`ArrowStreamUDTFSerializer` 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.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 24
scenario/UDTF combinations.
ASV benchmark comparison (`bench_eval_type.ArrowTableUDFTimeBench`, `-a
repeat=3`, 3 runs per side, averaged). before = `upstream/master`, after = this
PR.
```text
scenario udtf before (ms) after (ms) diff
------------------- --------------- ----------- ----------- --------
sm_batch_few_col identity_udtf 155.7ms 155.3ms -0.2%
sm_batch_few_col explode_udtf 157.3ms 164.7ms +4.7%
sm_batch_few_col filter_udtf 124.7ms 123.0ms -1.3%
sm_batch_few_col stringify_udtf 155.7ms 155.7ms +0.0%
sm_batch_many_col identity_udtf 51.8ms 51.4ms -0.8%
sm_batch_many_col explode_udtf 52.6ms 52.3ms -0.7%
sm_batch_many_col filter_udtf 44.9ms 43.1ms -4.1%
sm_batch_many_col stringify_udtf 51.7ms 51.4ms -0.5%
lg_batch_few_col identity_udtf 397.7ms 387.0ms -2.7%
lg_batch_few_col explode_udtf 393.7ms 393.3ms -0.1%
lg_batch_few_col filter_udtf 302.0ms 300.7ms -0.4%
lg_batch_few_col stringify_udtf 386.7ms 388.3ms +0.4%
lg_batch_many_col identity_udtf 207.0ms 203.7ms -1.6%
lg_batch_many_col explode_udtf 207.7ms 205.0ms -1.3%
lg_batch_many_col filter_udtf 175.3ms 169.0ms -3.6%
lg_batch_many_col stringify_udtf 206.7ms 207.0ms +0.2%
pure_ints identity_udtf 392.7ms 401.3ms +2.2%
pure_ints explode_udtf 399.0ms 399.7ms +0.2%
pure_ints filter_udtf 310.3ms 307.3ms -1.0%
pure_ints stringify_udtf 392.0ms 394.7ms +0.7%
pure_strings identity_udtf 410.7ms 419.0ms +2.0%
pure_strings explode_udtf 416.0ms 427.3ms +2.7%
pure_strings filter_udtf 329.0ms 324.3ms -1.4%
pure_strings stringify_udtf 412.0ms 409.0ms -0.7%
```
The `sm_batch_few_col / explode_udtf` cell is a noise artifact from one
noisy ASV run (a 180ms outlier; the other two runs measured 159ms/155ms, in
line with before). Re-running it in isolation with min-of-30 direct timing
shows no regression (min 153.0ms before vs 153.3ms after, median 154.6ms vs
154.9ms).
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #56458 from Yicong-Huang/refactor/arrow-table-udf.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
(cherry picked from commit 13d5f1ccd1b772a937bec93187bdc126576d64a6)
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/worker.py | 249 ++++++++++++++++++++++-------------------------
1 file changed, 118 insertions(+), 131 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 18cfc15f8f89..9e629138d537 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -83,7 +83,6 @@ from pyspark.sql.pandas.serializers import (
TransformWithStateInPandasInitStateSerializer,
TransformWithStateInPySparkRowSerializer,
TransformWithStateInPySparkRowInitStateSerializer,
- ArrowStreamUDTFSerializer,
)
from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type
from pyspark.sql.types import (
@@ -868,10 +867,12 @@ def read_udtf(pickleSer, udtf_info, eval_type,
runner_conf, eval_conf):
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
)
else:
- ser = ArrowStreamUDTFSerializer()
+ # Pure Arrow stream I/O; output struct wrapping is handled in the
+ # func below.
+ ser = ArrowStreamSerializer(write_start_stream=True)
elif eval_type == PythonEvalType.SQL_ARROW_UDTF:
# Pure Arrow stream I/O; table-arg flattening and output coercion
- # are handled in the mapper below.
+ # are handled in the func below.
ser = ArrowStreamSerializer(write_start_stream=True)
else:
# Each row is a group so do not batch but send one by one.
@@ -1616,148 +1617,130 @@ def read_udtf(pickleSer, udtf_info, eval_type,
runner_conf, eval_conf):
eval_type == PythonEvalType.SQL_ARROW_TABLE_UDF
and not runner_conf.use_legacy_pandas_udtf_conversion
):
+ import pyarrow as pa
- def wrap_arrow_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)
+ def verify_result(result: pa.Table, method_name: str) -> pa.Table:
+ if not isinstance(result, pa.Table):
+ raise PySparkTypeError(
+ errorClass="INVALID_ARROW_UDTF_RETURN_TYPE",
+ messageParameters={
+ "return_type": type(result).__name__,
+ "value": str(result),
+ "func": method_name,
+ },
+ )
- def verify_result(result):
- if not isinstance(result, pa.Table):
- raise PySparkTypeError(
- errorClass="INVALID_ARROW_UDTF_RETURN_TYPE",
+ # Validate the output schema when the result dataframe has either
output
+ # rows or columns. Note that we avoid using `df.empty` here
because the
+ # result dataframe may contain an empty row. For example, when a
UDTF is
+ # defined as follows: def eval(self): yield tuple().
+ if result.num_rows > 0 or result.num_columns > 0:
+ if result.num_columns != return_type_size:
+ raise PySparkRuntimeError(
+ errorClass="UDTF_RETURN_SCHEMA_MISMATCH",
messageParameters={
- "return_type": type(result).__name__,
- "value": str(result),
- "func": f.__name__,
+ "expected": str(return_type_size),
+ "actual": str(result.num_columns),
+ "func": method_name,
},
)
- # Validate the output schema when the result dataframe has
either output
- # rows or columns. Note that we avoid using `df.empty` here
because the
- # result dataframe may contain an empty row. For example, when
a UDTF is
- # defined as follows: def eval(self): yield tuple().
- if result.num_rows > 0 or result.num_columns > 0:
- 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__,
- },
- )
-
- # Verify the type and the schema of the result.
- verify_arrow_result(
- result,
- assign_cols_by_name=False,
- expected_cols_and_types=[
- (field.name, field.type) for field in arrow_return_type
- ],
- )
- return result
+ # Verify the type and the schema of the result.
+ verify_arrow_result(
+ result,
+ assign_cols_by_name=False,
+ expected_cols_and_types=[(field.name, field.type) for field in
arrow_return_type],
+ )
+ return result
- # 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:
+ def check_return_value(res: Any, method_name: str) -> Iterator:
+ # Check whether the result of an arrow UDTF is iterable before
+ # using it to construct a pandas DataFrame.
+ if res is not None:
+ if not isinstance(res, Iterable):
raise PySparkRuntimeError(
- errorClass="UDTF_EXEC_ERROR",
- messageParameters={"method_name": f.__name__, "error":
str(e)},
+ errorClass="UDTF_RETURN_NOT_ITERABLE",
+ messageParameters={
+ "type": type(res).__name__,
+ "func": method_name,
+ },
)
+ for row in res:
+ if not isinstance(row, tuple) and return_type_size == 1:
+ row = (row,)
+ if check_output_row_against_schema is not None:
+ if row is not None:
+ check_output_row_against_schema(row)
+ yield row
- def check_return_value(res):
- # Check whether the result of an arrow UDTF is iterable before
- # using it to construct a pandas DataFrame.
- 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__,
- },
- )
- for row in res:
- if not isinstance(row, tuple) and return_type_size ==
1:
- row = (row,)
- if check_output_row_against_schema is not None:
- if row is not None:
- check_output_row_against_schema(row)
- yield row
-
- def convert_to_arrow(data: Iterable):
- data = list(check_return_value(data))
- if len(data) == 0:
- # Return one empty RecordBatch to match the left side of
the lateral join
- return [
- pa.RecordBatch.from_pylist(data,
schema=pa.schema(list(arrow_return_type)))
- ]
+ def convert_rows_to_arrow(data: Iterable, method_name: str) ->
list[pa.RecordBatch]:
+ data = list(check_return_value(data, method_name))
+ if len(data) == 0:
+ # Return one empty RecordBatch to match the left side of the
lateral join
+ return [pa.RecordBatch.from_pylist(data,
schema=pa.schema(list(arrow_return_type)))]
+
+ def raise_conversion_error(original_exception):
+ raise PySparkRuntimeError(
+ errorClass="UDTF_ARROW_DATA_CONVERSION_ERROR",
+ messageParameters={
+ "data": str(data),
+ "schema": return_type.simpleString(),
+ "arrow_schema": str(arrow_return_type),
+ },
+ ) from original_exception
- def raise_conversion_error(original_exception):
+ try:
+ table = LocalDataToArrowConversion.convert(
+ data, return_type, runner_conf.use_large_var_types
+ )
+ except PySparkValueError as e:
+ if e.getErrorClass() == "AXIS_LENGTH_MISMATCH":
raise PySparkRuntimeError(
- errorClass="UDTF_ARROW_DATA_CONVERSION_ERROR",
+ errorClass="UDTF_RETURN_SCHEMA_MISMATCH",
messageParameters={
- "data": str(data),
- "schema": return_type.simpleString(),
- "arrow_schema": str(arrow_return_type),
+ "expected":
e.getMessageParameters()["expected_length"], # type: ignore[index]
+ "actual":
e.getMessageParameters()["actual_length"], # type: ignore[index]
+ "func": method_name,
},
- ) from original_exception
+ ) from e
+ # Fall through to general conversion error
+ raise_conversion_error(e)
+ except Exception as e:
+ raise_conversion_error(e)
+
+ return verify_result(table, method_name).to_batches()
+ def evaluate_rows(
+ method: Callable, *args: list, num_rows: int = 1
+ ) -> Iterator[pa.RecordBatch]:
+ rows = itertools.repeat((), num_rows) if len(args) == 0 else
zip(*args)
+ for row in rows:
+ # Wrap the exception thrown from the UDTF in a
PySparkRuntimeError.
try:
- table = LocalDataToArrowConversion.convert(
- data, return_type, runner_conf.use_large_var_types
- )
- except PySparkValueError as e:
- if e.getErrorClass() == "AXIS_LENGTH_MISMATCH":
- raise PySparkRuntimeError(
- errorClass="UDTF_RETURN_SCHEMA_MISMATCH",
- messageParameters={
- "expected":
e.getMessageParameters()["expected_length"], # type: ignore[index]
- "actual":
e.getMessageParameters()["actual_length"], # type: ignore[index]
- "func": f.__name__,
- },
- ) from e
- # Fall through to general conversion error
- raise_conversion_error(e)
+ res = method(*row)
+ except SkipRestOfInputTableException:
+ raise
except Exception as e:
- raise_conversion_error(e)
-
- return verify_result(table).to_batches()
-
- def evaluate(*args: list, num_rows=1):
- if len(args) == 0:
- for _ in range(num_rows):
- for batch in convert_to_arrow(func()):
- yield batch, arrow_return_type
-
- else:
- for row in zip(*args):
- for batch in convert_to_arrow(func(*row)):
- yield batch, arrow_return_type
-
- return evaluate
+ raise PySparkRuntimeError(
+ errorClass="UDTF_EXEC_ERROR",
+ messageParameters={"method_name": method.__name__,
"error": str(e)},
+ )
+ for batch in convert_rows_to_arrow(res, method.__name__):
+ yield ArrowBatchTransformer.wrap_struct(batch)
- 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_arrow_udtf(eval_func_kwargs_support, return_type)
-
- if hasattr(udtf, "terminate"):
- terminate = wrap_arrow_udtf(getattr(udtf, "terminate"),
return_type)
- else:
- terminate = None
-
- cleanup = getattr(udtf, "cleanup") if hasattr(udtf, "cleanup") else
None
+ terminate = getattr(udtf, "terminate", None)
+ cleanup = getattr(udtf, "cleanup", None)
- def mapper(_, it):
+ def func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
+ """Apply Arrow table UDF"""
try:
converters = [
ArrowTableToRowsConversion._create_converter(
@@ -1767,28 +1750,32 @@ def read_udtf(pickleSer, udtf_info, eval_type,
runner_conf, eval_conf):
)
for f in eval_conf.input_type
]
- for a in it:
+ for batch in data:
+ # Convert each input column to a list of Python values per
row,
+ # then call eval once per input row.
pylist = [
(
[conv(v) for v in column.to_pylist()]
if conv is not None
else column.to_pylist()
)
- for column, conv in zip(a.columns, converters)
+ for column, conv in zip(batch.columns, converters)
]
- # The eval function yields an iterator. Each element
produced by this
- # iterator is a tuple in the form of (pyarrow.RecordBatch,
arrow_return_type).
- yield from eval(*[pylist[o] for o in args_kwargs_offsets],
num_rows=a.num_rows)
+ yield from evaluate_rows(
+ eval_method,
+ *[pylist[o] for o in args_kwargs_offsets],
+ num_rows=batch.num_rows,
+ )
if terminate is not None:
- yield from terminate()
+ yield from evaluate_rows(terminate)
except SkipRestOfInputTableException:
if terminate is not None:
- yield from terminate()
+ yield from evaluate_rows(terminate)
finally:
if cleanup is not None:
cleanup()
- return mapper, None, ser, ser
+ return func, None, ser, ser
elif eval_type == PythonEvalType.SQL_ARROW_UDTF:
import pyarrow as pa
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
