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 e3ca71ecfd28 [SPARK-57400][PYTHON] Refactor
SQL_TRANSFORM_WITH_STATE_PANDAS_UDF
e3ca71ecfd28 is described below
commit e3ca71ecfd28b937a01a1e2d8a35a0bdf80c9ba2
Author: Yicong Huang <[email protected]>
AuthorDate: Wed Jun 24 00:25:36 2026 +0000
[SPARK-57400][PYTHON] Refactor SQL_TRANSFORM_WITH_STATE_PANDAS_UDF
### What changes were proposed in this pull request?
This PR refactors `SQL_TRANSFORM_WITH_STATE_PANDAS_UDF` so that the worker
uses the plain `ArrowStreamSerializer` for pure Arrow stream I/O, moving the
per-eval-type logic (regrouping rows by grouping key, re-chunking into pandas
DataFrames bounded by
`arrow_max_records_per_batch`/`arrow_max_bytes_per_batch`, and converting
result DataFrames back to Arrow) from `TransformWithStateInPandasSerializer`
into `read_udfs()` in `worker.py`. The serializer class is kept for now since
`Transfo [...]
### 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.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests
(`pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state`). No
behavior change: replaying identical worker input through the old and new code
paths produces byte-identical worker output across 12 scenario/UDF combinations.
ASV comparison (`bench_eval_type.TransformWithStatePandasUDFTimeBench` /
`TransformWithStatePandasUDFPeakmemBench`, `-a repeat=3`): before =
`upstream/master`, after = this PR. Values are from one representative run per
side; the conclusion is consistent across the 2 runs performed per side.
```text
time_worker
scenario udf before after diff
--------------- ------------ ------------ ------------ ------
few_groups_sm identity_udf 754+-5ms 743+-3ms -1.5%
few_groups_sm sort_udf 781+-20ms 760+-0.6ms -2.7%
few_groups_sm count_udf 726+-3ms 722+-1ms -0.6%
few_groups_lg identity_udf 7.01+-0.03s 6.99+-0.01s -0.3%
few_groups_lg sort_udf 7.21+-0.03s 7.10+-0.02s -1.5%
few_groups_lg count_udf 6.67+-0.04s 6.60+-0.01s -1.0%
many_groups_sm identity_udf 6.24+-0.03s 6.13+-0.01s -1.8%
many_groups_sm sort_udf 6.64+-0.01s 6.69+-0.02s +0.8%
many_groups_sm count_udf 5.76+-0.07s 5.89+-0.1s +2.3%
many_groups_lg identity_udf 3.54+-0.02s 3.53+-0.04s -0.3%
many_groups_lg sort_udf 3.65+-0.03s 3.63+-0.01s -0.5%
many_groups_lg count_udf 3.40+-0.04s 3.39+-0.03s -0.3%
wide_cols identity_udf 7.45+-0.07s 7.35+-0.03s -1.3%
wide_cols sort_udf 7.51+-0.03s 7.42+-0.03s -1.2%
wide_cols count_udf 6.85+-0.03s 6.80+-0.01s -0.7%
mixed_cols identity_udf 3.15+-0.01s 3.16+-0.01s +0.3%
mixed_cols sort_udf 3.36+-0.06s 3.28+-0s -2.4%
mixed_cols count_udf 2.85+-0.02s 2.91+-0.04s +2.1%
nested_struct identity_udf 7.47+-0.03s 7.33+-0.01s -1.9%
nested_struct sort_udf 8.63+-0.2s 8.36+-0.07s -3.1%
nested_struct count_udf 5.27+-0.03s 5.16+-0.01s -2.1%
```
```text
peakmem_worker
scenario udf before after
--------------- ------------ ------- -------
few_groups_sm identity_udf 104M 106M
few_groups_sm sort_udf 108M 107M
few_groups_sm count_udf 97.9M 98.1M
few_groups_lg identity_udf 200M 200M
few_groups_lg sort_udf 201M 189M
few_groups_lg count_udf 161M 161M
many_groups_sm identity_udf 130M 130M
many_groups_sm sort_udf 131M 131M
many_groups_sm count_udf 113M 113M
many_groups_lg identity_udf 136M 133M
many_groups_lg sort_udf 132M 134M
many_groups_lg count_udf 113M 113M
wide_cols identity_udf 240M 242M
wide_cols sort_udf 241M 241M
wide_cols count_udf 205M 205M
mixed_cols identity_udf 182M 182M
mixed_cols sort_udf 182M 182M
mixed_cols count_udf 182M 182M
nested_struct identity_udf 210M 210M
nested_struct sort_udf 210M 210M
nested_struct count_udf 210M 210M
```
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #56464 from Yicong-Huang/SPARK-57400/refactor/tws-pandas.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/worker.py | 175 +++++++++++++++++++++++++++++++++++------------
1 file changed, 130 insertions(+), 45 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 2be915fa358c..5a82d11ebe9b 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -79,7 +79,6 @@ from pyspark.sql.pandas.serializers import (
ArrowStreamPandasUDTFSerializer,
ArrowStreamCoGroupSerializer,
ApplyInPandasWithStateSerializer,
- TransformWithStateInPandasSerializer,
TransformWithStateInPandasInitStateSerializer,
TransformWithStateInPySparkRowSerializer,
TransformWithStateInPySparkRowInitStateSerializer,
@@ -503,18 +502,6 @@ def verify_arrow_result(
)
-def wrap_grouped_transform_with_state_pandas_udf(f, return_type, runner_conf):
- def wrapped(stateful_processor_api_client, mode, key, value_series_gen):
- result_iter = f(stateful_processor_api_client, mode, key,
value_series_gen)
-
- # TODO(SPARK-49100): add verification that elements in result_iter are
- # indeed of type pd.DataFrame and confirm to assigned cols
-
- return result_iter
-
- return lambda p, m, k, v: [(wrapped(p, m, k, v), return_type)]
-
-
def wrap_grouped_transform_with_state_pandas_init_state_udf(f, return_type,
runner_conf):
def wrapped(stateful_processor_api_client, mode, key, value_series_gen):
# Split the generator into two using itertools.tee
@@ -848,9 +835,7 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE:
return args_offsets, wrap_grouped_map_pandas_udf_with_state(func,
return_type, runner_conf)
elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF:
- return args_offsets, wrap_grouped_transform_with_state_pandas_udf(
- func, return_type, runner_conf
- )
+ return func, args_offsets, return_type
elif eval_type ==
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF:
return args_offsets,
wrap_grouped_transform_with_state_pandas_init_state_udf(
func, return_type, runner_conf
@@ -2103,16 +2088,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
prefers_large_var_types=runner_conf.use_large_var_types,
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
)
- elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF:
- ser = TransformWithStateInPandasSerializer(
- timezone=runner_conf.timezone,
- safecheck=runner_conf.safecheck,
- assign_cols_by_name=runner_conf.assign_cols_by_name,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
-
arrow_max_records_per_batch=runner_conf.arrow_max_records_per_batch,
-
arrow_max_bytes_per_batch=runner_conf.arrow_max_bytes_per_batch,
-
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
- )
elif eval_type ==
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF:
ser = TransformWithStateInPandasInitStateSerializer(
timezone=runner_conf.timezone,
@@ -3374,37 +3349,147 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
return func, None, ser, ser
if eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF:
- # We assume there is only one UDF here because grouped map doesn't
- # support combining multiple UDFs.
- assert num_udfs == 1
+ import pyarrow as pa
+ import pandas as pd
+
+ assert num_udfs == 1, "One TRANSFORM_WITH_STATE_PANDAS UDF expected
here."
+ udf, arg_offsets, return_type = udfs[0]
# See TransformWithStateInPandasExec for how arg_offsets are used to
# distinguish between grouping attributes and data attributes
- arg_offsets, f = udfs[0]
parsed_offsets = extract_key_value_indexes(arg_offsets)
- ser.key_offsets = parsed_offsets[0][0]
+ assert len(parsed_offsets) == 1, (
+ "Expected one pair of offsets for TRANSFORM_WITH_STATE_PANDAS UDF."
+ )
+
+ key_offsets = parsed_offsets[0][0]
+ value_offsets = parsed_offsets[0][1]
+ output_schema = StructType([StructField("_0", return_type)])
+
stateful_processor_api_client = StatefulProcessorApiClient(
eval_conf.state_server_socket_port, eval_conf.grouping_key_schema
)
- def mapper(a):
- mode = a[0]
+ arrow_max_records_per_batch = runner_conf.arrow_max_records_per_batch
+ arrow_max_records_per_batch = (
+ arrow_max_records_per_batch if arrow_max_records_per_batch > 0
else 2**31 - 1
+ )
+ arrow_max_bytes_per_batch = runner_conf.arrow_max_bytes_per_batch
- if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
- key = a[1]
+ def transform_with_state_func(
+ split_index: int,
+ batches: Iterator[pa.RecordBatch],
+ ) -> Iterator[pa.RecordBatch]:
+ """Apply transformWithStateInPandas UDF.
+
+ Data chunks for the same grouping key appear sequentially in the
+ input batches but may span batch boundaries, so rows are regrouped
+ by key and re-chunked into pandas DataFrames bounded by
+ arrow_max_records_per_batch and arrow_max_bytes_per_batch. The UDF
+ is invoked once per grouping key with a lazy iterator of chunks,
+ then once for PROCESS_TIMER and once for COMPLETE.
+ """
+ total_bytes = 0
+ total_rows = 0
+ average_arrow_row_size = 0.0
+
+ def row_stream():
+ nonlocal total_bytes, total_rows, average_arrow_row_size
+ for batch in batches:
+ # Short circuit batch size stats if the batch size is
+ # unlimited as computing batch size is computationally
+ # expensive.
+ if arrow_max_bytes_per_batch != 2**31 - 1 and
batch.num_rows > 0:
+ total_bytes += sum(
+ buf.size
+ for col in batch.columns
+ for buf in col.buffers()
+ if buf is not None
+ )
+ total_rows += batch.num_rows
+ average_arrow_row_size = total_bytes / total_rows
+ data_pandas = ArrowBatchTransformer.to_pandas(
+ batch,
+ timezone=runner_conf.timezone,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ )
+ for row in pd.concat(data_pandas,
axis=1).itertuples(index=False):
+ batch_key = tuple(row[o] for o in key_offsets)
+ yield (batch_key, row)
+
+ def generate_data_batches():
+ """
+ Deserialize ArrowRecordBatches and return a generator of
+ (grouping key, pandas.DataFrame) chunks.
+
+ This function must avoid materializing multiple Arrow
+ RecordBatches into memory at the same time. And data chunks
+ from the same grouping key should appear sequentially.
+ """
+ for batch_key, group_rows in itertools.groupby(row_stream(),
key=lambda x: x[0]):
+ rows = []
+ for _, row in group_rows:
+ rows.append(row)
+ if (
+ len(rows) >= arrow_max_records_per_batch
+ or len(rows) * average_arrow_row_size >=
arrow_max_bytes_per_batch
+ ):
+ yield (batch_key, pd.DataFrame(rows))
+ rows = []
+ if rows:
+ yield (batch_key, pd.DataFrame(rows))
+
+ def convert_results(result_iter):
+ for result in result_iter:
+ if isinstance(return_type, StructType) and not
isinstance(result, pd.DataFrame):
+ raise PySparkValueError(
+ "Invalid return type. Please make sure that the
UDF returns a "
+ "pandas.DataFrame when the specified return type
is StructType."
+ )
+ yield PandasToArrowConversion.convert(
+ [result],
+ output_schema,
+ timezone=runner_conf.timezone,
+ safecheck=runner_conf.safecheck,
+ arrow_cast=True,
+ assign_cols_by_name=runner_conf.assign_cols_by_name,
+
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
+ )
- def values_gen():
- for x in a[2]:
- retVal = x[1].iloc[:, parsed_offsets[0][1]]
- yield retVal
+ for key, group in itertools.groupby(generate_data_batches(),
key=lambda x: x[0]):
+ # This must be a generator expression - do not materialize.
+ values_gen = (df.iloc[:, value_offsets] for _, df in group)
+ yield from convert_results(
+ udf(
+ stateful_processor_api_client,
+ TransformWithStateInPandasFuncMode.PROCESS_DATA,
+ key,
+ values_gen,
+ )
+ )
- # This must be generator comprehension - do not materialize.
- return f(stateful_processor_api_client, mode, key,
values_gen())
- else:
- # mode == PROCESS_TIMER or mode == COMPLETE
- return f(stateful_processor_api_client, mode, None, iter([]))
+ yield from convert_results(
+ udf(
+ stateful_processor_api_client,
+ TransformWithStateInPandasFuncMode.PROCESS_TIMER,
+ None,
+ iter([]),
+ )
+ )
- elif eval_type ==
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF:
+ yield from convert_results(
+ udf(
+ stateful_processor_api_client,
+ TransformWithStateInPandasFuncMode.COMPLETE,
+ None,
+ iter([]),
+ )
+ )
+
+ # profiling is not supported for UDF
+ return transform_with_state_func, None, ser, ser
+
+ if eval_type ==
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF:
# We assume there is only one UDF here because grouped map doesn't
# support combining multiple UDFs.
assert num_udfs == 1
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
