diff --git a/fe/fe-core/src/main/java/org/apache/doris/common/profile/SummaryProfile.java b/fe/fe-core/src/main/java/org/apache/doris/common/profile/SummaryProfile.java index a51653bbfd35ce..bf286d263a203b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/common/profile/SummaryProfile.java +++ b/fe/fe-core/src/main/java/org/apache/doris/common/profile/SummaryProfile.java @@ -21,6 +21,8 @@ import org.apache.doris.common.io.Text; import org.apache.doris.common.util.TimeUtils; import org.apache.doris.persist.gson.GsonUtils; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.StmtExecutor; import org.apache.doris.system.Backend; import org.apache.doris.thrift.TNetworkAddress; import org.apache.doris.thrift.TUnit; @@ -240,6 +242,8 @@ public class SummaryProfile { @SerializedName(value = "nereidsCollectTablePartitionFinishTime") private long nereidsCollectTablePartitionFinishTime = -1; + @SerializedName(value = "nereidsCollectTablePartitionTime") + private long nereidsCollectTablePartitionTime = 0; @SerializedName(value = "nereidsAnalysisFinishTime") private long nereidsAnalysisFinishTime = -1; @SerializedName(value = "nereidsRewriteFinishTime") @@ -553,6 +557,10 @@ public void setNereidsCollectTablePartitionFinishTime() { this.nereidsCollectTablePartitionFinishTime = TimeUtils.getStartTimeMs(); } + public void addCollectTablePartitionTime(long elapsed) { + nereidsCollectTablePartitionTime += elapsed; + } + public void setNereidsAnalysisTime() { this.nereidsAnalysisFinishTime = TimeUtils.getStartTimeMs(); } @@ -829,7 +837,9 @@ public String getPrettyNereidsRewriteTime() { public String getPrettyNereidsCollectTablePartitionTime() { - return getPrettyTime(nereidsCollectTablePartitionFinishTime, nereidsRewriteFinishTime, TUnit.TIME_MS); + long totalTime = nereidsCollectTablePartitionFinishTime + - nereidsRewriteFinishTime + nereidsCollectTablePartitionTime; + return RuntimeProfile.printCounter(totalTime, TUnit.TIME_MS); } public String getPrettyNereidsOptimizeTime() { @@ -999,4 +1009,15 @@ public void write(DataOutput output) throws IOException { public void setAssignedWeightPerBackend(Map assignedWeightPerBackend) { this.assignedWeightPerBackend = assignedWeightPerBackend; } + + public static SummaryProfile getSummaryProfile(ConnectContext connectContext) { + ConnectContext ctx = connectContext == null ? ConnectContext.get() : connectContext; + if (ctx != null) { + StmtExecutor executor = ctx.getExecutor(); + if (executor != null) { + return executor.getSummaryProfile(); + } + } + return null; + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVRelationManager.java b/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVRelationManager.java index 29cba417f1b25f..ba498636d73619 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVRelationManager.java +++ b/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVRelationManager.java @@ -45,6 +45,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import java.util.BitSet; import java.util.Collection; import java.util.List; import java.util.Map; @@ -85,7 +86,7 @@ public Set getAvailableMTMVs(Set candidateMTMVs, ConnectContext ctx, boolean forceConsistent, BiPredicate predicate) { Set res = Sets.newLinkedHashSet(); Map, Set> queryUsedPartitions = PartitionCompensator.getQueryUsedPartitions( - ctx.getStatementContext()); + ctx.getStatementContext(), new BitSet()); for (MTMV mtmv : candidateMTMVs) { if (predicate.test(ctx, mtmv)) { continue; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/CascadesContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/CascadesContext.java index d6924e77ed3979..5944ce693da2f7 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/CascadesContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/CascadesContext.java @@ -24,7 +24,6 @@ import org.apache.doris.nereids.jobs.JobContext; import org.apache.doris.nereids.jobs.executor.Analyzer; import org.apache.doris.nereids.jobs.executor.TableCollectAndHookInitializer; -import org.apache.doris.nereids.jobs.executor.TablePartitionCollector; import org.apache.doris.nereids.jobs.rewrite.RewriteBottomUpJob; import org.apache.doris.nereids.jobs.rewrite.RewriteTopDownJob; import org.apache.doris.nereids.jobs.rewrite.RootPlanTreeRewriteJob.RootRewriteJobContext; @@ -229,10 +228,6 @@ public TableCollectAndHookInitializer newTableCollector() { return new TableCollectAndHookInitializer(this); } - public TablePartitionCollector newTablePartitionCollector() { - return new TablePartitionCollector(this); - } - public Analyzer newAnalyzer() { return new Analyzer(this); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/NereidsPlanner.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/NereidsPlanner.java index db4b0e524dfe65..6270b82874dd13 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/NereidsPlanner.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/NereidsPlanner.java @@ -369,21 +369,6 @@ protected void collectAndLockTable(boolean showPlanProcess) { } } - protected void collectTableUsedPartitions(boolean showPlanProcess) { - if (LOG.isDebugEnabled()) { - LOG.debug("Start to collect table used partition"); - } - keepOrShowPlanProcess(showPlanProcess, () -> cascadesContext.newTablePartitionCollector().execute()); - NereidsTracer.logImportantTime("EndCollectTablePartitions"); - if (LOG.isDebugEnabled()) { - LOG.debug("Start to collect table used partition"); - } - if (statementContext.getConnectContext().getExecutor() != null) { - statementContext.getConnectContext().getExecutor().getSummaryProfile() - .setNereidsCollectTablePartitionFinishTime(); - } - } - protected void analyze(boolean showPlanProcess) { if (LOG.isDebugEnabled()) { LOG.debug("Start analyze plan"); @@ -416,9 +401,6 @@ protected void rewrite(boolean showPlanProcess) { if (statementContext.getConnectContext().getExecutor() != null) { statementContext.getConnectContext().getExecutor().getSummaryProfile().setNereidsRewriteTime(); } - // collect partitions table used, this is for query rewrite by materialized view - // this is needed before init hook - collectTableUsedPartitions(showPlanProcess); cascadesContext.getStatementContext().getPlannerHooks().forEach(hook -> hook.afterRewrite(this)); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/TablePartitionCollector.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/TablePartitionCollector.java deleted file mode 100644 index e67b94d1314e89..00000000000000 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/TablePartitionCollector.java +++ /dev/null @@ -1,46 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you 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 -// -// http://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 org.apache.doris.nereids.jobs.executor; - -import org.apache.doris.nereids.CascadesContext; -import org.apache.doris.nereids.jobs.rewrite.RewriteJob; -import org.apache.doris.nereids.rules.RuleType; -import org.apache.doris.nereids.rules.rewrite.QueryPartitionCollector; - -import java.util.List; - -/** - * Collect partitions which query used, this is useful for optimizing get available mvs, - * should collect after RBO - */ -public class TablePartitionCollector extends AbstractBatchJobExecutor { - public TablePartitionCollector(CascadesContext cascadesContext) { - super(cascadesContext); - } - - @Override - public List getJobs() { - return buildCollectorJobs(); - } - - private static List buildCollectorJobs() { - return jobs( - custom(RuleType.COLLECT_PARTITIONS, QueryPartitionCollector::new) - ); - } -} 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 index 4dd7ea882a58a3..3a6f08f8cc8a1f 100644 --- 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 @@ -22,6 +22,8 @@ import org.apache.doris.common.AnalysisException; import org.apache.doris.common.Id; import org.apache.doris.common.Pair; +import org.apache.doris.common.profile.SummaryProfile; +import org.apache.doris.common.util.TimeUtils; import org.apache.doris.mtmv.BaseTableInfo; import org.apache.doris.nereids.CascadesContext; import org.apache.doris.nereids.StatementContext; @@ -298,7 +300,7 @@ protected List doRewrite(StructInfo queryStructInfo, CascadesContext casca MTMV mtmv = ((AsyncMaterializationContext) materializationContext).getMtmv(); BaseTableInfo relatedTableInfo = mtmv.getMvPartitionInfo().getRelatedTableInfo(); Map, Set> queryUsedPartitions = PartitionCompensator.getQueryUsedPartitions( - cascadesContext.getConnectContext().getStatementContext()); + cascadesContext.getStatementContext(), queryStructInfo.getRelationIdBitSet()); Set relateTableUsedPartitions = queryUsedPartitions.get(relatedTableInfo.toList()); if (relateTableUsedPartitions == null) { materializationContext.recordFailReason(queryStructInfo, @@ -420,6 +422,18 @@ protected List doRewrite(StructInfo queryStructInfo, CascadesContext casca logicalProperties, queryPlan.getLogicalProperties())); continue; } + // need to collect table partition again, because the rewritten plan would contain new relation + // and the rewritten plan would part in rewritten later , the table used partition info is needed + // for later rewrite + long startTimeMs = TimeUtils.getStartTimeMs(); + try { + MaterializedViewUtils.collectTableUsedPartitions(rewrittenPlan, cascadesContext); + } finally { + SummaryProfile summaryProfile = SummaryProfile.getSummaryProfile(cascadesContext.getConnectContext()); + if (summaryProfile != null) { + summaryProfile.addCollectTablePartitionTime(TimeUtils.getElapsedTimeMs(startTimeMs)); + } + } trySetStatistics(materializationContext, cascadesContext); rewriteResults.add(rewrittenPlan); recordIfRewritten(queryStructInfo.getOriginalPlan(), materializationContext, cascadesContext); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/InitMaterializationContextHook.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/InitMaterializationContextHook.java index 1ee227437491f1..de703ba9fd6fc9 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/InitMaterializationContextHook.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/InitMaterializationContextHook.java @@ -32,6 +32,7 @@ import org.apache.doris.nereids.CascadesContext; import org.apache.doris.nereids.NereidsPlanner; import org.apache.doris.nereids.PlannerHook; +import org.apache.doris.nereids.StatementContext; import org.apache.doris.nereids.parser.NereidsParser; import org.apache.doris.qe.ConnectContext; @@ -60,6 +61,15 @@ public class InitMaterializationContextHook implements PlannerHook { @Override public void afterRewrite(NereidsPlanner planner) { + CascadesContext cascadesContext = planner.getCascadesContext(); + // collect partitions table used, this is for query rewrite by materialized view + // this is needed before init hook, because compare partition version in init hook would use this + MaterializedViewUtils.collectTableUsedPartitions(cascadesContext.getRewritePlan(), cascadesContext); + StatementContext statementContext = cascadesContext.getStatementContext(); + if (statementContext.getConnectContext().getExecutor() != null) { + statementContext.getConnectContext().getExecutor().getSummaryProfile() + .setNereidsCollectTablePartitionFinishTime(); + } initMaterializationContext(planner.getCascadesContext()); } @@ -142,12 +152,17 @@ private List createAsyncMaterializationContext(CascadesC // so regenerate the struct info table bitset StructInfo mvStructInfo = mtmvCache.getStructInfo(); BitSet tableBitSetInCurrentCascadesContext = new BitSet(); - mvStructInfo.getRelations().forEach(relation -> tableBitSetInCurrentCascadesContext.set( - cascadesContext.getStatementContext().getTableId(relation.getTable()).asInt())); + BitSet relationIdBitSetInCurrentCascadesContext = new BitSet(); + mvStructInfo.getRelations().forEach(relation -> { + tableBitSetInCurrentCascadesContext.set( + cascadesContext.getStatementContext().getTableId(relation.getTable()).asInt()); + relationIdBitSetInCurrentCascadesContext.set(relation.getRelationId().asInt()); + }); asyncMaterializationContext.add(new AsyncMaterializationContext(materializedView, mtmvCache.getLogicalPlan(), mtmvCache.getOriginalPlan(), ImmutableList.of(), ImmutableList.of(), cascadesContext, - mtmvCache.getStructInfo().withTableBitSet(tableBitSetInCurrentCascadesContext))); + mtmvCache.getStructInfo().withTableBitSet(tableBitSetInCurrentCascadesContext, + relationIdBitSetInCurrentCascadesContext))); } catch (Exception e) { LOG.warn(String.format("MaterializationContext init mv cache generate fail, current queryId is %s", cascadesContext.getConnectContext().getQueryIdentifier()), e); 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 index 8a61d56646d812..b14261178f24d8 100644 --- 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 @@ -34,6 +34,7 @@ import org.apache.doris.nereids.rules.analysis.BindRelation; import org.apache.doris.nereids.rules.expression.ExpressionNormalization; import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext; +import org.apache.doris.nereids.rules.rewrite.QueryPartitionCollector; import org.apache.doris.nereids.trees.expressions.Alias; import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.Expression; @@ -300,6 +301,15 @@ public static List extractNondeterministicFunction(Plan plan) { return nondeterministicFunctions; } + /** + * Collect table used partitions, this is used for mv rewrite partition union + * can not cumulative, if called multi times, should clean firstly + */ + public static void collectTableUsedPartitions(Plan plan, CascadesContext cascadesContext) { + // the recorded partition is based on relation id + plan.accept(new QueryPartitionCollector(), cascadesContext); + } + /** * Check the query if Contains query operator * Such sql as following should return true diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensator.java index 98629f86028b81..0ee88ce82af070 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensator.java @@ -40,6 +40,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import java.util.BitSet; import java.util.Collection; import java.util.HashMap; import java.util.HashSet; @@ -164,8 +165,12 @@ public static boolean needUnionRewrite(MaterializationContext materializationCon /** * Get query used partitions * this is calculated from tableUsedPartitionNameMap and tables in statementContext - * */ - public static Map, Set> getQueryUsedPartitions(StatementContext statementContext) { + * + * @param customRelationIdSet if union compensate occurs, the new query used partitions is changed, + * so need to get used partitions by relation id set + */ + public static Map, Set> getQueryUsedPartitions(StatementContext statementContext, + BitSet customRelationIdSet) { // get table used partitions // if table is not in statementContext().getTables() which means the table is partition prune as empty relation Multimap, Pair>> tableUsedPartitionNameMap = statementContext @@ -174,7 +179,7 @@ public static Map, Set> getQueryUsedPartitions(StatementCon // if value is null, means query all partitions // if value is not empty, means query some partitions Map, Set> queryUsedRelatedTablePartitionsMap = new HashMap<>(); - outer: + tableLoop: for (Map.Entry, TableIf> queryUsedTableEntry : statementContext.getTables().entrySet()) { Set usedPartitionSet = new HashSet<>(); Collection>> tableUsedPartitions = @@ -185,11 +190,20 @@ public static Map, Set> getQueryUsedPartitions(StatementCon continue; } for (Pair> partitionPair : tableUsedPartitions) { - if (ALL_PARTITIONS.equals(partitionPair)) { - queryUsedRelatedTablePartitionsMap.put(queryUsedTableEntry.getKey(), null); - continue outer; + if (!customRelationIdSet.isEmpty()) { + if (ALL_PARTITIONS.equals(partitionPair)) { + continue; + } + if (customRelationIdSet.get(partitionPair.key().asInt())) { + usedPartitionSet.addAll(partitionPair.value()); + } + } else { + if (ALL_PARTITIONS.equals(partitionPair)) { + queryUsedRelatedTablePartitionsMap.put(queryUsedTableEntry.getKey(), null); + continue tableLoop; + } + usedPartitionSet.addAll(partitionPair.value()); } - usedPartitionSet.addAll(partitionPair.value()); } } queryUsedRelatedTablePartitionsMap.put(queryUsedTableEntry.getKey(), usedPartitionSet); 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 index 29240153b20d29..3b3f47e1f84cd8 100644 --- 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 @@ -100,6 +100,7 @@ public class StructInfo { // So if the cascadesContext currently is different form the cascadesContext which generated it. // Should regenerate the tableBitSet by current cascadesContext and call withTableBitSet method private final BitSet tableBitSet; + private final BitSet relationIdBitSet; // this is for LogicalCompatibilityContext later private final Map relationIdStructInfoNodeMap; // this recorde the predicates which can pull up, not shuttled @@ -137,6 +138,7 @@ private StructInfo(Plan originalPlan, ObjectId originalPlanId, HyperGraph hyperG shuttledExpressionsToExpressionsMap, Map> expressionToShuttledExpressionToMap, BitSet tableIdSet, + BitSet relationIdSet, SplitPredicate splitPredicate, EquivalenceClass equivalenceClass, List planOutputShuttledExpressions) { @@ -148,6 +150,7 @@ private StructInfo(Plan originalPlan, ObjectId originalPlanId, HyperGraph hyperG this.bottomPlan = bottomPlan; this.relations = relations; this.tableBitSet = tableIdSet; + this.relationIdBitSet = relationIdSet; this.relationIdStructInfoNodeMap = relationIdStructInfoNodeMap; this.predicates = predicates; this.splitPredicate = splitPredicate; @@ -164,17 +167,19 @@ 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.expressionToShuttledExpressionToMap, - this.tableBitSet, null, null, this.planOutputShuttledExpressions); + this.tableBitSet, this.relationIdBitSet, null, null, + this.planOutputShuttledExpressions); } /** * Construct StructInfo with new tableBitSet */ - public StructInfo withTableBitSet(BitSet tableBitSet) { + public StructInfo withTableBitSet(BitSet tableBitSet, BitSet relationIdBitSet) { return new StructInfo(this.originalPlan, this.originalPlanId, this.hyperGraph, this.valid, this.topPlan, this.bottomPlan, this.relations, this.relationIdStructInfoNodeMap, this.predicates, this.shuttledExpressionsToExpressionsMap, this.expressionToShuttledExpressionToMap, - tableBitSet, this.splitPredicate, this.equivalenceClass, this.planOutputShuttledExpressions); + tableBitSet, relationIdBitSet, this.splitPredicate, this.equivalenceClass, + this.planOutputShuttledExpressions); } private static boolean collectStructInfoFromGraph(HyperGraph hyperGraph, @@ -185,6 +190,7 @@ private static boolean collectStructInfoFromGraph(HyperGraph hyperGraph, List relations, Map relationIdStructInfoNodeMap, BitSet hyperTableBitSet, + BitSet relationBitSet, CascadesContext cascadesContext) { // Collect relations from hyper graph which in the bottom plan firstly @@ -194,8 +200,11 @@ private static boolean collectStructInfoFromGraph(HyperGraph hyperGraph, List nodeRelations = new ArrayList<>(); nodePlan.accept(RELATION_COLLECTOR, nodeRelations); relations.addAll(nodeRelations); - nodeRelations.forEach(relation -> hyperTableBitSet.set( - cascadesContext.getStatementContext().getTableId(relation.getTable()).asInt())); + nodeRelations.forEach(relation -> { + hyperTableBitSet.set( + cascadesContext.getStatementContext().getTableId(relation.getTable()).asInt()); + relationBitSet.set(relation.getRelationId().asInt()); + }); // plan relation collector and set to map StructInfoNode structInfoNode = (StructInfoNode) node; // record expressions in node @@ -314,12 +323,14 @@ public static StructInfo of(Plan originalPlan, @Nullable Plan topPlan, @Nullable Map>> shuttledHashConjunctsToConjunctsMap = new LinkedHashMap<>(); BitSet tableBitSet = new BitSet(); + BitSet relationBitSet = new BitSet(); Map> expressionToShuttledExpressionToMap = new HashMap<>(); boolean valid = collectStructInfoFromGraph(hyperGraph, topPlan, shuttledHashConjunctsToConjunctsMap, expressionToShuttledExpressionToMap, relationList, relationIdStructInfoNodeMap, tableBitSet, + relationBitSet, cascadesContext); valid = valid && hyperGraph.getNodes().stream().allMatch(n -> ((StructInfoNode) n).getExpressions() != null); @@ -338,8 +349,7 @@ public static StructInfo of(Plan originalPlan, @Nullable Plan topPlan, @Nullable return new StructInfo(originalPlan, originalPlanId, hyperGraph, valid, topPlan, bottomPlan, relationList, relationIdStructInfoNodeMap, predicates, shuttledHashConjunctsToConjunctsMap, expressionToShuttledExpressionToMap, - tableBitSet, null, null, - planOutputShuttledExpressions); + tableBitSet, relationBitSet, null, null, planOutputShuttledExpressions); } public List getRelations() { @@ -443,6 +453,10 @@ public BitSet getTableBitSet() { return tableBitSet; } + public BitSet getRelationIdBitSet() { + return relationIdBitSet; + } + public List getPlanOutputShuttledExpressions() { return planOutputShuttledExpressions; } @@ -752,10 +766,11 @@ public static Pair addFilterOnTableScan(Plan queryPlan, Map { + Plan filterAddedPlan = MaterializedViewUtils.rewriteByRules(parentCascadesContext, context -> { Rewriter.getWholeTreeRewriter(context).execute(); return context.getRewritePlan(); - }, queryPlanWithUnionFilter, queryPlan), true); + }, queryPlanWithUnionFilter, queryPlan); + return Pair.of(filterAddedPlan, true); } /** diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/QueryPartitionCollector.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/QueryPartitionCollector.java index 2ad993b361d43a..cfe5d6863c387b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/QueryPartitionCollector.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/QueryPartitionCollector.java @@ -20,17 +20,14 @@ import org.apache.doris.catalog.TableIf; import org.apache.doris.common.Pair; import org.apache.doris.datasource.ExternalTable; -import org.apache.doris.nereids.jobs.JobContext; +import org.apache.doris.nereids.CascadesContext; import org.apache.doris.nereids.rules.exploration.mv.PartitionCompensator; -import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.RelationId; import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation; import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan; import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan.SelectedPartitions; import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; -import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter; -import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter; -import org.apache.doris.qe.ConnectContext; +import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanVisitor; import com.google.common.collect.Multimap; import org.apache.logging.log4j.LogManager; @@ -43,29 +40,18 @@ /** * Used to collect query partitions, only collect once * */ -public class QueryPartitionCollector extends DefaultPlanRewriter implements CustomRewriter { +public class QueryPartitionCollector extends DefaultPlanVisitor { public static final Logger LOG = LogManager.getLogger(QueryPartitionCollector.class); @Override - public Plan rewriteRoot(Plan plan, JobContext jobContext) { - - ConnectContext connectContext = ConnectContext.get(); - if (connectContext != null && connectContext.getSessionVariable().internalSession) { - return plan; - } - plan.accept(this, connectContext); - return plan; - } - - @Override - public Plan visitLogicalCatalogRelation(LogicalCatalogRelation catalogRelation, ConnectContext context) { + public Void visitLogicalCatalogRelation(LogicalCatalogRelation catalogRelation, CascadesContext context) { TableIf table = catalogRelation.getTable(); if (table.getDatabase() == null) { LOG.error("QueryPartitionCollector visitLogicalCatalogRelation database is null, table is " + table.getName()); - return catalogRelation; + return null; } Multimap, Pair>> tableUsedPartitionNameMap = context.getStatementContext() .getTableUsedPartitionNameMap(); @@ -90,6 +76,6 @@ public Plan visitLogicalCatalogRelation(LogicalCatalogRelation catalogRelation, // not support get partition scene, we consider query all partitions from table tableUsedPartitionNameMap.put(table.getFullQualifiers(), PartitionCompensator.ALL_PARTITIONS); } - return catalogRelation; + return null; } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensatorTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensatorTest.java index 76246e52f9d542..25c0a679d8d388 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensatorTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensatorTest.java @@ -28,6 +28,7 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.util.BitSet; import java.util.List; import java.util.Map; import java.util.Set; @@ -107,9 +108,11 @@ public void testGetQueryTableUsedPartition() { + "left outer join orders_list_partition\n" + "on l1.l_shipdate = o_orderdate\n", nereidsPlanner -> { + MaterializedViewUtils.collectTableUsedPartitions(nereidsPlanner.getRewrittenPlan(), + nereidsPlanner.getCascadesContext()); Map, Set> queryUsedPartitions = PartitionCompensator.getQueryUsedPartitions( - nereidsPlanner.getCascadesContext().getStatementContext()); + nereidsPlanner.getCascadesContext().getStatementContext(), new BitSet()); List itmeQualifier = ImmutableList.of( "internal", "partition_compensate_test", "lineitem_list_partition"); @@ -131,6 +134,8 @@ public void testGetAllTableUsedPartition() { + "left outer join orders_list_partition\n" + "on l1.l_shipdate = o_orderdate\n", nereidsPlanner -> { + MaterializedViewUtils.collectTableUsedPartitions(nereidsPlanner.getRewrittenPlan(), + nereidsPlanner.getCascadesContext()); List qualifier = ImmutableList.of( "internal", "partition_compensate_test", "lineitem_list_partition"); @@ -140,7 +145,7 @@ public void testGetAllTableUsedPartition() { Map, Set> queryUsedPartitions = PartitionCompensator.getQueryUsedPartitions( - nereidsPlanner.getCascadesContext().getStatementContext()); + nereidsPlanner.getCascadesContext().getStatementContext(), new BitSet()); Set queryTableUsedPartition = queryUsedPartitions.get(qualifier); // if tableUsedPartitionNameMap contain any PartitionCompensator.ALL_PARTITIONS // consider query all partitions from table @@ -161,6 +166,8 @@ public void testGetAllTableUsedPartitionList() { + "left outer join orders_list_partition\n" + "on l1.l_shipdate = o_orderdate\n", nereidsPlanner -> { + MaterializedViewUtils.collectTableUsedPartitions(nereidsPlanner.getRewrittenPlan(), + nereidsPlanner.getCascadesContext()); List qualifier = ImmutableList.of( "internal", "partition_compensate_test", "lineitem_list_partition"); @@ -171,7 +178,7 @@ public void testGetAllTableUsedPartitionList() { Map, Set> queryUsedPartitions = PartitionCompensator.getQueryUsedPartitions( - nereidsPlanner.getCascadesContext().getStatementContext()); + nereidsPlanner.getCascadesContext().getStatementContext(), new BitSet()); Set queryTableUsedPartition = queryUsedPartitions.get(qualifier); // if tableUsedPartitionNameMap contain only PartitionCompensator.ALL_PARTITIONS // consider query all partitions from table diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/util/PlanChecker.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/util/PlanChecker.java index 48803d1d252c7f..71d0f0101b0413 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/util/PlanChecker.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/util/PlanChecker.java @@ -52,6 +52,7 @@ import org.apache.doris.nereids.rules.RuleType; import org.apache.doris.nereids.rules.exploration.mv.InitConsistentMaterializationContextHook; import org.apache.doris.nereids.rules.exploration.mv.InitMaterializationContextHook; +import org.apache.doris.nereids.rules.exploration.mv.MaterializedViewUtils; import org.apache.doris.nereids.rules.rewrite.OneRewriteRuleFactory; import org.apache.doris.nereids.trees.plans.GroupPlan; import org.apache.doris.nereids.trees.plans.Plan; @@ -255,7 +256,7 @@ public Rule build() { public PlanChecker rewrite() { Rewriter.getWholeTreeRewriter(cascadesContext).execute(); - cascadesContext.newTablePartitionCollector().execute(); + MaterializedViewUtils.collectTableUsedPartitions(cascadesContext.getRewritePlan(), cascadesContext); InitMaterializationContextHook.INSTANCE.initMaterializationContext(this.cascadesContext); cascadesContext.toMemo(); return this;