optimize partial pushdown: extend attribute when child already contains it

zml1206 · zml1206 · commit 76d9b6d51a8c · 2025-01-08T19:31:10.000+08:00
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -1132,7 +1132,6 @@ object CollapseProject extends Rule[LogicalPlan] with AliasHelper {
         r.copy(child = p.copy(projectList = buildCleanedProjectList(l1, p.projectList)))
       case Project(l1, s @ Sample(_, _, _, _, p2 @ Project(l2, _))) if isRenaming(l1, l2) =>
         s.copy(child = p2.copy(projectList = buildCleanedProjectList(l1, p2.projectList)))
-      case p @ Project(_, child) if child.output == p.output => child
     }
   }
 
@@ -1812,7 +1811,10 @@ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelpe
       maxIterationsSetting = SQLConf.OPTIMIZER_MAX_ITERATIONS.key)
     val batches = Seq(
       Batch("RewriteWithExpression", fixedPoint, RewriteWithExpression),
-      Batch("Optimize after RewriteWithExpression", fixedPoint, CollapseProject, ColumnPruning)
+      Batch("Optimize after RewriteWithExpression", fixedPoint,
+        CollapseProject,
+        ColumnPruning,
+        RemoveNoopOperators)
     )
   }
 
@@ -1830,9 +1832,9 @@ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelpe
       val aliasMap = getAliasMap(project)
       var newProjectList = fields ++ getWithAttributes(condition)
       val newCondition = rewriteCondition(condition, aliasMap)
-      val rewriteAlias = getWithAlias(newCondition).toSet
+      val exprIdSet = getWithAlias(newCondition).map(_.exprId).toSet
       newProjectList = newProjectList.map {
-        case a: Alias if rewriteAlias.contains(a) => a.toAttribute
+        case a: Alias if exprIdSet.contains(a.exprId) => a.toAttribute
         case e => e
       }
       project.copy(child = Filter(newCondition, grandChild), projectList = newProjectList)
@@ -1859,8 +1861,8 @@ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelpe
         // propagate the attributes directly need add the groupingExpressions may cause regression.
         // So Aggregate only need inline common expression from parent for original project
         // inheritance.
-        val newAggregateExpressions =
-          aggregate.aggregateExpressions ++ getWithAlias(pushDown.reduce(And))
+        val newAggregateExpressions = aggregate.aggregateExpressions ++
+          getWithAlias(pushDown.reduce(And)).map(replaceAliasButKeepName(_, aliasMap))
         val replaced = removeOriginAlias(rewriteCondition(pushDown.reduce(And), aliasMap))
         val newAggregate = aggregate.copy(child = Filter(replaced, aggregate.child),
           aggregateExpressions = newAggregateExpressions)
@@ -2038,8 +2040,7 @@ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelpe
       cond: Expression,
       aliasMap: AttributeMap[Alias]): Expression = {
     if (!SQLConf.get.getConf(SQLConf.ALWAYS_INLINE_COMMON_EXPR)) {
-      val replaceWithMap = cond.collect {case a: Attribute => a }
-        .distinct
+      val replaceWithMap = cond.references.toSeq
         .filter(attr => aliasMap.contains(attr))
         .map(attr => attr -> aliasMap(attr))
         .filter(m => !CollapseProject.isCheap(m._2))
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala
@@ -90,13 +90,11 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
     // 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.
     assert(p.output.length <= newPlan.output.length)
-    newPlan.output.diff(p.output).forall(_.name.startsWith("_"))
-    def hasOriginAlias(expr: Expression): Boolean = {
-      expr match {
-        case w: With =>
-          if (w.defs.exists(_.originAlias.nonEmpty)) true else false
-        case e => e.children.exists(hasOriginAlias)
-      }
+
+    def hasOriginAlias(expr: Expression): Boolean = expr match {
+      case w: With =>
+        if (w.defs.exists(_.originAlias.nonEmpty)) true else false
+      case e => e.children.exists(hasOriginAlias)
     }
     // If this iteration contains attribute that require propagate, the column cannot be pruning.
     val needPropagate = p.expressions.exists(hasOriginAlias)
@@ -148,6 +146,11 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
               // TODO: we should calculate the ref count and also inline the common expression
               //       if it's ref count is 1.
               refToExpr(id) = child
+            } else if (originAlias.nonEmpty &&
+              inputPlans.head.output.contains(originAlias.get.toAttribute)) {
+              // originAlias only exists in Project or Filter. If the child already contains this
+              // attribute, extend it.
+              refToExpr(id) = originAlias.get.toAttribute
             } else {
               val commonExprs = commonExprsPerChild(childPlanIndex)
               val existingCommonExpr = commonExprs.find(_._2 == id.id)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala
@@ -52,7 +52,8 @@ class FilterPushdownSuite extends PlanTest {
 
     val batches = Batch("ColumnPruning and CollapseProject", FixedPoint(10),
       ColumnPruning,
-      CollapseProject) :: Nil
+      CollapseProject,
+      RemoveNoopOperators) :: Nil
   }
 
   val attrA = $"a".int
@@ -1556,29 +1557,46 @@ class FilterPushdownSuite extends PlanTest {
   }
 
   test("SPARK-50589: avoid extra expression duplication when push filter") {
+//    withSQLConf(SQLConf.USE_COMMON_EXPR_ID_FOR_ALIAS.key -> "false") {
+//      // through project
+//      val originalQuery1 = testRelation
+//        .select($"a" + $"b" as "add", $"a" - $"b" as "sub")
+//        .where($"add" < 10 && $"add" + $"add" > 10 && $"sub" > 0)
+//        .analyze
+//      val optimized1 = Optimize.execute(originalQuery1)
+//      comparePlans(optimized1, originalQuery1)
+//
+//      // through aggregate
+//      val originalQuery2 = testRelation
+//        .groupBy($"a")($"a", $"a" + $"a" as "add", abs($"a") as "abs", count(1) as "ct")
+//        .where($"add" < 10 && $"add" + $"add" > 10 && $"abs" > 5)
+//      val optimized2 = Optimize.execute(originalQuery2.analyze)
+//      val correctAnswer2 = testRelation
+//        .select($"a", $"a" + $"a" as "_common_expr_0")
+//        .where($"_common_expr_0" < 10 &&
+//          $"_common_expr_0" + $"_common_expr_0" > 10 &&
+//          abs($"a") > 5)
+//        .select($"a")
+//        .groupBy($"a")($"a", $"a" + $"a" as "add", abs($"a") as "abs", count(1) as "ct")
+//        .analyze
+//      comparePlans(optimized2, correctAnswer2)
+//    }
     withSQLConf(SQLConf.USE_COMMON_EXPR_ID_FOR_ALIAS.key -> "false") {
-      // through project
-      val originalQuery1 = testRelation
-        .select($"a" + $"b" as "add", $"a" - $"b" as "sub")
-        .where($"add" < 10 && $"add" + $"add" > 10 && $"sub" > 0)
-        .analyze
-      val optimized1 = Optimize.execute(originalQuery1)
-      comparePlans(optimized1, originalQuery1)
-
-      // through aggregate
-      val originalQuery2 = testRelation
-        .groupBy($"a")($"a", $"a" + $"a" as "add", abs($"a") as "abs", count(1) as "ct")
-        .where($"add" < 10 && $"add" + $"add" > 10 && $"abs" > 5)
-      val optimized2 = Optimize.execute(originalQuery2.analyze)
-      val correctAnswer2 = testRelation
+      // partial push down
+      val originalQuery3 = testRelation
+        .groupBy($"a")($"a", count(1) as "ct")
+        .select($"a" + $"a" as "add", $"ct")
+        .where($"add" + $"add" > 10 && $"add" > $"ct")
+      val optimized3 = Optimize.execute(originalQuery3.analyze)
+      val correctAnswer3 = testRelation
         .select($"a", $"a" + $"a" as "_common_expr_0")
-        .where($"_common_expr_0" < 10 &&
-          $"_common_expr_0" + $"_common_expr_0" > 10 &&
-          abs($"a") > 5)
+        .where($"_common_expr_0" + $"_common_expr_0" > 10)
         .select($"a")
-        .groupBy($"a")($"a", $"a" + $"a" as "add", abs($"a") as "abs", count(1) as "ct")
+        .groupBy($"a")(count(1) as "ct", $"a" + $"a" as "add")
+        .where($"add" > $"ct")
+        .select($"add", $"ct")
         .analyze
-      comparePlans(optimized2, correctAnswer2)
-    }
+      comparePlans(optimized3, correctAnswer3)
     }
+  }
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InferFiltersFromConstraintsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InferFiltersFromConstraintsSuite.scala
@@ -161,10 +161,9 @@ class InferFiltersFromConstraintsSuite extends PlanTest {
       .join(t2, Inner, Some("t.a".attr === "t2.a".attr && "t.int_col".attr === "t2.a".attr))
       .analyze
     val correctAnswer = t1
-      .select($"a", Coalesce(Seq($"a", $"b")) as "int_col")
+      .select($"a", Coalesce(Seq($"a", $"b")) as "int_col").as("t")
       .where(IsNotNull($"a") && IsNotNull($"int_col") &&
         $"a" === $"int_col")
-      .select($"a", $"int_col").as("t")
       .join(t2.where(IsNotNull($"a")), Inner,
         Some("t.a".attr === "t2.a".attr && "t.int_col".attr === "t2.a".attr))
       .analyze
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala
@@ -28,13 +28,9 @@ import org.apache.spark.sql.catalyst.rules.RuleExecutor
 class RewriteWithExpressionSuite extends PlanTest {
 
   object Optimizer extends RuleExecutor[LogicalPlan] {
-    val batches =
-      Batch("Rewrite With expression", FixedPoint(5),
-        PullOutGroupingExpressions,
-        RewriteWithExpression) ::
-      Batch("Optimize after RewriteWithExpression", FixedPoint(5),
-        CollapseProject,
-        ColumnPruning) :: Nil
+    val batches = Batch("Rewrite With expression", FixedPoint(5),
+      PullOutGroupingExpressions,
+      RewriteWithExpression) :: Nil
   }
 
   private val testRelation = LocalRelation($"a".int, $"b".int)
@@ -72,7 +68,7 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(plan),
       testRelation
-        .select((a + a).as("_common_expr_0"))
+        .select((testRelation.output :+ (a + a).as("_common_expr_0")): _*)
         .select(($"_common_expr_0" * $"_common_expr_0").as("col"))
         .analyze
     )
@@ -90,8 +86,8 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(testRelation.select(outerExpr.as("col"))),
       testRelation
-        .select(b, (a + a).as("_common_expr_0"))
-        .select(($"_common_expr_0" + $"_common_expr_0" + b).as("_common_expr_1"))
+        .select(star(), (a + a).as("_common_expr_0"))
+        .select(a, b, ($"_common_expr_0" + $"_common_expr_0" + b).as("_common_expr_1"))
         .select(($"_common_expr_1" * $"_common_expr_1").as("col"))
         .analyze
     )
@@ -109,7 +105,8 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(testRelation.select(outerExpr.as("col"))),
       testRelation
-        .select((b + b).as("_common_expr_1"), (a + a).as("_common_expr_0"))
+        .select(star(), (b + b).as("_common_expr_1"))
+        .select(star(), (a + a).as("_common_expr_0"))
         .select(finalExpr.as("col"))
         .analyze
     )
@@ -130,10 +127,10 @@ class RewriteWithExpressionSuite extends PlanTest {
       Optimizer.execute(testRelation.select(outerExpr1.as("col"))),
       testRelation
         // The first Project contains the common expression of the outer With
-        .select(a, rewrittenOuterExpr)
+        .select(star(), rewrittenOuterExpr)
         // The second Project contains the common expression of the inner With, which references
         // the common expression of the outer With.
-        .select($"_common_expr_0", (a + a + $"_common_expr_0").as("_common_expr_1"))
+        .select(star(), (a + a + $"_common_expr_0").as("_common_expr_1"))
         // The final Project contains the final result expression, which references both common
         // expressions.
         .select(($"_common_expr_0" + ($"_common_expr_1" + $"_common_expr_1")).as("col"))
@@ -148,7 +145,11 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(testRelation.select(outerExpr2.as("col"))),
       testRelation
-        .select(rewrittenOuterExpr, (a + a).as("_common_expr_2"))
+        // The first Project contains the common expression of the outer With
+        .select(star(), rewrittenOuterExpr)
+        // The second Project contains the common expression of the inner With, which does not
+        // reference the common expression of the outer With.
+        .select(star(), (a + a).as("_common_expr_2"))
         // The final Project contains the final result expression, which references both common
         // expressions.
         .select(($"_common_expr_0" +
@@ -243,7 +244,7 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(plan2),
       testRelation
-        .select(a, (a + a).as("_common_expr_0"))
+        .select((testRelation.output :+ (a + a).as("_common_expr_0")): _*)
         .select(Coalesce(Seq(($"_common_expr_0" * $"_common_expr_0"), a)).as("col"))
         .analyze
     )
@@ -267,8 +268,24 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(plan),
       testRelation
-        .select(a, (a + 1).as("_common_expr_0"))
-        .select(Seq(
+        .select(testRelation.output :+ (a + 1).as("_common_expr_0"): _*)
+        .select(testRelation.output :+
+          ($"_common_expr_0" * $"_common_expr_0").as("_groupingexpression"): _*)
+        .select(testRelation.output ++ Seq($"_groupingexpression",
+          (a + 1).as("_common_expr_1")): _*)
+        .groupBy($"_groupingexpression")(
+          $"_groupingexpression",
+          count($"_common_expr_1" * $"_common_expr_1" - 3).as("_aggregateexpression")
+        )
+        .select(($"_groupingexpression" + 2).as("col1"), $"_aggregateexpression".as("col2"))
+        .analyze
+    )
+    // Running CollapseProject after the rule cleans up the unnecessary projections.
+    comparePlans(
+      CollapseProject(Optimizer.execute(plan)),
+      testRelation
+        .select(testRelation.output :+ (a + 1).as("_common_expr_0"): _*)
+        .select(testRelation.output ++ Seq(
           ($"_common_expr_0" * $"_common_expr_0").as("_groupingexpression"),
           (a + 1).as("_common_expr_1")): _*)
         .groupBy($"_groupingexpression")(
@@ -295,9 +312,9 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(plan),
       testRelation
-        .select(a, (b + 2).as("_common_expr_0"))
-        .groupBy(a)(a, max($"_common_expr_0" * $"_common_expr_0").as("_aggregateexpression"),
-          (a + 1).as("_common_expr_1"))
+        .select(testRelation.output :+ (b + 2).as("_common_expr_0"): _*)
+        .groupBy(a)(a, max($"_common_expr_0" * $"_common_expr_0").as("_aggregateexpression"))
+        .select(a, $"_aggregateexpression", (a + 1).as("_common_expr_1"))
         .select(
           (a + 3).as("col1"),
           ($"_common_expr_1" * $"_common_expr_1").as("col2"),
@@ -319,7 +336,6 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(plan),
       testRelation
-        .select(a)
         .groupBy(a)(a, count(a - 1).as("_aggregateexpression"))
         .select(
           (a - 1).as("col1"),
@@ -355,9 +371,9 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(plan),
       testRelation
-        .select(a, (a + 1).as("_common_expr_0"))
-        .groupBy(a)(max($"_common_expr_0" * $"_common_expr_0").as("_aggregateexpression"),
-          (a - 1).as("_common_expr_1"))
+        .select(testRelation.output :+ (a + 1).as("_common_expr_0"): _*)
+        .groupBy(a)(a, max($"_common_expr_0" * $"_common_expr_0").as("_aggregateexpression"))
+        .select($"a", $"_aggregateexpression", (a - 1).as("_common_expr_1"))
         .select(($"_common_expr_1" * $"_aggregateexpression" + $"_common_expr_1").as("col"))
         .analyze
     )
@@ -388,20 +404,22 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(plan),
       testRelation
-        .select(a, (b + 2).as("_common_expr_0"))
+        .select(a, b, (b + 2).as("_common_expr_0"))
         .window(
           Seq(windowExpr(count(a), windowSpec(Seq($"_common_expr_0" * $"_common_expr_0"), Nil,
             frame)).as("col2")),
           Seq($"_common_expr_0" * $"_common_expr_0"),
           Nil
         )
-        .select(a, $"col2", (a + 1).as("_common_expr_1"))
+        .select(a, b, $"col2")
+        .select(a, b, $"col2", (a + 1).as("_common_expr_1"))
         .window(
           Seq(windowExpr(sum($"_common_expr_1" * $"_common_expr_1"),
             windowSpec(Seq(a), Nil, frame)).as("col3")),
           Seq(a),
           Nil
         )
+        .select(a, b, $"col2", $"col3")
         .select((a - 1).as("col1"), $"col2", $"col3")
         .analyze
     )
@@ -420,7 +438,8 @@ class RewriteWithExpressionSuite extends PlanTest {
       testRelation
         .select(a)
         .window(Seq(winExpr.as("_we0")), Seq(a), Nil)
-        .select(($"_we0" * $"_we0").as("col"))
+        .select(a, $"_we0", ($"_we0" * $"_we0").as("col"))
+        .select($"col")
         .analyze
     )
   }
@@ -445,7 +464,7 @@ class RewriteWithExpressionSuite extends PlanTest {
     val plan = testRelation.having($"b")(avg("a").as("a"))(expr).analyze
     comparePlans(
       Optimizer.execute(plan),
-      testRelation.select($"b").groupBy($"b")(avg("a").as("a")).where($"a" === 1).analyze
+      testRelation.groupBy($"b")(avg("a").as("a")).where($"a" === 1).analyze
     )
   }
 
@@ -459,7 +478,7 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(plan),
       testRelation
-        .select((a + a).as("_common_expr_0"))
+        .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"))
