Copilot commented on code in PR #13254:
URL: https://github.com/apache/gluten/pull/13254#discussion_r4314006576
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/ImplementJoinAggregate.scala:
##########
@@ -62,42 +62,26 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
}
plan match {
- case agg: Aggregate if containsWrapperAggregate(agg) =>
- planJoinAggregate(agg).toSeq
+ case agg: Aggregate =>
+ agg.getTagValue(JoinAggregateTags.phase)
+ .flatMap(phase => planJoinAggregate(agg, phase))
+ .toSeq
Review Comment:
`ImplementJoinAggregate` now attempts to plan *any* `Aggregate` that has a
`JoinAggregateTags.phase`, but `planPhase(..., packOutputBuffers = true)`
assumes the aggregate expressions are joinagg wrappers (so it can pack/unpack
buffer structs). If an aggregate is tagged but contains non-wrapper aggregate
functions (non-empty `aggExpressions` without `JoinAggregateFunctionWrapper`),
`buildBufferAggregate` will repack *Spark buffer attributes* into structs and
rewrite outputs, which will produce an incorrect schema/results. A concrete
hardening is to validate inputs before planning: allow (a) grouping-only
aggregates (`aggExpressions.isEmpty`), or (b) wrapper-only aggregates (every
`AggregateExpression` uses `JoinAggregateFunctionWrapper` and targets match the
tag); otherwise return `None` so Spark’s normal strategies handle it.
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/ImplementJoinAggregate.scala:
##########
@@ -62,42 +62,26 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
}
plan match {
- case agg: Aggregate if containsWrapperAggregate(agg) =>
- planJoinAggregate(agg).toSeq
+ case agg: Aggregate =>
+ agg.getTagValue(JoinAggregateTags.phase)
+ .flatMap(phase => planJoinAggregate(agg, phase))
+ .toSeq
case _ =>
Nil
}
}
- private def planJoinAggregate(agg: Aggregate): Option[SparkPlan] = agg match
{
+ private def planJoinAggregate(agg: Aggregate, phase: TargetPhase):
Option[SparkPlan] = agg match {
case PhysicalAggregation(groupingExpressions, aes, resultExpressions,
child) =>
Review Comment:
`ImplementJoinAggregate` now attempts to plan *any* `Aggregate` that has a
`JoinAggregateTags.phase`, but `planPhase(..., packOutputBuffers = true)`
assumes the aggregate expressions are joinagg wrappers (so it can pack/unpack
buffer structs). If an aggregate is tagged but contains non-wrapper aggregate
functions (non-empty `aggExpressions` without `JoinAggregateFunctionWrapper`),
`buildBufferAggregate` will repack *Spark buffer attributes* into structs and
rewrite outputs, which will produce an incorrect schema/results. A concrete
hardening is to validate inputs before planning: allow (a) grouping-only
aggregates (`aggExpressions.isEmpty`), or (b) wrapper-only aggregates (every
`AggregateExpression` uses `JoinAggregateFunctionWrapper` and targets match the
tag); otherwise return `None` so Spark’s normal strategies handle it.
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/ImplementJoinAggregate.scala:
##########
@@ -62,42 +62,26 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
}
plan match {
- case agg: Aggregate if containsWrapperAggregate(agg) =>
- planJoinAggregate(agg).toSeq
+ case agg: Aggregate =>
+ agg.getTagValue(JoinAggregateTags.phase)
+ .flatMap(phase => planJoinAggregate(agg, phase))
+ .toSeq
case _ =>
Nil
}
}
- private def planJoinAggregate(agg: Aggregate): Option[SparkPlan] = agg match
{
+ private def planJoinAggregate(agg: Aggregate, phase: TargetPhase):
Option[SparkPlan] = agg match {
case PhysicalAggregation(groupingExpressions, aes, resultExpressions,
child) =>
// KEEP: For compatibility with Spark version before 3.5, which widens
// the aggregate expressions
// to `Seq[NamedExpression]` instead of `Seq[AggregateExpression]`.
val aggExpressions = aes.map(_.asInstanceOf[AggregateExpression])
- // A single logical aggregate must lower either entirely as pushed-phase
wrappers or entirely
- // as final-phase wrappers. Mixed-phase aggregates are rejected here.
val grouping = groupingExpressions.collect { case ne: NamedExpression =>
ne }
if (grouping.size != groupingExpressions.size) {
return None
}
- val wrappers = aggExpressions.flatMap {
- case ae @ AggregateExpression(wrapper: JoinAggregateFunctionWrapper,
_, _, _, _) =>
- Some((ae, wrapper))
- case _ =>
- None
- }
- if (wrappers.isEmpty) {
- return None
- }
-
- val phases = wrappers.map(_._2.targetPhase).distinct
- if (phases.size != 1) {
- return None
- }
- val phase = phases.head
-
val childPlan = planLater(child)
phase match {
Review Comment:
`ImplementJoinAggregate` now attempts to plan *any* `Aggregate` that has a
`JoinAggregateTags.phase`, but `planPhase(..., packOutputBuffers = true)`
assumes the aggregate expressions are joinagg wrappers (so it can pack/unpack
buffer structs). If an aggregate is tagged but contains non-wrapper aggregate
functions (non-empty `aggExpressions` without `JoinAggregateFunctionWrapper`),
`buildBufferAggregate` will repack *Spark buffer attributes* into structs and
rewrite outputs, which will produce an incorrect schema/results. A concrete
hardening is to validate inputs before planning: allow (a) grouping-only
aggregates (`aggExpressions.isEmpty`), or (b) wrapper-only aggregates (every
`AggregateExpression` uses `JoinAggregateFunctionWrapper` and targets match the
tag); otherwise return `None` so Spark’s normal strategies handle it.
##########
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 {
Review Comment:
`pushPartialAggToJoinSide`, `JoinSide`, `JoinLeft`, and `JoinRight` were
made fully public to support cross-package tests. This expands the public
surface area of a rule implementation detail and makes it easier for other code
to take a dependency on internal rewrite mechanics. If the intent is
“test-visible but not public API”, consider narrowing visibility (e.g.,
`private[gluten]` or `private[joinagg]`) for the method and the `JoinSide`
types while keeping the tests compiling (tests under `org.apache.gluten.*` can
still access `private[gluten]`).
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/ImplementJoinAggregate.scala:
##########
@@ -62,42 +62,26 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
}
plan match {
- case agg: Aggregate if containsWrapperAggregate(agg) =>
- planJoinAggregate(agg).toSeq
+ case agg: Aggregate =>
+ agg.getTagValue(JoinAggregateTags.phase)
+ .flatMap(phase => planJoinAggregate(agg, phase))
+ .toSeq
case _ =>
Nil
}
}
- private def planJoinAggregate(agg: Aggregate): Option[SparkPlan] = agg match
{
+ private def planJoinAggregate(agg: Aggregate, phase: TargetPhase):
Option[SparkPlan] = agg match {
case PhysicalAggregation(groupingExpressions, aes, resultExpressions,
child) =>
// KEEP: For compatibility with Spark version before 3.5, which widens
// the aggregate expressions
// to `Seq[NamedExpression]` instead of `Seq[AggregateExpression]`.
val aggExpressions = aes.map(_.asInstanceOf[AggregateExpression])
Review Comment:
`ImplementJoinAggregate` now attempts to plan *any* `Aggregate` that has a
`JoinAggregateTags.phase`, but `planPhase(..., packOutputBuffers = true)`
assumes the aggregate expressions are joinagg wrappers (so it can pack/unpack
buffer structs). If an aggregate is tagged but contains non-wrapper aggregate
functions (non-empty `aggExpressions` without `JoinAggregateFunctionWrapper`),
`buildBufferAggregate` will repack *Spark buffer attributes* into structs and
rewrite outputs, which will produce an incorrect schema/results. A concrete
hardening is to validate inputs before planning: allow (a) grouping-only
aggregates (`aggExpressions.isEmpty`), or (b) wrapper-only aggregates (every
`AggregateExpression` uses `JoinAggregateFunctionWrapper` and targets match the
tag); otherwise return `None` so Spark’s normal strategies handle it.
##########
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 {
Review Comment:
`pushPartialAggToJoinSide`, `JoinSide`, `JoinLeft`, and `JoinRight` were
made fully public to support cross-package tests. This expands the public
surface area of a rule implementation detail and makes it easier for other code
to take a dependency on internal rewrite mechanics. If the intent is
“test-visible but not public API”, consider narrowing visibility (e.g.,
`private[gluten]` or `private[joinagg]`) for the method and the `JoinSide`
types while keeping the tests compiling (tests under `org.apache.gluten.*` can
still access `private[gluten]`).
##########
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)
}
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] = {
Review Comment:
`pushPartialAggToJoinSide`, `JoinSide`, `JoinLeft`, and `JoinRight` were
made fully public to support cross-package tests. This expands the public
surface area of a rule implementation detail and makes it easier for other code
to take a dependency on internal rewrite mechanics. If the intent is
“test-visible but not public API”, consider narrowing visibility (e.g.,
`private[gluten]` or `private[joinagg]`) for the method and the `JoinSide`
types while keeping the tests compiling (tests under `org.apache.gluten.*` can
still access `private[gluten]`).
--
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]
