diff --git a/python/benchmarks/bench_eval_type.py b/python/benchmarks/bench_eval_type.py index da59c664df566..1632d9a857250 100644 --- a/python/benchmarks/bench_eval_type.py +++ b/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. diff --git a/python/pyspark/sql/pandas/serializers.py b/python/pyspark/sql/pandas/serializers.py index 48516f47be548..0b3b143bcf2cd 100644 --- a/python/pyspark/sql/pandas/serializers.py +++ b/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