fix: preserve Spark AVG buffer and overflow semantics - #5420
Open
sunchao wants to merge 6 commits into
Open
Conversation
sunchao
force-pushed
the
dev/chao/codex/oss-avg-empty-partial-state
branch
from
August 26, 2026 06:16
0b26695 to
e5aa8d7
Compare
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.
Which issue does this PR close?
Closes #5418.
Rationale for this change
An
AVGquery can currently returnNULL, or report decimal overflow, where Spark would return a valid average. The problem is in partial aggregation: each partition produces a(sum, count)buffer, and a later aggregate merges those buffers before dividing.An empty partition can erase a valid average
Suppose a table has four input partitions, and this filter leaves one matching row with integer
quantity = 2:The answer should be
2.0. Before this fix, the partial buffers differ as follows:(sum, count)(sum, count)(2.0, 1)(2.0, 1)(0.0, 0)(NULL, 0)When native partial aggregation feeds a Spark final aggregate, Spark adds the partial sums without replacing nulls with zero. The empty partitions turn the final sum into
NULL, losing the valid average from the matching row.For decimal AVG, a null sum also represents an overflowed partial. Replacing every null sum with zero would hide overflow, so empty input and overflow must remain distinct.
A valid decimal average can require a wider temporary sum
Consider two values of type
DECIMAL(38, 38):The average fits, but
DECIMAL(38, 38)has no room for the sum's integer digit. Spark's generated global aggregation and expanding window paths can retain the wider temporary sum until division. Comet checks precision earlier and can discard it before calculating the valid average.Partition assignment matters too. In the
AVG(DISTINCT ...)regression, Spark's shuffle with two partitions sends the distinctDECIMAL(38, 38)values0.6,0.2, and0.3to the same partition. Materializing their partial sum,1.1, overflows. Different native hashing can spread the values across partitions and hide that overflow. Comet must preserve Spark's partition assignments to match this behavior.What changes are included in this PR?
The PR makes empty partials contribute Spark's zero sum and zero count, while preserving decimal overflow markers through state export, merging, and shuffle. An empty partition can then contribute no rows without changing a valid average. An entirely empty or all-null aggregate still returns
NULL, and an overflowed partial still leads to an ANSI error or a legacy/TRY null result.Changing the buffer cannot recover a wider sum that native arithmetic has already discarded. The PR therefore keeps global
AVG/TRY_AVGand expanding decimal AVG windows in Spark when input precision isp >= 28. At that boundary, Spark's extra ten digits of sum precision have already reached the maximum of 38. This deliberately trades native execution in those cases for Spark's established behavior.Fallback must be safe across the aggregate pipeline. A Spark final aggregate cannot consume an incompatible native partial buffer. The included prerequisite from #5421 restores the feeding aggregate/exchange chain to Spark when necessary, while retaining eligible native work below it.
Partitioning follows the same compatibility principle. Decimal hash keys with precision greater than 18 use Spark's partition assignments: Comet's columnar shuffle in
automode when eligible, or Spark shuffle innativemode. Wide decimal payload columns and range keys remain eligible for native shuffle.How are these changes tested?
The regressions compare Spark and Comet results and assert the intended native, mixed, or fallback plan. They cover the examples above, empty/all-null input, overflow propagation, precision boundaries, ANSI/TRY behavior, windows, and AQE. Controls verify that supported cases still execute natively.
With JDK 17 and a fresh native build, the expression and shuffle suites passed 652 and 35 tests respectively. All five Spark profiles compiled and passed 738 execution and TPC-DS plan checks:
Additional Spark 4.1.3 validation passed the full
CometAggregateSuite(98 tests) and 194 standalone planner/shuffle comparisons, including exchange reuse, typed aggregate buffers, and exact decimal partition assignments.All 645 plan checks passed with the committed snapshots.
cargo fmt, Spotless, Scalastyle, and whitespace checks also passed.