Skip to content

fix: preserve aggregate result identity during exchange reuse - #5470

Draft
sunchao wants to merge 1 commit into
apache:mainfrom
sunchao:codex/aggregate-result-identity
Draft

fix: preserve aggregate result identity during exchange reuse#5470
sunchao wants to merge 1 commit into
apache:mainfrom
sunchao:codex/aggregate-result-identity

Conversation

@sunchao

@sunchao sunchao commented Aug 26, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

No issue is automatically closed. This is a standalone correctness fix for aggregate exchange reuse.

Rationale for this change

Spark can reuse a shuffle when another branch of the query has already computed the same data. That decision depends on the identity of the plan below the exchange. Comet currently compares an aggregate's grouping keys and aggregate functions without comparing its final result expressions. Two aggregates can therefore look equivalent even when they return different values.

For example, suppose a Parquet table t(k, v) contains (0, 2) and (0, 3). This excerpt shows the query shape covered by the regression, including a shuffle above each final aggregate:

import org.apache.spark.sql.functions.col

def branch(result: String) =
  spark.sql(s"SELECT $result AS c, k FROM t GROUP BY k")
    .repartition(2, col("c"), col("k"))

val plus = branch("COUNT(*) + 1")
val minus = branch("COUNT(*) - 1")
plus.unionAll(minus).collect()

The expected (c, k) rows are (3, 0) and (1, 0). Both branches count the same rows, but their final arithmetic differs. If Spark incorrectly reuses the first shuffle for the second, the query can finish successfully with (3, 0) twice.

The fix must also preserve legitimate reuse. Repeating COUNT(*) + 1 under a different column alias still computes the same values, even though Spark assigns fresh expression IDs. Treating those plans as different would avoid incorrect sharing at the cost of unnecessary computation and shuffling.

What changes are included in this PR?

Aggregate identity now includes the final result expressions, so equality and hashing distinguish computations such as COUNT(*) + 1 and COUNT(*) - 1 before Spark decides whether to share an exchange. The native aggregate already evaluates these projections; this change makes the plan's identity describe the result it actually produces.

To retain valid reuse, Comet also preserves Spark's original aggregate result attributes and uses them when normalizing plan identity. This is especially important for distinct aggregates, whose rewritten expressions do not always retain the original result IDs. Equivalent computations can still compare equal across fresh IDs and aliases, while different projections remain separate. This applies to both hash and object hash aggregate conversion. The scope is the aggregate implementation and its existing test suite; native execution kernels and dependencies are unchanged.

How are these changes tested?

Four regressions in CometAggregateSuite cover count, distinct count, distinct count combined with sum, and collect-set size. Each executes the two different projections and checks their results, then checks that equivalent projections with different aliases still reuse an exchange. These tests use Parquet input, native Comet shuffle, exchange reuse enabled, and AQE disabled; they do not establish coverage for execution with AQE enabled.

All four regressions passed at head b448894b in the macOS Spark 4.0 execution job, whose complete shard reported 696 successful tests and zero failures. This provides JVM runtime validation beyond the earlier local Maven attempt, which stopped during dependency resolution. These results validate the PR head; a separate runtime comparison against the base revision is not claimed.

Full CI is still in progress at the time of this update. The macOS scans shard is being rerun after a native HDFS thread-cleanup crash with the same pattern as an earlier failure on main.

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.

1 participant