[Enhancement] eliminate unnecessary project (#49724)

Signed-off-by: before-Sunrise <[email protected]> (cherry picked from commit 8921afd) # Conflicts: # fe/fe-core/src/test/java/com/starrocks/planner/MaterializedViewManualTest.java # fe/fe-core/src/test/java/com/starrocks/sql/plan/EliminateConstantValueTest.java # fe/fe-core/src/test/java/com/starrocks/sql/plan/JoinTest.java # fe/fe-core/src/test/java/com/starrocks/sql/plan/PruneComplexSubfieldTest.java # fe/fe-core/src/test/java/com/starrocks/sql/plan/SelectConstTest.java
StarRocks · Sep 4, 2024 · d56bcda · d56bcda
1 parent 3cd0732
commit d56bcda
Show file tree

Hide file tree

Showing 22 changed files with 773 additions and 108 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Optimizer.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Optimizer.java
@@ -82,6 +82,7 @@
 import com.starrocks.sql.optimizer.rule.tree.CloneDuplicateColRefRule;
 import com.starrocks.sql.optimizer.rule.tree.ExchangeSortToMergeRule;
 import com.starrocks.sql.optimizer.rule.tree.ExtractAggregateColumn;
+import com.starrocks.sql.optimizer.rule.tree.InlineCteProjectPruneRule;
 import com.starrocks.sql.optimizer.rule.tree.JoinLocalShuffleRule;
 import com.starrocks.sql.optimizer.rule.tree.PhysicalDistributionAggOptRule;
 import com.starrocks.sql.optimizer.rule.tree.PreAggregateTurnOnRule;
@@ -777,6 +778,7 @@ private OptExpression physicalRuleRewrite(TaskContext rootTaskContext, OptExpres
         result = new AddDecodeNodeForDictStringRule().rewrite(result, rootTaskContext);
         // Put before ScalarOperatorsReuseRule
         result = new PruneSubfieldsForComplexType().rewrite(result, rootTaskContext);
+        result = new InlineCteProjectPruneRule().rewrite(result, rootTaskContext);
         // This rule should be last
         result = new ScalarOperatorsReuseRule().rewrite(result, rootTaskContext);
         // Reorder predicates

diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/operator/Projection.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/operator/Projection.java
@@ -75,6 +75,9 @@ public List<ColumnRefOperator> getOutputColumns() {
     public ColumnRefSet getUsedColumns() {
         final ColumnRefSet usedColumns = new ColumnRefSet();
         columnRefMap.values().stream().forEach(e -> usedColumns.union(e.getUsedColumns()));
+        commonSubOperatorMap.values().stream().forEach(e -> usedColumns.union(e.getUsedColumns()));
+        // remove some of columnRefMap's used columns which are from commonSubOperatorMap's output column
+        commonSubOperatorMap.keySet().stream().forEach(e -> usedColumns.union(e.getUsedColumns()));
         return usedColumns;
     }
 

diff --git a/...e-core/src/main/java/com/starrocks/sql/optimizer/rule/tree/InlineCteProjectPruneRule.java b/...e-core/src/main/java/com/starrocks/sql/optimizer/rule/tree/InlineCteProjectPruneRule.java
@@ -0,0 +1,65 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.starrocks.sql.optimizer.rule.tree;
+
+import com.starrocks.sql.optimizer.OptExpression;
+import com.starrocks.sql.optimizer.OptExpressionVisitor;
+import com.starrocks.sql.optimizer.operator.Projection;
+import com.starrocks.sql.optimizer.operator.physical.PhysicalNoCTEOperator;
+import com.starrocks.sql.optimizer.operator.physical.PhysicalOperator;
+import com.starrocks.sql.optimizer.task.TaskContext;
+
+public class InlineCteProjectPruneRule implements TreeRewriteRule {
+    private static final InlineCteProjectPruneRuleVisitor HANDLER = new InlineCteProjectPruneRuleVisitor();
+
+    @Override
+    public OptExpression rewrite(OptExpression root, TaskContext taskContext) {
+        root.getOp().accept(HANDLER, root, taskContext);
+        return root;
+    }
+
+    private static class InlineCteProjectPruneRuleVisitor extends OptExpressionVisitor<Void, TaskContext> {
+        @Override
+        public Void visit(OptExpression opt, TaskContext context) {
+            for (OptExpression input : opt.getInputs()) {
+                input.getOp().accept(this, input, context);
+            }
+            return null;
+        }
+
+        @Override
+        public Void visitPhysicalNoCTE(OptExpression opt, TaskContext context) {
+            PhysicalNoCTEOperator op = (PhysicalNoCTEOperator) opt.getOp();
+            if (op.getProjection() == null) {
+                visit(opt, context);
+                return null;
+            }
+            Projection parentProjection = op.getProjection();
+
+            PhysicalOperator child = (PhysicalOperator) opt.getInputs().get(0).getOp();
+            if (child.getProjection() == null) {
+                visit(opt, context);
+                return null;
+            }
+
+            Projection childProjection = child.getProjection();
+            if (!parentProjection.getUsedColumns().containsAny(childProjection.getOutputColumns())) {
+                child.setProjection(null);
+            }
+
+            visit(opt, context);
+            return null;
+        }
+    }
+}
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/transformer/ExpressionMapping.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/transformer/ExpressionMapping.java
@@ -23,6 +23,7 @@
 import com.starrocks.sql.common.ErrorType;
 import com.starrocks.sql.common.StarRocksPlannerException;
 import com.starrocks.sql.optimizer.operator.scalar.ColumnRefOperator;
+import com.starrocks.sql.optimizer.operator.scalar.ScalarOperator;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -50,13 +51,28 @@ public class ExpressionMapping {
     private ColumnRefOperator[] fieldMappings;
     private RelationId outerScopeRelationId;
 
+    //  record columnRefOp which is generated by const expr in project node
+    //  if this columnRefOp is referenced by upper node, we should replace it with const expr in upper node
+    private Map<ColumnRefOperator, ScalarOperator> columnRefToConstOperators = new HashMap<>();
+
     public ExpressionMapping(Scope scope, List<ColumnRefOperator> fieldMappings) {
         this.scope = scope;
         this.fieldMappings = new ColumnRefOperator[fieldMappings.size()];
         fieldMappings.toArray(this.fieldMappings);
+        this.columnRefToConstOperators = new HashMap<>();
+    }
+
+    public ExpressionMapping(Scope scope, List<ColumnRefOperator> fieldMappings,
+                             Map<ColumnRefOperator, ScalarOperator> columnRefToConstOperators) {
+        this.scope = scope;
+        this.fieldMappings = new ColumnRefOperator[fieldMappings.size()];
+        fieldMappings.toArray(this.fieldMappings);
+        this.columnRefToConstOperators =
+                columnRefToConstOperators == null ? new HashMap<>() : columnRefToConstOperators;
     }
 
-    public ExpressionMapping(Scope scope, List<ColumnRefOperator> fieldMappings, ExpressionMapping outer) {
+    public ExpressionMapping(Scope scope, List<ColumnRefOperator> fieldMappings, ExpressionMapping outer,
+                             Map<ColumnRefOperator, ScalarOperator> columnRefToConstOperators) {
         this.scope = scope;
         List<ColumnRefOperator> fieldsList = new ArrayList<>(fieldMappings);
         if (outer != null) {
@@ -69,6 +85,9 @@ public ExpressionMapping(Scope scope, List<ColumnRefOperator> fieldMappings, Exp
         }
         this.fieldMappings = new ColumnRefOperator[fieldsList.size()];
         fieldsList.toArray(this.fieldMappings);
+
+        this.columnRefToConstOperators =
+                columnRefToConstOperators == null ? new HashMap<>() : columnRefToConstOperators;
     }
 
     public ExpressionMapping(Scope scope) {
@@ -128,6 +147,17 @@ public void put(Expr expression, ColumnRefOperator columnRefOperator) {
         expressionToColumns.put(expression, columnRefOperator);
     }
 
+    public void putConstOperator(ColumnRefOperator columnRefOperator, ScalarOperator constOperator) {
+        columnRefToConstOperators.put(columnRefOperator, constOperator);
+    }
+
+    public ScalarOperator getConstOperator(ColumnRefOperator columnRefOperator) {
+        if (columnRefToConstOperators == null) {
+            return null;
+        }
+        return columnRefToConstOperators.get(columnRefOperator);
+    }
+
     public void putWithSymbol(Expr expression, Expr resolveExpr, ColumnRefOperator columnRefOperator) {
         if (resolveExpr instanceof SlotRef) {
             if (expression instanceof SlotRef
@@ -164,4 +194,12 @@ public Map<Expr, ColumnRefOperator> getExpressionToColumns() {
     public void addExpressionToColumns(Map<Expr, ColumnRefOperator> expressionToColumns) {
         this.expressionToColumns.putAll(expressionToColumns);
     }
+
+    public Map<ColumnRefOperator, ScalarOperator> getColumnRefToConstOperators() {
+        return columnRefToConstOperators;
+    }
+
+    public void addColumnRefToConstOperators(Map<ColumnRefOperator, ScalarOperator> columnRefToConstOperators) {
+        this.columnRefToConstOperators.putAll(columnRefToConstOperators);
+    }
 }
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/transformer/OptExprBuilder.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/transformer/OptExprBuilder.java
@@ -18,9 +18,11 @@
 import com.starrocks.sql.optimizer.OptExpression;
 import com.starrocks.sql.optimizer.operator.Operator;
 import com.starrocks.sql.optimizer.operator.scalar.ColumnRefOperator;
+import com.starrocks.sql.optimizer.operator.scalar.ScalarOperator;
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 
 /**
@@ -49,6 +51,13 @@ public ExpressionMapping getExpressionMapping() {
         return expressionMapping;
     }
 
+    public Map<ColumnRefOperator, ScalarOperator> getColumnRefToConstOperators() {
+        if (expressionMapping == null) {
+            return null;
+        }
+        return expressionMapping.getColumnRefToConstOperators();
+    }
+
     public void setExpressionMapping(ExpressionMapping expressionMapping) {
         this.expressionMapping = expressionMapping;
     }

diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/transformer/QueryTransformer.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/transformer/QueryTransformer.java
@@ -84,7 +84,8 @@ public QueryTransformer(ColumnRefFactory columnRefFactory, ConnectContext sessio
 
     public LogicalPlan plan(SelectRelation queryBlock, ExpressionMapping outer) {
         OptExprBuilder builder = planFrom(queryBlock.getRelation(), cteContext);
-        builder.setExpressionMapping(new ExpressionMapping(builder.getScope(), builder.getFieldMappings(), outer));
+        builder.setExpressionMapping(new ExpressionMapping(builder.getScope(), builder.getFieldMappings(), outer,
+                builder.getColumnRefToConstOperators()));
 
         Map<Expr, SlotRef> generatedExprToColumnRef = queryBlock.getGeneratedExprToColumnRef();
         ExpressionMapping expressionMapping = builder.getExpressionMapping();
@@ -244,6 +245,8 @@ private OptExprBuilder projectForOrder(OptExprBuilder subOpt,
         }
 
         outputTranslations.addExpressionToColumns(subOpt.getExpressionMapping().getExpressionToColumns());
+        outputTranslations.addColumnRefToConstOperators(subOpt.getColumnRefToConstOperators());
+
         LogicalProjectOperator projectOperator = new LogicalProjectOperator(projections);
         return new OptExprBuilder(projectOperator, Lists.newArrayList(subOpt), outputTranslations);
     }
@@ -253,7 +256,8 @@ private OptExprBuilder project(OptExprBuilder subOpt, Iterable<Expr> expressions
     }
 
     private OptExprBuilder project(OptExprBuilder subOpt, Iterable<Expr> expressions, long limit) {
-        ExpressionMapping outputTranslations = new ExpressionMapping(subOpt.getScope(), subOpt.getFieldMappings());
+        ExpressionMapping outputTranslations = new ExpressionMapping(subOpt.getScope(), subOpt.getFieldMappings(),
+                subOpt.getColumnRefToConstOperators());
 
         Map<ColumnRefOperator, ScalarOperator> projections = Maps.newHashMap();
         for (Expr expression : expressions) {
@@ -268,9 +272,14 @@ private OptExprBuilder project(OptExprBuilder subOpt, Iterable<Expr> expressions
             ColumnRefOperator columnRefOperator = getOrCreateColumnRefOperator(expression, scalarOperator, projections);
             projections.put(columnRefOperator, scalarOperator);
             outputTranslations.put(expression, columnRefOperator);
+            if (scalarOperator.isConstant()) {
+                outputTranslations.putConstOperator(columnRefOperator, scalarOperator);
+            }
         }
 
         outputTranslations.addExpressionToColumns(subOpt.getExpressionMapping().getExpressionToColumns());
+        outputTranslations.addColumnRefToConstOperators(subOpt.getColumnRefToConstOperators());
+
         LogicalProjectOperator projectOperator = new LogicalProjectOperator(projections, limit);
         return new OptExprBuilder(projectOperator, Lists.newArrayList(subOpt), outputTranslations);
     }
@@ -436,7 +445,8 @@ private OptExprBuilder aggregate(OptExprBuilder subOpt,
             subOpt = project(subOpt, inputs);
         }
         ExpressionMapping groupingTranslations =
-                new ExpressionMapping(subOpt.getScope(), subOpt.getFieldMappings());
+                new ExpressionMapping(subOpt.getScope(), subOpt.getFieldMappings(),
+                        subOpt.getColumnRefToConstOperators());
 
         List<ColumnRefOperator> groupByColumnRefs = new ArrayList<>();