fix: revert unsafe partial aggregates after final fallback

[sunchao]

Which issue does this PR close?

Closes #5419.

Rationale for this change

Comet can leave a native partial aggregate underneath a Spark final aggregate after an enabled native shuffle turns out to be ineligible. The early eligibility check runs before child conversion, so it cannot always see that boundary. Spark then consumes an intermediate buffer whose native representation it does not support.

collect_list and collect_set demonstrate the problem independently of the decimal AVG fix in #5420. With Comet shuffle enabled in native-only mode but native hash partitioning disabled, this grouped query must fall back across the Spark shuffle:

SELECT group_id, sort_array(collect_list(value)), COUNT(*)
FROM records
GROUP BY group_id;

The native partial produces an Arrow array buffer, whereas Spark's final typed aggregate expects a serialized binary buffer. Leaving those stages in different engines can crash. This example uses ordinary integer grouping keys, so it does not depend on #5420's later wide-decimal shuffle restriction. The committed tests exercise both collection functions with AQE on and off.

Decimal AVG supplied the original wrong-result example: with four input partitions and only one surviving value of CAST(200 AS DECIMAL(20, 2)), an empty native partial could poison Spark's final sum and produce NULL instead of 200.000000. #5420 repairs that separate native state defect. It does not make the typed collection buffers interchangeable or remove the need to repair unsupported aggregate boundaries before AQE materializes them.

What changes are included in this PR?

A post-conversion check repairs the feeding aggregate/exchange chain whenever a Spark final would consume an incompatible native partial. The earlier tagging check remains, and the repaired partial is tagged so stage-only replanning cannot convert it again. Materialized and reused stages are not rewritten. Native scans, filters, and projections below the partial remain native.

The compatibility policy now distinguishes supportsNativePartialToSparkFinal from supportsSparkPartialToNativeFinal. COUNT permits NativePartialToSparkFinal: its native buffer is a non-null Long that Spark can consume. COUNT still rejects SparkPartialToNativeFinal to preserve the AQE and correlated-subquery rewrites behind that restriction.

AVG explicitly rejects NativePartialToSparkFinal until the native emitted-state repair in #5420 is present. A never-updated native non-decimal AVG partial emits (null, 0) rather than Spark's (0.0, 0) state. Allowing a compatible COUNT in the same aggregate must not admit that AVG buffer: with four input partitions and one surviving BIGINT value of 1, SELECT COUNT(*), AVG(value) ... WHERE id = 1 otherwise returns (1, NULL) instead of (1, 1.0). The new policy keeps the whole mixed aggregate in Spark while preserving native filters below it. AVG's existing non-decimal SparkPartialToNativeFinal support and fully native aggregate execution remain unchanged. Other exclusions, including decimal SUM/AVG and typed array buffers, remain conservative; ordinary-input results alone do not establish their overflow, null, mode, or multi-stage compatibility.

The planner repair remains a prerequisite for #5420's native AVG and wide-decimal partitioning fixes. The conservative AVG NativePartialToSparkFinal restriction in this prerequisite is separate from the native state repair: the repaired buffers in #5420 need their own admission validation. No native accumulator code changes in this PR.

How are these changes tested?

Array-key regression follow-up at c078ef136 (2026-08-29): added six cases from the additional #5419 reproducer, covering percentile, collect_list, and a safe SUM control with AQE off/on. Native hash partitioning remains enabled; the array key itself forces Spark shuffle fallback. The tests compare the reported queries with Spark and inspect both the initial and materialized plans, requiring compatible aggregate buffers while preserving native input projections and SUM partials.

With only the post-conversion revertUnsafePartialAggregates call temporarily bypassed, the four percentile/collect-list cases reproduced the reported EOFException/NullPointerException; both SUM controls passed. After restoring the existing repair, all 141 tests passed: the complete CometAggregateSuite (108), CometExecRuleSuite (31), and CometShuffleFallbackStickinessSuite (2). This run rebuilt the branch's unchanged native source and the Spark 4.1.3 / Scala 2.13.17 / JDK 17 JVM reactor. The selected native library and bundled copy have matching hashes. Spotless, Scalastyle, and git diff --check passed. This follow-up changes only tests; no production-code change or local cross-version run is claimed. Hosted CI for this commit is pending.

The review follow-ups add eight execution regressions: COUNT across a Spark shuffle, mixed COUNT+AVG with empty partial partitions, and both collection functions with native hash partitioning ineligible, each with AQE disabled and enabled. COUNT-only cases retain a native partial for grouped, all-null, and empty input. COUNT+AVG cases pin four Parquet input partitions, check one surviving value, an all-null value, and no surviving rows, and exercise both early tagging and post-conversion repair. They require the partial's unsafe-buffer tag, retain native filters, and check the fully native control before and after AQE materialization. Collection tests likewise compare results with Spark and verify the fully native control. Planner and existing AVG expectations now reflect the separate compatibility directions.

For the preceding update 46ad7eca3, I rebuilt this branch's native code from f1e5868d8 (unchanged native source in this update) and the complete Spark 4.1.3 / Scala 2.13 / JDK 17 JVM reactor. Spotless, Scalastyle, and git diff --check passed. Before the AVG admission guard, both new AQE variants failed with (1, NULL) instead of (1, 1.0) and a native Partial feeding Spark Final. With the guard, all 31 planner tests and 18 selected aggregate execution tests passed (49 total), including COUNT-only native partials, both AVG directions, the mixed empty-partition regressions, fully native controls, and the earlier typed-buffer repair cases.

The preceding 44-test follow-up reused a library containing #5420's native AVG repair and did not establish standalone COUNT+AVG safety for this prerequisite. The new before/after runs use this branch's own native implementation and verify the loaded library hash. These are focused tests, not a full Comet-suite or all-Spark-version run; hosted CI for that revision subsequently passed.

Earlier validation of the original planner repair passed 121 tests across CometAggregateSuite, CometExecRuleSuite, and CometShuffleFallbackStickinessSuite on Spark 4.0.4 / Java 17, with two existing version-gated cancellations. A separate Spark 4.0.2 replay verified the decimal AVG example with AQE on/off, and applying #5420's accumulator fix as well matched Spark in 20 synthetic cases. Those are historical results, not reruns of this follow-up.

[sunchao]

cc @andygrove @comphead discovered this correctness issue when running TPC-DS benchmark internally with spark.comet.shuffle.enabled=false

[andygrove]

I used an LLM (Claude Code) to help with this review, including building both PRs locally and running the sweeps below. I have gone over the results myself and I agree with them, but flagging the tooling up front.

Thanks for digging into this. The gap you identified is real, and the reasoning about canAggregateBeConverted skipping the child-native check matches what the comment in operators.scala says the tagging pass is supposed to cover. Before getting into the mechanics I want to raise something about the framing.

I swept every aggregate function through the Comet-Partial plus Spark-Final configuration, using spark.comet.exec.shuffle.enabled=false so the Final stays in Spark, over 8 rows in 4 Parquet files with a filter that leaves three partials empty. On Spark 4.1 with JDK 17, on current main, only two of the 26 aggregates I tried are wrong across that boundary:

avg(bigint)          comet=[null]                      spark=[1.0]
avg(decimal(20,2))   comet=threw ARITHMETIC_OVERFLOW   spark=[200.000000]

With this PR applied the decimal case is fixed and AVG(bigint) is still wrong, exactly as before. CometAverage.supportsMixedPartialFinal returns true for non-decimal input, so the allowlist waves through what looks like the same defect with the same root cause. That seems like the thing to fix first, and it would be worth a test either way, since it is the same bug this PR is named after.

Everything else the PR newly blocks was already correct in my sweep. COUNT, decimal SUM, FIRST, LAST, and the whole stddev/variance/covariance/correlation family all matched Spark on main, and all of them lose their native Partial with this change. I tried to justify that cost by hunting for a case where the guard earns it: decimal SUM overflowing DECIMAL(38,0), long SUM overflow, TRY_SUM, TRY_AVG, grouped and filtered variants, under ANSI on and off. All of those matched with the native Partial retained, including ANSI throwing and legacy returning null. The comment on CometSum.supportsMixedPartialFinal about overflow detection not surviving the split did not reproduce in this direction, which makes sense to me because that comment describes the Spark-partial to Comet-final direction rather than this one. The same goes for the COUNT exclusion, which the docstring attributes to PropagateEmptyRelationAfterAQE and Spark 4.0 count-bug decorrelation. Both of those only bite when the Final becomes a CometHashAggregateExec.

The other thing I noticed is that #5420 already contains revertUnsafePartialAggregates verbatim, so this PR is a subset of it. When I keep only #5420's native avg.rs and avg_decimal.rs changes and drop the planner guard entirely, both AVG cases are fixed and every aggregate keeps its native Partial.

So on the evidence I have, the native change in #5420 fixes strictly more than this guard does and costs no native aggregation, while this guard fixes less and gives up native partials for a fair number of aggregates. Could you say what this PR adds once #5420 lands? The description asserts the safeguard is still needed after that, but I was not able to construct a case, and none of the four new tests exercises one. If the value is defense in depth against buffer mismatches we have not found yet, rather than a bug that is broken today, that seems like a reasonable position to me, but I think it should be argued on that basis and weighed against the measured cost. It also affects the ordering of the two PRs.

Two smaller things while I was in there. supportsMixedPartialFinal is a single direction-agnostic flag, but most of the exclusions behind it are justified only in the Spark-partial to Comet-final direction. Would it be worth splitting it into two predicates so this pass consults only the one that applies? And where revertChain returns None, the unsafe boundary is left in the plan with no signal, which is the same failure mode the PR exists to prevent. A plain warning there would be too noisy, since I see the None path taken benignly on q10 and q35 where the Partial is already Spark, but gating it on findCometPartialAgg finding a Comet Partial that revertChain failed to reach would make the gap between those two traversals detectable.

One last note, unrelated to this PR. The sweep also turned up that PERCENTILE, APPROX_PERCENTILE, COLLECT_LIST and COLLECT_SET all trip the strict check with spark.comet.exec.shuffle.enabled=false, reporting "Comet did not convert ObjectHashAggregate but recorded no fallback reason". The shuffle-enabled guard on ObjectHashAggregateExec declines without calling withFallbackReason. I will file that separately.

[sunchao]

Thanks @andygrove . You’re right that the AVG examples alone don’t demonstrate what #5421 adds after #5420’s native fixes.

I’ve now reproduced a separate failure. I kept #5420’s native library and removed only the call to revertUnsafePartialAggregates. On Spark 4.1.3 / JDK 17, COLLECT_LIST and COLLECT_SET crash in Spark’s Collect.deserialize when grouping by a stored DECIMAL(19,0) key, with Comet shuffle enabled, mode native, and native hash partitioning enabled. This happens with AQE both off and on. Restoring the guard makes all four cases match Spark.

The decimal key makes the native shuffle fall back, leaving a Comet partial feeding a Spark final. That partial produces an array buffer, while Spark expects serialized binary state. The shuffle-disabled sweep doesn’t exercise this path because it prevents the ObjectHashAggregate partial from converting in the first place. The existing AVG regression passes with or without the guard.

So there is an independent correctness case for the planner repair. That still doesn’t justify every exclusion in the current compatibility flag. I agree we should distinguish the two directions and preserve native partials where their buffers are compatible.

The duplication is because #5420 currently includes #5421 as a prerequisite. I’ll clarify that relationship and update #5421’s tests and description around this separate failure.

[andygrove]

Note on this review: this was generated by an LLM (Claude Code) at my request while I worked through a review backlog. I have not verified the individual findings myself. Please treat everything below as suggestions to evaluate rather than as authoritative review feedback, and push back on anything that is wrong or already handled.

I read the existing thread and do not want to relitigate the "what does this add on top of #5420" question, which is already being worked out between you two. Two things I did not see raised there.

modes.distinct == Seq(Final) misses the distinct-rewrite final stage

Both tagUnsafePartialAggregates and the new revertUnsafePartialAggregates gate on agg.aggregateExpressions.map(_.mode).distinct == Seq(Final). Spark's AggUtils.planAggregateWithOneDistinct produces a fourth stage whose modes are Final for the regular aggregates and Complete for the distinct one, so a query like

SELECT AVG(CAST(v AS DECIMAL(20,2))), COUNT(DISTINCT k) FROM t GROUP BY g

lands on Seq(Final, Complete) and matches neither guard. The Complete half reads raw input so it is harmless on its own, but the Final half still consumes the buffer produced by stage three, which is exactly the boundary this PR exists to protect. Is that shape actually safe for some reason I am missing, or should the predicate be "contains Final" rather than "is exactly Final"? The description mentions distinct aggregations with intermediate merge stages as an affected case, which is what made me look.

The new CometExecRuleSuite test does cover distinct=true, but it asserts zero CometHashAggregateExec in the fallback configurations, and I think that outcome comes from the pre-existing tagging plus the missingCometProducer cascade rather than from revertChain reaching the stage-four aggregate. It would be worth adding an assertion that distinguishes those two mechanisms, otherwise the test passes whether or not the new pass handles the shape.

Re-running transform over a partially converted subtree

The transformUp body calls transform(agg.withNewChildren(Seq(child))). The child returned by revertChain is Spark operators down to the point where native work resumes, so the subtree handed back to transform still contains CometFilterExec, CometScanExec, and friends from the first pass. Is transform idempotent over already-converted operators? I would expect it to be, but the comment above the call only explains why rebuilding is necessary, not that re-entry is safe. A note there would help, and if there is any doubt, a test that puts a second unsafe Final above the first would exercise the nested case where transformUp triggers this twice on overlapping subtrees.

Related to that, is there a bound on how much work this can redo? With transformUp, N stacked unsafe finals means N re-conversions of progressively larger subtrees. Probably irrelevant for real plans, but worth knowing it is bounded.

[sunchao]

Updated in f1e5868d8.

I split the compatibility checks by direction and restored COUNT's safe native-partial-to-Spark-final path. The reverse direction stays disabled for the documented AQE/count-bug reasons. The new execution test covers grouped, all-null, and empty inputs with AQE on/off, and both early and post-conversion fallback. I have not enabled the remaining functions based only on the ordinary-input sweep; their buffer layout, overflow modes, and merge behavior still need targeted validation.

The PR now commits an independent collect_list/collect_set reproduction with Comet shuffle enabled but native hash partitioning disabled, using integer keys so it does not depend on #5420. Each AQE mode checks Spark-equivalent results, zero native aggregate stages, the unsafe-partial tag, and preservation of the native filter; re-enabling eligible native shuffle checks the two-native-stage control. These buffers are native arrays versus Spark's serialized binary state, so fixing AVG's empty sum/count in #5420 does not fix this boundary. The description now leads with that concrete case. #5421 remains the prerequisite, and I am synchronizing this update into #5420.

For the DISTINCT concern, Spark's one-DISTINCT rewrite has intermediate PartialMerge and mixed Partial/PartialMerge stages, but its final stage contains Final-mode expressions for both ordinary and distinct functions. I checked this across the supported Spark sources and in an actual four-stage plan; there was no Final+Complete final stage. The existing regression also exercised the new repair: the early pass added zero unsafe tags, conversion produced one native aggregate, and post-conversion repair removed it. Reapplying conversion preserved the same tree. I therefore kept the Final-only condition and did not broaden it to address an unobserved mode combination.

revertChain(None) intentionally stops at an already-Spark producer, materialized/reused stage, or unrelated boundary. An unconditional warning/assertion would flag those safe cases. Existing Comet nodes pass through conversion, and repair invokes conversion rather than recursively invoking the repair pass, so repeated-tree work is bounded by the finite plan. The current tests do not show an idempotence or planning-time defect. The separate pre-existing missing ObjectHashAggregate fallback explanation is tracked in #5500; it is not introduced by this buffer-compatibility change.

Validation for this update: full Spark 4.1.3 JVM reactor, Spotless, and Scalastyle passed, along with 44 planner/execution tests. Local execution reused the existing OSS native library and is not a fresh native build. CI for the new commit remains to be confirmed.

[sunchao]

Updated in 46ad7eca3.

Fixed the new COUNT+AVG empty-partial report by declining AVG NativePartialToSparkFinal until its emitted state is repaired. COUNT-only interoperability, AVG SparkPartialToNativeFinal, and fully native execution remain enabled.

Both AQE variants reproduced (1, NULL) instead of (1, 1.0) against this branch's own rebuilt native code before the guard. The final native/JVM builds, style checks, and 49 focused planner/execution tests pass. The description now explicitly explains why the earlier run with #5420's already-repaired native library did not establish this standalone boundary. The review thread has the detailed fix and test scope; no full-suite or multi-version rerun is claimed.

[ziting-openai]

Reviewed commit 46ad7ec: no P1 or significant P2 findings. Source review covered the changed execution and compatibility boundaries and the focused regression cases; runtime tests were not run.

Posted by Codex on behalf of ziting-openai using spark-pr-review-memo.

[ziting-openai]

Re-reviewed commit c078ef136add779662b0e74d6bb7035869799284: no P1 or significant P2 findings. The production code is unchanged from the previously approved head. The six added AQE-off/on cases cover the reported array-key fallback for percentile and collect_list, retain a safe native SUM partial, and inspect both initial and materialized plans.

Source review only; runtime tests and cross-version builds were not rerun. Hosted CI was still in progress at review time.

Posted by Codex on behalf of ziting-openai using spark-pr-review-memo.

[ziting-openai]

I'm sorry for the unsolicited AI reviews on this PR. The approvals I posted on August 29 (first review, second review) were the result of a misconfigured AI agent that should have been limited to OpenAI-owned repositories. Please disregard those approvals. I've corrected the configuration to prevent further reviews on non-OpenAI repositories. Neither approval opened an inline review thread, and I don't have permission to dismiss the approvals myself. Sorry for the noise.