[SPARK-58225][PYTHON] Delete unused TransformWithStateInPandasSerializer

python/benchmarks/bench_eval_type.py

@@ -1807,9 +1807,9 @@ class WindowAggPandasUDFPeakmemBench(_WindowAggPandasBenchMixin, _PeakmemBenchBa
 # Stateful streaming with Pandas. UDF signature is
 # ``(api_client, mode, key, pdfs)`` and returns ``Iterator[pandas.DataFrame]``.
 # The input wire stream is a single plain Arrow stream pre-sorted by the
-# grouping key column at offset 0; ``TransformWithStateInPandasSerializer``
-# chunks rows into one ``(mode, key, pdfs)`` tuple per group, then emits a
-# phantom ``PROCESS_TIMER`` and ``COMPLETE`` call with an empty pdf iterator.
+# grouping key column at offset 0; ``read_udfs`` in worker.py chunks rows into
+# one ``(mode, key, pdfs)`` tuple per group, then emits a phantom
+# ``PROCESS_TIMER`` and ``COMPLETE`` call with an empty pdf iterator.
 # ``StatefulProcessorApiClient.__init__`` opens a real TCP socket to the JVM
 # state server; the stub listener below satisfies that connect. The benchmark
 # UDFs never invoke any state API method, so no protocol exchange is needed.

python/pyspark/sql/pandas/serializers.py

@@ -836,150 +836,6 @@ def serialize_batches():
         return ArrowStreamSerializer.dump_stream(self, batches, stream)
 
 
-class TransformWithStateInPandasSerializer(ArrowStreamPandasUDFSerializer):
-    """
-    Serializer used by Python worker to evaluate UDF for
-    :meth:`pyspark.sql.GroupedData.transformWithStateInPandasSerializer`.
-
-    Parameters
-    ----------
-    timezone : str
-        A timezone to respect when handling timestamp values
-    safecheck : bool
-        If True, conversion from Arrow to Pandas checks for overflow/truncation
-    assign_cols_by_name : bool
-        If True, then Pandas DataFrames will get columns by name
-    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,
-        *,
-        timezone,
-        safecheck,
-        assign_cols_by_name,
-        prefer_int_ext_dtype,
-        arrow_max_records_per_batch,
-        arrow_max_bytes_per_batch,
-        int_to_decimal_coercion_enabled,
-    ):
-        super().__init__(
-            timezone=timezone,
-            safecheck=safecheck,
-            assign_cols_by_name=assign_cols_by_name,
-            df_for_struct=False,
-            struct_in_pandas="dict",
-            ndarray_as_list=False,
-            prefer_int_ext_dtype=prefer_int_ext_dtype,
-            arrow_cast=True,
-            input_type=None,
-            int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled,
-        )
-        self.arrow_max_records_per_batch = (
-            arrow_max_records_per_batch if arrow_max_records_per_batch > 0 else 2**31 - 1
-        )
-        self.arrow_max_bytes_per_batch = arrow_max_bytes_per_batch
-        self.key_offsets = None
-        self.average_arrow_row_size = 0
-        self.total_bytes = 0
-        self.total_rows = 0
-
-    def _update_batch_size_stats(self, batch):
-        """
-        Update batch size statistics for adaptive batching.
-        """
-        # Short circuit batch size calculation if the batch size is
-        # unlimited as computing batch size is computationally expensive.
-        if self.arrow_max_bytes_per_batch != 2**31 - 1 and batch.num_rows > 0:
-            batch_bytes = sum(
-                buf.size for col in batch.columns for buf in col.buffers() if buf is not None
-            )
-            self.total_bytes += batch_bytes
-            self.total_rows += batch.num_rows
-            self.average_arrow_row_size = self.total_bytes / self.total_rows
-
-    def load_stream(self, stream):
-        """
-        Read ArrowRecordBatches from stream, deserialize them to populate a list of data chunk, and
-        convert the data into Rows.
-
-        Please refer the doc of inner function `generate_data_batches` for more details how
-        this function works in overall.
-        """
-        import pandas as pd
-        from pyspark.sql.streaming.stateful_processor_util import (
-            TransformWithStateInPandasFuncMode,
-        )
-
-        def generate_data_batches(batches):
-            """
-            Deserialize ArrowRecordBatches and return a generator of Rows.
-
-            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.
-            """
-
-            def row_stream():
-                for batch in batches:
-                    self._update_batch_size_stats(batch)
-                    data_pandas = ArrowBatchTransformer.to_pandas(
-                        batch,
-                        timezone=self._timezone,
-                        schema=self._input_type,
-                        struct_in_pandas=self._struct_in_pandas,
-                        ndarray_as_list=self._ndarray_as_list,
-                        prefer_int_ext_dtype=self._prefer_int_ext_dtype,
-                        df_for_struct=self._df_for_struct,
-                    )
-                    for row in pd.concat(data_pandas, axis=1).itertuples(index=False):
-                        batch_key = tuple(row[s] for s in self.key_offsets)
-                        yield (batch_key, row)
-
-            for batch_key, group_rows in groupby(row_stream(), key=lambda x: x[0]):
-                rows = []
-                for _, row in group_rows:
-                    rows.append(row)
-                    if (
-                        len(rows) >= self.arrow_max_records_per_batch
-                        or len(rows) * self.average_arrow_row_size >= self.arrow_max_bytes_per_batch
-                    ):
-                        yield (batch_key, pd.DataFrame(rows))
-                        rows = []
-                if rows:
-                    yield (batch_key, pd.DataFrame(rows))
-
-        _batches = super(ArrowStreamPandasSerializer, self).load_stream(stream)
-        data_batches = generate_data_batches(_batches)
-
-        for k, g in groupby(data_batches, key=lambda x: x[0]):
-            yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, g)
-
-        yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None)
-
-        yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None)
-
-    def dump_stream(self, iterator, stream):
-        """
-        Read through an iterator of (iterator of pandas DataFrame), serialize them to Arrow
-        RecordBatches, and write batches to stream.
-        """
-
-        def flatten_iterator():
-            # iterator: iter[list[(iter[pandas.DataFrame], pdf_type)]]
-            for packed in iterator:
-                iter_pdf_with_type = packed[0]
-                iter_pdf = iter_pdf_with_type[0]
-                pdf_type = iter_pdf_with_type[1]
-                for pdf in iter_pdf:
-                    yield [(pdf, pdf_type)]
-
-        super().dump_stream(flatten_iterator(), stream)
-
-
 class TransformWithStateInPySparkRowSerializer(ArrowStreamUDFSerializer):
     """
     Serializer used by Python worker to evaluate UDF for