Yicong-Huang opened a new pull request, #57260: URL: https://github.com/apache/spark/pull/57260
### What changes were proposed in this pull request? This PR refactors `SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF` so that the worker uses the plain `ArrowStreamSerializer` for pure Arrow stream I/O, moving the per-eval-type logic (flattening the `inputData`/`initState` struct columns, regrouping rows by grouping key, re-chunking into pandas DataFrames bounded by `arrow_max_records_per_batch`/`arrow_max_bytes_per_batch`, splitting each group into a data iterator and an init-state iterator, and converting result DataFrames back to Arrow) from `TransformWithStateInPandasInitStateSerializer` into `read_udfs()` in `worker.py`. This mirrors the earlier refactor of the non-init variant in [SPARK-57400](https://issues.apache.org/jira/browse/SPARK-57400). The `TransformWithStateInPandasInitStateSerializer` class is kept for now and deleted in a follow-up cleanup. ### Why are the changes needed? Part of [SPARK-55388](https://issues.apache.org/jira/browse/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.TransformWithStatePandasInitStateUDFTimeBench` / `TransformWithStatePandasInitStateUDFPeakmemBench`, `-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 negatives track high before-side variance, e.g. `many_groups_lg/count` before `5.05+-0.9s`), showing no regression. ```text time_worker scenario udf before after diff --------------- ------------ ------------ ------------ ------ few_groups_sm identity_udf 854+-7ms 850+-8ms -0.5% few_groups_sm sort_udf 875+-10ms 848+-20ms -3.1% few_groups_sm count_udf 908+-30ms 881+-10ms -3.0% few_groups_lg identity_udf 7.70+-0.05s 7.78+-0.1s +1.0% few_groups_lg sort_udf 7.81+-0.04s 7.73+-0.01s -1.0% few_groups_lg count_udf 7.21+-0.1s 7.26+-0.09s +0.7% many_groups_sm identity_udf 8.07+-0.1s 8.13+-0.01s +0.7% many_groups_sm sort_udf 8.36+-0.03s 8.45+-0.02s +1.1% many_groups_sm count_udf 9.39+-0.09s 9.44+-0.05s +0.5% many_groups_lg identity_udf 4.21+-0.06s 4.24+-0.06s +0.7% many_groups_lg sort_udf 4.45+-0.2s 4.31+-0.05s -3.1% many_groups_lg count_udf 5.05+-0.9s 4.46+-0.1s -11.7% wide_cols identity_udf 9.14+-0.6s 8.33+-0.07s -8.9% wide_cols sort_udf 8.70+-0.1s 8.79+-0.3s +1.0% wide_cols count_udf 8.07+-0.1s 7.92+-0.2s -1.9% mixed_cols identity_udf 3.76+-0.1s 3.67+-0.08s -2.4% mixed_cols sort_udf 3.81+-0.2s 3.84+-0.09s +0.8% mixed_cols count_udf 3.51+-0.2s 3.57+-0.07s +1.7% nested_struct identity_udf 8.74+-0.2s 8.76+-0.1s +0.2% nested_struct sort_udf 9.55+-0.8s 8.79+-0.2s -8.0% nested_struct count_udf 6.33+-0.2s 6.23+-0.1s -1.6% ``` ```text peakmem_worker scenario udf before after --------------- ------------ ------- ------- few_groups_sm identity_udf 115M 118M few_groups_sm sort_udf 118M 116M few_groups_sm count_udf 107M 107M few_groups_lg identity_udf 249M 249M few_groups_lg sort_udf 249M 249M few_groups_lg count_udf 249M 249M many_groups_sm identity_udf 176M 176M many_groups_sm sort_udf 178M 180M many_groups_sm count_udf 162M 162M many_groups_lg identity_udf 152M 152M many_groups_lg sort_udf 152M 152M many_groups_lg count_udf 152M 152M wide_cols identity_udf 365M 368M wide_cols sort_udf 369M 377M wide_cols count_udf 343M 343M mixed_cols identity_udf 182M 183M mixed_cols sort_udf 182M 183M mixed_cols count_udf 182M 183M nested_struct identity_udf 211M 211M nested_struct sort_udf 211M 211M nested_struct count_udf 211M 211M ``` ### Was this patch authored or co-authored using generative AI tooling? No. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
