[GLUTEN-12959][VL] Keep the final stage of a grouping-only aggregate non-flushable - #12960
Conversation
…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>
|
cc @zhztheplayer @yikaifei for review when you have time, thanks! |
| /** | ||
| * 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`. | ||
| */ |
There was a problem hiding this comment.
nit: Can we only keep the 1st paragraph of the comment?
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.
|
🟢 Context Maintainer scanned commit |
|
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. |
|
Hi @zhztheplayer @kevinwilfong @LuciferYang can we please merge this....all the checks have been successfully passed. |
Fixes #12959.
What changes are proposed in this pull request?
This PR fixes a correctness issue where queries doing a
LEFT JOINagainst aSELECT DISTINCTsubquery 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:For grouping-only queries (like
SELECT DISTINCTor aGROUP BYwithout aggregate functions),aggregateExpressionsis empty. BecauseSeq.empty.forall(...)evaluates vacuously totrue, the rule fails to distinguish between partial and final stages, and incorrectly transforms the final distinct aggregation into aFlushableHashAggregateExecTransformer.When this final distinct output is repartitioned again (e.g. feeding into a
JOINon 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
isAggInputAlreadyDistributedWithAggKeys, introduced in [GLUTEN-4421][VL] Disable flushable aggregate when input is already partitioned by grouping keys #4443.AggUtils.planAggregateWithOneDistinct#12098 when the rule was narrowed to target only theAggUtils.planAggregateWithOneDistinctpipeline, which left the plainSELECT DISTINCTunprotected.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.In Spark's
AggUtils.planAggregateWithoutDistinct,requiredChildDistributionExpressionsis:Nonefor the partial aggregateSome(groupingAttributes)for the final aggregateThis 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?
flushable aggregate rule - distinct feeding a join keeps final agg regulartoVeloxAggregateFunctionsFlushSuite.SELECT DISTINCTrepartitioned for aJOINon a subset of its keys leaves the final distinct aggregate as aRegularHashAggregateExecTransformer../dev/format-scala-code.sh.