Copilot commented on code in PR #13254:
URL: https://github.com/apache/gluten/pull/13254#discussion_r4313383483
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/PushAggregateThroughJoin.scala:
##########
@@ -593,19 +619,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
Review Comment:
This PR makes previously-private internals (`JoinSide`,
`JoinLeft/JoinRight`, and `pushPartialAggToJoinSide`) part of the public
surface area of `PushAggregateThroughJoin` (even if nested). If this is
intended only for tests, consider reducing visibility (e.g.,
`private[joinagg]`) and placing the unit tests under
`org.apache.gluten.extension.joinagg` so they can access package-private
members, or annotate the method as test-only and keep the `JoinSide` helpers
non-public to avoid expanding the supported API unintentionally.
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/PushAggregateThroughJoin.scala:
##########
@@ -125,135 +130,170 @@ 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)
}
pushed match {
- case Some(newPlan) =>
- pushCount += 1
- changed = true
- newPlan
+ case Some(result) =>
+ pushCount += result.pushedAggregates.size
+ nextCandidates ++= result.pushedAggregates
+ result.plan
case None =>
partialAgg
}
}
+ candidates = nextCandidates.result()
+ depth += 1
}
successfulPushCount += pushCount
- current
+ if (pushCount == 0) {
+ // Keep the original aggregate when no join push succeeds.
+ None
+ } else {
+ Some(current)
+ }
}
- private def pushOnce(partialAgg: Aggregate): Option[LogicalPlan] = {
+ private def pushOnce(partialAgg: Aggregate): Option[PushResult] = {
+ val hasBuffers =
collectPartialWrapperAliases(partialAgg.aggregateExpressions).nonEmpty
+ if (!isDeterministicJoinPath(partialAgg.child)) {
+ // Skip pushdown: grouping rows before evaluating rand(), for example,
changes the results.
+ return None
+ }
extractJoin(partialAgg.child).flatMap {
case (join, wrapperRequiredAttrs, rebuild) =>
- Seq(JoinLeft, JoinRight).iterator
- .flatMap {
- side =>
- val maybePushedJoin =
- pushPartialAggToJoinSide(partialAgg, join,
wrapperRequiredAttrs, side)
- maybePushedJoin match {
- case Some(pushedJoin) =>
- Some(rebuild(pushedJoin, partialAgg.output.collect { case a:
Attribute => a }))
- case None => None
- }
+ val candidates: Seq[(JoinSide, Aggregate)] = Seq(JoinLeft,
JoinRight).flatMap {
+ side =>
+ pushPartialAggToJoinSide(
+ partialAgg.groupingExpressions,
+ partialAgg.aggregateExpressions,
+ join,
+ wrapperRequiredAttrs,
+ side).map(agg => (side, agg))
+ }
+ val selected = if (!hasBuffers) {
+ // Without aggregate functions, DISTINCT can be pushed to both join
sides simultaneously.
+ candidates
+ } else {
+ // When both sides qualify, choose one so the same aggregate buffer
is not produced on both
+ // sides.
+ //
+ // SUM(fact.amount) qualifies only on the fact side because it
requires fact.amount.
+ // COUNT(*) requires no columns, so both sides qualify; take(1)
prefers the left side.
+ // TODO: Use statistics to choose a side when both qualify, rather
than preferring the left.
+ candidates.take(1)
+ }
+ if (selected.isEmpty) {
+ None
+ } else {
+ val pushedJoin = selected.foldLeft(join) {
+ case (currentJoin, (side, pushedAgg)) => side.replace(currentJoin,
pushedAgg)
}
- .toSeq
- .headOption
+ Some(PushResult(
+ rebuild(pushedJoin, partialAgg.output.collect { case a: Attribute
=> a }),
+ selected.map(_._2)))
+ }
}
}
- private def pushPartialAggToJoinSide(
- partialAgg: Aggregate,
+ // Public for testing.
+ def pushPartialAggToJoinSide(
+ groupingExpressions: Seq[Expression],
+ aggregateExpressions: Seq[NamedExpression],
join: Join,
wrapperRequiredAttrs: Seq[Attribute],
- side: JoinSide): Option[Join] = {
- // A pushed wrapper aggregate may move to a join side only when all of its
aggregate inputs
- // come from that side. The pushed grouping then has to include:
- // - grouping attrs already present on this side
- // - join keys / join-condition attrs required to preserve join semantics
- // - attrs referenced by non-pushable expressions still above the pushed
aggregate
- // - attrs propagated from extracted Project / Filter nodes when those
attrs are required to
- // preserve derived grouping expressions
- //
- // The pushed grouping must *not* include pure measure inputs of the
pushed aggregate such as
- // `ss_net_profit`, otherwise the pre-aggregation becomes over-constrained
and ineffective.
- val wrapperAliases =
collectPartialWrapperAliases(partialAgg.aggregateExpressions)
- if (wrapperAliases.isEmpty) {
+ side: JoinSide): Option[Aggregate] = {
+ val wrapperAliases = collectPartialWrapperAliases(aggregateExpressions)
Review Comment:
This PR makes previously-private internals (`JoinSide`,
`JoinLeft/JoinRight`, and `pushPartialAggToJoinSide`) part of the public
surface area of `PushAggregateThroughJoin` (even if nested). If this is
intended only for tests, consider reducing visibility (e.g.,
`private[joinagg]`) and placing the unit tests under
`org.apache.gluten.extension.joinagg` so they can access package-private
members, or annotate the method as test-only and keep the `JoinSide` helpers
non-public to avoid expanding the supported API unintentionally.
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/PushAggregateThroughJoin.scala:
##########
@@ -593,19 +619,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
Review Comment:
This PR makes previously-private internals (`JoinSide`,
`JoinLeft/JoinRight`, and `pushPartialAggToJoinSide`) part of the public
surface area of `PushAggregateThroughJoin` (even if nested). If this is
intended only for tests, consider reducing visibility (e.g.,
`private[joinagg]`) and placing the unit tests under
`org.apache.gluten.extension.joinagg` so they can access package-private
members, or annotate the method as test-only and keep the `JoinSide` helpers
non-public to avoid expanding the supported API unintentionally.
##########
gluten-substrait/src/test/scala/org/apache/gluten/execution/PushAggregateThroughJoinRuleSuite.scala:
##########
@@ -0,0 +1,253 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.execution
+
+import org.apache.gluten.config.GlutenConfig
+import org.apache.gluten.extension.joinagg.{JoinAggregateFunctionWrapper,
JoinAggregateTags, PushAggregateThroughJoin}
+import
org.apache.gluten.extension.joinagg.JoinAggregateFunctionWrapper.{FinalPhase,
PartialMergePhase, PartialPhase}
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, And,
AttributeReference, EqualTo, GreaterThan, Literal, Rand}
+import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
+import org.apache.spark.sql.catalyst.plans.{Inner, PlanTest}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join,
JoinHint, LocalRelation, Project}
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.IntegerType
+
+class PushAggregateThroughJoinRuleSuite extends PlanTest with
SharedSparkSession {
+ private val joinAggregateRule = PushAggregateThroughJoin(spark)
+
+ test("pushPartialAggToJoinSide accepts only an eligible side") {
+ val leftKey = AttributeReference("left_key", IntegerType)()
+ val value = AttributeReference("value", IntegerType)()
+ val rightKey = AttributeReference("right_key", IntegerType)()
+ val join = Join(
+ LocalRelation(Seq(leftKey, value)),
+ LocalRelation(Seq(rightKey)),
+ Inner,
+ Some(EqualTo(leftKey, rightKey)),
+ JoinHint.NONE)
+ val leftPartialSum = Alias(
+
JoinAggregateFunctionWrapper.wrapperPartial(Sum(value)).toAggregateExpression(),
+ "left_partial_sum")()
+ val rightPartialSum = Alias(
+
JoinAggregateFunctionWrapper.wrapperPartial(Sum(rightKey)).toAggregateExpression(),
+ "right_partial_sum")()
+
+ val leftPushed = joinAggregateRule.pushPartialAggToJoinSide(
+ Seq(leftKey),
+ Seq(leftKey, leftPartialSum),
+ join,
+ Nil,
+ joinAggregateRule.JoinLeft)
+ assert(leftPushed.exists(_.child eq join.left))
+ assert(leftPushed.get.aggregateExpressions.last == leftPartialSum)
+
assert(leftPushed.get.getTagValue(JoinAggregateTags.phase).contains(PartialPhase))
+ assert(joinAggregateRule
+ .pushPartialAggToJoinSide(
+ Seq(leftKey),
+ Seq(leftKey, leftPartialSum),
+ join,
+ Nil,
+ joinAggregateRule.JoinRight)
+ .isEmpty)
+ val rightPushed = joinAggregateRule.pushPartialAggToJoinSide(
+ Seq(leftKey),
+ Seq(leftKey, rightPartialSum),
+ join,
+ Nil,
+ joinAggregateRule.JoinRight)
+ assert(rightPushed.exists(_.child eq join.right))
+ }
+
+ test("pushPartialAggToJoinSide supports grouping-only pushdown on either
side") {
+ val leftKey = AttributeReference("left_key", IntegerType)()
+ val rightKey = AttributeReference("right_key", IntegerType)()
+ val join = Join(
+ LocalRelation(Seq(leftKey)),
+ LocalRelation(Seq(rightKey)),
+ Inner,
+ Some(EqualTo(leftKey, rightKey)),
+ JoinHint.NONE)
+ val keys = Seq(leftKey, rightKey)
+
+ Seq(joinAggregateRule.JoinLeft, joinAggregateRule.JoinRight).foreach {
+ side =>
+ val pushed = joinAggregateRule
+ .pushPartialAggToJoinSide(keys, keys, join, Nil, side)
+ .get
+ assert(pushed.child eq side.plan(join))
+ assert(pushed.groupingExpressions == side.plan(join).output)
+ assert(pushed.aggregateExpressions == side.plan(join).output)
+
assert(pushed.getTagValue(JoinAggregateTags.phase).contains(PartialPhase))
+ }
+ }
+
+ test("grouping-only pushdown rule is idempotent") {
+ Seq(false, true).foreach {
+ partialMergeEnabled =>
+ withSQLConf(
+ GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_ENABLED.key -> "true",
+ GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_MAX_DEPTH.key -> "1",
+ GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_ENABLED.key ->
+ partialMergeEnabled.toString
+ ) {
+ val leftKey = AttributeReference("left_key", IntegerType)()
+ val rightKey = AttributeReference("right_key", IntegerType)()
+ val join = Join(
+ LocalRelation(Seq(leftKey)),
+ LocalRelation(Seq(rightKey)),
+ Inner,
+ Some(EqualTo(leftKey, rightKey)),
+ JoinHint.NONE)
+ val keys = Seq(leftKey, rightKey)
+ val plan = Aggregate(
+ keys,
+ Seq(Alias(leftKey, "left_output")(), Alias(rightKey,
"right_output")()),
+ Project(keys, join))
+ val rule = PushAggregateThroughJoin(spark)
+
+ val once = rule(plan)
+ assert(once.collect { case _: Aggregate => 1 }.size == 3)
+ assert(once.output == plan.output)
+ assert(once.resolved)
+ assert(once.collect { case p if p.missingInput.nonEmpty => p
}.isEmpty)
+ assert(rule.getSuccessfulPushCount == 2)
+
+ val phases = once.collect {
+ case agg: Aggregate => agg.getTagValue(JoinAggregateTags.phase)
+ }
+ assert(phases == Seq(Some(FinalPhase), Some(PartialPhase),
Some(PartialPhase)))
Review Comment:
This assertion is brittle because it depends on the traversal order returned
by `TreeNode.collect` (and implicitly on join child ordering). To avoid
unnecessary test flakes during Spark/tree rewrite changes, assert the multiset
of phases instead (e.g., counts per phase, or `phases.toSet` plus a size/count
check), while still verifying that exactly one `FinalPhase` and two
`PartialPhase` aggregates are present.
--
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]
