[SPARK-59704][PYTHON] Migrate Arrow aggregate and window UDF eval types to the EvalTypeHandler pipeline - #58956
Closed
Yicong-Huang wants to merge 4 commits into
Closed
Yicong-Huang wants to merge 4 commits into
Yicong-Huang wants to merge 4 commits into
Conversation
gaogaotiantian
approved these changes
Sep 22, 2026
Yicong-Huang
added a commit
that referenced
this pull request
Sep 23, 2026
…es to the EvalTypeHandler pipeline ### What changes were proposed in this pull request? This PR migrates the three non-incremental Arrow-native aggregate/window UDF eval types off the hand-written `if/elif` chain in `worker.py`'s `read_udfs` and onto the `EvalTypeHandler` pipeline introduced in SPARK-59425 and extended in SPARK-59622. The migrated types are `SQL_GROUPED_AGG_ARROW_UDF` (252), `SQL_GROUPED_AGG_ARROW_ITER_UDF` (254), and `SQL_WINDOW_AGG_ARROW_UDF` (253). Each becomes a `GroupedEvalTypeHandler` subclass in `python/pyspark/eval_handlers/_arrow.py` (`ArrowGroupedAggUDFHandler`, `ArrowGroupedAggIterUDFHandler`, `ArrowWindowAggUDFHandler`), declaring its `eval_type` so it self-registers and is dispatched via `get_eval_type_handler`. The handlers follow the existing lazy-pyarrow convention: string forward-ref bases, `require_minimum_pyarrow_version()` in `__init__`, and `import pyarrow as pa` inside `run`. The window handler keeps the unbounded (compute once, repeat per row) vs bounded (per-row frame slicing) dispatch verbatim. The repeated pre-19.0.0 `concat_batches` fallback is lifted into a shared `_concat_group_batches` helper. The corresponding branches and the migrated types' entries in the serializer-selection lists are removed from `worker.py`; the incremental variants are out of scope and stay in `worker.py`. This is a pure code move with no behavior change: the UDF logic, output coercion, and stream draining are identical to the removed branches. ### Why are the changes needed? Consolidating the Arrow eval types onto one registry-dispatched pipeline removes the growing per-eval-type `if/elif` chain in `read_udfs`, so the execution model for each type lives in one cohesive handler class rather than being split between a serializer-selection block and a function-construction block. Follows the umbrella SPARK-59415. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing per-eval-type tests, classic and Spark Connect: `pyspark.sql.tests.arrow.test_arrow_udf_grouped_agg`, `pyspark.sql.tests.arrow.test_arrow_udf_window`, and their `connect.arrow.test_parity_*` counterparts all pass. Worker microbenchmarks for the three eval types (including the window bounded and unbounded frames) show no regression versus the pre-refactor baseline. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Isaac This pull request and its description were written by Isaac. Closes #58956 from Yicong-Huang/SPARK-59415-arrow-agg-window. Authored-by: Yicong Huang <17627829+Yicong-Huang@users.noreply.github.com> Signed-off-by: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> (cherry picked from commit de8be6d) Signed-off-by: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com>
Contributor
Author
Yicong-Huang
pushed a commit
that referenced
this pull request
Sep 23, 2026
…atch interfaces ### What changes were proposed in this pull request? - Add a general `ArrowBatchTransformer.concat_batches` classmethod with compatibility for PyArrow versions before 19.0.0. - Keep element-wise list flattening and shape reconstruction local to the worker while wrapping flattened inputs in RecordBatches for batch-level pandas/Arrow conversion. - Reuse the shared concat helper in the grouped aggregate and window Arrow handlers migrated by [#58956](#58956), while preserving existing UDF contracts. ### Why are the changes needed? - The element-wise worker directly called the private Arrow-to-pandas array converter, coupling it to column-level conversion internals. - A RecordBatch conversion boundary allows lower-level implementations to change without worker edits and avoids duplicating the PyArrow compatibility fallback. ### Does this PR introduce _any_ user-facing change? No. This is an internal refactor that preserves existing UDF inputs, outputs, validation, and iterator behavior. ### How was this patch tested? - Added `test_concat_batches` for multi-batch concatenation and the single-batch identity fast path. - Ran the conversion, higher-order-function UDF, and Arrow eval-handler suites under pandas 2/PyArrow 24 and pandas 3/PyArrow 25; both profiles passed 125 tests and 117 subtests. Ruff, formatting, and mypy also passed. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Codex (GPT-5) Closes #58903 from Spenserrrr/arrow-worker-batch-interface. Authored-by: Spenser Sun <hsun112358@gmail.com> Signed-off-by: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com>
Yicong-Huang
pushed a commit
that referenced
this pull request
Sep 23, 2026
…atch interfaces ### What changes were proposed in this pull request? - Add a general `ArrowBatchTransformer.concat_batches` classmethod with compatibility for PyArrow versions before 19.0.0. - Keep element-wise list flattening and shape reconstruction local to the worker while wrapping flattened inputs in RecordBatches for batch-level pandas/Arrow conversion. - Reuse the shared concat helper in the grouped aggregate and window Arrow handlers migrated by [#58956](#58956), while preserving existing UDF contracts. ### Why are the changes needed? - The element-wise worker directly called the private Arrow-to-pandas array converter, coupling it to column-level conversion internals. - A RecordBatch conversion boundary allows lower-level implementations to change without worker edits and avoids duplicating the PyArrow compatibility fallback. ### Does this PR introduce _any_ user-facing change? No. This is an internal refactor that preserves existing UDF inputs, outputs, validation, and iterator behavior. ### How was this patch tested? - Added `test_concat_batches` for multi-batch concatenation and the single-batch identity fast path. - Ran the conversion, higher-order-function UDF, and Arrow eval-handler suites under pandas 2/PyArrow 24 and pandas 3/PyArrow 25; both profiles passed 125 tests and 117 subtests. Ruff, formatting, and mypy also passed. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Codex (GPT-5) Closes #58903 from Spenserrrr/arrow-worker-batch-interface. Authored-by: Spenser Sun <hsun112358@gmail.com> Signed-off-by: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> (cherry picked from commit 89062b8) Signed-off-by: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This PR migrates the three non-incremental Arrow-native aggregate/window UDF eval types off the hand-written
if/elifchain inworker.py'sread_udfsand onto theEvalTypeHandlerpipeline introduced in SPARK-59425 and extended in SPARK-59622. The migrated types areSQL_GROUPED_AGG_ARROW_UDF(252),SQL_GROUPED_AGG_ARROW_ITER_UDF(254), andSQL_WINDOW_AGG_ARROW_UDF(253).Each becomes a
GroupedEvalTypeHandlersubclass inpython/pyspark/eval_handlers/_arrow.py(ArrowGroupedAggUDFHandler,ArrowGroupedAggIterUDFHandler,ArrowWindowAggUDFHandler), declaring itseval_typeso it self-registers and is dispatched viaget_eval_type_handler. The handlers follow the existing lazy-pyarrow convention: string forward-ref bases,require_minimum_pyarrow_version()in__init__, andimport pyarrow as painsiderun. The window handler keeps the unbounded (compute once, repeat per row) vs bounded (per-row frame slicing) dispatch verbatim. The repeated pre-19.0.0concat_batchesfallback is lifted into a shared_concat_group_batcheshelper. The corresponding branches and the migrated types' entries in the serializer-selection lists are removed fromworker.py; the incremental variants are out of scope and stay inworker.py.This is a pure code move with no behavior change: the UDF logic, output coercion, and stream draining are identical to the removed branches.
Why are the changes needed?
Consolidating the Arrow eval types onto one registry-dispatched pipeline removes the growing per-eval-type
if/elifchain inread_udfs, so the execution model for each type lives in one cohesive handler class rather than being split between a serializer-selection block and a function-construction block. Follows the umbrella SPARK-59415.Does this PR introduce any user-facing change?
No.
How was this patch tested?
Existing per-eval-type tests, classic and Spark Connect:
pyspark.sql.tests.arrow.test_arrow_udf_grouped_agg,pyspark.sql.tests.arrow.test_arrow_udf_window, and theirconnect.arrow.test_parity_*counterparts all pass. Worker microbenchmarks for the three eval types (including the window bounded and unbounded frames) show no regression versus the pre-refactor baseline.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Isaac
This pull request and its description were written by Isaac.