[SPARK-58189][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE

[Yicong-Huang]

What changes were proposed in this pull request?

This PR refactors SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE (applyInPandasWithState) so that the worker uses the plain ArrowStreamSerializer for pure Arrow stream I/O, moving the per-eval-type logic from ApplyInPandasWithStateSerializer into read_udfs() in worker.py: reconstructing each GroupState from the trailing state-info column, slicing the data columns by each chunk's (startOffset, numRows), regrouping chunks by grouping key (via the reused state object), invoking the UDF once per key with a lazy iterator of pandas DataFrames and its state, validating the returned DataFrames, and serializing the output plus updated state back into count/data/state Arrow RecordBatches bounded by arrow_max_records_per_batch. This mirrors the earlier refactor of SQL_TRANSFORM_WITH_STATE_PANDAS_UDF in SPARK-57400. The ApplyInPandasWithStateSerializer class is kept for now and deleted in a follow-up cleanup.

Why are the changes needed?

Part of SPARK-55388. Keeping serializers as pure Arrow stream I/O and concentrating eval-type-specific logic in worker.py makes the per-eval-type data flow explicit.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

Existing tests. No behavior change.

ASV comparison (bench_eval_type.ApplyInPandasWithStateUDFTimeBench / ApplyInPandasWithStateUDFPeakmemBench, -a repeat=3): before = upstream/master, after = this PR. Values are from one representative run per side; the conclusion is consistent across runs. All deltas are within run-to-run noise (the larger swings track high before-side variance, e.g. few_groups_sm/count before 131ms+-228ms), showing no regression. Peak memory is unchanged.

time_worker

scenario         udf            before          after            diff
---------------  -------------  --------------  --------------  ------
few_groups_sm    identity_udf   1.05s+-25.2ms   1.04s+-13.8ms    -1.0%
few_groups_sm    sort_udf       1.07s+-37.3ms   1.08s+-7.33ms    +0.9%
few_groups_sm    count_udf      131ms+-228ms    129ms+-255ms     -1.5%
few_groups_lg    identity_udf   7.21s+-78.4ms   7.19s+-49.9ms    -0.3%
few_groups_lg    sort_udf       7.57s+-75.1ms   7.33s+-9.87ms    -3.2%
few_groups_lg    count_udf      800ms+-8.91ms   781ms+-8.37ms    -2.4%
many_groups_sm   identity_udf   7.12s+-106ms    7.00s+-34.9ms    -1.7%
many_groups_sm   sort_udf       7.85s+-58.7ms   7.83s+-24.1ms    -0.3%
many_groups_sm   count_udf      3.26s+-16.5ms   3.24s+-32.8ms    -0.6%
many_groups_lg   identity_udf   4.83s+-136ms    4.75s+-72.3ms    -1.7%
many_groups_lg   sort_udf       5.33s+-85ms     5.00s+-42.5ms    -6.2%
many_groups_lg   count_udf      951ms+-6.36ms   935ms+-2.69ms    -1.7%
wide_cols        identity_udf   6.13s+-30ms     5.94s+-29.8ms    -3.1%
wide_cols        sort_udf       6.36s+-404ms    6.30s+-78.9ms    -0.9%
wide_cols        count_udf      913ms+-25.6ms   940ms+-22.3ms    +3.0%
mixed_cols       identity_udf   4.57s+-72.5ms   4.84s+-49.7ms    +5.9%
mixed_cols       sort_udf       4.91s+-86.6ms   4.81s+-68.8ms    -2.0%
mixed_cols       count_udf      739ms+-20.2ms   709ms+-5.46ms    -4.1%
nested_struct    identity_udf   9.62s+-183ms    9.27s+-22.4ms    -3.6%
nested_struct    sort_udf       12.1s+-62.1ms   11.0s+-154ms     -9.1%
nested_struct    count_udf      3.95s+-21ms     3.53s+-27.1ms   -10.6%

peakmem_worker

scenario         udf            before   after
---------------  -------------  -------  -------
few_groups_sm    identity_udf   19M      19M
few_groups_sm    sort_udf       19M      19M
few_groups_sm    count_udf      2M       2M
few_groups_lg    identity_udf   156M     156M
few_groups_lg    sort_udf       156M     156M
few_groups_lg    count_udf      3M       3M
many_groups_sm   identity_udf   62M      63M
many_groups_sm   sort_udf       62M      63M
many_groups_sm   count_udf      20M      20M
many_groups_lg   identity_udf   61M      61M
many_groups_lg   sort_udf       61M      61M
many_groups_lg   count_udf      6M       6M
wide_cols        identity_udf   163M     163M
wide_cols        sort_udf       163M     163M
wide_cols        count_udf      6M       6M
mixed_cols       identity_udf   66M      66M
mixed_cols       sort_udf       66M      66M
mixed_cols       count_udf      4M       4M
nested_struct    identity_udf   103M     103M
nested_struct    sort_udf       103M     103M
nested_struct    count_udf      10M      10M

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

No.

[uros-b]

Thank you @Yicong-Huang and @HyukjinKwon!