[BugFix] fix project prune rule (StarRocks#13278)

murphyatwork · Nov 16, 2022 · 7e4bfc5 · 7e4bfc5
1 parent f18128b
commit 7e4bfc5
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 4 deletions.
diff --git a/...rc/main/java/com/starrocks/sql/optimizer/rule/transformation/PruneProjectColumnsRule.java b/...rc/main/java/com/starrocks/sql/optimizer/rule/transformation/PruneProjectColumnsRule.java
@@ -5,8 +5,10 @@
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.starrocks.catalog.FunctionSet;
+import com.starrocks.sql.optimizer.ExpressionContext;
 import com.starrocks.sql.optimizer.OptExpression;
 import com.starrocks.sql.optimizer.OptimizerContext;
+import com.starrocks.sql.optimizer.Utils;
 import com.starrocks.sql.optimizer.base.ColumnRefSet;
 import com.starrocks.sql.optimizer.operator.OperatorType;
 import com.starrocks.sql.optimizer.operator.logical.LogicalProjectOperator;
@@ -15,9 +17,11 @@
 import com.starrocks.sql.optimizer.operator.scalar.ColumnRefOperator;
 import com.starrocks.sql.optimizer.operator.scalar.ScalarOperator;
 import com.starrocks.sql.optimizer.rule.RuleType;
+import org.apache.commons.collections.CollectionUtils;
 
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 public class PruneProjectColumnsRule extends TransformationRule {
 
@@ -48,6 +52,20 @@ public List<OptExpression> transform(OptExpression input, OptimizerContext conte
             }
         }));
 
+        if (newMap.isEmpty()) {
+            List<ColumnRefOperator> outputColumns =
+                    projectOperator.getOutputColumns(new ExpressionContext(input.inputAt(0))).getStream().
+                            mapToObj(context.getColumnRefFactory()::getColumnRef).collect(Collectors.toList());
+            if (CollectionUtils.isNotEmpty(outputColumns)) {
+                ColumnRefOperator smallestColumn = Utils.findSmallestColumnRef(outputColumns);
+                ScalarOperator expr = projectOperator.getColumnRefMap().get(smallestColumn);
+                if (!smallestColumn.equals(expr) && !expr.isVariable()) {
+                    newMap.put(smallestColumn, expr);
+                    requiredInputColumns.union(smallestColumn);
+                }
+            }
+        }
+
         // Change the requiredOutputColumns in context
         requiredOutputColumns.union(requiredInputColumns);
 

diff --git a/fe/fe-core/src/test/java/com/starrocks/sql/plan/AggregateTest.java b/fe/fe-core/src/test/java/com/starrocks/sql/plan/AggregateTest.java
@@ -708,10 +708,14 @@ public void testMultiCountDistinct() throws Exception {
     public void testVarianceStddevAnalyze() throws Exception {
         String sql = "select stddev_pop(1222) from (select 1) t;";
         String plan = getFragmentPlan(sql);
-        assertContains(plan, "  1:AGGREGATE (update finalize)\n" +
+        assertContains(plan, "  2:AGGREGATE (update finalize)\n" +
                 "  |  output: stddev_pop(1222)\n" +
-                "  |  group by: ");
-        assertContains(plan, "  0:UNION\n" +
+                "  |  group by: \n" +
+                "  |  \n" +
+                "  1:Project\n" +
+                "  |  <slot 2> : 1\n" +
+                "  |  \n" +
+                "  0:UNION\n" +
                 "     constant exprs: \n" +
                 "         NULL");
     }

diff --git a/fe/fe-core/src/test/java/com/starrocks/sql/plan/NestLoopJoinTest.java b/fe/fe-core/src/test/java/com/starrocks/sql/plan/NestLoopJoinTest.java
@@ -34,6 +34,22 @@ public void testJoinColumnsPrune() throws Exception {
                 "  |  other join predicates: 3: v3 < 6: v3\n" +
                 "  |  \n" +
                 "  |----2:EXCHANGE\n"));
+
+        // Prune should make the HASH JOIN(LEFT ANTI) could output the left table, but not join slot
+        sql = "select distinct('const') from t0, t1, " +
+                " (select * from t2 where cast(v7 as string) like 'ss%' ) sub1 " +
+                "left anti join " +
+                " (select * from t3 where cast(v10 as string) like 'ss%' ) sub2" +
+                " on substr(cast(sub1.v7 as string), 1) = substr(cast(sub2.v10 as string), 1)";
+
+        PlanTestBase.connectContext.getSessionVariable().setJoinImplementationMode("auto");
+        assertPlanContains(sql, " 11:Project\n" +
+                "  |  <slot 14> : 14: substr\n" +
+                "  |  \n" +
+                "  10:HASH JOIN\n" +
+                "  |  join op: LEFT ANTI JOIN (BROADCAST)\n" +
+                "  |  colocate: false, reason: \n" +
+                "  |  equal join conjunct: 14: substr = 15: substr");
     }
 
     @Test

diff --git a/fe/fe-core/src/test/java/com/starrocks/sql/plan/SelectConstTest.java b/fe/fe-core/src/test/java/com/starrocks/sql/plan/SelectConstTest.java
@@ -93,7 +93,7 @@ public void testSubquery() throws Exception {
         assertPlanContains("select * from t0 where exists (select 9,10)", "  1:UNION\n" +
                 "     constant exprs: \n" +
                 "         NULL");
-        assertPlanContains("select * from t0 where not exists (select 9)", "  1:UNION\n" +
+        assertPlanContains("select * from t0 where not exists (select 9)", "  0:UNION\n" +
                 "     constant exprs: \n" +
                 "         NULL");
         assertPlanContains("select * from t0 where v3 = (select 6)", "  5:Project\n" +