From 9191d4ab973d918411915aeecce9a0893a79a5c3 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Thu, 23 Jul 2026 19:13:36 +0000 Subject: [PATCH] [SPARK-58297][PYTHON] Remove TransformWithStateInPySparkRowSerializer --- python/benchmarks/bench_eval_type.py | 4 +- python/pyspark/sql/pandas/serializers.py | 103 ----------------------- 2 files changed, 2 insertions(+), 105 deletions(-) diff --git a/python/benchmarks/bench_eval_type.py b/python/benchmarks/bench_eval_type.py index 1a352760af7b..3228d915f44e 100644 --- a/python/benchmarks/bench_eval_type.py +++ b/python/benchmarks/bench_eval_type.py @@ -2162,7 +2162,7 @@ class TransformWithStatePandasInitStateUDFPeakmemBench( # Stateful streaming with plain PySpark Rows. UDF signature is # ``(api_client, mode, key, rows)`` and returns ``Iterator[Row]``. The input # wire stream is a single plain Arrow stream pre-sorted by the grouping key -# column at offset 0; ``TransformWithStateInPySparkRowSerializer`` walks the +# column at offset 0; the worker walks the # batch row by row, materializing each into a ``Row`` (all columns, including # the key) via ``.as_py()``, groups consecutive rows by key, and yields one # ``(mode, key, rows)`` tuple per group, then a phantom ``PROCESS_TIMER`` and @@ -2181,7 +2181,7 @@ class _TransformWithStateRowBenchMixin: Each scenario emits one plain Arrow stream pre-sorted by the leading int key column. Unlike the Pandas variant, the key column is NOT projected out: UDFs receive an iterator of ``Row`` objects carrying every column (key - included), mirroring ``TransformWithStateInPySparkRowSerializer``. Row-by-row + included), mirroring the worker's row handling. Row-by-row materialization and re-encoding is ~10x slower than the columnar Pandas path, so row counts are scaled down accordingly to stay under ASV's 60s per-sample timeout. diff --git a/python/pyspark/sql/pandas/serializers.py b/python/pyspark/sql/pandas/serializers.py index a84debfab32e..e62d688c6a25 100644 --- a/python/pyspark/sql/pandas/serializers.py +++ b/python/pyspark/sql/pandas/serializers.py @@ -19,7 +19,6 @@ Serializers for PyArrow and pandas conversions. See `pyspark.serializers` for more details. """ -from itertools import groupby from typing import IO, TYPE_CHECKING, Iterable, Iterator, List, Optional, Tuple from pyspark.errors import PySparkRuntimeError, PySparkValueError @@ -29,7 +28,6 @@ write_int, UTF8Deserializer, ) -from pyspark.sql import Row from pyspark.sql.conversion import ( ArrowBatchTransformer, PandasToArrowConversion, @@ -339,104 +337,3 @@ def load_stream(self, stream): def __repr__(self): return "ArrowStreamPandasSerializer" - - -class TransformWithStateInPySparkRowSerializer(ArrowStreamUDFSerializer): - """ - Serializer used by Python worker to evaluate UDF for - :meth:`pyspark.sql.GroupedData.transformWithState`. - - Parameters - ---------- - arrow_max_records_per_batch : int - Limit of the number of records that can be written to a single ArrowRecordBatch in memory. - """ - - def __init__(self, *, arrow_max_records_per_batch): - super().__init__() - self.arrow_max_records_per_batch = ( - arrow_max_records_per_batch if arrow_max_records_per_batch > 0 else 2**31 - 1 - ) - self.key_offsets = None - - def load_stream(self, stream): - """ - Read ArrowRecordBatches from stream, deserialize them to populate a list of data chunks, - and convert the data into a list of pandas.Series. - - Please refer the doc of inner function `generate_data_batches` for more details how - this function works in overall. - """ - from pyspark.sql.streaming.stateful_processor_util import ( - TransformWithStateInPandasFuncMode, - ) - import itertools - - def generate_data_batches(batches): - """ - Deserialize ArrowRecordBatches and return a generator of Row. - - The deserialization logic assumes that Arrow RecordBatches contain the data with the - ordering that data chunks for same grouping key will appear sequentially. - - 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 in batches: - DataRow = Row(*batch.schema.names) - - # Iterate row by row without converting the whole batch - num_cols = batch.num_columns - for row_idx in range(batch.num_rows): - # build the key for this row - row_key = tuple(batch[o][row_idx].as_py() for o in self.key_offsets) - row = DataRow(*(batch.column(i)[row_idx].as_py() for i in range(num_cols))) - yield row_key, row - - _batches = super(ArrowStreamUDFSerializer, self).load_stream(stream) - data_batches = generate_data_batches(_batches) - - for k, g in groupby(data_batches, key=lambda x: x[0]): - chained = itertools.chain(g) - chained_values = map(lambda x: x[1], chained) - yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, chained_values) - - yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None) - - yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None) - - def dump_stream(self, iterator, stream): - """ - Read through an iterator of (iterator of Row), serialize them to Arrow - RecordBatches, and write batches to stream. - """ - import pyarrow as pa - - from pyspark.sql.pandas.types import to_arrow_type - - def flatten_iterator(): - # iterator: iter[list[(iter[Row], spark_type)]] - for packed in iterator: - iter_row_with_type = packed[0] - iter_row = iter_row_with_type[0] - spark_type = iter_row_with_type[1] - - # Convert spark type to arrow type - # TODO: WE need to make this configurable, currently using default values. - arrow_type = to_arrow_type( - spark_type, - timezone="UTC", - prefers_large_types=False, - ) - - rows_as_dict = [] - for row in iter_row: - row_as_dict = row.asDict(True) - rows_as_dict.append(row_as_dict) - - pdf_schema = pa.schema(list(arrow_type)) - record_batch = pa.RecordBatch.from_pylist(rows_as_dict, schema=pdf_schema) - - yield (record_batch, arrow_type) - - return ArrowStreamUDFSerializer.dump_stream(self, flatten_iterator(), stream)