fix ut

zml1206 · Jan 6, 2025 · 6b55ff2 · 6b55ff2
1 parent 906aff4
commit 6b55ff2
Show file tree

Hide file tree

Showing 6 changed files with 434 additions and 399 deletions.
diff --git a/...t/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala b/...t/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala
@@ -105,8 +105,7 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(testRelation.select(outerExpr.as("col"))),
       testRelation
-        .select(star(), (b + b).as("_common_expr_1"))
-        .select(star(), (a + a).as("_common_expr_0"))
+        .select(star(), (b + b).as("_common_expr_1"), (a + a).as("_common_expr_0"))
         .select(finalExpr.as("col"))
         .analyze
     )
@@ -145,11 +144,7 @@ class RewriteWithExpressionSuite extends PlanTest {
     comparePlans(
       Optimizer.execute(testRelation.select(outerExpr2.as("col"))),
       testRelation
-        // 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"))
+        .select(star(), rewrittenOuterExpr, (a + a).as("_common_expr_2"))
         // The final Project contains the final result expression, which references both common
         // expressions.
         .select(($"_common_expr_0" +
@@ -267,22 +262,6 @@ class RewriteWithExpressionSuite extends PlanTest {
     )
     comparePlans(
       Optimizer.execute(plan),
-      testRelation
-        .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(
@@ -313,8 +292,8 @@ class RewriteWithExpressionSuite extends PlanTest {
       Optimizer.execute(plan),
       testRelation
         .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"))
+        .groupBy(a)(a, max($"_common_expr_0" * $"_common_expr_0").as("_aggregateexpression"),
+          (a + 1).as("_common_expr_1"))
         .select(
           (a + 3).as("col1"),
           ($"_common_expr_1" * $"_common_expr_1").as("col2"),
@@ -372,8 +351,8 @@ class RewriteWithExpressionSuite extends PlanTest {
       Optimizer.execute(plan),
       testRelation
         .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"))
+        .groupBy(a)(a, max($"_common_expr_0" * $"_common_expr_0").as("_aggregateexpression"),
+          (a - 1).as("_common_expr_1"))
         .select(($"_common_expr_1" * $"_aggregateexpression" + $"_common_expr_1").as("col"))
         .analyze
     )
@@ -411,15 +390,13 @@ class RewriteWithExpressionSuite extends PlanTest {
           Seq($"_common_expr_0" * $"_common_expr_0"),
           Nil
         )
-        .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
     )
@@ -438,8 +415,7 @@ class RewriteWithExpressionSuite extends PlanTest {
       testRelation
         .select(a)
         .window(Seq(winExpr.as("_we0")), Seq(a), Nil)
-        .select(a, $"_we0", ($"_we0" * $"_we0").as("col"))
-        .select($"col")
+        .select(($"_we0" * $"_we0").as("col"))
         .analyze
     )
   }

diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_count_if.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_count_if.explain
@@ -1,4 +1,3 @@
-Project [_aggregateexpression#0L AS count_if((a > 0))#0L]
-+- Aggregate [count(if ((_common_expr_0#0 = false)) null else _common_expr_0#0) AS _aggregateexpression#0L]
-   +- Project [id#0L, a#0, b#0, d#0, e#0, f#0, g#0, (a#0 > 0) AS _common_expr_0#0]
-      +- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]
+Aggregate [count(if ((_common_expr_0#0 = false)) null else _common_expr_0#0) AS count_if((a > 0))#0L]
++- Project [id#0L, a#0, b#0, d#0, e#0, f#0, g#0, (a#0 > 0) AS _common_expr_0#0]
+   +- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]