Skip to content

fix: preserve Spark AVG buffer and overflow semantics - #5420

Open
sunchao wants to merge 6 commits into
apache:mainfrom
sunchao:dev/chao/codex/oss-avg-empty-partial-state
Open

fix: preserve Spark AVG buffer and overflow semantics#5420
sunchao wants to merge 6 commits into
apache:mainfrom
sunchao:dev/chao/codex/oss-avg-empty-partial-state

Conversation

@sunchao

@sunchao sunchao commented Aug 22, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #5418.

Rationale for this change

An AVG query can currently return NULL, 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:

SELECT AVG(quantity) FROM input WHERE id = 1;

The answer should be 2.0. Before this fix, the partial buffers differ as follows:

Partition Spark's (sum, count) Comet's (sum, count)
The partition containing the matching row (2.0, 1) (2.0, 1)
Each of the three empty partitions (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):

Input values:   0.6, 0.6
Temporary sum:  1.2
Average:        0.6

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 distinct DECIMAL(38, 38) values 0.6, 0.2, and 0.3 to 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_AVG and expanding decimal AVG windows in Spark when input precision is p >= 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 auto mode when eligible, or Spark shuffle in native mode. 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:

Spark Scala Execution checks Plan checks
3.4.3 2.12 12 129
3.5.9 2.13 7 129
4.0.4 2.13 8 129
4.1.3 2.13 53 129
4.2.0 2.13 13 129

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.

@sunchao sunchao changed the title fix: use zero sums for empty AVG partial buffers fix: preserve Spark AVG buffer and overflow semantics Aug 26, 2026
@sunchao
sunchao force-pushed the dev/chao/codex/oss-avg-empty-partial-state branch from 0b26695 to e5aa8d7 Compare August 26, 2026 06:16
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.

Incorrect AVG result from empty native partial buffers

1 participant