Copilot commented on code in PR #13033:
URL: https://github.com/apache/gluten/pull/13033#discussion_r4023408143
##########
backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala:
##########
@@ -79,74 +72,30 @@ case class FlushableHashAggregateRule(session:
SparkSession) extends Rule[SparkP
}
/**
- * Walks the plan downward, applying func to each
RegularHashAggregateExecTransformer or
- * SortHashAggregateExecTransformer that is eligible for flushable
conversion. An aggregate is
- * eligible when all expressions are Partial/PartialMerge, it is not the
final stage of a
- * grouping-only aggregate, it is not the protected PartialMerge aggregate
directly below a
- * distinct-partial aggregate, and no aggregate function disallows flushing.
+ * An aggregate is eligible when all expressions are Partial/PartialMerge,
it is not the final
+ * stage of a grouping-only aggregate, it is not the protected PartialMerge
aggregate directly
+ * below a distinct-partial aggregate, and no aggregate function disallows
flushing.
*/
- private def replaceEligibleAggregates(
- plan: SparkPlan,
- protectedAggs: mutable.Map[Int, HashAggregateExecTransformer]):
SparkPlan = {
- def toFlushableAgg(agg: HashAggregateExecTransformer):
FlushableHashAggregateExecTransformer = {
- FlushableHashAggregateExecTransformer(
- agg.requiredChildDistributionExpressions,
- agg.groupingExpressions,
- agg.aggregateExpressions,
- agg.aggregateAttributes,
- agg.initialInputBufferOffset,
- agg.resultExpressions,
- agg.child
- )
- }
-
- def transformDown: SparkPlan => SparkPlan = {
- case agg: RegularHashAggregateExecTransformer if
isGroupingOnlyFinalAgg(agg) =>
- // Final stage of a grouping-only aggregate. It must fully aggregate.
Skip.
- agg
- case agg: RegularHashAggregateExecTransformer
- if !agg.aggregateExpressions.forall(p => p.mode == Partial || p.mode
== PartialMerge) =>
- // Not an intermediate agg. Skip.
- agg
- case agg: RegularHashAggregateExecTransformer
- if protectedAggs.contains(agg.id) =>
- // This is the PartialMerge aggregate directly below a
distinct-partial aggregate in
- // Spark's one-distinct pipeline. Keep it non-flushable so the
distinct step continues to
- // see globally de-duplicated (grouping + distinct) keys.
- agg
- case agg: RegularHashAggregateExecTransformer
- if aggregatesNotSupportFlush(agg.aggregateExpressions) =>
- // Aggregate uses a function that is unsafe to flush. Skip.
- agg
- case agg: RegularHashAggregateExecTransformer =>
- // All guards passed; replace with the flushable variant.
- toFlushableAgg(agg)
- case agg: SortHashAggregateExecTransformer if
isGroupingOnlyFinalAgg(agg) =>
- // See the RegularHashAggregateExecTransformer branch above.
- agg
- case agg: SortHashAggregateExecTransformer
- if !agg.aggregateExpressions.forall(p => p.mode == Partial || p.mode
== PartialMerge) =>
- // Not an intermediate agg. Skip.
- agg
- case agg: SortHashAggregateExecTransformer if
protectedAggs.contains(agg.id) =>
- // See the RegularHashAggregateExecTransformer branch above.
- agg
- case agg: SortHashAggregateExecTransformer
- if aggregatesNotSupportFlush(agg.aggregateExpressions) =>
- // Aggregate uses a function that is unsafe to flush. Skip.
- agg
- case agg: SortHashAggregateExecTransformer =>
- // All guards passed; replace with the flushable variant.
- toFlushableAgg(agg)
- case exchange: ShuffleExchangeLike =>
- // Stop at the next exchange. This rule is applied from an exchange
boundary and should not
- // continue rewriting into a different shuffle region.
- exchange
- case other => other.withNewChildren(other.children.map(transformDown))
- }
+ private def isEligible(
+ agg: HashAggregateExecTransformer,
+ protectedAggIds: Set[Int]): Boolean = {
+ !isGroupingOnlyFinalAgg(agg) &&
+ agg.aggregateExpressions.forall(p => p.mode == Partial || p.mode ==
PartialMerge) &&
+ !protectedAggIds.contains(agg.id) &&
Review Comment:
`protectedAggIds` is collected from the original tree, but this is now a
bottom-up transform. When the partial aggregate below the exchange is replaced,
Spark rebuilds the protected `PartialMerge` parent through
`withNewChildInternal`/`copy(child = ...)`, so that copied node gets a
different `id`; `contains(agg.id)` then misses it and the protected node is
converted as well. This breaks the one-distinct invariant that the surrounding
tests are intended to preserve. Carry the protection through the traversal (or
use a stable structural marker) rather than matching copied plan IDs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
