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/06/11 18:23:51 UTC
(spark) branch branch-4.x updated: [SPARK-57020][PYTHON][TEST] Add ASV microbenchmark for SQL_TRANSFORM_WITH_STATE_PANDAS_UDF
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 b872c047a536 [SPARK-57020][PYTHON][TEST] Add ASV microbenchmark for SQL_TRANSFORM_WITH_STATE_PANDAS_UDF
b872c047a536 is described below
commit b872c047a536981206d285505e1f296778b00e8f
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Thu Jun 11 18:23:19 2026 +0000
[SPARK-57020][PYTHON][TEST] Add ASV microbenchmark for SQL_TRANSFORM_WITH_STATE_PANDAS_UDF
### What changes were proposed in this pull request?
Add an ASV micro-benchmark for `SQL_TRANSFORM_WITH_STATE_PANDAS_UDF` to `bench_eval_type.py`.
A stub TCP listener (`_StubStateServer`) satisfies `StatefulProcessorApiClient`'s socket connect; the benchmark UDFs never call any state API so no protocol exchange beyond connect is needed.
Scenarios cover few/many groups, small/large group sizes, wide columns, mixed value types (string/binary/boolean) and a nested struct column. UDFs: `identity_udf`, `sort_udf`, `count_udf`. The input pdfs are value-only (the grouping key is projected out before the UDF, mirroring `worker.py`'s `values_gen`), so `identity_udf`/`sort_udf` pass values through and `count_udf` reconstructs the grouping key from the `key` arg to keep it in the output, matching the common transformWithState o [...]
### Why are the changes needed?
Part of [SPARK-55724](https://issues.apache.org/jira/browse/SPARK-55724). Establishes a performance baseline before refactoring `SQL_TRANSFORM_WITH_STATE_PANDAS_UDF`.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
`COLUMNS=120 ./python/asv run --python=same --bench "TransformWithStatePandas" -a "repeat=(3,5,5.0)"` (one of two stable runs):
`TransformWithStatePandasUDFTimeBench`:
```text
============== ============ ========== ==========
scenario identity_udf sort_udf count_udf
============== ============ ========== ==========
few_groups_sm 742±20ms 768±9ms 724±2ms
few_groups_lg 7.04±0.04s 7.23±0.1s 6.63±0.01s
many_groups_sm 6.14±0.02s 6.65±0.01s 5.73±0.01s
many_groups_lg 3.50±0.02s 3.68±0.04s 3.37±0.02s
wide_cols 7.51±0.1s 7.51±0.01s 6.82±0.03s
mixed_cols 3.17±0.02s 3.36±0.05s 2.87±0.03s
nested_struct 7.47±0.01s 8.62±0.07s 5.19±0.04s
============== ============ ========== ==========
```
`TransformWithStatePandasUDFPeakmemBench`:
```text
============== ============ ======== =========
scenario identity_udf sort_udf count_udf
============== ============ ======== =========
few_groups_sm 486M 488M 477M
few_groups_lg 569M 581M 541M
many_groups_sm 513M 512M 492M
many_groups_lg 517M 514M 492M
wide_cols 620M 621M 585M
mixed_cols 561M 561M 561M
nested_struct 589M 589M 589M
============== ============ ======== =========
```
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #56192 from Yicong-Huang/SPARK-57020/bench/tws-pandas.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
(cherry picked from commit 8cced6f5bd834f027aadb5b6684566879d89831e)
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
python/benchmarks/bench_eval_type.py | 195 +++++++++++++++++++++++++++++++++++
1 file changed, 195 insertions(+)
diff --git a/python/benchmarks/bench_eval_type.py b/python/benchmarks/bench_eval_type.py
index 17b112e5f89f..745885ddb1b7 100644
--- a/python/benchmarks/bench_eval_type.py
+++ b/python/benchmarks/bench_eval_type.py
@@ -26,9 +26,11 @@ by constructing the complete binary protocol that ``worker.py``'s
import io
import os
import json
+import socket
import struct
import sys
import tempfile
+import threading
from typing import Any, Callable, Iterator
import numpy as np
@@ -1791,3 +1793,196 @@ class WindowAggPandasUDFTimeBench(_WindowAggPandasBenchMixin, _TimeBenchBase):
class WindowAggPandasUDFPeakmemBench(_WindowAggPandasBenchMixin, _PeakmemBenchBase):
pass
+
+
+# -- SQL_TRANSFORM_WITH_STATE_PANDAS_UDF ---------------------------------------
+# 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.
+# ``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.
+
+
+class _StubStateServer:
+ """Stub TCP listener so ``StatefulProcessorApiClient`` init succeeds.
+
+ One instance per benchmark process; the port is reused across all scenarios
+ and ASV iterations. The accept loop stashes connections to keep them open
+ until the worker process tears them down (the worker never closes its end
+ explicitly, but Python GCs the socket on ``main`` return).
+ """
+
+ _instance: "_StubStateServer | None" = None
+
+ @classmethod
+ def get_port(cls) -> int:
+ if cls._instance is None:
+ cls._instance = cls()
+ return cls._instance.port
+
+ def __init__(self) -> None:
+ self._sock = socket.socket()
+ self._sock.bind(("127.0.0.1", 0))
+ self._sock.listen(128)
+ self.port = self._sock.getsockname()[1]
+ self._connections: list[socket.socket] = []
+ self._thread = threading.Thread(target=self._accept_loop, daemon=True)
+ self._thread.start()
+
+ def _accept_loop(self) -> None:
+ while True:
+ try:
+ conn, _ = self._sock.accept()
+ except OSError:
+ break
+ self._connections.append(conn)
+
+
+class _TransformWithStatePandasBenchMixin:
+ """Provides ``_write_scenario`` for SQL_TRANSFORM_WITH_STATE_PANDAS_UDF.
+
+ Each scenario emits one plain Arrow stream pre-sorted by the leading int
+ key column. UDFs receive an iterator of value-only Pandas DataFrames per
+ group (the grouping key is projected out before the UDF, mirroring
+ ``worker.py``'s ``values_gen``) plus phantom ``PROCESS_TIMER``/``COMPLETE``
+ calls (empty iterator).
+ """
+
+ # Per-scenario value-column type pool. ``mixed_cols`` exercises the
+ # string/binary/boolean encode paths and ``nested_struct`` exercises the
+ # struct (dict) conversion path; the rest stay numeric to keep encode cost
+ # dominated by row volume rather than per-value Python work.
+ _MIXED_POOL = MockDataFactory.MIXED_TYPES
+ _NESTED_POOL = [
+ MockDataFactory.TYPE_REGISTRY["int"],
+ MockDataFactory.make_struct_type(num_fields=3, base_types=MockDataFactory.MIXED_TYPES),
+ ]
+
+ # Each scenario: (num_groups, rows_per_group, num_value_cols, value_pool).
+ # Row counts are scaled so identity_udf (full pdf passthrough -> ~equal
+ # input and output volume) stays under ASV's 60s per-sample timeout.
+ _scenario_configs = {
+ "few_groups_sm": (50, 5_000, 5, MockDataFactory.NUMERIC_TYPES),
+ "few_groups_lg": (50, 50_000, 5, MockDataFactory.NUMERIC_TYPES),
+ "many_groups_sm": (2_000, 500, 5, MockDataFactory.NUMERIC_TYPES),
+ "many_groups_lg": (500, 2_000, 5, MockDataFactory.NUMERIC_TYPES),
+ "wide_cols": (200, 5_000, 20, MockDataFactory.NUMERIC_TYPES),
+ "mixed_cols": (200, 5_000, 5, _MIXED_POOL),
+ "nested_struct": (200, 5_000, 4, _NESTED_POOL),
+ }
+
+ @staticmethod
+ def _build_scenario(name):
+ """Build a single TWS Pandas scenario.
+
+ Returns ``(batches, schema)`` where ``batches`` is a plain list of Arrow
+ RecordBatches with rows pre-sorted by the leading int32 key column.
+ """
+ np.random.seed(42)
+ num_groups, rows_per_group, num_value_cols, value_pool = (
+ _TransformWithStatePandasBenchMixin._scenario_configs[name]
+ )
+ total_rows = num_groups * rows_per_group
+ key_array = pa.array(
+ np.repeat(np.arange(num_groups, dtype=np.int32), rows_per_group),
+ type=pa.int32(),
+ )
+ value_arrays = [
+ value_pool[i % len(value_pool)][0](total_rows) for i in range(num_value_cols)
+ ]
+ names = ["col_0"] + [f"col_{i + 1}" for i in range(num_value_cols)]
+ full_batch = pa.RecordBatch.from_arrays([key_array] + value_arrays, names=names)
+ batch_size = MockDataFactory.MAX_RECORDS_PER_BATCH
+ batches = [
+ full_batch.slice(offset, min(batch_size, total_rows - offset))
+ for offset in range(0, total_rows, batch_size)
+ ]
+ schema = StructType(
+ [StructField("col_0", IntegerType())]
+ + [
+ StructField(f"col_{i + 1}", value_pool[i % len(value_pool)][1])
+ for i in range(num_value_cols)
+ ]
+ )
+ return batches, schema
+
+ def _tws_pandas_identity(api_client, mode, key, pdfs):
+ from pyspark.sql.streaming.stateful_processor_util import (
+ TransformWithStateInPandasFuncMode,
+ )
+
+ if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
+ yield from pdfs
+
+ def _tws_pandas_sort(api_client, mode, key, pdfs):
+ from pyspark.sql.streaming.stateful_processor_util import (
+ TransformWithStateInPandasFuncMode,
+ )
+
+ if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
+ for pdf in pdfs:
+ yield pdf.sort_values(pdf.columns[0])
+
+ def _tws_pandas_count(api_client, mode, key, pdfs):
+ import pandas as pd
+ from pyspark.sql.streaming.stateful_processor_util import (
+ TransformWithStateInPandasFuncMode,
+ )
+
+ if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
+ total = sum(len(pdf) for pdf in pdfs)
+ # The input pdfs are value-only, so an aggregating UDF reconstructs
+ # the grouping key from the ``key`` arg to keep it in the output --
+ # the common shape for transformWithState (see the
+ # transformWithStateInPandas docstring example).
+ yield pd.DataFrame({"col_0": [key[0]], "col_1": [total]})
+
+ # ret_type=None means "use all value columns of the input schema" (identity
+ # and sort are pure value-column passthroughs, so their output omits the
+ # key, which the input pdf never carries). count_udf re-emits the key, so it
+ # declares an explicit output schema of (key, count).
+ _udfs = {
+ "identity_udf": (_tws_pandas_identity, None),
+ "sort_udf": (_tws_pandas_sort, None),
+ "count_udf": (
+ _tws_pandas_count,
+ StructType([StructField("col_0", IntegerType()), StructField("col_1", IntegerType())]),
+ ),
+ }
+ params = [list(_scenario_configs), list(_udfs)]
+ param_names = ["scenario", "udf"]
+
+ _NUM_KEY_COLS = 1
+
+ def _write_scenario(self, scenario, udf_name, buf):
+ batches, schema = self._build_scenario(scenario)
+ udf_func, ret_type = self._udfs[udf_name]
+ if ret_type is None:
+ ret_type = StructType(schema.fields[self._NUM_KEY_COLS :])
+ n_value_cols = len(schema.fields) - self._NUM_KEY_COLS
+ arg_offsets = MockUDFFactory.make_grouped_arg_offsets(self._NUM_KEY_COLS, n_value_cols)
+ grouping_key_schema = StructType(schema.fields[: self._NUM_KEY_COLS])
+ MockProtocolWriter.write_worker_input(
+ PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF,
+ lambda b: MockProtocolWriter.write_udf_payload(udf_func, ret_type, arg_offsets, b),
+ lambda b: MockProtocolWriter.write_data_payload(iter(batches), b),
+ buf,
+ eval_conf={
+ "state_server_socket_port": str(_StubStateServer.get_port()),
+ "grouping_key_schema": grouping_key_schema.json(),
+ },
+ )
+
+
+class TransformWithStatePandasUDFTimeBench(_TransformWithStatePandasBenchMixin, _TimeBenchBase):
+ pass
+
+
+class TransformWithStatePandasUDFPeakmemBench(
+ _TransformWithStatePandasBenchMixin, _PeakmemBenchBase
+):
+ pass
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org