[improvement](mtmv) improve mv rewrite performance by reuse the shuttled expression (#37197) (#37935)

seawinde · web-flow · commit dbd2b0abf75a · 2024-07-17T00:52:58.000+08:00
## Proposed changes chrry-pick 2.1 pr: #37197 commitId: 701c7db
diff --git a/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVCache.java b/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVCache.java
@@ -108,7 +108,7 @@ public Plan visitLogicalResultSink(LogicalResultSink<? extends Plan> logicalResu
             return childContext.getRewritePlan();
         }, mvPlan, originPlan);
         // Construct structInfo once for use later
-        Optional<StructInfo> structInfoOptional = MaterializationContext.constructStructInfo(mvPlan,
+        Optional<StructInfo> structInfoOptional = MaterializationContext.constructStructInfo(mvPlan, originPlan,
                 planner.getCascadesContext(),
                 new BitSet());
         return new MTMVCache(mvPlan, originPlan, planner.getCascadesContext().getMemo().getRoot().getStatistics(),
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewRule.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewRule.java
@@ -145,8 +145,8 @@ protected List<StructInfo> getValidQueryStructInfos(Plan queryPlan, CascadesCont
             BitSet materializedViewTableSet) {
         List<StructInfo> validStructInfos = new ArrayList<>();
         // For every materialized view we should trigger refreshing struct info map
-        List<StructInfo> uncheckedStructInfos = MaterializedViewUtils.extractStructInfo(queryPlan, cascadesContext,
-                materializedViewTableSet);
+        List<StructInfo> uncheckedStructInfos = MaterializedViewUtils.extractStructInfo(queryPlan, queryPlan,
+                cascadesContext, materializedViewTableSet);
         uncheckedStructInfos.forEach(queryStructInfo -> {
             boolean valid = checkQueryPattern(queryStructInfo, cascadesContext) && queryStructInfo.isValid();
             if (!valid) {
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/AsyncMaterializationContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/AsyncMaterializationContext.java
@@ -50,6 +50,7 @@ public class AsyncMaterializationContext extends MaterializationContext {
 
     private static final Logger LOG = LogManager.getLogger(AsyncMaterializationContext.class);
     private final MTMV mtmv;
+    private List<String> materializationQualifier;
 
     /**
      * MaterializationContext, this contains necessary info for query rewriting by mv
@@ -72,7 +73,10 @@ Plan doGenerateScanPlan(CascadesContext cascadesContext) {
 
     @Override
     List<String> getMaterializationQualifier() {
-        return this.mtmv.getFullQualifiers();
+        if (this.materializationQualifier == null) {
+            this.materializationQualifier = this.mtmv.getFullQualifiers();
+        }
+        return this.materializationQualifier;
     }
 
     @Override
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializationContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializationContext.java
@@ -36,7 +36,6 @@
 import org.apache.doris.nereids.trees.plans.physical.PhysicalPlan;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalRelation;
 import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanVisitor;
-import org.apache.doris.nereids.util.ExpressionUtils;
 import org.apache.doris.statistics.ColumnStatistic;
 import org.apache.doris.statistics.Statistics;
 
@@ -119,32 +118,31 @@ public MaterializationContext(Plan plan, Plan originalPlan, Plan scanPlan,
                 this.exprToScanExprMapping.put(originalPlanOutput.get(slotIndex), scanPlanOutput.get(slotIndex));
             }
         }
-        this.planOutputShuttledExpressions = ExpressionUtils.shuttleExpressionWithLineage(originalPlanOutput,
-                originalPlan, new BitSet());
-        // materialization output expression shuttle, this will be used to expression rewrite
-        this.shuttledExprToScanExprMapping = ExpressionMapping.generate(
-                this.planOutputShuttledExpressions,
-                scanPlanOutput);
         // Construct materialization struct info, catch exception which may cause planner roll back
-        if (structInfo == null) {
-            Optional<StructInfo> structInfoOptional = constructStructInfo(plan, cascadesContext, new BitSet());
-            if (!structInfoOptional.isPresent()) {
-                this.available = false;
-            }
-            this.structInfo = structInfoOptional.orElseGet(() -> null);
-        } else {
-            this.structInfo = structInfo;
+        this.structInfo = structInfo == null
+                ? constructStructInfo(plan, originalPlan, cascadesContext, new BitSet()).orElseGet(() -> null)
+                : structInfo;
+        this.available = this.structInfo != null;
+        if (available) {
+            this.planOutputShuttledExpressions = this.structInfo.getPlanOutputShuttledExpressions();
+            // materialization output expression shuttle, this will be used to expression rewrite
+            this.shuttledExprToScanExprMapping = ExpressionMapping.generate(
+                    this.planOutputShuttledExpressions,
+                    scanPlanOutput);
         }
     }
 
     /**
      * Construct materialized view Struct info
+     * @param plan maybe remove unnecessary plan node, and the logical output maybe wrong
+     * @param originalPlan original plan, the output is right
      */
-    public static Optional<StructInfo> constructStructInfo(Plan plan, CascadesContext cascadesContext,
-            BitSet expectedTableBitSet) {
+    public static Optional<StructInfo> constructStructInfo(Plan plan, Plan originalPlan,
+            CascadesContext cascadesContext, BitSet expectedTableBitSet) {
         List<StructInfo> viewStructInfos;
         try {
-            viewStructInfos = MaterializedViewUtils.extractStructInfo(plan, cascadesContext, expectedTableBitSet);
+            viewStructInfos = MaterializedViewUtils.extractStructInfo(plan, originalPlan,
+                    cascadesContext, expectedTableBitSet);
             if (viewStructInfos.size() > 1) {
                 // view struct info should only have one, log error and use the first struct info
                 LOG.warn(String.format("view strut info is more than one, materialization plan is %s",
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializedViewUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializedViewUtils.java
@@ -167,8 +167,10 @@ public static boolean containTableQueryOperator(Plan analyzedPlan) {
 
     /**
      * Extract struct info from plan, support to get struct info from logical plan or plan in group.
+     * @param plan maybe remove unnecessary plan node, and the logical output maybe wrong
+     * @param originalPlan original plan, the output is right
      */
-    public static List<StructInfo> extractStructInfo(Plan plan, CascadesContext cascadesContext,
+    public static List<StructInfo> extractStructInfo(Plan plan, Plan originalPlan, CascadesContext cascadesContext,
             BitSet materializedViewTableSet) {
         // If plan belong to some group, construct it with group struct info
         if (plan.getGroupExpression().isPresent()) {
@@ -188,7 +190,7 @@ public static List<StructInfo> extractStructInfo(Plan plan, CascadesContext casc
                         continue;
                     }
                     StructInfo structInfo = structInfoMap.getStructInfo(cascadesContext,
-                            queryTableSet, ownerGroup, plan);
+                            queryTableSet, ownerGroup, originalPlan);
                     if (structInfo != null) {
                         structInfosBuilder.add(structInfo);
                     }
@@ -197,7 +199,7 @@ public static List<StructInfo> extractStructInfo(Plan plan, CascadesContext casc
             }
         }
         // if plan doesn't belong to any group, construct it directly
-        return ImmutableList.of(StructInfo.of(plan, cascadesContext));
+        return ImmutableList.of(StructInfo.of(plan, originalPlan, cascadesContext));
     }
 
     /**
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/Predicates.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/Predicates.java
@@ -41,7 +41,6 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
 
 /**
  * This record the predicates which can be pulled up or some other type predicates.
@@ -70,11 +69,6 @@ public Predicates merge(Collection<Expression> predicates) {
         return new Predicates(mergedPredicates);
     }
 
-    public Expression composedExpression() {
-        return ExpressionUtils.and(pulledUpPredicates.stream().map(Expression.class::cast)
-                .collect(Collectors.toList()));
-    }
-
     /**
      * Split the expression to equal, range and residual predicate.
      */
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/StructInfo.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/StructInfo.java
@@ -106,15 +106,16 @@ public class StructInfo {
     // this is for LogicalCompatibilityContext later
     private final Map<RelationId, StructInfoNode> relationIdStructInfoNodeMap;
     // this recorde the predicates which can pull up, not shuttled
-    private Predicates predicates;
+    private final Predicates predicates;
     // split predicates is shuttled
-    private final SplitPredicate splitPredicate;
-    private final EquivalenceClass equivalenceClass;
+    private SplitPredicate splitPredicate;
+    private EquivalenceClass equivalenceClass;
     // Key is the expression shuttled and the value is the origin expression
     // this is for building LogicalCompatibilityContext later.
     private final Map<ExpressionPosition, Map<Expression, Expression>> shuttledExpressionsToExpressionsMap;
     // Record the exprId and the corresponding expr map, this is used by expression shuttled
     private final Map<ExprId, Expression> namedExprIdAndExprMapping;
+    private final List<? extends Expression> planOutputShuttledExpressions;
 
     /**
      * The construct method for StructInfo
@@ -125,30 +126,25 @@ private StructInfo(Plan originalPlan, ObjectId originalPlanId, HyperGraph hyperG
             @Nullable Predicates predicates,
             Map<ExpressionPosition, Map<Expression, Expression>> shuttledExpressionsToExpressionsMap,
             Map<ExprId, Expression> namedExprIdAndExprMapping,
-            BitSet tableIdSet) {
+            BitSet tableIdSet,
+            SplitPredicate splitPredicate,
+            EquivalenceClass equivalenceClass,
+            List<? extends Expression> planOutputShuttledExpressions) {
         this.originalPlan = originalPlan;
         this.originalPlanId = originalPlanId;
         this.hyperGraph = hyperGraph;
-        this.valid = valid
-                && hyperGraph.getNodes().stream().allMatch(n -> ((StructInfoNode) n).getExpressions() != null);
+        this.valid = valid;
         this.topPlan = topPlan;
         this.bottomPlan = bottomPlan;
         this.relations = relations;
         this.tableBitSet = tableIdSet;
         this.relationIdStructInfoNodeMap = relationIdStructInfoNodeMap;
         this.predicates = predicates;
-        if (predicates == null) {
-            // collect predicate from top plan which not in hyper graph
-            Set<Expression> topPlanPredicates = new LinkedHashSet<>();
-            topPlan.accept(PREDICATE_COLLECTOR, topPlanPredicates);
-            this.predicates = Predicates.of(topPlanPredicates);
-        }
-        Pair<SplitPredicate, EquivalenceClass> derivedPredicates =
-                predicatesDerive(this.predicates, topPlan, tableBitSet);
-        this.splitPredicate = derivedPredicates.key();
-        this.equivalenceClass = derivedPredicates.value();
+        this.splitPredicate = splitPredicate;
+        this.equivalenceClass = equivalenceClass;
         this.shuttledExpressionsToExpressionsMap = shuttledExpressionsToExpressionsMap;
         this.namedExprIdAndExprMapping = namedExprIdAndExprMapping;
+        this.planOutputShuttledExpressions = planOutputShuttledExpressions;
     }
 
     /**
@@ -157,7 +153,8 @@ private StructInfo(Plan originalPlan, ObjectId originalPlanId, HyperGraph hyperG
     public StructInfo withPredicates(Predicates predicates) {
         return new StructInfo(this.originalPlan, this.originalPlanId, this.hyperGraph, this.valid, this.topPlan,
                 this.bottomPlan, this.relations, this.relationIdStructInfoNodeMap, predicates,
-                this.shuttledExpressionsToExpressionsMap, this.namedExprIdAndExprMapping, this.tableBitSet);
+                this.shuttledExpressionsToExpressionsMap, this.namedExprIdAndExprMapping, this.tableBitSet,
+                null, null, this.planOutputShuttledExpressions);
     }
 
     /**
@@ -166,7 +163,8 @@ public StructInfo withPredicates(Predicates predicates) {
     public StructInfo withTableBitSet(BitSet tableBitSet) {
         return new StructInfo(this.originalPlan, this.originalPlanId, this.hyperGraph, this.valid, this.topPlan,
                 this.bottomPlan, this.relations, this.relationIdStructInfoNodeMap, this.predicates,
-                this.shuttledExpressionsToExpressionsMap, this.namedExprIdAndExprMapping, tableBitSet);
+                this.shuttledExpressionsToExpressionsMap, this.namedExprIdAndExprMapping, tableBitSet,
+                this.splitPredicate, this.equivalenceClass, this.planOutputShuttledExpressions);
     }
 
     private static boolean collectStructInfoFromGraph(HyperGraph hyperGraph,
@@ -252,11 +250,10 @@ private static boolean collectStructInfoFromGraph(HyperGraph hyperGraph,
     }
 
     // derive some useful predicate by predicates
-    private Pair<SplitPredicate, EquivalenceClass> predicatesDerive(Predicates predicates, Plan originalPlan,
-            BitSet tableBitSet) {
+    private static Pair<SplitPredicate, EquivalenceClass> predicatesDerive(Predicates predicates, Plan originalPlan) {
         // construct equivalenceClass according to equals predicates
         List<Expression> shuttledExpression = ExpressionUtils.shuttleExpressionWithLineage(
-                        new ArrayList<>(predicates.getPulledUpPredicates()), originalPlan, tableBitSet).stream()
+                        new ArrayList<>(predicates.getPulledUpPredicates()), originalPlan, new BitSet()).stream()
                 .map(Expression.class::cast)
                 .collect(Collectors.toList());
         SplitPredicate splitPredicate = Predicates.splitPredicates(ExpressionUtils.and(shuttledExpression));
@@ -328,9 +325,19 @@ public static StructInfo of(Plan originalPlan, @Nullable Plan topPlan, @Nullable
                 relationIdStructInfoNodeMap,
                 tableBitSet,
                 cascadesContext);
+        valid = valid
+                && hyperGraph.getNodes().stream().allMatch(n -> ((StructInfoNode) n).getExpressions() != null);
+        // collect predicate from top plan which not in hyper graph
+        Set<Expression> topPlanPredicates = new LinkedHashSet<>();
+        topPlan.accept(PREDICATE_COLLECTOR, topPlanPredicates);
+        Predicates predicates = Predicates.of(topPlanPredicates);
+        // this should use the output of originalPlan to make sure the output right order
+        List<? extends Expression> planOutputShuttledExpressions =
+                ExpressionUtils.shuttleExpressionWithLineage(originalPlan.getOutput(), originalPlan, new BitSet());
         return new StructInfo(originalPlan, originalPlanId, hyperGraph, valid, topPlan, bottomPlan,
-                relationList, relationIdStructInfoNodeMap, null, shuttledHashConjunctsToConjunctsMap,
-                namedExprIdAndExprMapping, tableBitSet);
+                relationList, relationIdStructInfoNodeMap, predicates, shuttledHashConjunctsToConjunctsMap,
+                namedExprIdAndExprMapping, tableBitSet, null, null,
+                planOutputShuttledExpressions);
     }
 
     /**
@@ -350,10 +357,6 @@ public Predicates getPredicates() {
         return predicates;
     }
 
-    public EquivalenceClass getEquivalenceClass() {
-        return equivalenceClass;
-    }
-
     public Plan getOriginalPlan() {
         return originalPlan;
     }
@@ -362,8 +365,28 @@ public HyperGraph getHyperGraph() {
         return hyperGraph;
     }
 
+    /**
+     * lazy init for performance
+     */
     public SplitPredicate getSplitPredicate() {
-        return splitPredicate;
+        if (this.splitPredicate == null && this.predicates != null) {
+            Pair<SplitPredicate, EquivalenceClass> derivedPredicates = predicatesDerive(this.predicates, topPlan);
+            this.splitPredicate = derivedPredicates.key();
+            this.equivalenceClass = derivedPredicates.value();
+        }
+        return this.splitPredicate;
+    }
+
+    /**
+     * lazy init for performance
+     */
+    public EquivalenceClass getEquivalenceClass() {
+        if (this.equivalenceClass == null && this.predicates != null) {
+            Pair<SplitPredicate, EquivalenceClass> derivedPredicates = predicatesDerive(this.predicates, topPlan);
+            this.splitPredicate = derivedPredicates.key();
+            this.equivalenceClass = derivedPredicates.value();
+        }
+        return this.equivalenceClass;
     }
 
     public boolean isValid() {
@@ -416,6 +439,10 @@ public BitSet getTableBitSet() {
         return tableBitSet;
     }
 
+    public List<? extends Expression> getPlanOutputShuttledExpressions() {
+        return planOutputShuttledExpressions;
+    }
+
     /**
      * Judge the source graph logical is whether the same as target
      * For inner join should judge only the join tables,
diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/exploration/mv/HyperGraphAggTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/exploration/mv/HyperGraphAggTest.java
@@ -89,9 +89,9 @@ void testIJWithAgg() {
     }
 
     LogicalCompatibilityContext constructContext(Plan p1, Plan p2) {
-        StructInfo st1 = MaterializedViewUtils.extractStructInfo(p1,
+        StructInfo st1 = MaterializedViewUtils.extractStructInfo(p1, p1,
                 null, new BitSet()).get(0);
-        StructInfo st2 = MaterializedViewUtils.extractStructInfo(p2,
+        StructInfo st2 = MaterializedViewUtils.extractStructInfo(p2, p2,
                 null, new BitSet()).get(0);
         RelationMapping rm = RelationMapping.generate(st1.getRelations(), st2.getRelations()).get(0);
         SlotMapping sm = SlotMapping.generate(rm);
diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/sqltest/SqlTestBase.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/sqltest/SqlTestBase.java
@@ -100,9 +100,9 @@ protected void runBeforeEach() throws Exception {
     }
 
     protected LogicalCompatibilityContext constructContext(Plan p1, Plan p2, CascadesContext context) {
-        StructInfo st1 = MaterializedViewUtils.extractStructInfo(p1,
+        StructInfo st1 = MaterializedViewUtils.extractStructInfo(p1, p1,
                 context, new BitSet()).get(0);
-        StructInfo st2 = MaterializedViewUtils.extractStructInfo(p2,
+        StructInfo st2 = MaterializedViewUtils.extractStructInfo(p2, p2,
                 context, new BitSet()).get(0);
         RelationMapping rm = RelationMapping.generate(st1.getRelations(), st2.getRelations()).get(0);
         SlotMapping sm = SlotMapping.generate(rm);

Original file line number	Diff line number	Diff line change
`@@ -167,8 +167,10 @@ public static boolean containTableQueryOperator(Plan analyzedPlan) {`
`167`	`167`
`168`	`168`	`/**`
`169`	`169`	`* Extract struct info from plan, support to get struct info from logical plan or plan in group.`
	`170`	`+ * @param plan maybe remove unnecessary plan node, and the logical output maybe wrong`
	`171`	`+ * @param originalPlan original plan, the output is right`
`170`	`172`	`*/`
`171`		`- public static List<StructInfo> extractStructInfo(Plan plan, CascadesContext cascadesContext,`
	`173`	`+ public static List<StructInfo> extractStructInfo(Plan plan, Plan originalPlan, CascadesContext cascadesContext,`
`172`	`174`	`BitSet materializedViewTableSet) {`
`173`	`175`	`// If plan belong to some group, construct it with group struct info`
`174`	`176`	`if (plan.getGroupExpression().isPresent()) {`
`@@ -188,7 +190,7 @@ public static List<StructInfo> extractStructInfo(Plan plan, CascadesContext casc`
`188`	`190`	`continue;`
`189`	`191`	`}`
`190`	`192`	`StructInfo structInfo = structInfoMap.getStructInfo(cascadesContext,`
`191`		`- queryTableSet, ownerGroup, plan);`
	`193`	`+ queryTableSet, ownerGroup, originalPlan);`
`192`	`194`	`if (structInfo != null) {`
`193`	`195`	`structInfosBuilder.add(structInfo);`
`194`	`196`	`}`
`@@ -197,7 +199,7 @@ public static List<StructInfo> extractStructInfo(Plan plan, CascadesContext casc`
`197`	`199`	`}`
`198`	`200`	`}`
`199`	`201`	`// if plan doesn't belong to any group, construct it directly`
`200`		`- return ImmutableList.of(StructInfo.of(plan, cascadesContext));`
	`202`	`+ return ImmutableList.of(StructInfo.of(plan, originalPlan, cascadesContext));`
`201`	`203`	`}`
`202`	`204`
`203`	`205`	`/**`