Copilot commented on code in PR #13254:
URL: https://github.com/apache/gluten/pull/13254#discussion_r4313541789
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/PushAggregateThroughJoin.scala:
##########
@@ -125,135 +130,171 @@ case class PushAggregateThroughJoin(spark: SparkSession)
// Any join-specific key widening or subtree rebuild is handled later by
pushOnce /
// pushPartialAggToJoinSide. This keeps the split step reusable even when
the current child is
// not yet a join-shaped subtree.
+ if (hasDistinctAggExpr(agg.aggregateExpressions)) {
+ return None
+ }
val pushableSpecs = collectPushableSpecs(agg.aggregateExpressions)
- if (pushableSpecs.isEmpty) {
- None
- } else {
- val groupingRefAttrs =
dedupeAttrs(agg.groupingExpressions.flatMap(referencedAttrsInOrder))
- // Aggregate subexpressions that are not rewritten remain above the
lower aggregate after the
- // split. Their referenced attrs must therefore still be produced by the
lower aggregate, but
- // at this stage we only derive that requirement from the aggregate
itself, not from any
- // particular child shape.
- val nonPushableAggAttrs = dedupeAttrs(agg.aggregateExpressions.flatMap {
- case Alias(expr, _) if !isPushableExpr(expr) =>
referencedAttrsInOrder(expr)
- case expr if !isPushableExpr(expr) => referencedAttrsInOrder(expr)
- case _ => Nil
- })
- val partialGroupingKeys = dedupeAttrs(groupingRefAttrs ++
nonPushableAggAttrs)
- val partialAliases = pushableSpecs.zipWithIndex.map {
- case (spec, idx) =>
- Alias(
- JoinAggregateFunctionWrapper
- .wrapperPartial(spec.aggregate, spec.wrapperKey)
- .toAggregateExpression(),
- s"_pushed_${spec.aggregate.prettyName}_buffer_$idx"
- )()
- }
+ if (
+ pushableSpecs.isEmpty &&
+ !isGroupingOnly(agg.groupingExpressions, agg.aggregateExpressions)
+ ) {
+ return None
+ }
+ val groupingRefAttrs =
dedupeAttrs(agg.groupingExpressions.flatMap(referencedAttrsInOrder))
+ // Aggregate subexpressions that are not rewritten remain above the lower
aggregate after the
+ // split. Their referenced attrs must therefore still be produced by the
lower aggregate, but
+ // at this stage we only derive that requirement from the aggregate
itself, not from any
+ // particular child shape.
+ val nonPushableAggAttrs = dedupeAttrs(agg.aggregateExpressions.flatMap {
+ case Alias(expr, _) if !isPushableExpr(expr) =>
referencedAttrsInOrder(expr)
+ case expr if !isPushableExpr(expr) => referencedAttrsInOrder(expr)
+ case _ => Nil
+ })
+ val partialGroupingKeys = dedupeAttrs(groupingRefAttrs ++
nonPushableAggAttrs)
+ val partialAliases = pushableSpecs.zipWithIndex.map {
+ case (spec, idx) =>
+ Alias(
+ JoinAggregateFunctionWrapper
+ .wrapperPartial(spec.aggregate, spec.wrapperKey)
+ .toAggregateExpression(),
+ s"_pushed_${spec.aggregate.prettyName}_buffer_$idx"
+ )()
+ }
- val partialAgg = Aggregate(
+ val partialAgg = withPhase(
+ Aggregate(
groupingExpressions = partialGroupingKeys,
aggregateExpressions = partialGroupingKeys ++ partialAliases,
- child = agg.child
- )
-
- val partialOutputAttrs =
- partialAgg.output
- .slice(partialGroupingKeys.size, partialGroupingKeys.size +
partialAliases.size)
- .map(cleanAttr)
- val partialRefs = pushableSpecs.zip(partialOutputAttrs).map {
- case (spec, attr) => SidePartialRef(spec, attr)
- }
+ child = agg.child),
+ PartialPhase)
+
+ val partialOutputAttrs =
+ partialAgg.output
+ .slice(partialGroupingKeys.size, partialGroupingKeys.size +
partialAliases.size)
+ .map(cleanAttr)
+ val partialRefs = pushableSpecs.zip(partialOutputAttrs).map {
+ case (spec, attr) => SidePartialRef(spec, attr)
+ }
- rewriteAggregateExpressions(agg.aggregateExpressions, partialRefs).map {
- rewrittenAggExprs =>
- (agg.copy(aggregateExpressions = rewrittenAggExprs, child =
partialAgg), partialAgg)
- }
+ rewriteAggregateExpressions(agg.aggregateExpressions, partialRefs).map {
+ rewrittenAggExprs =>
+ val finalAgg = withPhase(
+ agg.copy(aggregateExpressions = rewrittenAggExprs, child =
partialAgg),
+ FinalPhase)
+ (finalAgg, partialAgg)
}
}
- private def pushPartialWrapperAggregate(lowerPartialAgg: Aggregate):
LogicalPlan = {
- // Push one join edge per iteration. This keeps the rewrite local and lets
`maxDepth` bound
- // how far a pushed aggregate is allowed to travel through a multi-join
subtree.
+ private def pushPartialAggregate(lowerPartialAgg: Aggregate):
Option[LogicalPlan] = {
+ // Push one join level per iteration; grouping-only aggregation can push
to both sides.
val partialMergeEnabled =
GlutenConfig.get.pushAggregateThroughJoinPartialMergeEnabled
var current: LogicalPlan = lowerPartialAgg
- var changed = true
+ var candidates = Seq(lowerPartialAgg)
var pushCount = 0
- while (changed && pushCount < maxDepth) {
- changed = false
+ var depth = 0
+ while (candidates.nonEmpty && depth < maxDepth) {
+ val nextCandidates = Seq.newBuilder[Aggregate]
current = current.transformUp {
- case partialAgg: Aggregate if
isPurePartialWrapperAggregate(partialAgg) =>
+ // Only revisit aggregates created by the previous push, not unrelated
aggregates.
+ case partialAgg: Aggregate if candidates.exists(_ eq partialAgg) =>
// The physical strategy owns the top merge; retain logical merges
only between joins.
val retainPartialMerge = partialMergeEnabled && pushCount > 0
val pushed = if (retainPartialMerge) {
val (partialMergeAgg, freshLowerPartialAgg) =
splitPartialAggregateForMerge(partialAgg)
pushOnce(freshLowerPartialAgg).map {
- pushedPlan => partialMergeAgg.copy(child = pushedPlan)
+ result =>
+ val mergeAgg = withPhase(
+ partialMergeAgg.copy(child = result.plan),
+ PartialMergePhase)
+ result.copy(plan = mergeAgg)
}
} else {
pushOnce(partialAgg)
}
Review Comment:
For grouping-only aggregates (i.e., DISTINCT pushdown), `partialMergeEnabled
&& pushCount > 0` will trigger `splitPartialAggregateForMerge`, creating an
extra `PartialMergePhase` layer even though there are no wrapper buffers to
merge. This adds redundant aggregates and can bloat plans for multi-join
depths. Consider gating `retainPartialMerge` on the presence of wrapper buffers
for the current `partialAgg` (e.g.,
`collectPartialWrapperAliases(partialAgg.aggregateExpressions).nonEmpty`) so
grouping-only pushdown never generates merge aggregates.
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/PushAggregateThroughJoin.scala:
##########
@@ -593,19 +620,19 @@ case class PushAggregateThroughJoin(spark: SparkSession)
expr.collect { case attr: Attribute => attr }
}
- sealed private trait JoinSide {
+ sealed trait JoinSide {
def plan(join: Join): LogicalPlan
def outputSet(join: Join): AttributeSet
def replace(join: Join, newPlan: LogicalPlan): Join
}
- private case object JoinLeft extends JoinSide {
+ case object JoinLeft extends JoinSide {
override def plan(join: Join): LogicalPlan = join.left
override def outputSet(join: Join): AttributeSet = join.left.outputSet
override def replace(join: Join, newPlan: LogicalPlan): Join =
join.copy(left = newPlan)
}
- private case object JoinRight extends JoinSide {
+ case object JoinRight extends JoinSide {
override def plan(join: Join): LogicalPlan = join.right
override def outputSet(join: Join): AttributeSet = join.right.outputSet
override def replace(join: Join, newPlan: LogicalPlan): Join =
join.copy(right = newPlan)
Review Comment:
`JoinSide`, `JoinLeft`, and `JoinRight` were widened from private to public.
Along with the now-public `pushPartialAggToJoinSide` (commented as 'Public for
testing'), this expands the surface area of `joinagg` internals and makes later
refactors harder. If tests are the only consumer, consider restricting
visibility using package-private scope (e.g., `private[joinagg]` /
`private[gluten]`) rather than making these symbols fully public.
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/PushAggregateThroughJoin.scala:
##########
@@ -125,135 +130,171 @@ case class PushAggregateThroughJoin(spark: SparkSession)
// Any join-specific key widening or subtree rebuild is handled later by
pushOnce /
// pushPartialAggToJoinSide. This keeps the split step reusable even when
the current child is
// not yet a join-shaped subtree.
+ if (hasDistinctAggExpr(agg.aggregateExpressions)) {
+ return None
+ }
val pushableSpecs = collectPushableSpecs(agg.aggregateExpressions)
- if (pushableSpecs.isEmpty) {
- None
- } else {
- val groupingRefAttrs =
dedupeAttrs(agg.groupingExpressions.flatMap(referencedAttrsInOrder))
- // Aggregate subexpressions that are not rewritten remain above the
lower aggregate after the
- // split. Their referenced attrs must therefore still be produced by the
lower aggregate, but
- // at this stage we only derive that requirement from the aggregate
itself, not from any
- // particular child shape.
- val nonPushableAggAttrs = dedupeAttrs(agg.aggregateExpressions.flatMap {
- case Alias(expr, _) if !isPushableExpr(expr) =>
referencedAttrsInOrder(expr)
- case expr if !isPushableExpr(expr) => referencedAttrsInOrder(expr)
- case _ => Nil
- })
- val partialGroupingKeys = dedupeAttrs(groupingRefAttrs ++
nonPushableAggAttrs)
- val partialAliases = pushableSpecs.zipWithIndex.map {
- case (spec, idx) =>
- Alias(
- JoinAggregateFunctionWrapper
- .wrapperPartial(spec.aggregate, spec.wrapperKey)
- .toAggregateExpression(),
- s"_pushed_${spec.aggregate.prettyName}_buffer_$idx"
- )()
- }
+ if (
+ pushableSpecs.isEmpty &&
+ !isGroupingOnly(agg.groupingExpressions, agg.aggregateExpressions)
+ ) {
+ return None
+ }
+ val groupingRefAttrs =
dedupeAttrs(agg.groupingExpressions.flatMap(referencedAttrsInOrder))
+ // Aggregate subexpressions that are not rewritten remain above the lower
aggregate after the
+ // split. Their referenced attrs must therefore still be produced by the
lower aggregate, but
+ // at this stage we only derive that requirement from the aggregate
itself, not from any
+ // particular child shape.
+ val nonPushableAggAttrs = dedupeAttrs(agg.aggregateExpressions.flatMap {
+ case Alias(expr, _) if !isPushableExpr(expr) =>
referencedAttrsInOrder(expr)
+ case expr if !isPushableExpr(expr) => referencedAttrsInOrder(expr)
+ case _ => Nil
+ })
+ val partialGroupingKeys = dedupeAttrs(groupingRefAttrs ++
nonPushableAggAttrs)
+ val partialAliases = pushableSpecs.zipWithIndex.map {
+ case (spec, idx) =>
+ Alias(
+ JoinAggregateFunctionWrapper
+ .wrapperPartial(spec.aggregate, spec.wrapperKey)
+ .toAggregateExpression(),
+ s"_pushed_${spec.aggregate.prettyName}_buffer_$idx"
+ )()
+ }
- val partialAgg = Aggregate(
+ val partialAgg = withPhase(
+ Aggregate(
groupingExpressions = partialGroupingKeys,
aggregateExpressions = partialGroupingKeys ++ partialAliases,
- child = agg.child
- )
-
- val partialOutputAttrs =
- partialAgg.output
- .slice(partialGroupingKeys.size, partialGroupingKeys.size +
partialAliases.size)
- .map(cleanAttr)
- val partialRefs = pushableSpecs.zip(partialOutputAttrs).map {
- case (spec, attr) => SidePartialRef(spec, attr)
- }
+ child = agg.child),
+ PartialPhase)
+
+ val partialOutputAttrs =
+ partialAgg.output
+ .slice(partialGroupingKeys.size, partialGroupingKeys.size +
partialAliases.size)
+ .map(cleanAttr)
+ val partialRefs = pushableSpecs.zip(partialOutputAttrs).map {
+ case (spec, attr) => SidePartialRef(spec, attr)
+ }
- rewriteAggregateExpressions(agg.aggregateExpressions, partialRefs).map {
- rewrittenAggExprs =>
- (agg.copy(aggregateExpressions = rewrittenAggExprs, child =
partialAgg), partialAgg)
- }
+ rewriteAggregateExpressions(agg.aggregateExpressions, partialRefs).map {
+ rewrittenAggExprs =>
+ val finalAgg = withPhase(
+ agg.copy(aggregateExpressions = rewrittenAggExprs, child =
partialAgg),
+ FinalPhase)
+ (finalAgg, partialAgg)
}
}
- private def pushPartialWrapperAggregate(lowerPartialAgg: Aggregate):
LogicalPlan = {
- // Push one join edge per iteration. This keeps the rewrite local and lets
`maxDepth` bound
- // how far a pushed aggregate is allowed to travel through a multi-join
subtree.
+ private def pushPartialAggregate(lowerPartialAgg: Aggregate):
Option[LogicalPlan] = {
+ // Push one join level per iteration; grouping-only aggregation can push
to both sides.
val partialMergeEnabled =
GlutenConfig.get.pushAggregateThroughJoinPartialMergeEnabled
var current: LogicalPlan = lowerPartialAgg
- var changed = true
+ var candidates = Seq(lowerPartialAgg)
var pushCount = 0
- while (changed && pushCount < maxDepth) {
- changed = false
+ var depth = 0
+ while (candidates.nonEmpty && depth < maxDepth) {
+ val nextCandidates = Seq.newBuilder[Aggregate]
current = current.transformUp {
- case partialAgg: Aggregate if
isPurePartialWrapperAggregate(partialAgg) =>
+ // Only revisit aggregates created by the previous push, not unrelated
aggregates.
+ case partialAgg: Aggregate if candidates.exists(_ eq partialAgg) =>
Review Comment:
Using `candidates.exists(_ eq partialAgg)` makes membership checks O(n) for
every visited `Aggregate` during `transformUp`. Since grouping-only pushdown
can produce multiple candidates per depth (potentially growing quickly),
consider switching `candidates` to an identity-based set (e.g.,
`java.util.IdentityHashMap`-backed set) to make membership checks O(1) and
avoid quadratic behavior in large join trees.
--
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]
