Skip to content

[GLUTEN-12959][VL] Keep the final stage of a grouping-only aggregate non-flushable - #12960

Merged
zhztheplayer merged 2 commits into
apache:mainfrom
Dhruv-meesho:fix-flushable-grouping-only-final-agg
Sep 4, 2026
Merged

[GLUTEN-12959][VL] Keep the final stage of a grouping-only aggregate non-flushable#12960
zhztheplayer merged 2 commits into
apache:mainfrom
Dhruv-meesho:fix-flushable-grouping-only-final-agg

Conversation

@Dhruv-meesho

@Dhruv-meesho Dhruv-meesho commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

Fixes #12959.

What changes are proposed in this pull request?

This PR fixes a correctness issue where queries doing a LEFT JOIN against a SELECT DISTINCT subquery can produce extra duplicate rows and non-deterministic results when Velox triggers early aggregation abandonment.

The Problem

In FlushableHashAggregateRule.scala, the rule determines if an aggregate is an intermediate stage using:

agg.aggregateExpressions.forall(p => p.mode == Partial || p.mode == PartialMerge)

For grouping-only queries (like SELECT DISTINCT or a GROUP BY without aggregate functions), aggregateExpressions is empty. Because Seq.empty.forall(...) evaluates vacuously to true, the rule fails to distinguish between partial and final stages, and incorrectly transforms the final distinct aggregation into a FlushableHashAggregateExecTransformer.

When this final distinct output is repartitioned again (e.g. feeding into a JOIN on a strict subset of the distinct keys), it falls below an exchange, entering the rule's scope. If Velox later abandons the aggregation under memory pressure, it emits duplicate keys. Because it's already the final aggregation stage, there is no subsequent aggregate to collapse them, resulting in leaked duplicate keys and inflated join outputs.

History & Context

The Fix

Rather than reverting the performance gains from #12098 by restoring the general distribution check, this PR specifically identifies the final grouping-only aggregate using requiredChildDistributionExpressions.isDefined.

private def isGroupingOnlyFinalAgg(agg: HashAggregateExecTransformer): Boolean = {
  agg.aggregateExpressions.isEmpty && agg.requiredChildDistributionExpressions.isDefined
}

In Spark's AggUtils.planAggregateWithoutDistinct, requiredChildDistributionExpressions is:

  • None for the partial aggregate
  • Some(groupingAttributes) for the final aggregate

This matches the exact logic Spark uses in its own check for the equivalent runtime bypass (HashAggregateExec.adaptivePartialAggEnabled, introduced in SPARK-58511).

This ensures the optimization is preserved for valid aggregates with aggregate functions, while safely skipping the final stage for grouping-only aggregations.

How was this patch tested?

  • Added regression test: flushable aggregate rule - distinct feeding a join keeps final agg regular to VeloxAggregateFunctionsFlushSuite.
  • Validated that a SELECT DISTINCT repartitioned for a JOIN on a subset of its keys leaves the final distinct aggregate as a RegularHashAggregateExecTransformer.
  • Verified the fix query output matches vanilla Spark exactly.
  • Verified style check passed via ./dev/format-scala-code.sh.

…non-flushable

FlushableHashAggregateRule decides an aggregate is an intermediate one via
`aggregateExpressions.forall(_.mode == Partial || PartialMerge)`. For a
grouping-only aggregate -- `SELECT DISTINCT`, or a `GROUP BY` with no aggregate
functions -- `aggregateExpressions` is empty, so `forall` is vacuously true and
the check cannot tell the partial stage from the final one. The final stage is
therefore converted to FlushableHashAggregateExecTransformer, which is
serialized as `allowFlush=1` and mapped to `AggregationNode::Step::kPartial` in
SubstraitToVeloxPlan. Velox may then abandon aggregation and emit duplicate
grouping keys, and as no further aggregate follows, the duplicates reach the
consumer.

The rule only visits aggregates below a shuffle, so a plain `SELECT DISTINCT`
whose result is collected is unaffected. The problem shows up when the distinct
output is repartitioned again, e.g. when it feeds a join on a strict subset of
the distinct keys: the extra rows become extra join output rows.

This case used to be covered by `isAggInputAlreadyDistributedWithAggKeys`, added
in apache#4443 for apache#4421, and removed in apache#12098 when the rule was narrowed to protect
only `AggUtils.planAggregateWithOneDistinct`.

Rather than restoring that broader guard and giving back the gains of apache#12098,
this only skips aggregates for which the existing mode check is uninformative:
no aggregate functions, and a required child distribution, which Spark sets for
the final stage but not the partial one (see AggUtils.planAggregateWithoutDistinct).
Spark makes the same check in HashAggregateExec.adaptivePartialAggEnabled, its
guard for the equivalent runtime bypass. Aggregates that do have aggregate
functions keep the behavior introduced by apache#12098.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@github-actions github-actions Bot added the VELOX label Sep 2, 2026
@Dhruv-meesho

Copy link
Copy Markdown
Contributor Author

cc @zhztheplayer @yikaifei for review when you have time, thanks!

@zhztheplayer zhztheplayer left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Thanks for the fix.

Comment on lines +73 to +88
/**
* Returns true if the aggregate applies no aggregate functions and is the final (or complete)
* stage, e.g. the last step of `SELECT DISTINCT`, or of a `GROUP BY` without aggregate functions.
*
* Such an aggregate must fully aggregate. Flushing lets Velox abandon aggregation and emit
* duplicate grouping keys, and since no further aggregate follows, the duplicates reach the
* consumer. When the consumer is a join, they turn into extra join output rows.
*
* The mode check used by the other guards cannot detect this case: `aggregateExpressions` is
* empty here, so `forall(_.mode == Partial | PartialMerge)` is vacuously true and says nothing
* about which stage this is. Spark distinguishes the stages with
* `requiredChildDistributionExpressions`, which is `None` for a partial aggregate and
* `Some(groupingAttributes)` for the final one -- see `AggUtils.planAggregateWithoutDistinct`.
* Spark makes the same check in its own guard for the equivalent runtime bypass, see
* `HashAggregateExec.adaptivePartialAggEnabled`.
*/

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

nit: Can we only keep the 1st paragraph of the comment?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Trimmed, thanks. @zhztheplayer

The final stage of a grouping-only aggregate is regular again, so the approved
plans for the 13 queries apache#12098 touched need updating: q8, q14, q14a, q14b,
q36a, q38, q70a, q86a, q87 on spark40 and spark41. The partial stage below the
shuffle still flushes.

Also trimmed the isGroupingOnlyFinalAgg comment to the first paragraph as asked
in review.
@Dhruv-meesho

Copy link
Copy Markdown
Contributor Author

🟢 Context Maintainer scanned commit 67ed775d6f9f — 13 pre-existing drifts on the baseline, no new drift on this commit.

@github-actions github-actions Bot added the CORE works for Gluten Core label Sep 3, 2026
@Dhruv-meesho

Dhruv-meesho commented Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the quick review!

The spark40 / spark41 shard-3 failures were the TPC-DS plan stability checks, on the same 13 queries #12098 updated — expected, since this puts those final aggregates back to regular. I've updated the approved plans for both spark versions; the partial stage under the shuffle is untouched and still flushes.

Also trimmed the comment to the first paragraph as suggested.
Looks like the new workflow runs are waiting for approval — Can you please trigger them when you get a chance?
Thanks!

cc:- @ArnavBalyan @zhztheplayer

@Dhruv-meesho

Copy link
Copy Markdown
Contributor Author

Hi @zhztheplayer @kevinwilfong @LuciferYang can we please merge this....all the checks have been successfully passed.

@zhztheplayer
zhztheplayer merged commit 89a79ac into apache:main Sep 4, 2026
55 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CORE works for Gluten Core VELOX

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[VL] Flushable aggregate rule converts the final stage of a grouping-only aggregate, producing duplicate rows

3 participants