Skip to content

fix: revert unsafe partial aggregates after final fallback - #5421

Open
sunchao wants to merge 1 commit into
apache:mainfrom
sunchao:dev/chao/codex/oss-unsafe-partial-aggregate-fallback
Open

fix: revert unsafe partial aggregates after final fallback#5421
sunchao wants to merge 1 commit into
apache:mainfrom
sunchao:dev/chao/codex/oss-unsafe-partial-aggregate-fallback

Conversation

@sunchao

@sunchao sunchao commented Aug 22, 2026

Copy link
Copy Markdown
Member

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:

SELECT AVG(amount)
FROM decimal_avg_probe
WHERE id = 1;

Only one partition contains the matching row. Spark correctly returns 200.000000, but Comet returns NULL with adaptive execution either enabled or disabled. The problematic plan is:

Spark HashAggregate [Final AVG]
  Spark ShuffleExchange
    CometHashAggregate [Partial AVG]  <-- unsupported buffer boundary
      CometProject
        CometFilter
          Comet scan conversion

Decimal AVG is 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:

Spark HashAggregate [Final AVG]
  Spark ShuffleExchange
    Spark HashAggregate [Partial AVG]
      CometColumnarToRow
        CometProject
          CometFilter
            Comet scan conversion

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 MIN and MAX, 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 AVG buffer defect is tracked in #5418 and fixed by #5420. Correcting that representation does not eliminate the need for this planner safeguard: decimal AVG and 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 main and pass after this change. They cover the decimal AVG wrong-result case with AQE disabled and enabled, along with ordinary and distinct aggregate chains when native shuffle cannot be used.

The complete CometAggregateSuite, CometExecRuleSuite, and CometShuffleFallbackStickinessSuite passed 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 mixed MIN/MAX, and fully native aggregate chains.

An independent standalone-JAR replay against stock Spark 4.0.2 confirmed the corrected decimal AVG result 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, and git diff --check also passed. GitHub CI reports 63 passing checks, with nine inapplicable checks skipped.

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
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.

Unsafe native partial aggregates survive child-triggered final fallback

1 participant