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/07/24 03:42:32 UTC
(spark) branch branch-4.x updated: [SPARK-58297][PYTHON] Remove TransformWithStateInPySparkRowSerializer
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 9b0b266c9d07 [SPARK-58297][PYTHON] Remove TransformWithStateInPySparkRowSerializer
9b0b266c9d07 is described below
commit 9b0b266c9d077bb15a6b40d8b794019e9f7043be
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Fri Jul 24 03:42:07 2026 +0000
[SPARK-58297][PYTHON] Remove TransformWithStateInPySparkRowSerializer
### What changes were proposed in this pull request?
Delete `TransformWithStateInPySparkRowSerializer` from `python/pyspark/sql/pandas/serializers.py`. Its row-by-row deserialization and re-encoding logic was moved into `read_udfs()` in `worker.py` by [SPARK-58191](https://issues.apache.org/jira/browse/SPARK-58191) (#57337), which switched `SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF` to use the plain `ArrowStreamSerializer` for pure Arrow stream I/O. The class is no longer referenced by any eval type. The now-unused `itertools.groupby` and [...]
Part of [SPARK-55384](https://issues.apache.org/jira/browse/SPARK-55384).
### Why are the changes needed?
The class is dead code after the refactor in SPARK-58191. Removing it keeps serializers as pure I/O and avoids a stale, unused class.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests. This is a dead-code deletion with no remaining references (verified repo-wide, including dynamic lookups); behavior is unchanged since the eval type already runs through `read_udfs()`.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #57470 from Yicong-Huang/SPARK-58297.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
(cherry picked from commit fd3d03c76726a82c6345d120fec0b6c9175af1a3)
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
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 @@ from pyspark.serializers import (
write_int,
UTF8Deserializer,
)
-from pyspark.sql import Row
from pyspark.sql.conversion import (
ArrowBatchTransformer,
PandasToArrowConversion,
@@ -339,104 +337,3 @@ class ArrowStreamPandasSerializer(ArrowStreamSerializer):
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)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org