[SPARK-44571][SQL] Merge subplans with one row result

### What changes were proposed in this pull request?

This PR renames `MergeScalarSubqueries` rule to `MergeSubplans` and extends plan merging capabilities to non-grouping aggregate subplans, which are very similar to scalar subqueries in terms they return one row result.

Consider the following query that joins 2 non-grouping aggregates:
```
Join Inner
:- Aggregate [min(a) AS min(a)]
:  +- Relation [a, b, c]
+- Aggregate [sum(b) AS sum(b), avg(cast(c as double)) AS avg(c)]
   +- Relation [a, b, c]
```
with the improved rule the plan is optimized to:
```
WithCTE
:- CTERelationDef 0
:  +- Project [named_struct(min(a), min(a), sum(b), sum(b), avg(c), avg(c)) AS mergedValue]
:     +- Aggregate [min(a) AS min(a), sum(b) AS sum(b), avg(cast(c as double)) AS avg(c)]
:        +- Relation [a, b, c]
+- Join Inner
   :- Project [scalar-subquery [].min(a) AS min(a)]
   :  :  +- CTERelationRef 0
   :  +- OneRowRelation
   +- Project [scalar-subquery [].sum(b) AS sum(b), scalar-subquery [].avg(c) AS avg(c)]
      :  :- CTERelationRef 0
      :  +- CTERelationRef 0
      +- OneRowRelation
```
so as to scan `Relation` only once.

Please note that the above plan where the 2 aggregations are part of a "join group" could be rewritten as one aggregate without the need to introduce a CTE and keeping the join. But there are more complex cases when the proposed CTE based approach is the only viable option. Such cases include when the aggregates reside at different parts of plan, maybe even in diffrent subquery expressions.

E.g. the following query:
```
Join Inner
:- Project [scalar-subquery [] AS scalarsubquery()]
:  :  +- Aggregate [min(a) AS min(a)]
:  :     +- Relation [a, b, c]
:  +- OneRowRelation
+- Aggregate [sum(b) AS sum(b), avg(cast(c as double)) AS avg(c)]
   +- Relation [a, b, c]
```
can be optimized to:
```
WithCTE
:- CTERelationDef 0
:  +- Project [named_struct(min(a), min(a), sum(b), sum(b), avg(c), avg(c)) AS mergedValue]
:     +- Aggregate [min(a) AS min(a), sum(b) AS sum(b), avg(cast(c as double)) AS avg(c)]
:        +- Relation [a, b, c]
+- Join Inner
   :- Project [scalar-subquery [].min(a) AS scalarsubquery()]
   :  :  +- CTERelationRef 0
   :  +- OneRowRelation
   +- Project [scalar-subquery [].sum(b) AS sum(b), scalar-subquery [].avg(c) AS avg(c)]
      :  :- CTERelationRef 0
      :  +- CTERelationRef 0
      +- OneRowRelation
```

### Why are the changes needed?

To improve plan merging logic to further reduce redundant IO.

Please also note that TPCDS q28 and q88 contain non-grouping aggregates, but this PR can't deal with them yet. Those queries will improve once [SPARK-40193](https://issues.apache.org/jira/browse/SPARK-40193) / #37630 lands in Spark.

### Does this PR introduce _any_ user-facing change?

Yes.

### How was this patch tested?

Existing and new UTs.

### Was this patch authored or co-authored using generative AI tooling?

No.

Closes #53019 from peter-toth/SPARK-44571-merge-subplans-with-one-row-result.

Authored-by: Peter Toth <peter.toth@gmail.com>
Signed-off-by: Peter Toth <peter.toth@gmail.com>

Author: peter-toth

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala renamed to sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala

@@ -20,9 +20,9 @@ package org.apache.spark.sql.catalyst.optimizer
 import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.logical.{CTERelationDef, CTERelationRef, LogicalPlan, Project, Subquery, WithCTE}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, CTERelationDef, CTERelationRef, LeafNode, LogicalPlan, OneRowRelation, Project, Subquery, WithCTE}
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.catalyst.trees.TreePattern.{SCALAR_SUBQUERY, SCALAR_SUBQUERY_REFERENCE, TreePattern}
+import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, NO_GROUPING_AGGREGATE_REFERENCE, SCALAR_SUBQUERY, SCALAR_SUBQUERY_REFERENCE, TreePattern}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.DataType
 
@@ -100,7 +100,7 @@ import org.apache.spark.sql.types.DataType
  * :  +- ReusedSubquery Subquery scalar-subquery#242, [id=#125]
  * +- *(1) Scan OneRowRelation[]
  */
-object MergeScalarSubqueries extends Rule[LogicalPlan] {
+object MergeSubplans extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = {
     plan match {
       // Subquery reuse needs to be enabled for this optimization.
@@ -117,26 +117,24 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
   }
 
   private def extractCommonScalarSubqueries(plan: LogicalPlan) = {
-    // Collect `ScalarSubquery` plans by level into `PlanMerger`s and insert references in place of
-    // `ScalarSubquery`s.
+    // Collect subplans by level into `PlanMerger`s and insert references in place of them.
     val planMergers = ArrayBuffer.empty[PlanMerger]
-    val planWithReferences = insertReferences(plan, planMergers)._1
+    val planWithReferences = insertReferences(plan, true, planMergers)._1
 
     // Traverse level by level and convert merged plans to `CTERelationDef`s and keep non-merged
     // ones. While traversing replace references in plans back to `CTERelationRef`s or to original
-    // `ScalarSubquery`s. This is safe as a subquery plan at a level can reference only lower level
-    // other subqueries.
-    val subqueryPlansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
+    // plans. This is safe as a subplan at a level can reference only lower level ot other subplans.
+    val subplansByLevel = ArrayBuffer.empty[IndexedSeq[LogicalPlan]]
     planMergers.foreach { planMerger =>
       val mergedPlans = planMerger.mergedPlans()
-      subqueryPlansByLevel += mergedPlans.map { mergedPlan =>
-        val planWithoutReferences = if (subqueryPlansByLevel.isEmpty) {
+      subplansByLevel += mergedPlans.map { mergedPlan =>
+        val planWithoutReferences = if (subplansByLevel.isEmpty) {
           // Level 0 plans can't contain references
           mergedPlan.plan
         } else {
-          removeReferences(mergedPlan.plan, subqueryPlansByLevel)
+          removeReferences(mergedPlan.plan, subplansByLevel)
         }
-        if (mergedPlan.merged && mergedPlan.plan.output.size > 1) {
+        if (mergedPlan.merged) {
           CTERelationDef(
             Project(
               Seq(Alias(
@@ -151,38 +149,42 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
       }
     }
 
-    // Replace references back to `CTERelationRef`s or to original `ScalarSubquery`s in the main
-    // plan.
-    val newPlan = removeReferences(planWithReferences, subqueryPlansByLevel)
+    // Replace references back to `CTERelationRef`s or to original subplans.
+    val newPlan = removeReferences(planWithReferences, subplansByLevel)
 
     // Add `CTERelationDef`s to the plan.
-    val subqueryCTEs = subqueryPlansByLevel.flatMap(_.collect { case cte: CTERelationDef => cte })
-    if (subqueryCTEs.nonEmpty) {
-      WithCTE(newPlan, subqueryCTEs.toSeq)
+    val subplanCTEs = subplansByLevel.flatMap(_.collect { case cte: CTERelationDef => cte })
+    if (subplanCTEs.nonEmpty) {
+      WithCTE(newPlan, subplanCTEs.toSeq)
     } else {
       newPlan
     }
   }
 
-  // First traversal inserts `ScalarSubqueryReference`s to the plan and tries to merge subquery
-  // plans by each level.
+  // First traversal inserts `ScalarSubqueryReference`s and `NoGroupingAggregateReference`s to the
+  // plan and tries to merge subplans by each level. Levels are separated eiter by scalar subqueries
+  // or by non-grouping aggregate nodes. Nodes with the same level make sense to try merging.
   private def insertReferences(
       plan: LogicalPlan,
+      root: Boolean,
       planMergers: ArrayBuffer[PlanMerger]): (LogicalPlan, Int) = {
-    // The level of a subquery plan is maximum level of its inner subqueries + 1 or 0 if it has no
-    // inner subqueries.
-    var maxLevel = 0
-    val planWithReferences =
-      plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY)) {
+    if (!plan.containsAnyPattern(AGGREGATE, SCALAR_SUBQUERY)) {
+      return (plan, 0)
+    }
+
+    // Calculate the level propagated from subquery plans, which is the maximum level of the
+    // subqueries of the node + 1 or 0 if the node has no subqueries.
+    var levelFromSubqueries = 0
+    val nodeSubqueriesWithReferences =
+      plan.transformExpressionsWithPruning(_.containsPattern(SCALAR_SUBQUERY)) {
         case s: ScalarSubquery if !s.isCorrelated && s.deterministic =>
-          val (planWithReferences, level) = insertReferences(s.plan, planMergers)
+          val (planWithReferences, level) = insertReferences(s.plan, true, planMergers)
 
-          while (level >= planMergers.size) planMergers += new PlanMerger()
           // The subquery could contain a hint that is not propagated once we merge it, but as a
           // non-correlated scalar subquery won't be turned into a Join the loss of hints is fine.
-          val mergeResult = planMergers(level).merge(planWithReferences)
+          val mergeResult = getPlanMerger(planMergers, level).merge(planWithReferences, true)
 
-          maxLevel = maxLevel.max(level + 1)
+          levelFromSubqueries = levelFromSubqueries.max(level + 1)
 
           val mergedOutput = mergeResult.outputMap(planWithReferences.output.head)
           val outputIndex =
@@ -195,26 +197,96 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] {
             s.exprId)
         case o => o
       }
-    (planWithReferences, maxLevel)
+
+    // Calculate the level of the node, which is the maximum of the above calculated level
+    // propagated from subqueries and the level propagated from child nodes.
+    val (planWithReferences, level) = nodeSubqueriesWithReferences match {
+      case a: Aggregate if !root && a.groupingExpressions.isEmpty =>
+        val (childWithReferences, levelFromChild) = insertReferences(a.child, false, planMergers)
+        val aggregateWithReferences = a.withNewChildren(Seq(childWithReferences))
+
+        // Level is the maximum of the level from subqueries and the level from child.
+        val level = levelFromChild.max(levelFromSubqueries)
+
+        val mergeResult = getPlanMerger(planMergers, level).merge(aggregateWithReferences, false)
+
+        val mergedOutput = aggregateWithReferences.output.map(mergeResult.outputMap)
+        val outputIndices =
+          mergedOutput.map(a => mergeResult.mergedPlan.plan.output.indexWhere(_.exprId == a.exprId))
+        val aggregateReference = NonGroupingAggregateReference(
+          level,
+          mergeResult.mergedPlanIndex,
+          outputIndices,
+          a.output
+        )
+
+        // This is a non-grouping aggregate node so propagate the level of the node + 1 to its
+        // parent
+        (aggregateReference, level + 1)
+      case o =>
+        val (newChildren, levels) = o.children.map(insertReferences(_, false, planMergers)).unzip
+        // Level is the maximum of the level from subqueries and the level from the children.
+        (o.withNewChildren(newChildren), (levelFromSubqueries +: levels).max)
+    }
+
+    (planWithReferences, level)
+  }
+
+  private def getPlanMerger(planMergers: ArrayBuffer[PlanMerger], level: Int) = {
+    while (level >= planMergers.size) planMergers += new PlanMerger()
+    planMergers(level)
   }
 
-  // Second traversal replaces `ScalarSubqueryReference`s to either
-  // `GetStructField(ScalarSubquery(CTERelationRef to the merged plan)` if the plan is merged from
-  // multiple subqueries or `ScalarSubquery(original plan)` if it isn't.
+  // Second traversal replaces:
+  // - a `ScalarSubqueryReference` either to
+  //   `GetStructField(ScalarSubquery(CTERelationRef to the merged plan), merged output index)` if
+  //   the plan is merged from multiple subqueries or to `ScalarSubquery(original plan)` if it
+  //   isn't.
+  // - a `NoGroupingAggregateReference` either to
+  //   ```
+  //     Project(
+  //       Seq(
+  //         GetStructField(
+  //           ScalarSubquery(CTERelationRef to the merged plan),
+  //           merged output index 1),
+  //         GetStructField(
+  //           ScalarSubquery(CTERelationRef to the merged plan),
+  //           merged output index 2),
+  //         ...),
+  //       OneRowRelation)
+  //   ```
+  //   if the plan is merged from multiple subqueries or to `original plan` if it isn't.
   private def removeReferences(
       plan: LogicalPlan,
-      subqueryPlansByLevel: ArrayBuffer[IndexedSeq[LogicalPlan]]) = {
-    plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY_REFERENCE)) {
-      case ssr: ScalarSubqueryReference =>
-        subqueryPlansByLevel(ssr.level)(ssr.mergedPlanIndex) match {
+      subplansByLevel: ArrayBuffer[IndexedSeq[LogicalPlan]]) = {
+    plan.transformUpWithPruning(
+        _.containsAnyPattern(NO_GROUPING_AGGREGATE_REFERENCE, SCALAR_SUBQUERY_REFERENCE)) {
+      case ngar: NonGroupingAggregateReference =>
+        subplansByLevel(ngar.level)(ngar.mergedPlanIndex) match {
           case cte: CTERelationDef =>
-            GetStructField(
-              ScalarSubquery(
-                CTERelationRef(cte.id, _resolved = true, cte.output, cte.isStreaming),
-                exprId = ssr.exprId),
-              ssr.outputIndex)
-          case o => ScalarSubquery(o, exprId = ssr.exprId)
+            val projectList = ngar.outputIndices.zip(ngar.output).map { case (i, a) =>
+              Alias(
+                GetStructField(
+                  ScalarSubquery(
+                    CTERelationRef(cte.id, _resolved = true, cte.output, cte.isStreaming)),
+                  i),
+                a.name)(a.exprId)
+            }
+            Project(projectList, OneRowRelation())
+          case o => o
         }
+      case o => o.transformExpressionsUpWithPruning(_.containsPattern(SCALAR_SUBQUERY_REFERENCE)) {
+        case ssr: ScalarSubqueryReference =>
+          subplansByLevel(ssr.level)(ssr.mergedPlanIndex) match {
+            case cte: CTERelationDef =>
+              GetStructField(
+                ScalarSubquery(
+                  CTERelationRef(cte.id, _resolved = true, cte.output, cte.isStreaming),
+                  exprId = ssr.exprId),
+                ssr.outputIndex)
+            case o => ScalarSubquery(o, exprId = ssr.exprId)
+          }
+      }
     }
   }
 }
@@ -233,9 +305,26 @@ case class ScalarSubqueryReference(
     level: Int,
     mergedPlanIndex: Int,
     outputIndex: Int,
-    dataType: DataType,
+    override val dataType: DataType,
     exprId: ExprId) extends LeafExpression with Unevaluable {
   override def nullable: Boolean = true
 
   final override val nodePatterns: Seq[TreePattern] = Seq(SCALAR_SUBQUERY_REFERENCE)
 }
+
+/**
+ * Temporal reference to a non-grouping aggregate which is added to a `PlanMerger`.
+ *
+ * @param level The level of the replaced aggregate. It defines the `PlanMerger` instance into which
+ *              the aggregate is merged.
+ * @param mergedPlanIndex The index of the merged plan in the `PlanMerger`.
+ * @param outputIndices The indices of the output attributes of the merged plan.
+ * @param output The output of original aggregate.
+ */
+case class NonGroupingAggregateReference(
+    level: Int,
+    mergedPlanIndex: Int,
+    outputIndices: Seq[Int],
+    override val output: Seq[Attribute]) extends LeafNode {
+  final override val nodePatterns: Seq[TreePattern] = Seq(NO_GROUPING_AGGREGATE_REFERENCE)
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala

@@ -58,8 +58,8 @@ case class MergedPlan(plan: LogicalPlan, merged: Boolean)
  * 2. Merge a new plan with a cached plan by combining their outputs
  *
  * The merging process preserves semantic equivalence while combining outputs from multiple
- * plans into a single plan. This is primarily used by [[MergeScalarSubqueries]] to deduplicate
- * scalar subquery execution.
+ * plans into a single plan. This is primarily used by [[MergeSubplans]] to deduplicate subplan
+ * execution.
  *
  * Supported plan types for merging:
  * - [[Project]]: Merges project lists
@@ -88,16 +88,21 @@ class PlanMerger {
    * 3. If no merge is possible, add as a new cache entry
    *
    * @param plan The logical plan to merge or cache.
+   * @param subqueryPlan If the logical plan is a subquery plan.
    * @return A [[MergeResult]] containing:
    *         - The merged/cached plan to use
    *         - Its index in the cache
    *         - An attribute mapping for rewriting expressions
    */
-  def merge(plan: LogicalPlan): MergeResult = {
+  def merge(plan: LogicalPlan, subqueryPlan: Boolean): MergeResult = {
     cache.zipWithIndex.collectFirst(Function.unlift {
       case (mp, i) =>
         checkIdenticalPlans(plan, mp.plan).map { outputMap =>
-          val newMergePlan = MergedPlan(mp.plan, true)
+          // Identical subquery expression plans are not marked as `merged` as the
+          // `ReusedSubqueryExec` rule can handle them without extracting the plans to CTEs.
+          // But, when a non-subquery subplan is identical to a cached plan we need to mark the plan
+          // `merged` and so extract it to a CTE later.
+          val newMergePlan = MergedPlan(mp.plan, cache(i).merged || !subqueryPlan)
           cache(i) = newMergePlan
           MergeResult(newMergePlan, i, outputMap)
         }.orElse {

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala

@@ -150,6 +150,7 @@ object TreePattern extends Enumeration  {
   val LOCAL_RELATION: Value = Value
   val LOGICAL_QUERY_STAGE: Value = Value
   val NATURAL_LIKE_JOIN: Value = Value
+  val NO_GROUPING_AGGREGATE_REFERENCE: Value = Value
   val OFFSET: Value = Value
   val OUTER_JOIN: Value = Value
   val PARAMETERIZED_QUERY: Value = Value