You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ur...@apache.org on 2026/07/23 10:47:00 UTC
(spark) branch master updated: [SPARK-58157][PYTHON] Remove TransformWithStateInPySparkRowInitStateSerializer
This is an automated email from the ASF dual-hosted git repository.
uros-b 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 93c3367d4fa6 [SPARK-58157][PYTHON] Remove TransformWithStateInPySparkRowInitStateSerializer
93c3367d4fa6 is described below
commit 93c3367d4fa6f37aaf4c92a696a1be17c630f6e0
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Thu Jul 23 10:46:48 2026 +0000
[SPARK-58157][PYTHON] Remove TransformWithStateInPySparkRowInitStateSerializer
### What changes were proposed in this pull request?
Delete `TransformWithStateInPySparkRowInitStateSerializer` from `python/pyspark/sql/pandas/serializers.py`. Its init-state deserialization 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_INIT_STATE_UDF` to use the plain `ArrowStreamSerializer` for pure Arrow stream I/O. The class is no longer referenced by any eval type. A benchmark comment in `bench_eval [...]
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 #57438 from Yicong-Huang/SPARK-58157.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Uros Bojanic <22...@users.noreply.github.com>
---
python/benchmarks/bench_eval_type.py | 8 +--
python/pyspark/sql/pandas/serializers.py | 116 +------------------------------
2 files changed, 5 insertions(+), 119 deletions(-)
diff --git a/python/benchmarks/bench_eval_type.py b/python/benchmarks/bench_eval_type.py
index c7fe702930a1..1a352760af7b 100644
--- a/python/benchmarks/bench_eval_type.py
+++ b/python/benchmarks/bench_eval_type.py
@@ -2333,10 +2333,10 @@ class TransformWithStateRowUDFPeakmemBench(_TransformWithStateRowBenchMixin, _Pe
# Each batch carries either inputData or initState rows -- never both -- with the
# inactive column written as an all-null struct. Matching the JVM ``initData ++
# data`` ordering, all initial-state batches are emitted first, then all data
-# batches. ``TransformWithStateInPySparkRowInitStateSerializer`` walks each
-# batch row by row, materializing every column (key included) into a ``Row`` via
-# ``.as_py()``, then regroups consecutive rows by the leading key column so each
-# key surfaces as one ``(mode, key, (input_rows, init_rows))`` call. This is the
+# batches. ``read_udfs`` walks each batch row by row, materializing every column
+# (key included) into a ``Row`` via ``.as_py()``, then regroups consecutive rows
+# by the leading key column so each key surfaces as one
+# ``(mode, key, (input_rows, init_rows))`` call. This is the
# per-row Python object round trip the Row variant is built around, layered on
# top of the nested-struct init-state deserialization, in contrast to the
# columnar Pandas init-state variant above.
diff --git a/python/pyspark/sql/pandas/serializers.py b/python/pyspark/sql/pandas/serializers.py
index c035a20e37af..2823cefd5821 100644
--- a/python/pyspark/sql/pandas/serializers.py
+++ b/python/pyspark/sql/pandas/serializers.py
@@ -20,7 +20,7 @@ Serializers for PyArrow and pandas conversions. See `pyspark.serializers` for mo
"""
from itertools import groupby
-from typing import IO, TYPE_CHECKING, Any, Iterable, Iterator, List, Optional, Tuple
+from typing import IO, TYPE_CHECKING, Iterable, Iterator, List, Optional, Tuple
from pyspark.errors import PySparkRuntimeError, PySparkValueError
from pyspark.serializers import (
@@ -528,117 +528,3 @@ class TransformWithStateInPySparkRowSerializer(ArrowStreamUDFSerializer):
yield (record_batch, arrow_type)
return ArrowStreamUDFSerializer.dump_stream(self, flatten_iterator(), stream)
-
-
-class TransformWithStateInPySparkRowInitStateSerializer(TransformWithStateInPySparkRowSerializer):
- """
- Serializer used by Python worker to evaluate UDF for
- :meth:`pyspark.sql.GroupedData.transformWithStateInPySparkRowInitStateSerializer`.
- Parameters
- ----------
- Same as input parameters in TransformWithStateInPySparkRowSerializer.
- """
-
- def __init__(self, *, arrow_max_records_per_batch):
- super().__init__(arrow_max_records_per_batch=arrow_max_records_per_batch)
- self.init_key_offsets = None
-
- def load_stream(self, stream):
- import pyarrow as pa
- from pyspark.sql.streaming.stateful_processor_util import (
- TransformWithStateInPandasFuncMode,
- )
-
- def generate_data_batches(batches) -> Iterator[Tuple[Any, Optional[Any], Optional[Any]]]:
- """
- 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.
- See `TransformWithStateInPySparkPythonInitialStateRunner` for arrow batch schema sent
- from JVM.
- This function flattens the columns of input rows and initial state rows and feed them
- into the data generator.
- """
-
- def extract_rows(
- cur_batch, col_name, key_offsets
- ) -> Optional[Iterator[Tuple[Any, Any]]]:
- data_column = cur_batch.column(cur_batch.schema.get_field_index(col_name))
-
- # Check if the entire column is null
- if data_column.null_count == len(data_column):
- return None
-
- data_field_names = [
- data_column.type[i].name for i in range(data_column.type.num_fields)
- ]
- data_field_arrays = [
- data_column.field(i) for i in range(data_column.type.num_fields)
- ]
-
- DataRow = Row(*data_field_names)
-
- table = pa.Table.from_arrays(data_field_arrays, names=data_field_names)
-
- if table.num_rows == 0:
- return None
-
- def row_iterator():
- for row_idx in range(table.num_rows):
- key = tuple(table.column(o)[row_idx].as_py() for o in key_offsets)
- row = DataRow(
- *(table.column(i)[row_idx].as_py() for i in range(table.num_columns))
- )
- yield (key, row)
-
- return row_iterator()
-
- """
- The arrow batch is written in the schema:
- schema: StructType = new StructType()
- .add("inputData", dataSchema)
- .add("initState", initStateSchema)
- We'll parse batch into Tuples of (key, inputData, initState) and pass into the Python
- data generator. Each batch will have either init_data or input_data, not mix.
- """
- for batch in batches:
- # Detect which column has data - each batch contains only one type
- input_result = extract_rows(batch, "inputData", self.key_offsets)
- init_result = extract_rows(batch, "initState", self.init_key_offsets)
-
- assert not (input_result is not None and init_result is not None)
-
- if input_result is not None:
- for key, input_data_row in input_result:
- yield (key, input_data_row, None)
- elif init_result is not None:
- for key, init_state_row in init_result:
- yield (key, None, init_state_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]):
- input_rows = []
- init_rows = []
-
- for batch_key, input_row, init_row in g:
- if input_row is not None:
- input_rows.append(input_row)
- if init_row is not None:
- init_rows.append(init_row)
-
- total_len = len(input_rows) + len(init_rows)
- if total_len >= self.arrow_max_records_per_batch:
- ret_tuple = (iter(input_rows), iter(init_rows))
- yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, ret_tuple)
- input_rows = []
- init_rows = []
-
- if input_rows or init_rows:
- ret_tuple = (iter(input_rows), iter(init_rows))
- yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, ret_tuple)
-
- yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None)
-
- yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org