[SPARK-59564][SQL] Combine adjacent aggregation across a GroupPartitionsExec

sql/core/src/main/scala/org/apache/spark/sql/execution/CombineAdjacentAggregation.scala

@@ -20,6 +20,7 @@ package org.apache.spark.sql.execution
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Complete, Final, Partial, PartialMerge}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
+import org.apache.spark.sql.execution.datasources.v2.GroupPartitionsExec
 import org.apache.spark.sql.internal.SQLConf
 
 /**
@@ -43,6 +44,10 @@ import org.apache.spark.sql.internal.SQLConf
  *       Exchange
  *
  * It supports [[HashAggregateExec]], [[SortAggregateExec]] and [[ObjectHashAggregateExec]].
+ *
+ * A [[GroupPartitionsExec]] and the local sorts `EnsureRequirements` put between the two aggregates
+ * are looked through, so the pair is combined even when the final aggregate's distribution was
+ * satisfied without a shuffle. See `detachAggregate`.
  */
 object CombineAdjacentAggregation extends Rule[SparkPlan] {
   private case class CombinedAggregate(
@@ -55,34 +60,91 @@ object CombineAdjacentAggregation extends Rule[SparkPlan] {
     }
 
     plan.transformDown {
-      case finalAgg @ HashAggregateExec(_, _, _, _, _, _, _, _, partialAgg: HashAggregateExec) =>
-        combinedAggregate(partialAgg, finalAgg)
-          .map(combineHashAggregates(partialAgg, finalAgg, _))
-          .getOrElse(finalAgg)
+      case finalAgg: HashAggregateExec =>
+        detachAggregate(finalAgg.child, hasUpperSort = false) match {
+          case Some((partialAgg: HashAggregateExec, child)) =>
+            combinedAggregate(partialAgg, finalAgg)
+              .map(combineHashAggregates(partialAgg, finalAgg, _, child))
+              .getOrElse(finalAgg)
+          case _ => finalAgg
+        }
 
-      case finalAgg @ SortAggregateExec(_, _, _, _, _, _, _, _, partialAgg: SortAggregateExec)
-          if isPartialAgg(partialAgg, finalAgg) =>
-        finalAgg.copy(
-          groupingExpressions = partialAgg.groupingExpressions,
-          aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode = Complete)),
-          initialInputBufferOffset = 0,
-          child = partialAgg.child)
+      case finalAgg: SortAggregateExec =>
+        detachAggregate(finalAgg.child, hasUpperSort = false) match {
+          case Some((partialAgg: SortAggregateExec, child)) if isPartialAgg(partialAgg, finalAgg) =>
+            finalAgg.copy(
+              groupingExpressions = partialAgg.groupingExpressions,
+              aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode = Complete)),
+              initialInputBufferOffset = 0,
+              child = child)
+          case _ => finalAgg
+        }
 
-      case finalAgg @ ObjectHashAggregateExec(_, _, _, _, _, _, _, _,
-        partialAgg: ObjectHashAggregateExec)
-          if isPartialAgg(partialAgg, finalAgg) =>
-        finalAgg.copy(
-          groupingExpressions = partialAgg.groupingExpressions,
-          aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode = Complete)),
-          initialInputBufferOffset = 0,
-          child = partialAgg.child)
+      case finalAgg: ObjectHashAggregateExec =>
+        detachAggregate(finalAgg.child, hasUpperSort = false) match {
+          case Some((partialAgg: ObjectHashAggregateExec, child))
+              if isPartialAgg(partialAgg, finalAgg) =>
+            finalAgg.copy(
+              groupingExpressions = partialAgg.groupingExpressions,
+              aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode = Complete)),
+              initialInputBufferOffset = 0,
+              child = child)
+          case _ => finalAgg
+        }
     }
   }
 
+  /**
+   * Detaches the aggregate at the bottom of the chain `plan` starts and hands it back together with
+   * the subtree to leave where it was, or `None` when the chain bottoms out at no aggregate, or at
+   * one that cannot leave. The chain's `GroupPartitionsExec` and local sorts are crossed in place,
+   * so the combined aggregate reads whatever ends up on top of them.
+   *
+   * A sort crossed above the aggregate orders the rows the combined aggregate reads, by the
+   * grouping the two aggregates share, so the sort the aggregate reads goes with it: that crossed
+   * sort is what orders those rows. Where the aggregate holds no sort of its own, it stays, being
+   * the only cardinality reducer before that sort. With no sort crossed at all, the sort below the
+   * aggregate stays below it: the aggregate reads what it read, and the ordering it claims is the
+   * one it had.
+   *
+   * @param hasUpperSort whether a local sort has been crossed above `plan`, which is what makes the
+   *                     sort below the aggregate dead.
+   */
+  private def detachAggregate(
+      plan: SparkPlan,
+      hasUpperSort: Boolean): Option[(BaseAggregateExec, SparkPlan)] = plan match {
+    case aggregate: BaseAggregateExec =>
+      if (!hasUpperSort) {
+        Some((aggregate, aggregate.child))
+      } else {
+        aggregate.child match {
+          case sort: SortExec if !sort.global => Some((aggregate, sort.child))
+          case _ => None
+        }
+      }
+
+    case group: GroupPartitionsExec =>
+      detachAggregate(group.child, hasUpperSort) match {
+        case Some((aggregate, child)) =>
+          group.withKeyPositionsFor(child).map(regrouped => (aggregate, regrouped))
+        case _ => None
+      }
+
+    case sort: SortExec if !sort.global =>
+      detachAggregate(sort.child, hasUpperSort = true) match {
+        case Some((aggregate, child)) =>
+          Some((aggregate, sort.withNewChildren(Seq(child))))
+        case _ => None
+      }
+
+    case _ => None
+  }
+
   private def combineHashAggregates(
       partialAgg: HashAggregateExec,
       finalAgg: HashAggregateExec,
-      combined: CombinedAggregate): HashAggregateExec = {
+      combined: CombinedAggregate,
+      child: SparkPlan): HashAggregateExec = {
     // Keep the final aggregate's distribution requirement because the rule runs after
     // EnsureRequirements. The other child-facing metadata comes from the removed aggregate.
     finalAgg.copy(
@@ -91,7 +153,7 @@ object CombineAdjacentAggregation extends Rule[SparkPlan] {
       groupingExpressions = partialAgg.groupingExpressions,
       aggregateExpressions = combined.aggregateExpressions,
       initialInputBufferOffset = combined.initialInputBufferOffset,
-      child = partialAgg.child)
+      child = child)
   }
 
   private def combinedAggregate(

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala

@@ -551,6 +551,43 @@ case class GroupPartitionsExec(
     }
   }
 
+  /**
+   * This node reading `newChild`, with the key positions it projects moved to where the expressions
+   * they name sit there, or `None` when it may not be re-parented onto it. It is named for what it
+   * answers rather than for the node it returns, which is the node it is called on.
+   *
+   * `EnsureRequirements` computed `joinKeyPositions` against the child this node was planned for,
+   * whose partitioning is that child's projected down to the positions the operator above keeps, so
+   * the positions name a key space `newChild` need not share. The projected expressions are the
+   * planned child's own (`KeyedPartitioning.project` builds them that way), which makes moving them
+   * a lookup; a child that does not hold one is turned away.
+   *
+   * Moving them is all this does: what the operator above reads, the ordering it was planned
+   * against included, is the caller's to hold, since nothing here knows what that operator
+   * requires.
+   */
+  def withKeyPositionsFor(newChild: SparkPlan): Option[GroupPartitionsExec] = {
+    // The member of each child's partitioning this reads has to be the one `grouping` reads, since
+    // the positions are only meaningful for that member. `representativeOf` answers as the lookup
+    // there does: the first keyed member, nested collections included.
+    val childKeyed = PartitioningCollection.representativeOf(child.outputPartitioning)
+    val newChildKeyed = PartitioningCollection.representativeOf(newChild.outputPartitioning)
+    (childKeyed, newChildKeyed) match {
+      case (Some(childKp), Some(newChildKp)) =>
+        val plannedInNewChild = childKp.expressions.map(newChildKp.expressions.indexOf)
+        if (plannedInNewChild.exists(_ < 0)) {
+          return None
+        }
+        val positions = joinKeyPositions.fold(plannedInNewChild)(_.map(plannedInNewChild))
+        val regrouped = copy(
+          child = newChild,
+          joinKeyPositions = Option.when(positions != newChildKp.expressions.indices)(positions))
+        regrouped.copyTagsFrom(this)
+        Some(regrouped)
+      case _ => None
+    }
+  }
+
   override def simpleString(maxFields: Int): String = {
     s"$nodeName${planSummaryParts(maxFields).map(" " + _).mkString("")}"
   }