This is an automated email from the ASF dual-hosted git repository.
zhengruifeng 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 1d1bf06fa6c5 [SPARK-58130][PYTHON] Refactor legacy SQL_ARROW_TABLE_UDF
pandas path and remove ArrowStreamPandasUDTFSerializer
1d1bf06fa6c5 is described below
commit 1d1bf06fa6c51e5c8842d90162f7ea61c901d2e3
Author: Yicong Huang <[email protected]>
AuthorDate: Fri Jul 17 08:18:35 2026 +0800
[SPARK-58130][PYTHON] Refactor legacy SQL_ARROW_TABLE_UDF pandas path and
remove ArrowStreamPandasUDTFSerializer
### What changes were proposed in this pull request?
Following SPARK-57394, which refactored the non-legacy
`SQL_ARROW_TABLE_UDF` path onto `ArrowStreamSerializer`, this PR does the same
for the legacy pandas conversion path
(`spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled=true`):
- The pandas (de)serialization is moved out of
`ArrowStreamPandasUDTFSerializer` and into the `read_udtf` wrapper in
`worker.py`, so the legacy path uses a plain `ArrowStreamSerializer` for I/O
like the non-legacy path. Input Arrow batches are converted to pandas Series
via `ArrowBatchTransformer.to_pandas`, and the output pandas `DataFrame` is
converted back via `PandasToArrowConversion.convert(..., is_legacy=True)` --
the same parameters the serializer used, so behavior is unchanged.
- The legacy branch is rewritten flat to mirror the non-legacy path's
structure (`verify_result` / `check_return_value` / `evaluate_rows` / `func`),
dropping the nested `wrap_arrow_udtf` factory.
- With no remaining callers, `ArrowStreamPandasUDTFSerializer` is removed.
- A `conversion` axis (`arrow` / `legacy_pandas`) is added to the
`ArrowTableUDF` microbenchmark so both paths are covered.
The legacy flag and the parent serializers
(`ArrowStreamPandasUDFSerializer` / `ArrowStreamPandasSerializer`) are retained.
### Why are the changes needed?
This is part of the worker.py eval-type refactor (SPARK-55388) that
consolidates each eval type's serialization onto `ArrowStreamSerializer` and
keeps the (de)serialization logic self-contained in `read_udtf` / `read_udfs`.
It removes the last caller of `ArrowStreamPandasUDTFSerializer` and aligns the
legacy and non-legacy `SQL_ARROW_TABLE_UDF` paths structurally.
### Does this PR introduce _any_ user-facing change?
No. The legacy path's behavior is unchanged; the conversion logic is only
relocated.
### How was this patch tested?
Existing tests. No behavior change. The legacy path is fully covered by
`LegacyUDTFArrowTests` and `LegacyArrowUDTFParityTests` (flag=true), which pass
unchanged, along with the non-legacy `UDTFArrowTests` / `ArrowUDTFParityTests`.
Microbenchmark of the refactored legacy path (`ArrowTableUDFTimeBench` with
`conversion=legacy_pandas`, `time_worker`, 7 samples/param; representative run,
median +- half-range). No regression -- differences are within run-to-run noise
(overlapping intervals); the conclusion is consistent across runs.
```text
udtf scenario before(ms) after(ms) diff
identity_udtf sm_batch_few_col 770.0+-8 799.5+-81 +3.8%
identity_udtf lg_batch_few_col 1902.5+-40 1864.0+-142 -2.0%
filter_udtf pure_ints 2086.0+-117 2074.3+-16 -0.6%
filter_udtf pure_strings 2088.3+-19 2088.8+-18 +0.0%
stringify_udtf pure_ints 2132.8+-136 2099.6+-84 -1.6%
stringify_udtf pure_strings 2122.8+-19 1972.0+-119 -7.1%
```
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #57273 from Yicong-Huang/refactor-legacy-udtf.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 3dc79a4e1f61612a7716d4de92773d3b1f2963fe)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/benchmarks/bench_eval_type.py | 16 ++-
python/pyspark/sql/pandas/serializers.py | 48 -------
python/pyspark/worker.py | 211 ++++++++++++++++---------------
3 files changed, 121 insertions(+), 154 deletions(-)
diff --git a/python/benchmarks/bench_eval_type.py
b/python/benchmarks/bench_eval_type.py
index 8db72f335c63..8e0f6a312754 100644
--- a/python/benchmarks/bench_eval_type.py
+++ b/python/benchmarks/bench_eval_type.py
@@ -701,16 +701,25 @@ class _ArrowTableUDFBenchMixin:
"filter_udtf": (_ArrowTableUDFFilter, None, [0]),
"stringify_udtf": (_ArrowTableUDFStringify, StringType(), [0]),
}
- params = [list(_scenario_configs), list(_udtfs)]
- param_names = ["scenario", "udtf"]
+ # "arrow": non-legacy pure Arrow path; "legacy_pandas": legacy pandas
+ # conversion path (flag on), which goes through pandas Series/DataFrame.
+ _conversions = ["arrow", "legacy_pandas"]
+ params = [list(_scenario_configs), list(_udtfs), list(_conversions)]
+ param_names = ["scenario", "udtf", "conversion"]
- def _write_scenario(self, scenario, udtf_name, buf):
+ def _write_scenario(self, scenario, udtf_name, conversion, buf):
batches, schema = self._build_scenario(scenario)
handler, ret_type, arg_offsets = self._udtfs[udtf_name]
if ret_type is None:
ret_type = schema.fields[0].dataType
return_type = StructType([StructField("c0", ret_type)])
+ runner_conf = None
+ if conversion == "legacy_pandas":
+ runner_conf = {
+
"spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled": "true"
+ }
+
MockProtocolWriter.write_worker_input(
PythonEvalType.SQL_ARROW_TABLE_UDF,
lambda b: MockProtocolWriter.write_arrow_udtf_payload(
@@ -718,6 +727,7 @@ class _ArrowTableUDFBenchMixin:
),
lambda b: MockProtocolWriter.write_data_payload(iter(batches), b),
buf,
+ runner_conf=runner_conf,
eval_conf={"input_type": schema.json()},
)
diff --git a/python/pyspark/sql/pandas/serializers.py
b/python/pyspark/sql/pandas/serializers.py
index 1bbd4ed296cd..13de009fb77c 100644
--- a/python/pyspark/sql/pandas/serializers.py
+++ b/python/pyspark/sql/pandas/serializers.py
@@ -444,54 +444,6 @@ class
ArrowStreamPandasUDFSerializer(ArrowStreamPandasSerializer):
return "ArrowStreamPandasUDFSerializer"
-class ArrowStreamPandasUDTFSerializer(ArrowStreamPandasUDFSerializer):
- """
- Serializer used by Python worker to evaluate Arrow-optimized Python UDTFs.
- """
-
- def __init__(
- self,
- *,
- timezone,
- safecheck,
- input_type,
- prefer_int_ext_dtype,
- int_to_decimal_coercion_enabled,
- ):
- super().__init__(
- timezone=timezone,
- safecheck=safecheck,
- # The output pandas DataFrame's columns are unnamed.
- assign_cols_by_name=False,
- # Set to 'False' to avoid converting struct type inputs into a
pandas DataFrame.
- df_for_struct=False,
- # Defines how struct type inputs are converted. If set to "row",
struct type inputs
- # are converted into Rows. Without this setting, a struct type
input would be treated
- # as a dictionary. For example, for named_struct('name', 'Alice',
'age', 1),
- # if struct_in_pandas="dict", it becomes {"name": "Alice", "age":
1}
- # if struct_in_pandas="row", it becomes Row(name="Alice", age=1)
- struct_in_pandas="row",
- # When dealing with array type inputs, Arrow converts them into
numpy.ndarrays.
- # To ensure consistency across regular and arrow-optimized UDTFs,
we further
- # convert these numpy.ndarrays into Python lists.
- ndarray_as_list=True,
- prefer_int_ext_dtype=prefer_int_ext_dtype,
- # Enables explicit casting for mismatched return types of Arrow
Python UDTFs.
- arrow_cast=True,
- input_type=input_type,
- # Enable additional coercions for UDTF serialization
- int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled,
- # UDTF-specific: ignore unexpected complex type values in converter
- ignore_unexpected_complex_type_values=True,
- # Legacy UDTF pandas conversion: enables broader Arrow exception
- # handling to allow more implicit type coercions
- is_legacy=True,
- )
-
- def __repr__(self):
- return "ArrowStreamPandasUDTFSerializer"
-
-
class ApplyInPandasWithStateSerializer(ArrowStreamPandasUDFSerializer):
"""
Serializer used by Python worker to evaluate UDF for
applyInPandasWithState.
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 83ceeed1bd10..59c389d8f737 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -77,7 +77,6 @@ from pyspark.sql.functions import
SkipRestOfInputTableException
from pyspark.sql.pandas.serializers import (
ArrowStreamSerializer,
ArrowStreamGroupSerializer,
- ArrowStreamPandasUDTFSerializer,
ArrowStreamCoGroupSerializer,
ApplyInPandasWithStateSerializer,
TransformWithStateInPandasInitStateSerializer,
@@ -857,19 +856,10 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
# the UDTF logic to input rows.
def read_udtf(pickleSer, udtf_info, eval_type, runner_conf, eval_conf):
if eval_type == PythonEvalType.SQL_ARROW_TABLE_UDF:
- if runner_conf.use_legacy_pandas_udtf_conversion:
- # NOTE: if timezone is set here, that implies
respectSessionTimeZone is True
- ser = ArrowStreamPandasUDTFSerializer(
- timezone=runner_conf.timezone,
- safecheck=runner_conf.safecheck,
- input_type=eval_conf.input_type,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
-
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
- )
- else:
- # Pure Arrow stream I/O; output struct wrapping is handled in the
- # func below.
- ser = ArrowStreamSerializer(write_start_stream=True)
+ # Pure Arrow stream I/O for both the legacy pandas conversion path and
the
+ # non-legacy path; the pandas (de)serialization for the legacy path
and the
+ # output struct wrapping are both 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 func below.
@@ -1495,123 +1485,138 @@ def read_udtf(pickleSer, udtf_info, eval_type,
runner_conf, eval_conf):
eval_type == PythonEvalType.SQL_ARROW_TABLE_UDF
and runner_conf.use_legacy_pandas_udtf_conversion
):
+ import pandas as pd
- def wrap_arrow_udtf(f, return_type):
- import pandas as pd
+ return_type_size = len(return_type)
+ # The output pandas DataFrame is converted as a single struct column
named
+ # "_0" against this schema.
+ output_schema = StructType([StructField("_0", return_type)])
- return_type_size = len(return_type)
+ def verify_result(result: Any, method_name: str) -> Any:
+ if not isinstance(result, pd.DataFrame):
+ 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, pd.DataFrame):
- 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 len(result) > 0 or len(result.columns) > 0:
+ if len(result.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(len(result.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 len(result) > 0 or len(result.columns) > 0:
- if len(result.columns) != return_type_size:
- raise PySparkRuntimeError(
- errorClass="UDTF_RETURN_SCHEMA_MISMATCH",
- messageParameters={
- "expected": str(return_type_size),
- "actual": str(len(result.columns)),
- "func": f.__name__,
- },
- )
+ # Verify the type and the schema of the result.
+ verify_pandas_result(
+ result, return_type, assign_cols_by_name=False,
truncate_return_schema=False
+ )
+ return result
- # Verify the type and the schema of the result.
- verify_pandas_result(
- result, return_type, assign_cols_by_name=False,
truncate_return_schema=False
- )
- return result
+ 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_RETURN_NOT_ITERABLE",
+ messageParameters={
+ "type": type(res).__name__,
+ "func": method_name,
+ },
+ )
+ if check_output_row_against_schema is not None:
+ for row in res:
+ if row is not None:
+ check_output_row_against_schema(row)
+ yield row
+ else:
+ yield from res
+
+ def convert_df_to_arrow(result: "pd.DataFrame") -> "pa.RecordBatch":
+ # Convert the output pandas DataFrame into a single "_0" struct
column,
+ # applying the legacy pandas-to-Arrow coercions.
+ return PandasToArrowConversion.convert(
+ [result],
+ output_schema,
+ timezone=runner_conf.timezone,
+ safecheck=runner_conf.safecheck,
+ arrow_cast=True,
+ assign_cols_by_name=False,
+
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
+ ignore_unexpected_complex_type_values=True,
+ is_legacy=True,
+ )
- # Wrap the exception thrown from the UDTF in a PySparkRuntimeError.
- def func(*a: Any) -> Any:
+ def evaluate_rows(
+ method: Callable, *args: list, num_rows: int = 1
+ ) -> Iterator["pa.RecordBatch"]:
+ # Create tuples from the input pandas Series, each tuple
represents a row
+ # across all Series.
+ 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:
- return f(*a)
+ res = method(*row)
except SkipRestOfInputTableException:
raise
except Exception as e:
raise PySparkRuntimeError(
errorClass="UDTF_EXEC_ERROR",
- messageParameters={"method_name": f.__name__, "error":
str(e)},
+ messageParameters={"method_name": method.__name__,
"error": str(e)},
)
+ result = verify_result(
+ pd.DataFrame(list(check_return_value(res,
method.__name__))), method.__name__
+ )
+ yield convert_df_to_arrow(result)
- 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__,
- },
- )
- if check_output_row_against_schema is not None:
- for row in res:
- if row is not None:
- check_output_row_against_schema(row)
- yield row
- else:
- yield from res
-
- def evaluate(*args: pd.Series, num_rows=1):
- if len(args) == 0:
- for _ in range(num_rows):
- yield (
-
verify_result(pd.DataFrame(list(check_return_value(func())))),
- return_type,
- )
- else:
- # Create tuples from the input pandas Series, each tuple
- # represents a row across all Series.
- row_tuples = zip(*args)
- for row in row_tuples:
- yield (
-
verify_result(pd.DataFrame(list(check_return_value(func(*row))))),
- return_type,
- )
-
- return evaluate
-
- 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 legacy pandas Arrow table UDF"""
try:
- for a in it:
- # The eval function yields an iterator. Each element
produced by this
- # iterator is a tuple in the form of (pandas.DataFrame,
arrow_return_type).
- yield from eval(*[a[o] for o in args_kwargs_offsets],
num_rows=len(a[0]))
+ for batch in data:
+ # Deserialize the Arrow batch into a list of pandas Series
(one per
+ # input column), then call eval once per input row.
+ series_list = ArrowBatchTransformer.to_pandas(
+ batch,
+ timezone=runner_conf.timezone,
+ schema=eval_conf.input_type,
+ struct_in_pandas="row",
+ ndarray_as_list=True,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ df_for_struct=False,
+ )
+ yield from evaluate_rows(
+ eval_method,
+ *[series_list[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_TABLE_UDF
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
