fix: preserve aggregate result identity during exchange reuse - #5470
Draft
sunchao wants to merge 1 commit into
Draft
fix: preserve aggregate result identity during exchange reuse#5470sunchao wants to merge 1 commit into
sunchao wants to merge 1 commit into
Conversation
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?
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: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(*) + 1under 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(*) + 1andCOUNT(*) - 1before 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
CometAggregateSuitecover 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
b448894bin 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.