You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by yi...@apache.org on 2026/06/09 22:51:03 UTC
(spark) branch branch-4.x updated: [SPARK-56758][PYTHON] Refactor SQL_MAP_PANDAS_ITER_UDF
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 876efe803418 [SPARK-56758][PYTHON] Refactor SQL_MAP_PANDAS_ITER_UDF
876efe803418 is described below
commit 876efe803418d774b3bc93ef37df6d5ee10bfca3
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Tue Jun 9 20:52:40 2026 +0000
[SPARK-56758][PYTHON] Refactor SQL_MAP_PANDAS_ITER_UDF
### What changes were proposed in this pull request?
Consolidate the `SQL_MAP_PANDAS_ITER_UDF` (`mapInPandas`) execution path so that input transformation, UDF invocation, result verification, and output transformation live in one block in `read_udfs()`.
### Why are the changes needed?
Part of [SPARK-55388](https://issues.apache.org/jira/browse/SPARK-55388). The full data flow for `mapInPandas` is now visible in one place.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests. No behavior change.
ASV `MapPandasIterUDFTimeBench` comparison (`-a repeat=3`, before = `upstream/master`, after = this branch):
```text
scenario udf before (ms) after (ms) diff (%)
sm_batch_few_col identity_udf 311 310 -0.3
sm_batch_few_col sort_udf 365 367 +0.5
sm_batch_few_col filter_udf 347 341 -1.7
sm_batch_many_col identity_udf 213 212 -0.5
sm_batch_many_col sort_udf 231 233 +0.9
sm_batch_many_col filter_udf 216 217 +0.5
lg_batch_few_col identity_udf 850 765 -10.0
lg_batch_few_col sort_udf 1030 965 -6.3
lg_batch_few_col filter_udf 815 804 -1.3
lg_batch_many_col identity_udf 791 787 -0.5
lg_batch_many_col sort_udf 1290 1280 -0.8
lg_batch_many_col filter_udf 889 811 -8.8
pure_ints identity_udf 152 140 -7.9
pure_ints sort_udf 168 166 -1.2
pure_ints filter_udf 164 151 -7.9
pure_floats identity_udf 158 139 -12.0
pure_floats sort_udf 183 166 -9.3
pure_floats filter_udf 179 161 -10.1
pure_strings identity_udf 636 609 -4.2
pure_strings sort_udf 894 799 -10.6
pure_strings filter_udf 762 650 -14.7
pure_ts identity_udf 267 211 -21.0
pure_ts sort_udf 303 230 -24.1
pure_ts filter_udf 273 233 -14.7
mixed_types identity_udf 571 435 -23.8
mixed_types sort_udf 668 518 -22.5
mixed_types filter_udf 507 460 -9.3
```
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #55750 from Yicong-Huang/SPARK-56758.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
(cherry picked from commit 6c2325a0cde405246b80d77b407a0d487b0fb52f)
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
python/pyspark/worker.py | 125 ++++++++++++++++++++++++++++++++++-------------
1 file changed, 92 insertions(+), 33 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 076ef55dd21d..136e438f5354 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -972,7 +972,7 @@ def read_single_udf(pickleSer, udf_info, eval_type, runner_conf, udf_index):
elif eval_type == PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF:
return func, args_offsets, kwargs_offsets, return_type
elif eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF:
- return args_offsets, wrap_pandas_batch_iter_udf(func, return_type, runner_conf)
+ return func, None, None, return_type
elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF:
return func, None, None, None
elif eval_type in (
@@ -2284,6 +2284,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
)
elif eval_type in (
PythonEvalType.SQL_MAP_ARROW_ITER_UDF,
+ PythonEvalType.SQL_MAP_PANDAS_ITER_UDF,
PythonEvalType.SQL_SCALAR_ARROW_UDF,
PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
PythonEvalType.SQL_ARROW_BATCHED_UDF,
@@ -2293,10 +2294,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
else:
# Scalar Pandas UDF handles struct type arguments as pandas DataFrames instead of
# pandas Series. See SPARK-27240.
- df_for_struct = (
- eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
- or eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
- )
+ df_for_struct = eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
ser = ArrowStreamPandasUDFSerializer(
timezone=runner_conf.timezone,
@@ -2984,6 +2982,75 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
# profiling is not supported for UDF
return cogrouped_func, None, ser, ser
+ if eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF:
+ import pyarrow as pa
+ import pandas as pd
+
+ assert num_udfs == 1, "One MAP_PANDAS_ITER UDF expected here."
+ map_udf, _, _, return_type = udfs[0]
+ output_schema = StructType([StructField("_0", return_type)])
+ iter_type_label = (
+ "pandas.DataFrame" if isinstance(return_type, StructType) else "pandas.Series"
+ )
+ elem_type = pd.DataFrame if isinstance(return_type, StructType) else pd.Series
+
+ def func(
+ split_index: int,
+ data: Iterator[pa.RecordBatch],
+ ) -> Iterator[pa.RecordBatch]:
+ """Apply mapInPandas UDF."""
+
+ def dataframe_iter():
+ # Input batches have a single struct column (see
+ # MapInBatchEvaluatorFactory); convert lazily so peakmem stays
+ # bounded by one batch.
+ for batch in data:
+ yield ArrowBatchTransformer.to_pandas(
+ batch,
+ timezone=runner_conf.timezone,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ df_for_struct=True,
+ )[0]
+
+ # mapInPandas accepts any iterable (e.g. a list), not just an
+ # iterator, so the standard verify_return_type (which requires an
+ # Iterator) is intentionally not reused here.
+ result = map_udf(dataframe_iter())
+ if not isinstance(result, Iterator) and not hasattr(result, "__iter__"):
+ raise PySparkTypeError(
+ errorClass="UDF_RETURN_TYPE",
+ messageParameters={
+ "expected": "iterator of {}".format(iter_type_label),
+ "actual": type(result).__name__,
+ },
+ )
+
+ for df in result:
+ if not isinstance(df, elem_type):
+ raise PySparkTypeError(
+ errorClass="UDF_RETURN_TYPE",
+ messageParameters={
+ "expected": "iterator of {}".format(iter_type_label),
+ "actual": "iterator of {}".format(type(df).__name__),
+ },
+ )
+ verify_pandas_result(
+ df, return_type, assign_cols_by_name=True, truncate_return_schema=True
+ )
+ yield PandasToArrowConversion.convert(
+ [df],
+ output_schema,
+ timezone=runner_conf.timezone,
+ safecheck=runner_conf.safecheck,
+ arrow_cast=True,
+ prefers_large_types=runner_conf.use_large_var_types,
+ assign_cols_by_name=runner_conf.assign_cols_by_name,
+ int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
+ )
+
+ # profiling is not supported for UDF
+ return func, None, ser, ser
+
if (
eval_type == PythonEvalType.SQL_ARROW_BATCHED_UDF
and not runner_conf.use_legacy_pandas_udf_conversion
@@ -3243,15 +3310,8 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
# profiling is not supported for UDF
return func, None, ser, ser
- is_scalar_iter = eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
- is_map_pandas_iter = eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
-
- if is_scalar_iter or is_map_pandas_iter:
- # TODO: Better error message for num_udfs != 1
- if is_scalar_iter:
- assert num_udfs == 1, "One SCALAR_ITER UDF expected here."
- if is_map_pandas_iter:
- assert num_udfs == 1, "One MAP_PANDAS_ITER UDF expected here."
+ if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF:
+ assert num_udfs == 1, "One SCALAR_ITER UDF expected here."
arg_offsets, udf = udfs[0]
@@ -3279,31 +3339,30 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
# by consuming the input iterator in user side. Therefore,
# it's very unlikely the output length is higher than
# input length.
- if is_scalar_iter and num_output_rows > num_input_rows:
+ if num_output_rows > num_input_rows:
raise PySparkRuntimeError(
errorClass="OUTPUT_EXCEEDS_INPUT_ROWS", messageParameters={}
)
yield (result_batch, result_type)
- if is_scalar_iter:
- try:
- next(iterator)
- except StopIteration:
- pass
- else:
- raise PySparkRuntimeError(
- errorClass="INPUT_NOT_FULLY_CONSUMED",
- messageParameters={},
- )
+ try:
+ next(iterator)
+ except StopIteration:
+ pass
+ else:
+ raise PySparkRuntimeError(
+ errorClass="INPUT_NOT_FULLY_CONSUMED",
+ messageParameters={},
+ )
- if num_output_rows != num_input_rows:
- raise PySparkRuntimeError(
- errorClass="RESULT_ROWS_MISMATCH",
- messageParameters={
- "output_length": str(num_output_rows),
- "input_length": str(num_input_rows),
- },
- )
+ if num_output_rows != num_input_rows:
+ raise PySparkRuntimeError(
+ errorClass="RESULT_ROWS_MISMATCH",
+ messageParameters={
+ "output_length": str(num_output_rows),
+ "input_length": str(num_input_rows),
+ },
+ )
# profiling is not supported for UDF
return func, None, ser, ser
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org