[improvement](mtmv) Support to partition prune when query rewrite by …

…sync materialized view
seawinde · Jul 30, 2024 · 786ebd5 · 786ebd5
1 parent 6ab5535
commit 786ebd5
Show file tree

Hide file tree

Showing 7 changed files with 325 additions and 21 deletions.
diff --git a/...main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewRule.java b/...main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewRule.java
@@ -233,14 +233,7 @@ protected List<Plan> doRewrite(StructInfo queryStructInfo, CascadesContext casca
                 continue;
             }
             Plan rewrittenPlan;
-            Plan mvScan = materializationContext.getScanPlan();
-            if (materializationContext instanceof SyncMaterializationContext
-                    && queryStructInfo.getRelations().size() == 1
-                    && queryStructInfo.getRelations().get(0) instanceof LogicalOlapScan
-                    && !((LogicalOlapScan) queryStructInfo.getRelations().get(0)).getSelectedPartitionIds().isEmpty()) {
-                mvScan = materializationContext.getScanPlan(
-                        ((LogicalOlapScan) queryStructInfo.getRelations().get(0)).getSelectedPartitionIds());
-            }
+            Plan mvScan = materializationContext.getScanPlan(queryStructInfo);
             Plan queryPlan = queryStructInfo.getTopPlan();
             if (compensatePredicates.isAlwaysTrue()) {
                 rewrittenPlan = mvScan;

diff --git a/.../main/java/org/apache/doris/nereids/rules/exploration/mv/AsyncMaterializationContext.java b/.../main/java/org/apache/doris/nereids/rules/exploration/mv/AsyncMaterializationContext.java
@@ -110,7 +110,8 @@ public Optional<Pair<Id, Statistics>> getPlanStatistics(CascadesContext cascades
             return Optional.empty();
         }
         RelationId relationId = null;
-        Optional<LogicalOlapScan> logicalOlapScan = this.getScanPlan().collectFirst(LogicalOlapScan.class::isInstance);
+        Optional<LogicalOlapScan> logicalOlapScan = this.getScanPlan(null)
+                .collectFirst(LogicalOlapScan.class::isInstance);
         if (logicalOlapScan.isPresent()) {
             relationId = logicalOlapScan.get().getRelationId();
         }
@@ -130,7 +131,8 @@ boolean isFinalChosen(Relation relation) {
         );
     }
 
-    public Plan getScanPlan() {
+    @Override
+    public Plan getScanPlan(StructInfo queryInfo) {
         return scanPlan;
     }
 

diff --git a/...e/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializationContext.java b/...e/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializationContext.java
@@ -35,7 +35,6 @@
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.algebra.Relation;
 import org.apache.doris.nereids.trees.plans.commands.ExplainCommand.ExplainLevel;
-import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
 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;
@@ -277,14 +276,7 @@ public Plan getOriginalPlan() {
         return originalPlan;
     }
 
-    public Plan getScanPlan() {
-        return scanPlan;
-    }
-
-    public Plan getScanPlan(List<Long> selectPartitions) {
-        if (scanPlan instanceof LogicalOlapScan) {
-            return ((LogicalOlapScan) scanPlan).withSelectedPartitionIds(selectPartitions);
-        }
+    public Plan getScanPlan(StructInfo queryStructInfo) {
         return scanPlan;
     }
 

diff --git a/...c/main/java/org/apache/doris/nereids/rules/exploration/mv/SyncMaterializationContext.java b/...c/main/java/org/apache/doris/nereids/rules/exploration/mv/SyncMaterializationContext.java
@@ -99,13 +99,29 @@ String getStringInfo() {
     @Override
     Optional<Pair<Id, Statistics>> getPlanStatistics(CascadesContext cascadesContext) {
         RelationId relationId = null;
-        Optional<LogicalOlapScan> scanObj = this.getScanPlan().collectFirst(LogicalOlapScan.class::isInstance);
+        Optional<LogicalOlapScan> scanObj = this.getScanPlan(null)
+                .collectFirst(LogicalOlapScan.class::isInstance);
         if (scanObj.isPresent()) {
             relationId = scanObj.get().getRelationId();
         }
         return Optional.of(Pair.of(relationId, normalizeStatisticsColumnExpression(statistics)));
     }
 
+    @Override
+    public Plan getScanPlan(StructInfo queryStructInfo) {
+        if (queryStructInfo == null) {
+            return scanPlan;
+        }
+        if (queryStructInfo.getRelations().size() == 1
+                && queryStructInfo.getRelations().get(0) instanceof LogicalOlapScan
+                && !((LogicalOlapScan) queryStructInfo.getRelations().get(0)).getSelectedPartitionIds().isEmpty()
+                && scanPlan instanceof LogicalOlapScan) {
+            return ((LogicalOlapScan) scanPlan).withSelectedPartitionIds(
+                    ((LogicalOlapScan) queryStructInfo.getRelations().get(0)).getSelectedPartitionIds());
+        }
+        return scanPlan;
+    }
+
     /**
      * Calc the relation is chosen finally or not
      */

diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/mv/IdStatisticsMapTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/mv/IdStatisticsMapTest.java
@@ -75,7 +75,8 @@ public boolean isMVPartitionValid(MTMV mtmv, ConnectContext ctx, boolean isMVPar
                 .analyze()
                 .rewrite();
         // scan plan output will be refreshed after mv rewrite successfully, so need tmp store
-        Set<Slot> materializationScanOutput = c1.getMaterializationContexts().get(0).getScanPlan().getOutputSet();
+        Set<Slot> materializationScanOutput = c1.getMaterializationContexts().get(0)
+                .getScanPlan(null).getOutputSet();
         tmpPlanChecker
                 .optimize()
                 .printlnBestPlanTree();

diff --git a/regression-test/data/mv_p0/partition_prune/partition_prune.out b/regression-test/data/mv_p0/partition_prune/partition_prune.out
@@ -0,0 +1,43 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !query1_before --
+aa	bc	2024-07-03T01:15:30	1
+ad	be	2024-07-03T07:06:30	1
+af	bf	2024-07-04T10:01:30	1
+ag	bc	2024-07-04T10:01:35	1
+as	bd	2024-07-03T01:15:30	1
+
+-- !query1_after --
+aa	bc	2024-07-03T01:15:30	1
+ad	be	2024-07-03T07:06:30	1
+af	bf	2024-07-04T10:01:30	1
+ag	bc	2024-07-04T10:01:35	1
+as	bd	2024-07-03T01:15:30	1
+
+-- !query2_before --
+2024-07-03T01:00	aa	bc
+2024-07-03T06:00	as	bd
+2024-07-03T07:00	ad	be
+2024-07-04T10:00	af	bf
+2024-07-04T12:00	ag	bc
+
+-- !query2_after --
+2024-07-03T01:00	aa	bc
+2024-07-03T06:00	as	bd
+2024-07-03T07:00	ad	be
+2024-07-04T10:00	af	bf
+2024-07-04T12:00	ag	bc
+
+-- !query3_before --
+aa	bc	2024-07-03T01:00	2.1
+ad	be	2024-07-03T07:00	3.1
+af	bf	2024-07-04T10:00	4.1
+ag	bc	2024-07-04T12:00	5.1
+as	bd	2024-07-03T06:00	1.1
+
+-- !query3_after --
+aa	bc	2024-07-03T01:00	2.1
+ad	be	2024-07-03T07:00	3.1
+af	bf	2024-07-04T10:00	4.1
+ag	bc	2024-07-04T12:00	5.1
+as	bd	2024-07-03T06:00	1.1
+