Skip to content

feat: let a partial aggregate with a double AVG stop aggregating - #7421

Draft
andygrove wants to merge 2 commits into
apache:mainfrom
andygrove:avg-skip-partial-aggregation
Draft

andygrove wants to merge 2 commits into
apache:mainfrom
andygrove:avg-skip-partial-aggregation

Conversation

@andygrove

Copy link
Copy Markdown
Member

Which issue does this PR close?

Part of #7406.

Rationale for this change

With spark.comet.exec.aggregate.skipPartial.enabled, a native partial aggregate can stop aggregating once its input looks mostly distinct and send the rest of the task's rows on as partial states. Comet only allows this when every aggregate function is a single-argument COUNT, so a query such as TPC-H q17, which takes an AVG grouped by a nearly unique key, keeps aggregating. In #7289 this change cut q17 at SF1000 from about 277 s to 242 s and stopped an executor from being OOM-killed. That PR runs the query on Comet Native Executors, but the gate is shared, so the change helps Comet on Spark too. This splits it out of #7289.

What changes are included in this PR?

  • AvgGroupsAccumulator, the Float64 average that Comet uses for every non-decimal AVG (and for Spark's rewrite of a low-precision decimal AVG to avg(UnscaledValue(x))), implements convert_to_state. It used to return NotImplemented. Each row's state is its value and a count of one. A null or filtered-out row gets a zero sum and a zero count, which is the state state() emits for such a group.
  • configure_skip_partial_aggregation also admits a single-argument AVG whose result type is a double. A decimal AVG runs as AvgDecimal, whose convert_to_state is still not implemented, so it stays excluded.
  • The setting's description now lists AVG.

How are these changes tested?

  • A new unit test in avg.rs checks that convert_to_state gives the same states as aggregating each row into its own group, with a null row and a filtered-out row.
  • skip_partial_eligibility_is_fail_closed now covers a double AVG (admitted), an AVG with no child, and a decimal AVG (both rejected).
  • skip partial aggregation admits only supported native shuffle plans in CometAggregateSuite checks that a double AVG over nullable columns skips and still matches Spark, and that a DECIMAL(30, 2) AVG does not skip.

A partial aggregate that stops aggregating sends each row on as a partial
state, which DataFusion builds with GroupsAccumulator::convert_to_state.
Comet's Float64 AvgGroupsAccumulator, which runs Spark's average of a decimal
as avg(UnscaledValue(x)), returned NotImplemented. Each row's state is now its
value and a count of one, or a zero sum and count for a null or filtered-out
row, what state() emits for such a group.

(cherry picked from commit 5000f3a)
configure_skip_partial_aggregation only admitted single-argument COUNT. With
the double AvgGroupsAccumulator now converting rows to partial states, a
single-argument AVG whose result is a double is admitted too. A decimal AVG
runs as AvgDecimal, whose convert_to_state is not implemented, so it stays
excluded.
@github-actions github-actions Bot added enhancement New feature or request area:aggregation Hash aggregates, aggregate expressions labels Nov 2, 2026
@andygrove
andygrove marked this pull request as draft November 2, 2026 19:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:aggregation Hash aggregates, aggregate expressions enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant