fix: revert unsafe partial aggregates after final fallback - #5421
Open
sunchao wants to merge 1 commit into
Open
Conversation
sunchao
added a commit
to sunchao/arrow-datafusion-comet
that referenced
this pull request
Aug 26, 2026
sunchao
added a commit
to sunchao/arrow-datafusion-comet
that referenced
this pull request
Aug 26, 2026
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
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.
Why are the changes needed?
Closes #5419.
Comet can silently produce incorrect results when an aggregate starts in Comet but finishes in Spark after its shuffle falls back. Although the planner already tries to reject unsafe mixed-engine aggregation, its existing check runs before child operators are converted. A final aggregate can therefore appear convertible early in planning but later remain in Spark because its shuffle child is not native. Its partial aggregate has already been converted to Comet, leaving an execution boundary that the aggregate's intermediate buffer was never declared safe to cross.
For example, consider eight rows spread across four Parquet files, each with
amount = CAST(200 AS DECIMAL(20, 2)). With native shuffle disabled, run:Only one partition contains the matching row. Spark correctly returns
200.000000, but Comet returnsNULLwith adaptive execution either enabled or disabled. The problematic plan is:Decimal
AVGis explicitly ineligible for mixed Spark/Comet execution because its intermediate state has compatibility requirements beyond simply matching the output data type. In this example, an empty Comet partial contributes a null sum that poisons Spark's final aggregation. The same admission gap can affect other aggregates whose buffers are not declared safe for mixed execution, including distinct aggregations with intermediate merge stages. With AQE enabled, the plan must be repaired before a shuffle stage materializes; afterward, its incompatible buffers have already been produced.What changes were proposed in this PR?
Base the fallback decision on the execution plan that conversion actually produced, not solely on an early estimate of whether aggregate expressions look convertible. The existing pre-conversion check remains useful, but a new post-conversion reconciliation verifies that any aggregate left on Spark is not consuming an unsupported Comet partial. When such a boundary exists, its feeding partial aggregate and associated aggregation/shuffle chain are restored to their Spark implementations before execution begins.
The corrected plan keeps the incompatible aggregate buffer inside Spark while preserving native work below it:
This fallback is deliberately local rather than a blanket retreat to Spark. Distinct-aggregate chains and AQE stage replanning remain consistent, while existing materialized or reused stages are not rewritten. Native scans, filters, and projections below the aggregate stay native; supported mixed aggregates such as
MINandMAX, as well as aggregate pipelines that can run entirely in Comet, retain their existing native execution.This PR changes neither the mixed-execution eligibility policy nor native accumulator behavior. The separate empty-partial
AVGbuffer defect is tracked in #5418 and fixed by #5420. Correcting that representation does not eliminate the need for this planner safeguard: decimalAVGand other unsupported aggregate buffers still cannot be freely exchanged between Spark and Comet.How was this PR tested?
All four new targeted regressions fail on public
mainand pass after this change. They cover the decimalAVGwrong-result case with AQE disabled and enabled, along with ordinary and distinct aggregate chains when native shuffle cannot be used.The complete
CometAggregateSuite,CometExecRuleSuite, andCometShuffleFallbackStickinessSuitepassed on Spark 4.0.4 / Java 17: 121 tests passed, with no failures or aborted suites and two existing Spark-version-gated cancellations. The public-base native library was unchanged. Plan assertions verify AQE materialization, repeated whole-plan and stage-only planning, preserved native operators below the partial, safe mixedMIN/MAX, and fully native aggregate chains.An independent standalone-JAR replay against stock Spark 4.0.2 confirmed the corrected decimal
AVGresult with AQE on and off. The separate integer/narrow-decimal empty-partition defect from #5418 remains reproducible when this planner fix is tested alone; applying both fixes together matches Spark in 20/20 synthetic cases. Root-reactor Maven packaging, Spotless, Scala style checks, andgit diff --checkalso passed. GitHub CI reports 63 passing checks, with nine inapplicable checks skipped.