[SPARK-50679][SQL] Duplicated common expressions in different With should be projected only once

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

Sometimes we may need to duplicate the `With` expression with some shared common expressions. This PR improves the `With` rewriting rule to recognize it and only project the duplicated common expressions once.

### Why are the changes needed?

Produce better plans.

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

no

### How was this patch tested?

a new test

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

no

Closes apache#49303 from cloud-fan/with.

Authored-by: Wenchen Fan <wenchen@databricks.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>

Author: cloud-fan

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

@@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, Plan
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreePattern.{COMMON_EXPR_REF, WITH_EXPRESSION}
 import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.util.Utils
 
 /**
  * Rewrites the `With` expressions by adding a `Project` to pre-evaluate the common expressions, or
@@ -66,11 +67,19 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
   }
 
   private def applyInternal(p: LogicalPlan): LogicalPlan = {
-    val inputPlans = p.children.toArray
+    val inputPlans = p.children
+    val commonExprsPerChild = Array.fill(inputPlans.length)(mutable.ListBuffer.empty[(Alias, Long)])
     var newPlan: LogicalPlan = p.mapExpressions { expr =>
-      rewriteWithExprAndInputPlans(expr, inputPlans)
+      rewriteWithExprAndInputPlans(expr, inputPlans, commonExprsPerChild)
     }
-    newPlan = newPlan.withNewChildren(inputPlans.toIndexedSeq)
+    val newChildren = inputPlans.zip(commonExprsPerChild).map { case (inputPlan, commonExprs) =>
+      if (commonExprs.isEmpty) {
+        inputPlan
+      } else {
+        Project(inputPlan.output ++ commonExprs.map(_._1), inputPlan)
+      }
+    }
+    newPlan = newPlan.withNewChildren(newChildren)
     // Since we add extra Projects with extra columns to pre-evaluate the common expressions,
     // the current operator may have extra columns if it inherits the output columns from its
     // child, and we need to project away the extra columns to keep the plan schema unchanged.
@@ -85,17 +94,19 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
 
   private def rewriteWithExprAndInputPlans(
       e: Expression,
-      inputPlans: Array[LogicalPlan],
+      inputPlans: Seq[LogicalPlan],
+      commonExprsPerChild: Array[mutable.ListBuffer[(Alias, Long)]],
       isNestedWith: Boolean = false): Expression = {
     if (!e.containsPattern(WITH_EXPRESSION)) return e
     e match {
       // Do not handle nested With in one pass. Leave it to the next rule executor batch.
       case w: With if !isNestedWith =>
         // Rewrite nested With expressions first
-        val child = rewriteWithExprAndInputPlans(w.child, inputPlans, isNestedWith = true)
-        val defs = w.defs.map(rewriteWithExprAndInputPlans(_, inputPlans, isNestedWith = true))
+        val child = rewriteWithExprAndInputPlans(
+          w.child, inputPlans, commonExprsPerChild, isNestedWith = true)
+        val defs = w.defs.map(rewriteWithExprAndInputPlans(
+          _, inputPlans, commonExprsPerChild, isNestedWith = true))
         val refToExpr = mutable.HashMap.empty[CommonExpressionId, Expression]
-        val childProjections = Array.fill(inputPlans.length)(mutable.ArrayBuffer.empty[Alias])
 
         defs.zipWithIndex.foreach { case (CommonExpressionDef(child, id), index) =>
           if (id.canonicalized) {
@@ -106,10 +117,10 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
           if (CollapseProject.isCheap(child)) {
             refToExpr(id) = child
           } else {
-            val childProjectionIndex = inputPlans.indexWhere(
+            val childPlanIndex = inputPlans.indexWhere(
               c => child.references.subsetOf(c.outputSet)
             )
-            if (childProjectionIndex == -1) {
+            if (childPlanIndex == -1) {
               // When we cannot rewrite the common expressions, force to inline them so that the
               // query can still run. This can happen if the join condition contains `With` and
               // the common expression references columns from both join sides.
@@ -120,31 +131,33 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
               //       if it's ref count is 1.
               refToExpr(id) = child
             } else {
-              val aliasName = if (SQLConf.get.getConf(SQLConf.USE_COMMON_EXPR_ID_FOR_ALIAS)) {
-                s"_common_expr_${id.id}"
-              } else {
-                s"_common_expr_$index"
-              }
-              val alias = Alias(child, aliasName)()
-              val fakeProj = Project(Seq(alias), inputPlans(childProjectionIndex))
-              if (PlanHelper.specialExpressionsInUnsupportedOperator(fakeProj).nonEmpty) {
-                // We have to inline the common expression if it cannot be put in a Project.
-                refToExpr(id) = child
+              val commonExprs = commonExprsPerChild(childPlanIndex)
+              val existingCommonExpr = commonExprs.find(_._2 == id.id)
+              if (existingCommonExpr.isDefined) {
+                if (Utils.isTesting) {
+                  assert(existingCommonExpr.get._1.child.semanticEquals(child))
+                }
+                refToExpr(id) = existingCommonExpr.get._1.toAttribute
               } else {
-                childProjections(childProjectionIndex) += alias
-                refToExpr(id) = alias.toAttribute
+                val aliasName = if (SQLConf.get.getConf(SQLConf.USE_COMMON_EXPR_ID_FOR_ALIAS)) {
+                  s"_common_expr_${id.id}"
+                } else {
+                  s"_common_expr_$index"
+                }
+                val alias = Alias(child, aliasName)()
+                val fakeProj = Project(Seq(alias), inputPlans(childPlanIndex))
+                if (PlanHelper.specialExpressionsInUnsupportedOperator(fakeProj).nonEmpty) {
+                  // We have to inline the common expression if it cannot be put in a Project.
+                  refToExpr(id) = child
+                } else {
+                  commonExprs.append((alias, id.id))
+                  refToExpr(id) = alias.toAttribute
+                }
               }
             }
           }
         }
 
-        for (i <- inputPlans.indices) {
-          val projectList = childProjections(i)
-          if (projectList.nonEmpty) {
-            inputPlans(i) = Project(inputPlans(i).output ++ projectList, inputPlans(i))
-          }
-        }
-
         child.transformWithPruning(_.containsPattern(COMMON_EXPR_REF)) {
           // `child` may contain nested With and we only replace `CommonExpressionRef` that
           // references common expressions in the current `With`.
@@ -158,7 +171,7 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
 
       case c: ConditionalExpression =>
         val newAlwaysEvaluatedInputs = c.alwaysEvaluatedInputs.map(
-          rewriteWithExprAndInputPlans(_, inputPlans, isNestedWith))
+          rewriteWithExprAndInputPlans(_, inputPlans, commonExprsPerChild, isNestedWith))
         val newExpr = c.withNewAlwaysEvaluatedInputs(newAlwaysEvaluatedInputs)
         // Use transformUp to handle nested With.
         newExpr.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
@@ -171,7 +184,9 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
             }
         }
 
-      case other => other.mapChildren(rewriteWithExprAndInputPlans(_, inputPlans, isNestedWith))
+      case other => other.mapChildren(
+        rewriteWithExprAndInputPlans(_, inputPlans, commonExprsPerChild, isNestedWith)
+      )
     }
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala

@@ -404,17 +404,16 @@ class RewriteWithExpressionSuite extends PlanTest {
       Optimizer.execute(plan),
       testRelation
         .select(a, b, (b + 2).as("_common_expr_0"))
-        .select(a, b, $"_common_expr_0", (b + 2).as("_common_expr_1"))
         .window(
           Seq(windowExpr(count(a), windowSpec(Seq($"_common_expr_0" * $"_common_expr_0"), Nil,
             frame)).as("col2")),
-          Seq($"_common_expr_1" * $"_common_expr_1"),
+          Seq($"_common_expr_0" * $"_common_expr_0"),
           Nil
         )
         .select(a, b, $"col2")
-        .select(a, b, $"col2", (a + 1).as("_common_expr_2"))
+        .select(a, b, $"col2", (a + 1).as("_common_expr_1"))
         .window(
-          Seq(windowExpr(sum($"_common_expr_2" * $"_common_expr_2"),
+          Seq(windowExpr(sum($"_common_expr_1" * $"_common_expr_1"),
             windowSpec(Seq(a), Nil, frame)).as("col3")),
           Seq(a),
           Nil
@@ -467,4 +466,28 @@ class RewriteWithExpressionSuite extends PlanTest {
       testRelation.groupBy($"b")(avg("a").as("a")).where($"a" === 1).analyze
     )
   }
+
+  test("SPARK-50679: duplicated common expressions in different With") {
+    val a = testRelation.output.head
+    val exprDef = CommonExpressionDef(a + a)
+    val exprRef = new CommonExpressionRef(exprDef)
+    val expr1 = With(exprRef * exprRef, Seq(exprDef))
+    val expr2 = With(exprRef - exprRef, Seq(exprDef))
+    val plan = testRelation.select(expr1.as("c1"), expr2.as("c2")).analyze
+    comparePlans(
+      Optimizer.execute(plan),
+      testRelation
+        .select(star(), (a + a).as("_common_expr_0"))
+        .select(
+          ($"_common_expr_0" * $"_common_expr_0").as("c1"),
+          ($"_common_expr_0" - $"_common_expr_0").as("c2"))
+        .analyze
+    )
+
+    val wrongExprDef = CommonExpressionDef(a * a, exprDef.id)
+    val wrongExprRef = new CommonExpressionRef(wrongExprDef)
+    val expr3 = With(wrongExprRef + wrongExprRef, Seq(wrongExprDef))
+    val wrongPlan = testRelation.select(expr1.as("c1"), expr3.as("c3")).analyze
+    intercept[AssertionError](Optimizer.execute(wrongPlan))
+  }
 }