From 1512ec8fb5352ebfd72f2f58b70d8fe3f4494d51 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 10 Jul 2026 00:29:46 +0000 Subject: [PATCH] test: add ASV microbenchmark for SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE --- python/benchmarks/bench_eval_type.py | 271 ++++++++++++++++++++++++++- 1 file changed, 270 insertions(+), 1 deletion(-) diff --git a/python/benchmarks/bench_eval_type.py b/python/benchmarks/bench_eval_type.py index cec2ead32be08..4e3d83ae6088d 100644 --- a/python/benchmarks/bench_eval_type.py +++ b/python/benchmarks/bench_eval_type.py @@ -37,12 +37,13 @@ import pyarrow as pa from pyspark.cloudpickle import dumps as cloudpickle_dumps -from pyspark.serializers import write_int, write_long, SpecialLengths +from pyspark.serializers import CPickleSerializer, write_int, write_long, SpecialLengths from pyspark.sql.types import ( BinaryType, BooleanType, DoubleType, IntegerType, + LongType, StringType, StructField, StructType, @@ -2146,3 +2147,271 @@ class TransformWithStatePandasInitStateUDFPeakmemBench( _TransformWithStatePandasInitStateBenchMixin, _PeakmemBenchBase ): pass + + +# -- SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE ------------------------------------- +# Stateful grouped map with Pandas (applyInPandasWithState). The UDF signature is +# ``(key, pdfs, state)`` and returns ``Iterator[pandas.DataFrame]``, where +# ``state`` is a ``GroupState`` the UDF may read (``getOption``) and write +# (``update``/``remove``). Unlike TransformWithState, no state server socket is +# involved: ``ApplyInPandasWithStateSerializer`` reconstructs each ``GroupState`` +# entirely from a metadata column carried inline in the Arrow stream. +# +# The wire stream is a single plain Arrow IPC stream whose batch schema is the +# data columns followed by one trailing struct column (``__state``, matching the +# JVM ``ApplyInPandasWithStateWriter.STATE_METADATA_SCHEMA``): fields +# ``properties`` (GroupStateImpl.json()), ``keyRowAsUnsafe``, ``object`` (pickled +# state value), ``startOffset``, ``numRows``, ``isLastChunk``. Data and state +# columns must share a row count, so exactly one populated state row per data +# chunk sits at the top of each group's range and the remaining state rows are +# null structs (which the serializer treats as end-of-data padding). + + +class _ApplyInPandasWithStateBenchMixin: + """Provides ``_write_scenario`` for SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE. + + Each scenario emits one plain Arrow stream whose leading int32 column is the + grouping key (pre-sorted) and whose trailing struct column carries the state + metadata. One populated state row per chunk marks that chunk's ``[startOffset, + numRows)`` range; the rest of the state column is null-padded. A group that + overflows a batch is split into multiple chunks (mirroring the JVM writer). + """ + + _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 mirror the TransformWithState pandas scenarios so identity_udf + # (full pdf passthrough) 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), + } + + # State value schema (the type of the object each GroupState carries). A + # single long counter, which count_udf reads and updates per group key. + _STATE_VALUE_SCHEMA = StructType([StructField("count", LongType())]) + + # Field order must match ApplyInPandasWithStateWriter.STATE_METADATA_SCHEMA. + _STATE_METADATA_TYPE = pa.struct( + [ + pa.field("properties", pa.string()), + pa.field("keyRowAsUnsafe", pa.binary()), + pa.field("object", pa.binary()), + pa.field("startOffset", pa.int32()), + pa.field("numRows", pa.int32()), + pa.field("isLastChunk", pa.bool_()), + ] + ) + + # GroupStateImpl.json() for a fresh, defined state with no timeout. The + # serializer overwrites ``optionalValue`` with the unpickled ``object``, so + # its placeholder value here is irrelevant. + _STATE_PROPERTIES_JSON = json.dumps( + { + "optionalValue": None, + "batchProcessingTimeMs": 0, + "eventTimeWatermarkMs": -1, + "timeoutConf": "NoTimeout", + "hasTimedOut": False, + "watermarkPresent": False, + "defined": True, + "updated": False, + "removed": False, + "timeoutTimestamp": -1, + } + ) + + @classmethod + def _build_scenario(cls, name): + """Build a single applyInPandasWithState scenario. + + Returns ``(batches, schema)`` where each batch is an Arrow RecordBatch of + data columns plus a trailing ``__state`` struct column, pre-sorted by the + leading int32 key column with one populated state row per group. + """ + np.random.seed(42) + num_groups, rows_per_group, num_value_cols, value_pool = cls._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) + ] + data_names = ["col_0"] + [f"col_{i + 1}" for i in range(num_value_cols)] + data_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) + ] + ) + + # Pickle of the initial state value (count=0), matching what the JVM + # writer pickles via PythonSQLUtils.toPyRow. + pickled_object = CPickleSerializer().dumps(cls._STATE_VALUE_SCHEMA.toInternal((0,))) + + batch_size = MockDataFactory.MAX_RECORDS_PER_BATCH + # Plan the batch/chunk layout the JVM ApplyInPandasWithStateWriter would + # produce: bin-pack whole groups into batches of at most ``batch_size`` + # data rows, and split a group that overflows a batch into multiple chunks + # (one state row each). ``isLastChunk`` is True only on a group's final + # chunk, so the serializer reassembles a group split across batches into a + # single GroupState segment. + batch_plans = cls._plan_batches(num_groups, rows_per_group, batch_size) + + batches = [] + row_offset = 0 + for chunks in batch_plans: + batch_rows = sum(c[1] for c in chunks) + data_arrays = [key_array.slice(row_offset, batch_rows)] + [ + arr.slice(row_offset, batch_rows) for arr in value_arrays + ] + state_array = cls._build_state_column(chunks, batch_rows, pickled_object) + batches.append( + pa.RecordBatch.from_arrays( + data_arrays + [state_array], names=data_names + ["__state"] + ) + ) + row_offset += batch_rows + return batches, data_schema + + @staticmethod + def _plan_batches(num_groups, rows_per_group, batch_size): + """Simulate the JVM writer's bin-packing + chunking. + + Returns a list of batches; each batch is a list of chunks + ``(start_offset_in_batch, num_rows, is_last_chunk)``. A batch holds at most + ``batch_size`` data rows; a group that exceeds the remaining space is split + into contiguous chunks, with ``is_last_chunk`` set only on the chunk that + ends the group. + """ + batch_plans = [] + cur_chunks = [] + cur_rows = 0 + for _ in range(num_groups): + remaining = rows_per_group + chunk_start = cur_rows + chunk_rows = 0 + while remaining > 0: + take = min(batch_size - cur_rows, remaining) + cur_rows += take + chunk_rows += take + remaining -= take + if cur_rows == batch_size: + cur_chunks.append((chunk_start, chunk_rows, remaining == 0)) + batch_plans.append(cur_chunks) + cur_chunks = [] + cur_rows = 0 + chunk_start = 0 + chunk_rows = 0 + elif remaining == 0: + cur_chunks.append((chunk_start, chunk_rows, True)) + if cur_chunks: + batch_plans.append(cur_chunks) + return batch_plans + + @classmethod + def _build_state_column(cls, chunks, num_rows, pickled_object): + """Build the trailing ``__state`` struct column for one batch. + + One populated state row per chunk sits at the top of the batch; the + remaining rows are null structs -- the padding the serializer stops at + when scanning state rows. + """ + properties = [None] * num_rows + key_bytes = [None] * num_rows + objects = [None] * num_rows + start_offsets = [None] * num_rows + num_rows_col = [None] * num_rows + is_last_chunk = [None] * num_rows + mask = np.ones(num_rows, dtype=bool) # True == null + + for idx, (start_offset, chunk_rows, is_last) in enumerate(chunks): + properties[idx] = cls._STATE_PROPERTIES_JSON + key_bytes[idx] = b"" # opaque to the worker + objects[idx] = pickled_object + start_offsets[idx] = start_offset + num_rows_col[idx] = chunk_rows + is_last_chunk[idx] = is_last + mask[idx] = False + + return pa.StructArray.from_arrays( + [ + pa.array(properties, type=pa.string()), + pa.array(key_bytes, type=pa.binary()), + pa.array(objects, type=pa.binary()), + pa.array(start_offsets, type=pa.int32()), + pa.array(num_rows_col, type=pa.int32()), + pa.array(is_last_chunk, type=pa.bool_()), + ], + fields=list(cls._STATE_METADATA_TYPE), + mask=pa.array(mask), + ) + + def _apply_pandas_state_identity(key, pdfs, state): + yield from pdfs + + def _apply_pandas_state_sort(key, pdfs, state): + for pdf in pdfs: + yield pdf.sort_values(pdf.columns[0]) + + def _apply_pandas_state_count(key, pdfs, state): + import pandas as pd + + # Exercise the per-group state read/write path: read the running count, + # add this group's row total, write it back. + prev = state.getOption + prev_count = prev[0] if prev is not None else 0 + total = prev_count + sum(len(pdf) for pdf in pdfs) + state.update((total,)) + 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 pass the value-only pdf straight through). count_udf re-emits the + # key plus the count, so it declares an explicit (key, count) output schema. + _udfs = { + "identity_udf": (_apply_pandas_state_identity, None), + "sort_udf": (_apply_pandas_state_sort, None), + "count_udf": ( + _apply_pandas_state_count, + StructType([StructField("col_0", IntegerType()), StructField("col_1", LongType())]), + ), + } + params = [list(_scenario_configs), list(_udfs)] + param_names = ["scenario", "udf"] + + _NUM_KEY_COLS = 1 + + def _write_scenario(self, scenario, udf_name, buf): + batches, data_schema = self._build_scenario(scenario) + udf_func, ret_type = self._udfs[udf_name] + if ret_type is None: + ret_type = StructType(data_schema.fields[self._NUM_KEY_COLS :]) + n_value_cols = len(data_schema.fields) - self._NUM_KEY_COLS + arg_offsets = MockUDFFactory.make_grouped_arg_offsets(self._NUM_KEY_COLS, n_value_cols) + MockProtocolWriter.write_worker_input( + PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE, + 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_value_schema": self._STATE_VALUE_SCHEMA.json()}, + ) + + +class ApplyInPandasWithStateUDFTimeBench(_ApplyInPandasWithStateBenchMixin, _TimeBenchBase): + pass + + +class ApplyInPandasWithStateUDFPeakmemBench(_ApplyInPandasWithStateBenchMixin, _PeakmemBenchBase): + pass