feat: expose native aggregate spill and memory metrics - #5423
Merged
sunchao merged 2 commits intoAug 25, 2026
Conversation
comphead
reviewed
Aug 24, 2026
| else children.flatMap(_.leafNodes) | ||
| } | ||
|
|
||
| private[comet] def withoutAggregateMetrics(plan: SparkPlan): CometMetricNode = |
Contributor
There was a problem hiding this comment.
it would be useful to comment the reason aggregation metrics are excluded
Member
Author
There was a problem hiding this comment.
Good point. Range partitioning executes the native child once to sample partition boundaries and then again for the real shuffle. If aggregate metrics are updated in both passes, spill count, spilled bytes/rows, and peak native memory are counted twice. The sampling pass therefore suppresses only aggregate metrics; scan/input and other operator metrics remain intact, and the real shuffle execution reports aggregate metrics normally. This is why the filtering helper needs an explanatory comment.
sunchao
force-pushed
the
dev/chao/codex/comet-aggregate-spill-metrics
branch
2 times, most recently
from
August 24, 2026 18:58
e01bf08 to
836a01d
Compare
Member
Author
|
Merged, thanks @comphead for the review! |
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.
Why are the changes needed?
Grouped aggregation can consume substantial memory: every distinct grouping key adds hash-table entries and aggregate state, and memory pressure can force intermediate results to spill to disk. Those spills often explain why a query becomes unexpectedly slow even though it still completes successfully.
For example, consider a query that aggregates a large table by a high-cardinality key:
Comet can execute both the partial and final aggregates natively. DataFusion already records how much memory each grouped aggregate uses and whether it spills, but those measurements are not currently exposed on the corresponding Spark SQL operator. As a result, the Spark UI can show that an aggregate produced rows and spent time executing, while leaving the most important diagnostic question unanswered: was that time spent aggregating in memory, or repeatedly spilling intermediate state to disk?
For illustration, the operator currently looks roughly like this:
After this change, the same operator can also report:
The values above are illustrative. Together, these metrics make native aggregation behavior visible through the same Spark SQL interfaces users already rely on when investigating expensive stages, skewed grouping keys, shuffle partitioning, and memory pressure. The spilled bytes represent temporary aggregate-state spills; they are not a substitute for shuffle-write or scan-input metrics.
What changes were proposed in this PR?
This PR closes the observability gap by carrying DataFusion's existing grouped-aggregate measurements through Comet's existing native-to-Spark metric pipeline and exposing them on the corresponding
CometHashAggregateExec. Spill counts and spilled-row counts are presented as count metrics, while spilled bytes and peak aggregate memory are presented as size metrics. Both partial and final grouped aggregates retain their own operator-level measurements. No new spill accounting, memory manager, or parallel metric collection framework is introduced.The main implementation concern is preserving the meaning and ownership of those measurements as native execution is mapped back onto a Spark plan. A single Spark aggregate may be represented natively by an
AggregateExecwrapped in a projection that shares the same Spark plan ID. In that case, the aggregate's existing metrics need to be associated with the enclosing Spark operator rather than lost under an extra native-plan level; existing child relationships and output-row accounting must remain intact.The same care is needed for range-partitioned native shuffles. A plan such as:
may execute its input once to sample partition boundaries and again to produce the actual shuffle output. Reporting aggregate metrics from both passes would make a single logical aggregate appear to have used more memory or spilled more data than its material execution actually did. The sampling pass therefore suppresses only aggregate-specific metrics; scan input accounting and unrelated operator metrics continue to work normally. The real execution still reports the aggregate measurements.
Finally, the new metrics are exposed only where DataFusion actually provides the corresponding grouped-aggregation instrumentation. A global aggregate such as
SELECT SUM(amount) FROM ordersdoes not acquire fabricated spill or memory counters. This preserves the distinction between an observed value of zero, meaning that a supported metric was measured and no spill occurred, and an absent metric, meaning that the operator does not provide that measurement.How was this PR tested?
The Spark-side regression coverage executes a real native grouped aggregation and verifies that both partial and final aggregate operators expose correctly typed spill and memory metrics, including a nonzero measured native-memory value. It also runs a global aggregation and verifies that unsupported aggregate metrics remain absent rather than being fabricated as zeroes.
A separate native-shuffle regression exercises range partitioning through the actual sampling path. It verifies that sampling does not update grouped-aggregate metrics, that scan bytes and records are still reported during sampling, and that the subsequent real execution does publish the aggregate's memory usage. The latest focused run passed: 1 test, 0 failures.
Native unit coverage verifies absent-versus-zero semantics, propagation of nonzero spill and memory values, and correct metric ownership when a projection and aggregate share a Spark plan ID. Spotless formatting checks also passed.