fix: preserve Spark AVG buffer and overflow semantics

[sunchao]

Which issue does this PR close?

Closes #5418. Builds on the now-merged #5421. Wide-decimal shuffle routing and its benchmark are separate in #6005.

Rationale for this change

An empty native AVG partial exports (NULL, 0), while Spark expects (0, 0). When Spark merges that buffer with a nonempty partition, the null sum can erase the result. Decimal AVG must also distinguish empty input from overflow and keep overflow sticky through merging and shuffle.

What changes are included in this PR?

• Initialize AVG sums to zero and use the count for empty results. Scalar decimal AVG uses Some(0) for empty input and None for overflow.
• Preserve grouped decimal overflow through state export and shuffle; report ANSI errors when evaluating the affected group.
• Keep global and expanding-window decimal AVG/TRY_AVG in Spark for input precision at least 28, where Spark can retain a wider temporary sum until division and return a valid average. The global fallback applies to the whole aggregate operator, including sibling MIN/MAX/COUNT/SUM expressions.
• Keep grouped ANSI decimal AVG natively enabled, documenting that batch finalization can report overflow for a group Spark would not consume after LIMIT.
• Enable nondecimal AVG partial buffers to cross engines using fix: revert unsafe partial aggregates after final fallback #5421's directional compatibility policy. Decimal AVG buffers remain excluded from mixed execution. Match Spark ObjectHashAggregate decimal AVG buffer semantics #5509 tracks the separate ObjectHash AVG discrepancy.

Rebased onto main dd68a531c. The prerequisite is no longer included in this PR's diff: 11 files, +581/-216, with no benchmark or golden-file changes. Main's newer decimal SUM behavior and compatibility guards are preserved.

How are these changes tested?

Focused Rust and Scala regressions cover empty buffers crossing engines in both directions, sticky scalar overflow, grouped state round trips and deferred errors, JVM and native shuffle, DISTINCT, and the global/window precision boundary. The global AVG fixture asserts explicit expected results, including 0.6 at scale 38 after repartitioning. The grouped overflow fixture asserts the actual partial-buffer shuffle writer in each mode, covering legacy AVG, ANSI TRY_AVG, ANSI errors, valid groups, and all-null groups.

The follow-up at 7e9026df4 changes only documentation and the two existing regression tests. Current-source native compilation, Spark 4.1.3 / JDK 17 full-reactor test compilation, Spotless, Scalastyle, whitespace checks, and both affected tests pass. The loaded JNI library checksum matches the current-source build. Current-head CI is running.

Rebased-head CI at 71a1a1d80 passed native build/Rust tests, the Comet suites on Spark 3.4/3.5/4.0/4.1/4.2, TPC-H/TPC-DS result verification, and Spark 4.1 SQL tests. Those results precede this docs/test follow-up. The PR retains run-all-spark-profiles and run-spark-4.1-tests.

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

The AvgAccumulator::state fix is clearly correct and clearly valuable. Spark's Average.mergeExpressions is a plain sum.left + sum.right with no coalesce, so a partial that reports a null sum for an empty partition poisons the whole final. Getting decimal to distinguish "empty partial" from "overflowed partial" through the state round trip is the fiddly part and the new accumulator tests read well.

I have concerns about how much else is in here.

The decimal precision restriction on native shuffle

Changing CometShuffleExchangeExec so that hash partitioning keys with precision > 18 fall back is by far the widest-reaching change in this PR, and the justification given is that different partition assignments hide an overflow in the AVG(DISTINCT ...) regression. That is a very large hammer for making one test deterministic. Every query that hash-partitions on a DECIMAL(38, s) key now loses native shuffle, whether or not it involves AVG.

The removed TODO points at #3079 and says wider decimals need BigDecimal conversion before hashing. If native hashing genuinely disagrees with Spark for precision > 18, that is a correctness bug in its own right and it deserves its own PR, its own tests, and a note in the migration guide, rather than arriving as a side effect of an AVG fix. If it is not a correctness bug and this is purely about the test, could the test pin the partitioning some other way, for example by constructing the input so the overflow occurs regardless of assignment?

Either way, could you say what this costs? A before and after on a query that groups by a DECIMAL(38, 2) key would settle whether this is a minor or a major regression.

This PR contains #5421 verbatim

That is already being discussed on the other thread so I will not repeat it, except to note that the diff here is harder to review because of it. Once the ordering is settled, rebasing so this PR shows only its own changes would help a lot.

aggregateSupportLevel only covers the ungrouped case

The check is op.groupingExpressions.isEmpty, and the docs say grouped averages stay native. I believe the reasoning is that Spark's grouped path stores the buffer in an UnsafeRow at the declared precision, so it overflows at the same point Comet does, while the ungrouped codegen path keeps a wider local. Could you state that explicitly in the comment? And does the same hold for SortAggregateExec and for grouped ObjectHashAggregateExec after a spill, where the buffer round-trips through a different representation?

The relaxed ANSI condition in AvgDecimalGroupsAccumulator::evaluate

The guard changes from !is_not_null && count > 0 && ansi to !is_not_null && ansi. That is deliberate, and the comment explains why. But it means a group whose is_not_null is false and whose count is zero now throws in ANSI mode where it previously did not. I convinced myself that an all-null group keeps is_not_null true and so is unaffected, but I would rather see that pinned by a test: ANSI mode, grouped decimal AVG, with one group that is entirely null alongside a group that really does overflow. If that test already exists somewhere in the 381 new lines in CometAggregateSuite, pointing at it in the description would be enough.

Non-local return in CometWindowExec

The new block uses return None from inside three levels of nested match inside getSupportLevel. It matches the existing style just below it, so I am not asking you to change it. But four nested matches to reach one condition is hard to follow. Would a small predicate helper, along the lines of isMaxPrecisionDecimalAverage(windowExpr), make both this block and the existing one below it easier to read?

Splitting

Counting them up, this PR contains: the float AVG empty-partial fix, the decimal AVG overflow-marker plumbing, an ungrouped decimal AVG operator fallback, a window AVG fallback, a native shuffle restriction on wide decimals, and all of #5421. The first two are tightly coupled and obviously belong together. The shuffle restriction in particular looks separable and would be much easier to evaluate on its own.

[sunchao]

Updated in 6dca7f116.

I worked through the review points and added the wide-decimal routing benchmark, carried the directional buffer policy from #5421, and corrected the grouped-AVG explanation. The updated description includes current validation and measured costs, with earlier validation clearly identified as historical.

1. Scope beyond AVG. Yes, the hash-key check affects every native shuffle whose decimal key has precision greater than 18. In auto mode, eligible queries use columnar shuffle and can retain native aggregates. In native mode, they use Spark shuffle, which can also require aggregate fallback. Wide decimal payloads and range keys are not excluded by this hash-key check. The description and existing tuning guidance now make that distinction explicit.

2. Hash compatibility versus test determinism. This is a real partition-assignment mismatch, also tracked in Hash partitioning not falling back to Spark for high precision decimal partition keys #3079. Spark hashes wide decimals through BigInteger bytes, while the current native path uses the fixed-width little-endian integer representation. Decimal partial-sum overflow depends on which values meet in a partition, so matching the final numeric type is insufficient. The guard is not just a way to stabilize one test, and I kept it.

3. Splitting the changes. The planner repair is already separated as fix: revert unsafe partial aggregates after final fallback #5421. I kept the AVG state repair, admission restrictions and wide-key partitioning change together because the AVG(DISTINCT ...) regression needs compatible partial buffers and partition assignments. Splitting those into independently mergeable fixes would leave that regression incomplete. The unrelated, pre-existing ObjectHash case now has a separate issue, Match Spark ObjectHashAggregate decimal AVG buffer semantics #5509.

4. Focused shuffle tests. The existing wide decimal hash keys use Spark-compatible shuffle partitioning regression covers precision 18/19/38 across native, auto and jvm, together with payload/range/single-partition controls. It passed again in the current 47-test run. The new benchmark also checks actual input partition counts, results and routes rather than inferring execution from configuration.

5. Migration guide. The migration guide excludes correctness and native-admission changes from its scope. The tuning guide is the existing home for the shuffle-mode behavior, so I retained that explanation there and expanded the PR description rather than adding a migration flag or suggesting users restore incompatible hashing.

6. Measured performance. For 1,048,576 rows and 10,000 DECIMAL(38,2) groups, the four-process ABBA run measured 137→501.5 ms in native mode and 137→188.5 ms in auto mode, averaging the two run means for each variant. All eight variant/mode comparisons matched all Spark results and asserted their routes. This isolates the guard by removing only that check while keeping the same current native library. It is not the entire PR against its base. The description includes both raw run means and the resulting plans. The cost is material, especially when native mode restores Spark aggregation, but the faster incompatible route is not a safe alternative.

7. Relationship to fix: revert unsafe partial aggregates after final fallback #5421. fix: revert unsafe partial aggregates after final fallback #5421 remains the prerequisite and should land first. Its latest fix conservatively rejects native AVG partials before Spark Final because that branch still emits (NULL, 0) for never-updated partials. This PR repairs that state and explicitly enables non-decimal NativePartialToSparkFinal again. Decimal buffers remain excluded. The positive counterpart passed 12 mixed-route combinations plus six native controls with this branch's own rebuilt library. I have not treated fix: revert unsafe partial aggregates after final fallback #5421's temporary AVG restriction as appropriate for the repaired implementation.

8. Why the high-precision guard selects ungrouped AVG. The original wording was too broad. Ordinary Spark hash aggregation stores grouped decimal buffers in UnsafeRow, whose precision checks differ from generated global aggregation's local buffer. That explains the intended guard but does not prove all grouped implementations equivalent. The description now names the representation and qualifies the ObjectHash exception.

9. SortAggregateExec. Spark sort aggregation is already outside this native conversion path. The two ANSI controls in the earlier review probe remained in Spark and matched Spark results. No additional admission guard is needed for that path.

10. ObjectHashAggregateExec and spill. This concern is valid as a separate existing issue. SpecificInternalRow can preserve a wider temporary decimal sum, including after ObjectHash falls back to sort, while native grouped AVG can lose it earlier. I filed Match Spark ObjectHashAggregate decimal AVG buffer semantics #5509 with the 0.6, 0.6, -0.4 counterexample, an all-null group, default/forced-spill checks and exact native base/head evidence showing the discrepancy predates this PR. The description no longer claims that grouping alone establishes compatibility.

11. All-null and overflow groups. The existing Rust grouped_overflow_with_zeroed_null_count_is_not_empty test and Scala grouped decimal AVG preserves overflow across columnar shuffle regression cover the distinction, including the all-null group alongside overflow. I retained them instead of adding duplicate tests. Their earlier validation is identified separately from this follow-up's focused run.

12. Return in windowExprToProto. The return is in a regular match within the method, not a closure, so it is not a Scala non-local return. I did not extract a helper solely for that concern.

The current optimized native build, full Spark 4.1.3 JVM reactor/style checks, and all 47 focused Scala tests pass. The full historical cross-version suite was not rerun for this follow-up. CI on the pushed revision is still separate from these local results.

[sunchao]

Updated in ae98cef70.

I resolved the main-branch conflict with a history-preserving merge of 525c05b52. The only manual resolution combines the decimal and upstream shuffle-test imports, retaining both sets of tests. The AVG implementation and policy, wide-decimal guard, and benchmark source are unchanged.

The merged branch has a fresh default-release native build, a full Spark 4.1.3 JVM reactor/style pass, and 55 passing focused Scala tests, including the 18 AVG interoperability query combinations and upstream's new RSS/JNI/protocol coverage. An additional untimed wide-decimal run matches all 10,000 Spark results and verifies both expected shuffle/aggregate routes. Each run's loaded native library was checked against the merged build.

The description now distinguishes this validation from the earlier 6dca7f116 performance measurements. Those numbers remain pinned to the measured revision; I did not relabel them as merged-head timings or repeat the full cross-version suite.

[sunchao]

Fixed the grouped decimal AVG evaluation-order issue in 728c55e8d. Grouped decimal AVG in ANSI mode now keeps every stage evaluating that AVG in Spark. This prevents native batch emission from throwing for a group that Spark's LIMIT does not consume, while preserving errors when the overflowing group is consumed. The native empty/count-state and overflow repairs are unchanged.

The guard applies to every decimal precision and to the whole affected aggregate operator, including colocated aggregates, even without LIMIT. This intentionally gives up native aggregation for those ANSI queries. Grouped TRY_AVG, legacy decimal AVG, non-decimal AVG, and unrelated native work retain their existing eligibility. The description and generated-documentation source now state that scope; existing benchmark timings use legacy mode and do not measure this new fallback.

The full Spark 4.1.3 JVM reactor passed 139 tests (107 aggregate, 31 planner, and the AVG SQL fixture), with the loaded native library verified against the unchanged native tree. Both new LIMIT regressions fail against the unchanged planner and pass with the fix, with AQE off/on. Consumed-overflow and native-control checks also pass, as do Spotless, Scalastyle, and CI's syntactic Scalafix check. CI for the new commit is pending.

[sunchao]

Fixed both re-review findings in 97f605ede.

The six affected TPC-DS plan expectations were regenerated with the supported harness. They now preserve the intended grouped ANSI decimal AVG fallback on Spark 4.x; the shared and Spark 3.x expectations are unchanged. The operator guide also explains that the guard applies when AVG still has a DECIMAL input after Spark optimization, including other functions in the same aggregate operator and queries without LIMIT.

Validation passed all 129 plan checks on each of Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0 with generation disabled (645 total). Current JVM classes and the actual loaded native library were verified. No implementation or native code changed; no new native build or performance result is claimed. The PR description has been updated, and the new hosted CI run is pending.

[sunchao]

Corrected the generated AVG note in 754a1df66 to say that grouped ANSI fallback applies when the input remains DECIMAL after Spark optimization, matching the operator guide.

This changes only the documentation reason; aggregate admission and native code are unchanged. Full reactor compilation, Spotless, and fresh GenerateDocs output checks passed. The generated Average page now includes the qualification.

[andygrove]

Following up on my earlier question about aggregateSupportLevel only covering the ungrouped case. The answer in operators.scala went further than I expected. Grouped ANSI decimal AVG now falls back at every precision, and because ANSI is the default on Spark 4.x that removes the whole aggregate operator from Comet for a very common shape. The six regenerated Spark 4.x TPC-DS plans contain eighteen aggregate operators reverting to Spark, and q1, q18, q30, q65, q81 and q18a are all rooted at TakeOrderedAndProject, which consumes the entire aggregate output. The lazy-LIMIT argument does not apply to any of them, so the guard is paying a real cost in exactly the cases where Spark would raise the same error. The benchmark in the description is legacy mode, so the cost of this particular fallback is not measured anywhere, and Unsupported leaves users no way to opt back in.

Could the guard be narrowed? Two things I would like your read on. AvgDecimalGroupsAccumulator::merge_batch in avg_decimal.rs around line 592 raises the ANSI overflow error while it is still consuming input, before any group is emitted. If that deferred to evaluate, the remaining divergence from Spark would be batch granularity rather than whole-input granularity. Could the operator check also look at whether a limit can actually truncate this aggregate's output, instead of firing unconditionally? If neither is practical, could the fallback sit behind a config so a user who does not care about the LIMIT case can keep native aggregation?

Separately, in AvgDecimalAccumulator the is_empty field looks like it no longer does anything. Its only reader is the has_overflow expression at avg_decimal.rs line 382, and after this change I cannot construct a state where sum is None while is_not_null is true and count is greater than zero, because every path that clears sum also clears is_not_null. The flag has also stopped meaning what the comment above line 347 says, since merging a batch of empty partials with non-null zero counts now sets it to false. Is it still load bearing, or can it go?

The accumulator work itself reads well. The state round trip for decimal, the zeroed-null-count case and the precision boundary at 27 versus 28 are all pinned by tests in both ANSI modes, which was the part I was most worried about last time.

[sunchao]

Rebased onto current main in 45800b356.

Thanks, Andy. I addressed the accumulator and configuration points:

• Grouped merge overflow now records a sticky invalid state and raises during evaluation. Tests cover precision and i128 overflow, prefix emission, state round trips, and later updates in ANSI, legacy, and TRY modes.
• Removed the redundant scalar is_empty field.
• Changed grouped ANSI decimal AVG from Unsupported to Incompatible. Users can opt in with spark.comet.operator.HashAggregateExec.allowIncompatible=true or spark.comet.operator.ObjectHashAggregateExec.allowIncompatible=true, according to the operator. Decimal buffer boundaries and global/window restrictions still apply.

The default remains conservative. Deferring errors helps prefix emission, but native finalization can still evaluate an unconsumed group within a batch. I agree that queries which consume every group pay for a fallback they do not need. A top-level TakeOrderedAndProjectExec alone does not establish that guarantee: Spark can use take(limit) when its child already provides the requested ordering, and AQE can plan aggregate stages separately from their consumer. I documented that limitation and the opt-in.

Added --grouped-ansi-avg to compare default fallback and native opt-in on the same ANSI workload, with result and plan checks outside timing. Measurements remain pending: the configured local registry cannot resolve the rebased lockfile's DataFusion 55.1.0, preventing a fresh native build and benchmark run.

The previous hosted run passed 1,273 Rust tests and 862 Spark 4.1 execution tests. Its six failing jobs hit the same redundant string-interpolation prefix in the benchmark; I removed it. Those results cover the previous synthetic merge. The new revision needs its own hosted CI results. Full-reactor Spark 4.0.4 test compilation, syntactic Scalafix 0.14.6, Rust formatting, Spotless, Scalastyle, and whitespace checks pass locally.

[andygrove]

Two things before I can read this as an AVG fix.

The copy of #5421 in here predates its review. supportsNativePartialToSparkFinal still defaults to supportsSparkPartialToNativeFinal(fn), which is the inherited default that #5421 replaced with nine explicit per-function opt-ins in 6c14796f5 after @comphead raised it, and revertUnsafePartialAggregates is the version from before the unrepaired-producer diagnostic in 4b5f1704d. Merging as it stands would put both decisions back. Could you rebase onto #5421's head? operators.scala conflicts with main now as well.

The shuffle routing change is a separate fix riding along. CometShuffleExchangeExec.supportedShuffleType goes from case _: DecimalType => true with a TODO pointing at #3079 to d.precision <= 18, which resolves a hashing-compatibility question that has nothing to do with AVG, and your own numbers put it at 137 ms to 501.5 ms in native mode. Could it go in its own PR? That would let the routing change be benchmarked, reviewed and reverted independently of the accumulator work.

On the grouped ANSI decimal AVG guard, moving to Incompatible helps, but I do not think the trade has changed. This is the same divergence #5533 is about, which is that Comet evaluates a batch eagerly where Spark evaluates lazily and stops at a LIMIT, and over there the open question is whether to document the class or special-case the planner. Here we are picking the expensive side of it for a shape that is the default on Spark 4.x. All six regenerated TPC-DS plans are rooted at TakeOrderedAndProject, which consumes every group, so none of them benefits from the fallback and all of them pay for it. Would documenting this under "Known result-value divergences" in compatibility/index.md and keeping Compatible() be the better call? If the fallback stays, the opt-in deserves a line saying it is operator-wide, since spark.comet.operator.HashAggregateExec.allowIncompatible=true also accepts anything else that operator marks incompatible later.

The new expanding-frame guard in CometWindowExec only inspects Average. Spark's AggregateProcessor writes the window buffer into a SpecificInternalRow, which does not check decimal overflow the way UnsafeRowWriter does, so an expanding-frame decimal SUM has the same sticky-overflow divergence you are guarding AVG against. Should that be in the same match, or tracked alongside the decimal SUM point I left on #5421?

Last, this regenerates golden plans for the -spark4_0 and -spark3_5 tiers and the base tier, but only Spark 4.1 ran. Could you add run-all-spark-profiles?

The accumulator work itself reads well now. AvgAccumulator emitting (0.0, count) matches Literal.default(sumDataType) exactly, the scalar is_empty field is gone, and deferring the grouped merge overflow to evaluate is the right shape.

[andygrove]

On the window question above, I filed the decimal SUM divergence as #6002, which covers the ever-expanding frame alongside the ungrouped aggregate and ObjectHashAggregate. Happy for the CometWindowExec guard to stay AVG-only here and for SUM to be handled there, as long as this PR is not the last place it gets written down.

[sunchao]

Thanks, Andy. Rebased onto main 3b942e593 and pushed the follow-up in c59caf673.

• fix: revert unsafe partial aggregates after final fallback #5421: Synchronized the prerequisite with its reviewed head 4b5f1704d, including the independent default-false policies, nine explicit producer opt-ins, restoration explanations, and unrepaired-producer diagnostic. Repaired nondecimal AVG now has its own explicit producer opt-in. I also retained the newer prerequisite regressions and changed their deliberately unsafe AVG fixtures to decimal AVG, since ordinary AVG is now safe.
• Shuffle scope: Moved the wide-decimal restriction, shuffle-specific expectations, tuning guidance, and benchmark into fix: use Spark-compatible shuffle for wide decimal hash keys #6005. Its standalone COUNT benchmark is independent of AVG, and the historical AVG-workload measurements are identified as historical there. That draft includes fix: revert unsafe partial aggregates after final fallback #5421 because a new shuffle fallback must also restore incompatible collection buffers; a focused collect_list/collect_set regression covers the boundary. The AVG tests here select JVM shuffle explicitly where Spark's partition assignments or columnar buffer conversion are part of the test.
• Grouped ANSI policy: I retained the configurable Incompatible fallback as the conservative option. It still costs native coverage for queries that consume every group. The known-divergence section now describes the behavior, and the operator guide explicitly says that each opt-in accepts every incompatibility for that operator, including future additions. The remaining divergence is eager finalization of a group that Spark never consumes; deferring merge errors does not remove it. Grouped ANSI timings remain unmeasured.
• Window SUM: Kept the guard AVG-only as agreed and documented SUM(decimal) returns NULL or throws under ANSI when an intermediate sum overflows the buffer precision but the final result fits #6002 alongside the known decimal SUM divergence.
• Profiles: Added run-all-spark-profiles. The 11 shuffle-related golden changes moved to fix: use Spark-compatible shuffle for wide decimal hash keys #6005; the six grouped-ANSI plan expectations remain here. Both PR descriptions now reflect their final scope and validation.

Local Spark 4.1.3 validation passed full-reactor test compilation, Spotless, Scalastyle, Rust formatting, all 37 planner tests, and all 129 TPC-DS plan checks with generation disabled. The planner checks use a verified main CI library for initialization and do not execute the changed native AVG implementation. The configured local registry still lacks locked DataFusion 55.1.0, so fresh native AVG execution and cross-profile results are pending default-profile CI and additional-profile CI. Separately, #6005 passed seven focused execution tests, nine affected plan checks, and benchmark result validation using a verified CI library whose native source tree matches that branch exactly.

[sunchao]

Andy, I simplified the AVG patch further in 42a6bad39. This supersedes the grouped ANSI policy in my previous reply: grouped AVG now remains Compatible(), with the eager batch/LIMIT overflow difference documented under known divergences, as you suggested. That removes the blanket fallback, operator opt-in tests, benchmark, and all six associated golden snapshots.

The scalar decimal accumulator also uses one update loop, and the test coverage is consolidated around distinct failures: empty partials crossing engines in both directions, sticky overflow, grouped state round trips with a zeroed null-count payload, and the 27/28 precision boundary. The grouped columnar-shuffle test includes overflowing, valid, and all-null groups and checks legacy, TRY, and ANSI behavior.

The reviewed #5421 prerequisite remains intact, wide-decimal shuffle remains separate in #6005, and the window guard stays AVG-only with decimal SUM documented under #6002. Beyond the prerequisite on the same main baseline, the patch is now +571/-216, down from +1475/-224. No benchmark or golden-file changes remain here.

Full-reactor Spark 4.1.3 test compilation, formatting/style checks, all 37 planner tests, and all 129 TPC-DS plan checks pass. The planning checks use a verified main native library only for initialization. Current-head CI has run-all-spark-profiles; native execution results are pending because the configured local registry cannot resolve locked DataFusion 55.1.0.

[andygrove]

The compatibility argument reads well now, and I checked it against Average.scala rather than taking it on trust. UnsafeRow.setDecimal nulling the field when changePrecision fails is what makes the grouped case genuinely equivalent to Comet, and doProduceWithoutKeys keeping the buffer in Decimal mutable state is what makes the ungrouped case genuinely different. Deferring the grouped merge error to evaluate and dropping count > 0 from the ANSI check are both right. I had not appreciated that the second is a fix rather than a relaxation, since the old guard would have swallowed the error after a columnar shuffle zeroed the count payload.

Three things I would still like before this goes in.

The docs do not say the ungrouped fallback takes the whole operator with it. The new Aggregation section in docs/source/user-guide/latest/compatibility/operators.md reads as if only the average falls back, but your own change to final decimal avg shows SELECT MIN(a), MAX(a), COUNT(a), SUM(a), AVG(a) FROM $table dropping from 2 native aggregates to 0. A single wide-decimal AVG costs a user native execution for every aggregate sharing that operator. Could the section say so? That seems like the thing a reader of that page would most want to know.

The fixture that justifies the ungrouped fallback never asserts the value that justifies it. Every case in high-precision global decimal AVG and TRY_AVG fall back to Spark uses checkSparkAnswerAndNumOfAggregates(..., 0), and once the aggregate has fallen back Comet is Spark, so the answer comparison cannot fail. That pins the fallback without pinning the reason for it. Would you add a checkAnswer on the AVG(v38) cases asserting 0.6 at the result scale? I worked through DecimalDivideWithOverflowCheck and I believe Spark does return 0.6 where native returns NULL, but that is the entire justification for the guard and nothing in the suite records it. The repartition(2, col("ord")) case is the better one to pin, since it does not depend on the exchange being elided by coalesce(1).

The grouped overflow round trip only runs over columnar shuffle. grouped decimal AVG preserves overflow across columnar shuffle pins CometColumnarShuffle, which is the harder case and the right one to have. But spark.comet.shuffle.mode defaults to auto and prefers native, and the grouping key there is an int, so the default path for that shape is the one not covered. Could the test loop over jvm and native? Removing the partial_counts.null_count() handling leaves the sum's null bit as the only thing carrying overflow across the exchange, so it would be good to have both writers on record.

Ordering note rather than a finding. This still carries #5421 in full, now at 4b5f1704d, so the merge order still matters and #5421 needs to land first.

[sunchao]

Thanks, Andy. Addressed the three points in 7e9026df4:

• The aggregation guide now explicitly says the wide-decimal AVG/TRY_AVG fallback moves the whole operator to Spark, including sibling MIN/MAX/COUNT/SUM expressions, in both partial and final stages.
• The global fixture now uses explicit expected results while retaining the native-aggregate count assertions. It pins 0.6 at scale 38 for AVG/TRY_AVG, including the repartition(2, col("ord")) case. The existing collect_list case explicitly expects NULL: Spark's object aggregate materializes an overflowing partial buffer before division.
• The grouped overflow regression now loops over jvm and native, preserving the legacy AVG, ANSI TRY_AVG, ANSI error, valid-group, and all-null-group checks. It asserts the shuffle writer directly above the partial aggregate, so the input repartition cannot satisfy the assertion accidentally.

#5421 has merged, and the rebase removed it from this PR's diff. This follow-up changes two files. Both affected tests pass on Spark 4.1.3 / JDK 17 with the current-source native library; full-reactor compilation, Spotless, Scalastyle, and whitespace checks pass. Current-head CI is running with all Spark profiles and Spark 4.1 SQL tests enabled.

[sunchao]

Summary

• Prior state and problem: Empty native AVG partials could poison Spark’s final sum. Decimal overflow could be lost during updates, merging or shuffle.
• Design approach: Initialize sums to zero, preserve sticky decimal overflow, defer grouped ANSI errors until evaluation, and keep incompatible global/window decimal averages in Spark.
• Correctness / compatibility analysis: Compared Spark’s AVG expressions, decimal arithmetic, UnsafeRow buffers, generated aggregation and window processing across Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. The changes correctly distinguish empty input from overflow. Reviewed mixed execution, DISTINCT, precision boundaries and both shuffle writers.
• Key design decisions: Nondecimal AVG explicitly permits native partials before Spark final aggregation. Decimal buffers remain excluded from mixed execution. Grouped ANSI AVG retains the documented batch/LIMIT error-timing difference.
• Implementation sketch: Rust accumulator changes simplify scalar state by removing redundant flags. Scala admission checks reuse the existing compatibility policy. Regression tests cover state round trips, explicit decimal results and actual shuffle-writer selection.
• Behavioral changes worth calling out: Global and expanding-window decimal AVG/TRY_AVG fall back at input precision 28 and above. The aggregate fallback also moves sibling aggregates to Spark, reducing native coverage as documented. Performance was assessed structurally, without benchmark measurements.
• Suggested improvements: None meeting the requested P1/P2 bar. Previous review requests are addressed. No introduced P1/P2 issues found within this review, and no unresolved existing P1/P2 blockers remain.

Reviewed all 11 changed files across both commits at 7e9026df4968111be023dbe72610bc4e3df9ba0e, relative to dd68a531c413c287c1f5fbe12ade4edb488900e5. The PR remains open and non-draft. Read repository guidance and all existing reviews/discussion. Routed skills: review-comet-pr, review-comet-expression-pr, and review-comet-shuffle-pr.

Exact-head CI: run 35946785166 passed. GitHub reports 48 successful and 14 skipped checks, with none failed or pending. Logs confirm 1,627 Rust tests and 992 Spark 4.1 execution tests passed, including the new AVG regressions. Comet suites across Spark 3.4–4.2, Spark 4.1 SQL suites, and TPC-H/TPC-DS verification passed.

Local validation: All four focused AVG Rust tests passed using cargo test --locked --offline -p datafusion-comet-spark-expr agg_funcs::avg --lib. Whitespace checks passed and the working tree remains clean. JVM integration suites were not rerun locally. CI skipped macOS, Iceberg and upstream Spark SQL suites outside 4.1.

[andygrove]

Everything from my last review is in, and I checked each piece by removing what it pins. With the ungrouped guard turned off, the new checkAnswer calls fail for AVG(v38), TRY_AVG(v38), the COUNT(DISTINCT ord) case and the repartition(2, col("ord")) case, so the fixture now records why the fallback exists. With main merged in, CometAggregateSuite, CometWindowExecSuite, CometExecRuleSuite and CometCelebornShufflePlanningSuite all pass locally on 4.1 and 3.5. CI at this head ran the [exec] bucket, plan-stability suites included, on every profile from 3.4 to 4.2.

There are two things left, both inline. One is a single extra fixture row so the native shuffle leg actually depends on the null bit that carries overflow across the exchange. The other is new since my last review. #6041 now declines grouped maximum-precision SUM under ObjectHashAggregateExec, and doing the same for AVG would close #5509 in two lines. I'm happy to approve once those are in.

[andygrove]

Could group 1 get a third row, (1, 0, "-0.6")? As written, the native leg passes even with the partial_sums.is_null(idx) check removed from AvgDecimalGroupsAccumulator::merge_batch. Native shuffle keeps the payload of the null sum, which here is the overflowed 1.2, so the final merge overflows again without needing the null bit. Only the JVM leg's ANSI case catches that mutation. With the extra row the payload comes back to 0.6, and the native leg fails under the same mutation. Unmutated, the test still passes on both legs with the same expected results, because Spark's UnsafeRow buffer has already latched group 1 to null by then.

[andygrove]

Grouped max-precision AVG under ObjectHashAggregateExec still diverges, which is #5509. With main merged in, SELECT g, AVG(v), sort_array(collect_list(ord)) FROM t GROUP BY g over DECIMAL(38,38) values 0.6, 0.6, -0.4 returns NULL natively, or ARITHMETIC_OVERFLOW under ANSI, where Spark returns 0.26666666666666666666666666666666666667. The SUM version of the same query falls back and matches, because #6041 declines it in CometObjectHashAggregateExec.getSupportLevel, and its comment there already says decimal AVG has the same gap. Extending this check to ObjectHashAggregateExec with grouping keys fixed both ANSI modes for me, and all of CometAggregateSuite still passed. A hasMaxPrecisionDecimalAvg next to hasMaxPrecisionDecimalSum might be the tidiest way to write it. Could this PR decline that case too and close #5509? If you'd rather keep it separate, could #5509 go in the known result-value divergences list in compatibility/index.md, since it returns NULL without an error in legacy mode?

[sunchao]

Summary

• Prior state and problem: Empty native AVG partials exported (NULL, 0), which could poison Spark’s final sum. Decimal overflow could lose its invalid-state marker.
• Design approach: Initialize empty sums to zero, preserve sticky decimal overflow, defer grouped ANSI errors until evaluation, and restrict incompatible global/window aggregation.
• Correctness / compatibility analysis: Compared Spark’s AVG, decimal arithmetic, generated aggregation and row-buffer semantics across Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. No introduced P1/P2 issues found within this review.
• Key design decisions: Nondecimal AVG buffers can cross engines in both directions. Decimal buffers remain excluded. Grouped ANSI AVG retains the documented batch/LIMIT error-timing difference.
• Implementation sketch: Scalar state becomes simpler by removing redundant flags. Grouped state preserves overflow through shuffle. Existing planner compatibility checks enforce fallback without a new configuration or abstraction.
• Behavioral changes worth calling out: Global and expanding-window decimal AVG/TRY_AVG fall back at input precision 28 and above. Global fallback also moves sibling aggregates to Spark. This reduces native coverage as documented. Performance was assessed structurally, without benchmark measurements.
• Suggested improvements: No additional P1/P2 findings. The existing ObjectHash AVG thread remains unresolved. For grouped DECIMAL(38,38) values 0.6, 0.6, -0.4 alongside collect_list, Spark returns 0.26666666666666666666666666666666666667, while native execution returns NULL or an ANSI overflow. Spark’s object buffer and native overflow latching substantiate this concern. It predates this PR, is tracked by #5509, and is not duplicated as a new finding.

Reviewed all 11 changed files across both commits at 7e9026df4968111be023dbe72610bc4e3df9ba0e, relative to dd68a531c413c287c1f5fbe12ade4edb488900e5. Confirmed non-draft status and read existing reviews, comments and threads. Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-shuffle-pr.

Exact-head CI: run 35946785166 passed. GitHub reports 48 successful and 14 skipped checks, with none failed or pending. Logs confirm 1,627 Rust tests and 992 Spark 4.1 execution tests passed, including the AVG regressions. All supported Comet Spark profiles, Spark 4.1 SQL suites and TPC-H/TPC-DS verification passed.

Local validation: All four tests passed with cargo test --locked --offline -p datafusion-comet-spark-expr agg_funcs::avg --lib. Whitespace checks passed and the working tree is clean. JVM integration tests were not rerun locally. CI skipped macOS, Iceberg and upstream Spark SQL suites outside 4.1.

[andygrove]

This is a light fully automated review since there are so many PRs open.

avg() at native/spark-expr/src/agg_funcs/avg_decimal.rs:606 multiplies the sum by 10^(target_scale - sum_scale), which is 10^4 for most scales, in i128 before dividing, and both callers (lines 328 and 551) turn its None into NULL. Spark's DecimalDivideWithOverflowCheck divides with BigDecimal and only checks the final result. That leaves two gaps in grouped decimal AVG, which this PR keeps native at every precision. With v DECIMAL(38, 18) and one group holding two 9000000000000000 rows, the unscaled sum 1.8e34 times 10^4 passes i128::MAX, so native returns NULL in every mode, while Spark's sum fits its DECIMAL(38, 18) buffer and it returns 9000000000000000. For that type any group whose sum passes about 1.7e16 hits this, long before Spark's buffer overflows at 1e20. And with a single 100000000000000000 row the average does not fit DECIMAL(38, 22), so Spark raises NUMERIC_VALUE_OUT_OF_RANGE under ANSI where native returns NULL. Both predate this PR, but they sit in the evaluate paths it rewrites. Could avg() divide before scaling (or use i256), and raise SparkError::NumericValueOutOfRange in ANSI mode when the result really does not fit, with a grouped test covering both rows? If you'd rather keep this separate, could it get an issue and an entry next to the new grouped AVG line in the known result-value divergences list?

[sunchao]

Summary

• Prior state and problem: Empty native AVG partials exported (NULL, 0), which could poison Spark’s final sum. Decimal overflow state could be lost during updates, merging or shuffle.
• Design approach: Initialize empty sums to zero, preserve sticky decimal overflow, defer grouped ANSI errors until evaluation, and restrict incompatible global/window aggregation.
• Correctness / compatibility analysis: Compared Spark’s AVG, decimal division, aggregation buffers and window processing across Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. Reviewed empty/all-null input, mixed execution, DISTINCT, precision boundaries and both shuffle writers. No introduced P1/P2 issues found within this review.
• Key design decisions: Nondecimal AVG buffers can cross engines in both directions. Decimal buffers remain excluded. Grouped ANSI AVG retains the documented batch/LIMIT error-timing difference.
• Implementation sketch: Scalar decimal state removes redundant flags and uses Some(0) for empty input and None for overflow. Grouped state retains its overflow bitmap. Planner guards reuse existing compatibility checks without adding configuration or another abstraction.
• Behavioral changes worth calling out: Global and expanding-window decimal AVG/TRY_AVG fall back at input precision 28 and above. Global fallback also moves sibling aggregates to Spark. This reduces native coverage as documented. Performance was assessed structurally, without benchmark measurements.
• Suggested improvements: No new P1/P2 findings. Existing correctness concerns remain unresolved: the ObjectHash AVG thread, tracked by #5509, and the decimal division/ANSI comment. Disposable native probes reproduced both concerns using the base and head accumulator implementations. For grouped DECIMAL(38,38) values 0.6, 0.6, -0.4, native returns NULL or an ANSI overflow where Spark’s object buffer permits a valid average. For DECIMAL(38,18), two 9000000000000000 values incorrectly produce NULL during native scaling. A single 100000000000000000 also returns NULL instead of Spark’s ANSI range error. These are existing blockers, not duplicated new findings.

Reviewed all 11 changed files across both commits at 7e9026df4968111be023dbe72610bc4e3df9ba0e, relative to dd68a531c413c287c1f5fbe12ade4edb488900e5. Confirmed non-draft status and read existing reviews, issue comments, inline comments and threads. Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-shuffle-pr.

Exact-head CI: run 35946785166 passed, with 48 successful and 14 skipped checks and none failed or pending. Logs confirm 1,627 Rust tests and 992 Spark 4.1 execution tests passed, including the AVG regressions. All supported Comet Spark profiles, Spark 4.1 SQL suites and TPC-H/TPC-DS verification passed.

Local validation: All four focused AVG Rust tests passed. Base/head probes confirmed the existing concerns. Disposable changes were removed, restored-head tests passed again, and whitespace checks and working-tree checks were clean. JVM integration suites were not rerun locally. CI skipped macOS, Iceberg and upstream Spark SQL suites outside 4.1.