Skip to content

feat: expose native aggregate spill and memory metrics - #5423

Merged
sunchao merged 2 commits into
apache:mainfrom
sunchao:dev/chao/codex/comet-aggregate-spill-metrics
Aug 25, 2026
Merged

feat: expose native aggregate spill and memory metrics#5423
sunchao merged 2 commits into
apache:mainfrom
sunchao:dev/chao/codex/comet-aggregate-spill-metrics

Conversation

@sunchao

@sunchao sunchao commented Aug 23, 2026

Copy link
Copy Markdown
Member

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:

SELECT customer_id, SUM(amount)
FROM orders
GROUP BY customer_id;

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:

CometHashAggregate
  number of output rows: 2,000,000
  time spent in this operator: 48 s

After this change, the same operator can also report:

CometHashAggregate
  number of output rows: 2,000,000
  time spent in this operator: 48 s
  number of spills: 4
  total spilled bytes: 1.2 GiB
  number of spilled rows: 8,000,000
  peak native aggregate memory: 256 MiB

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 AggregateExec wrapped 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:

orders
  .groupBy("customer_id")
  .sum("amount")
  .repartitionByRange(col("customer_id"))

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 orders does 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.

@sunchao sunchao changed the title Expose native aggregate spill and memory metrics feat: Expose native aggregate spill and memory metrics Aug 23, 2026
@sunchao sunchao changed the title feat: Expose native aggregate spill and memory metrics feat: expose native aggregate spill and memory metrics Aug 23, 2026
else children.flatMap(_.leafNodes)
}

private[comet] def withoutAggregateMetrics(plan: SparkPlan): CometMetricNode =

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it would be useful to comment the reason aggregation metrics are excluded

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
sunchao force-pushed the dev/chao/codex/comet-aggregate-spill-metrics branch 2 times, most recently from e01bf08 to 836a01d Compare August 24, 2026 18:58

@comphead comphead left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @sunchao it looks good to me

@sunchao
sunchao merged commit f90a3aa into apache:main Aug 25, 2026
71 checks passed
@sunchao

sunchao commented Aug 25, 2026

Copy link
Copy Markdown
Member Author

Merged, thanks @comphead for the review!

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.

2 participants