Skip to content

[SPARK-58067][PYTHON][TESTS] Add ASV microbenchmark for SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF#57167

Open
Yicong-Huang wants to merge 2 commits into
apache:masterfrom
Yicong-Huang:SPARK-58067
Open

[SPARK-58067][PYTHON][TESTS] Add ASV microbenchmark for SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF#57167
Yicong-Huang wants to merge 2 commits into
apache:masterfrom
Yicong-Huang:SPARK-58067

Conversation

@Yicong-Huang

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

Add an ASV microbenchmark for the SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF eval type (served by TransformWithStateInPySparkRowInitStateSerializer), mirroring the existing SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF benchmark and the non-init Row benchmark (#57047). It reuses the existing bench_eval_type.py harness and the plain-Row scenario grid from _TransformWithStateRowBenchMixin, adding TransformWithStateRowInitStateUDF{Time,Peakmem}Bench over 7 scenarios x 3 UDFs. The UDF signature is (api_client, mode, key, rows, init_rows) where both rows and init_rows are iterators of Row. The input wire stream is the nested struct<inputData, initState> Arrow stream (init batches first, then data batches); the serializer materializes every column into a Row via .as_py() and regroups by the leading key, so each key surfaces as one init-only call followed by one data-only call -- the per-row Python object path this eval type is built around, layered on the init-state deserialization.

Stacked on #57047 (_TransformWithStateRowBenchMixin); will rebase onto master once that merges.

Why are the changes needed?

Establishes a no-regression baseline for the upcoming refactor that moves the Row init-state transformWithState logic out of the serializer into read_udfs().

Does this PR introduce any user-facing change?

No.

How was this patch tested?

Benchmark-only change. Ran asv run --bench 'TransformWithStateRowInitState' -a repeat=3 --python=same twice; results were stable across runs (time within ~2%, peakmem identical). One representative run:

[time_worker]
================ ============== ============= ============
--                                  udf
---------------- -----------------------------------------
    scenario      identity_udf   rebuild_udf   count_udf
================ ============== ============= ============
 few_groups_sm      475±3ms        529±2ms      391±1ms
 few_groups_lg     4.04±0.02s     4.60±0.01s    3.24±0s
 many_groups_sm    4.73±0.02s     4.96±0.02s   4.20±0.01s
 many_groups_lg    2.39±0.01s     2.62±0.03s    2.01±0s
   wide_cols       5.14±0.01s     5.62±0.02s    4.24±0s
   mixed_cols      1.96±0.01s     2.19±0.01s   1.62±0.01s
 nested_struct      3.31±0s       3.51±0.01s    2.68±0s
================ ============== ============= ============

[peakmem_worker]
================ ============== ============= ===========
    scenario      identity_udf   rebuild_udf   count_udf
================ ============== ============= ===========
 few_groups_sm       92.2M          92.2M        88.7M
 few_groups_lg        114M           114M         104M
 many_groups_sm       105M           105M         102M
 many_groups_lg      98.7M          98.7M        92.6M
   wide_cols          113M           113M         113M
   mixed_cols         101M           101M        97.1M
 nested_struct        103M           103M         103M
================ ============== ============= ===========

Was this patch authored or co-authored using generative AI tooling?

No.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant