From acd7ab379d497c8e62e91f5c26df6d58d6e6cd8f Mon Sep 17 00:00:00 2001 From: Kikyou1997 <33112463+Kikyou1997@users.noreply.github.com> Date: Sun, 28 Aug 2022 23:39:09 +0800 Subject: [PATCH] [feature](Nereids)add range partition prune (#11964) 1. Rewrite Filter(Project) to Project(Filter) to make sure when do partition prune the tree looks like this: Project(Filter(OlapScan)). 2. Enable the MergeConsecutiveProject MergeConsecutiveFilter rules. 3. prune range partition just like what Legacy Planner do. --- .../apache/doris/nereids/NereidsPlanner.java | 11 +- .../doris/nereids/jobs/batch/RewriteJob.java | 11 +- .../apache/doris/nereids/rules/RuleType.java | 2 + .../LogicalOlapScanToPhysicalOlapScan.java | 16 +- .../logical/PruneOlapScanPartition.java | 179 ++++++++++++++++++ .../rewrite/logical/SwapFilterAndProject.java | 56 ++++++ .../nereids/trees/expressions/Expression.java | 20 ++ .../trees/expressions/InPredicate.java | 13 ++ .../trees/expressions/SlotReference.java | 5 + .../trees/plans/logical/LogicalOlapScan.java | 62 +++--- .../trees/plans/logical/LogicalRelation.java | 16 +- .../apache/doris/planner/OlapScanNode.java | 2 +- .../org/apache/doris/planner/ScanNode.java | 8 +- .../nereids/jobs/RewriteTopDownJobTest.java | 4 +- .../logical/PruneOlapScanPartitionTest.java | 165 ++++++++++++++++ 15 files changed, 520 insertions(+), 50 deletions(-) create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/logical/PruneOlapScanPartition.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/logical/SwapFilterAndProject.java create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/logical/PruneOlapScanPartitionTest.java 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 5dcfc4f936..e617e2a79a 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 @@ -61,8 +61,7 @@ public class NereidsPlanner extends Planner { } @Override - public void plan(StatementBase queryStmt, - org.apache.doris.thrift.TQueryOptions queryOptions) throws UserException { + public void plan(StatementBase queryStmt, org.apache.doris.thrift.TQueryOptions queryOptions) throws UserException { if (!(queryStmt instanceof LogicalPlanAdapter)) { throw new RuntimeException("Wrong type of queryStmt, expected: "); } @@ -80,8 +79,8 @@ public class NereidsPlanner extends Planner { // set output exprs logicalPlanAdapter.setResultExprs(root.getOutputExprs()); - ArrayList columnLabelList = physicalPlan.getOutput().stream() - .map(NamedExpression::getName).collect(Collectors.toCollection(ArrayList::new)); + ArrayList columnLabelList = physicalPlan.getOutput().stream().map(NamedExpression::getName) + .collect(Collectors.toCollection(ArrayList::new)); logicalPlanAdapter.setColLabels(columnLabelList); } @@ -144,8 +143,8 @@ public class NereidsPlanner extends Planner { } private void deriveStats() { - cascadesContext - .pushJob(new DeriveStatsJob(getRoot().getLogicalExpression(), cascadesContext.getCurrentJobContext())); + cascadesContext.pushJob( + new DeriveStatsJob(getRoot().getLogicalExpression(), cascadesContext.getCurrentJobContext())); cascadesContext.getJobScheduler().executeJobPool(cascadesContext); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/batch/RewriteJob.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/batch/RewriteJob.java index 15f394d8de..8a3e5b4c83 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/batch/RewriteJob.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/batch/RewriteJob.java @@ -25,8 +25,10 @@ import org.apache.doris.nereids.rules.rewrite.logical.FindHashConditionForJoin; import org.apache.doris.nereids.rules.rewrite.logical.MergeConsecutiveFilters; import org.apache.doris.nereids.rules.rewrite.logical.MergeConsecutiveLimits; import org.apache.doris.nereids.rules.rewrite.logical.MergeConsecutiveProjects; +import org.apache.doris.nereids.rules.rewrite.logical.PruneOlapScanPartition; import org.apache.doris.nereids.rules.rewrite.logical.PushPredicateThroughJoin; import org.apache.doris.nereids.rules.rewrite.logical.ReorderJoin; +import org.apache.doris.nereids.rules.rewrite.logical.SwapFilterAndProject; import com.google.common.collect.ImmutableList; @@ -43,15 +45,18 @@ public class RewriteJob extends BatchRulesJob { public RewriteJob(CascadesContext cascadesContext) { super(cascadesContext); ImmutableList jobs = new ImmutableList.Builder() - .add(bottomUpBatch(ImmutableList.of(new MergeConsecutiveProjects()))) - .add(bottomUpBatch(ImmutableList.of(new MergeConsecutiveFilters()))) - .add(bottomUpBatch(ImmutableList.of(new MergeConsecutiveLimits()))) .add(topDownBatch(ImmutableList.of(new ExpressionNormalization()))) .add(topDownBatch(ImmutableList.of(new ReorderJoin()))) .add(topDownBatch(ImmutableList.of(new FindHashConditionForJoin()))) .add(topDownBatch(ImmutableList.of(new PushPredicateThroughJoin()))) .add(topDownBatch(ImmutableList.of(new AggregateDisassemble()))) + .add(topDownBatch(ImmutableList.of(new SwapFilterAndProject()))) + .add(bottomUpBatch(ImmutableList.of(new MergeConsecutiveProjects()))) + .add(topDownBatch(ImmutableList.of(new MergeConsecutiveFilters()))) + .add(bottomUpBatch(ImmutableList.of(new MergeConsecutiveLimits()))) + .add(topDownBatch(ImmutableList.of(new PruneOlapScanPartition()))) .build(); + rulesJob.addAll(jobs); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java index 4d991102e3..d1c6d92329 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java @@ -75,6 +75,8 @@ public enum RuleType { MERGE_CONSECUTIVE_LIMITS(RuleTypeClass.REWRITE), FIND_HASH_CONDITION_FOR_JOIN(RuleTypeClass.REWRITE), REWRITE_SENTINEL(RuleTypeClass.REWRITE), + OLAP_SCAN_PARTITION_PRUNE(RuleTypeClass.REWRITE), + SWAP_FILTER_AND_PROJECT(RuleTypeClass.REWRITE), // exploration rules LOGICAL_JOIN_COMMUTATIVE(RuleTypeClass.EXPLORATION), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java index 1555450624..ee48065739 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java @@ -44,14 +44,14 @@ public class LogicalOlapScanToPhysicalOlapScan extends OneImplementationRuleFact public Rule build() { return logicalOlapScan().then(olapScan -> new PhysicalOlapScan( - olapScan.getTable(), - olapScan.getQualifier(), - olapScan.getSelectedIndexId(), - olapScan.getSelectedTabletId(), - olapScan.getSelectedPartitionId(), - convertDistribution(olapScan), - Optional.empty(), - olapScan.getLogicalProperties()) + olapScan.getTable(), + olapScan.getQualifier(), + olapScan.getSelectedIndexId(), + olapScan.getSelectedTabletId(), + olapScan.getSelectedPartitionIds(), + convertDistribution(olapScan), + Optional.empty(), + olapScan.getLogicalProperties()) ).toRule(RuleType.LOGICAL_OLAP_SCAN_TO_PHYSICAL_OLAP_SCAN_RULE); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/logical/PruneOlapScanPartition.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/logical/PruneOlapScanPartition.java new file mode 100644 index 0000000000..506e2d6b17 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/logical/PruneOlapScanPartition.java @@ -0,0 +1,179 @@ +// 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.rules.rewrite.logical; + +import org.apache.doris.analysis.LiteralExpr; +import org.apache.doris.catalog.OlapTable; +import org.apache.doris.catalog.PartitionInfo; +import org.apache.doris.catalog.PartitionItem; +import org.apache.doris.catalog.PartitionType; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.rules.rewrite.OneRewriteRuleFactory; +import org.apache.doris.nereids.trees.expressions.ComparisonPredicate; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.GreaterThan; +import org.apache.doris.nereids.trees.expressions.GreaterThanEqual; +import org.apache.doris.nereids.trees.expressions.LessThan; +import org.apache.doris.nereids.trees.expressions.LessThanEqual; +import org.apache.doris.nereids.trees.expressions.Or; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; +import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.nereids.util.Utils; +import org.apache.doris.planner.ColumnBound; +import org.apache.doris.planner.ColumnRange; +import org.apache.doris.planner.PartitionPruner; +import org.apache.doris.planner.RangePartitionPrunerV2; +import org.apache.doris.planner.ScanNode.ColumnRanges; + +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; +import com.google.common.collect.Range; +import com.google.common.collect.Sets; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Used to prune partition of olap scan, should execute after SwapProjectAndFilter, MergeConsecutiveFilters, + * MergeConsecutiveProjects and all predicate push down related rules. + */ +public class PruneOlapScanPartition extends OneRewriteRuleFactory { + + @Override + public Rule build() { + return logicalFilter(logicalOlapScan()).when(p -> !p.child().isPartitionPruned()).thenApply(ctx -> { + LogicalFilter filter = ctx.root; + LogicalOlapScan scan = filter.child(); + Expression predicate = filter.getPredicates(); + OlapTable table = scan.getTable(); + Set partitionColumnNameSet = Utils.execWithReturnVal(table::getPartitionColumnNames); + PartitionInfo partitionInfo = table.getPartitionInfo(); + // TODO: 1. support grammar: SELECT * FROM tbl PARTITION(p1,p2) + // 2. support list partition + if (partitionColumnNameSet.isEmpty() || !partitionInfo.getType().equals(PartitionType.RANGE)) { + return ctx.root; + } + List expressionList = ExpressionUtils.extractConjunction(predicate); + // TODO: Process all partition column for now, better to process required column only. + Map columnNameToRange = Maps.newHashMap(); + for (String colName : partitionColumnNameSet) { + ColumnRange columnRange = createColumnRange(colName, expressionList); + columnNameToRange.put(colName, columnRange); + } + + Map keyItemMap = partitionInfo.getIdToItem(false); + PartitionPruner partitionPruner = new RangePartitionPrunerV2(keyItemMap, + partitionInfo.getPartitionColumns(), columnNameToRange); + Collection selectedPartitionId = Utils.execWithReturnVal(partitionPruner::prune); + LogicalOlapScan rewrittenScan = + scan.withSelectedPartitionId(new ArrayList<>(selectedPartitionId)); + return new LogicalFilter<>(filter.getPredicates(), rewrittenScan); + }).toRule(RuleType.OLAP_SCAN_PARTITION_PRUNE); + } + + private ColumnRange createColumnRange(String colName, List expressionList) { + ColumnRange result = ColumnRange.create(); + for (Expression expression : expressionList) { + List slotReferenceList = expression.collect(SlotReference.class::isInstance); + int slotReferenceListSize = new HashSet<>(slotReferenceList).size(); + if (slotReferenceListSize != 1 || !slotReferenceList.get(0).getName().equals(colName)) { + continue; + } + if (expression instanceof Or) { + List disjunctiveList = ExpressionUtils.extractDisjunction(expression); + if (disjunctiveList.isEmpty()) { + continue; + } + List> disjunctiveRanges = Lists.newArrayList(); + Set hasIsNull = Sets.newHashSet(); + boolean allMatch = disjunctiveList.stream().allMatch(e -> { + ColumnRanges ranges = exprToRanges(e, colName); + switch (ranges.type) { + case IS_NULL: + hasIsNull.add(true); + return true; + case CONVERT_SUCCESS: + disjunctiveRanges.addAll(ranges.ranges); + return true; + case CONVERT_FAILURE: + default: + return false; + } + }); + if (allMatch && !(disjunctiveRanges.isEmpty() && hasIsNull.isEmpty())) { + result.intersect(disjunctiveRanges); + result.setHasDisjunctiveIsNull(!hasIsNull.isEmpty()); + } + } else { + ColumnRanges ranges = exprToRanges(expression, colName); + switch (ranges.type) { + case IS_NULL: + result.setHasConjunctiveIsNull(true); + break; + case CONVERT_SUCCESS: + result.intersect(ranges.ranges); + break; + case CONVERT_FAILURE: + default: + break; + } + } + } + return result; + } + + private ColumnRanges exprToRanges(Expression expression, String colName) { + // TODO: process in/is null expression + if (!(expression instanceof ComparisonPredicate)) { + return ColumnRanges.createFailure(); + } + List> result = Lists.newArrayList(); + ComparisonPredicate comparisonPredicate = (ComparisonPredicate) expression; + Expression rightChild = comparisonPredicate.child(1); + if (rightChild == null || !rightChild.isConstant() || !(rightChild instanceof Literal)) { + return ColumnRanges.createFailure(); + } + LiteralExpr value = ((Literal) rightChild).toLegacyLiteral(); + if (expression instanceof EqualTo) { + ColumnBound bound = ColumnBound.of(value); + result.add(Range.closed(bound, bound)); + } else if (expression instanceof GreaterThanEqual) { + result.add(Range.atLeast(ColumnBound.of(value))); + } else if (expression instanceof GreaterThan) { + result.add(Range.greaterThan(ColumnBound.of(value))); + } else if (expression instanceof LessThan) { + result.add(Range.lessThan(ColumnBound.of(value))); + } else if (expression instanceof LessThanEqual) { + result.add(Range.atMost(ColumnBound.of(value))); + } + if (result.isEmpty()) { + return ColumnRanges.createFailure(); + } else { + return ColumnRanges.create(result); + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/logical/SwapFilterAndProject.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/logical/SwapFilterAndProject.java new file mode 100644 index 0000000000..d8c1e56c13 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/logical/SwapFilterAndProject.java @@ -0,0 +1,56 @@ +// 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.rules.rewrite.logical; + +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.rules.rewrite.OneRewriteRuleFactory; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionReplacer; +import org.apache.doris.nereids.trees.plans.GroupPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Rewrite filter -> project to project -> filter. + */ +public class SwapFilterAndProject extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalFilter(logicalProject()).thenApply(ctx -> { + LogicalFilter> filter = ctx.root; + LogicalProject project = filter.child(); + List namedExpressionList = project.getProjects(); + Map slotToAlias = new HashMap<>(); + namedExpressionList.stream().filter(Alias.class::isInstance).forEach(s -> { + slotToAlias.put(s.toSlot(), ((Alias) s).child()); + }); + Expression rewrittenPredicate = ExpressionReplacer.INSTANCE.visit(filter.getPredicates(), slotToAlias); + LogicalFilter rewrittenFilter = + new LogicalFilter(rewrittenPredicate, project.child()); + return new LogicalProject(project.getProjects(), rewrittenFilter); + }).toRule(RuleType.SWAP_FILTER_AND_PROJECT); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Expression.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Expression.java index ba68f8fca3..818ef3796e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Expression.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Expression.java @@ -146,4 +146,24 @@ public abstract class Expression extends AbstractTreeNode { public int hashCode() { return 0; } + + /** + * Return true if all the SlotRef in the expr tree is bound to the same column. + */ + public boolean boundToColumn(String column) { + for (Expression child : children) { + if (!child.boundToColumn(column)) { + return false; + } + } + return true; + } + + public Expression leftMostNode() { + Expression leftChild = this; + while (leftChild.children.size() > 0) { + leftChild = leftChild.child(0); + } + return leftChild; + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/InPredicate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/InPredicate.java index a3904b2367..8f6b5f14ab 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/InPredicate.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/InPredicate.java @@ -18,6 +18,7 @@ package org.apache.doris.nereids.trees.expressions; import org.apache.doris.nereids.exceptions.UnboundException; +import org.apache.doris.nereids.trees.expressions.literal.Literal; import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; import org.apache.doris.nereids.types.BooleanType; import org.apache.doris.nereids.types.DataType; @@ -103,4 +104,16 @@ public class InPredicate extends Expression { public List getOptions() { return options; } + + /** + * Return true when all children are Literal , otherwise, return false. + */ + public boolean isLiteralChildren() { + for (Expression expression : options) { + if (!(expression instanceof Literal)) { + return false; + } + } + return true; + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SlotReference.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SlotReference.java index a7f02be03d..e8e04e901a 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SlotReference.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SlotReference.java @@ -154,4 +154,9 @@ public class SlotReference extends Slot { public Slot withQualifier(List qualifiers) { return new SlotReference(exprId, name, dataType, nullable, qualifiers); } + + @Override + public boolean boundToColumn(String name) { + return this.name.equals(name); + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java index d17944b0bd..01031f1cfd 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java @@ -42,14 +42,20 @@ public class LogicalOlapScan extends LogicalRelation { private final long selectedIndexId; private final List selectedTabletId; - private final List selectedPartitionId; + private final boolean partitionPruned; public LogicalOlapScan(OlapTable table) { this(table, ImmutableList.of()); } public LogicalOlapScan(OlapTable table, List qualifier) { - this(table, qualifier, Optional.empty(), Optional.empty()); + this(table, qualifier, Optional.empty(), Optional.empty(), + table.getPartitionIds(), false); + } + + public LogicalOlapScan(Table table, List qualifier) { + this(table, qualifier, Optional.empty(), Optional.empty(), + ((OlapTable) table).getPartitionIds(), false); } /** @@ -58,27 +64,17 @@ public class LogicalOlapScan extends LogicalRelation { * @param table Doris table * @param qualifier table name qualifier */ - public LogicalOlapScan(Table table, List qualifier, - Optional groupExpression, Optional logicalProperties) { - super(PlanType.LOGICAL_OLAP_SCAN, table, qualifier, groupExpression, logicalProperties); + public LogicalOlapScan(Table table, List qualifier, Optional groupExpression, + Optional logicalProperties, List selectedPartitionIdList, + boolean partitionPruned) { + super(PlanType.LOGICAL_OLAP_SCAN, table, qualifier, + groupExpression, logicalProperties, selectedPartitionIdList); this.selectedIndexId = getTable().getBaseIndexId(); this.selectedTabletId = Lists.newArrayList(); - this.selectedPartitionId = getTable().getPartitionIds(); for (Partition partition : getTable().getAllPartitions()) { selectedTabletId.addAll(partition.getBaseIndex().getTabletIdsInOrder()); } - } - - public List getSelectedTabletId() { - return selectedTabletId; - } - - public long getSelectedIndexId() { - return selectedIndexId; - } - - public List getSelectedPartitionId() { - return selectedPartitionId; + this.partitionPruned = partitionPruned; } @Override @@ -89,11 +85,8 @@ public class LogicalOlapScan extends LogicalRelation { @Override public String toString() { - return "ScanOlapTable (" - + qualifiedName() - + ", output: " - + getOutput().stream().map(Objects::toString).collect(Collectors.joining(", ", "[", "]")) - + ")"; + return "ScanOlapTable (" + qualifiedName() + ", output: " + getOutput().stream().map(Objects::toString) + .collect(Collectors.joining(", ", "[", "]")) + ")"; } @Override @@ -109,16 +102,35 @@ public class LogicalOlapScan extends LogicalRelation { @Override public Plan withGroupExpression(Optional groupExpression) { - return new LogicalOlapScan(table, qualifier, groupExpression, Optional.of(logicalProperties)); + return new LogicalOlapScan(table, qualifier, groupExpression, Optional.of(logicalProperties), + selectedPartitionIds, partitionPruned); } @Override public LogicalOlapScan withLogicalProperties(Optional logicalProperties) { - return new LogicalOlapScan(table, qualifier, Optional.empty(), logicalProperties); + return new LogicalOlapScan(table, qualifier, Optional.empty(), logicalProperties, selectedPartitionIds, + partitionPruned); + } + + public LogicalOlapScan withSelectedPartitionId(List selectedPartitionId) { + return new LogicalOlapScan(table, qualifier, Optional.empty(), Optional.of(logicalProperties), + selectedPartitionId, true); } @Override public R accept(PlanVisitor visitor, C context) { return visitor.visitLogicalOlapScan(this, context); } + + public boolean isPartitionPruned() { + return partitionPruned; + } + + public List getSelectedTabletId() { + return selectedTabletId; + } + + public long getSelectedIndexId() { + return selectedIndexId; + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalRelation.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalRelation.java index 21a977d394..eee472ace7 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalRelation.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalRelation.java @@ -29,7 +29,9 @@ import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; import org.apache.doris.nereids.util.Utils; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; +import java.util.Collections; import java.util.List; import java.util.Objects; import java.util.Optional; @@ -42,8 +44,10 @@ public abstract class LogicalRelation extends LogicalLeaf implements Scan { protected final Table table; protected final List qualifier; + protected List selectedPartitionIds = Lists.newArrayList(); + public LogicalRelation(PlanType type, Table table, List qualifier) { - this(type, table, qualifier, Optional.empty(), Optional.empty()); + this(type, table, qualifier, Optional.empty(), Optional.empty(), Collections.emptyList()); } /** @@ -53,10 +57,13 @@ public abstract class LogicalRelation extends LogicalLeaf implements Scan { * @param qualifier qualified relation name */ public LogicalRelation(PlanType type, Table table, List qualifier, - Optional groupExpression, Optional logicalProperties) { + Optional groupExpression, + Optional logicalProperties, + List selectedPartitionIdList) { super(type, groupExpression, logicalProperties); this.table = Objects.requireNonNull(table, "table can not be null"); this.qualifier = ImmutableList.copyOf(Objects.requireNonNull(qualifier, "qualifier can not be null")); + this.selectedPartitionIds = selectedPartitionIdList; } public Table getTable() { @@ -115,4 +122,9 @@ public abstract class LogicalRelation extends LogicalLeaf implements Scan { public String qualifiedName() { return Utils.qualifiedName(qualifier, table.getName()); } + + public List getSelectedPartitionIds() { + return selectedPartitionIds; + } + } diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java index 5bc2d6834e..5b010be574 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java @@ -231,7 +231,7 @@ public class OlapScanNode extends ScanNode { this.tupleIds = tupleIds; } - // only used for UT + // only used for UT and Nereids public void setSelectedPartitionIds(Collection selectedPartitionIds) { this.selectedPartitionIds = selectedPartitionIds; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/ScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/planner/ScanNode.java index f8a90cfa99..c522823622 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/planner/ScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/planner/ScanNode.java @@ -352,8 +352,8 @@ public abstract class ScanNode extends PlanNode { return partitionColumnFilter; } - private static class ColumnRanges { - enum Type { + public static class ColumnRanges { + public enum Type { // Expression is `is null` predicate. IS_NULL, // Succeed to convert expression to ranges. @@ -362,8 +362,8 @@ public abstract class ScanNode extends PlanNode { CONVERT_FAILURE } - final Type type; - final List> ranges; + public final Type type; + public final List> ranges; private ColumnRanges(Type type, List> ranges) { this.type = type; diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/jobs/RewriteTopDownJobTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/jobs/RewriteTopDownJobTest.java index 06b9d2ecac..faf6863bb6 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/jobs/RewriteTopDownJobTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/jobs/RewriteTopDownJobTest.java @@ -42,6 +42,7 @@ import com.google.common.collect.Lists; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.util.Collections; import java.util.List; import java.util.Optional; @@ -97,7 +98,8 @@ public class RewriteTopDownJobTest { public LogicalBoundRelation(Table table, List qualifier, Optional groupExpression, Optional logicalProperties) { - super(PlanType.LOGICAL_BOUND_RELATION, table, qualifier, groupExpression, logicalProperties); + super(PlanType.LOGICAL_BOUND_RELATION, table, qualifier, groupExpression, logicalProperties, + Collections.emptyList()); } @Override diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/logical/PruneOlapScanPartitionTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/logical/PruneOlapScanPartitionTest.java new file mode 100644 index 0000000000..1cec7f2738 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/logical/PruneOlapScanPartitionTest.java @@ -0,0 +1,165 @@ +// 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.rules.rewrite.logical; + +import org.apache.doris.analysis.IntLiteral; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.OlapTable; +import org.apache.doris.catalog.PartitionItem; +import org.apache.doris.catalog.PartitionKey; +import org.apache.doris.catalog.RangePartitionInfo; +import org.apache.doris.catalog.RangePartitionItem; +import org.apache.doris.catalog.Type; +import org.apache.doris.common.jmockit.Deencapsulation; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.trees.expressions.And; +import org.apache.doris.nereids.trees.expressions.CompoundPredicate; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.GreaterThan; +import org.apache.doris.nereids.trees.expressions.GreaterThanEqual; +import org.apache.doris.nereids.trees.expressions.LessThan; +import org.apache.doris.nereids.trees.expressions.LessThanEqual; +import org.apache.doris.nereids.trees.expressions.Or; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.util.MemoTestUtils; + +import com.google.common.collect.BoundType; +import com.google.common.collect.Lists; +import com.google.common.collect.Range; +import mockit.Expectations; +import mockit.Mocked; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +class PruneOlapScanPartitionTest { + + @Test + public void testOlapScanPartitionWithSingleColumnCase(@Mocked OlapTable olapTable) throws Exception { + List columnNameList = new ArrayList<>(); + columnNameList.add(new Column("col1", Type.INT.getPrimitiveType())); + columnNameList.add(new Column("col2", Type.INT.getPrimitiveType())); + Map keyItemMap = new HashMap<>(); + PartitionKey k0 = new PartitionKey(); + k0.pushColumn(new IntLiteral(0), Type.INT.getPrimitiveType()); + PartitionKey k1 = new PartitionKey(); + k1.pushColumn(new IntLiteral(5), Type.INT.getPrimitiveType()); + keyItemMap.put(0L, new RangePartitionItem(Range.range(k0, BoundType.CLOSED, k1, BoundType.OPEN))); + PartitionKey k2 = new PartitionKey(); + k2.pushColumn(new IntLiteral(5), Type.INT.getPrimitiveType()); + PartitionKey k3 = new PartitionKey(); + k3.pushColumn(new IntLiteral(10), Type.INT.getPrimitiveType()); + keyItemMap.put(1L, new RangePartitionItem(Range.range(k2, BoundType.CLOSED, k3, BoundType.OPEN))); + RangePartitionInfo rangePartitionInfo = new RangePartitionInfo(columnNameList); + Deencapsulation.setField(rangePartitionInfo, "idToItem", keyItemMap); + new Expectations() {{ + olapTable.getPartitionInfo(); + result = rangePartitionInfo; + olapTable.getPartitionColumnNames(); + result = rangePartitionInfo.getPartitionColumns().stream().map(c -> c.getName().toLowerCase()) + .collect(Collectors.toSet()); + + }}; + LogicalOlapScan scan = new LogicalOlapScan(olapTable); + SlotReference slotRef = new SlotReference("col1", IntegerType.INSTANCE); + Expression expression = new LessThan(slotRef, new IntegerLiteral(4)); + LogicalFilter filter = new LogicalFilter<>(expression, scan); + + CascadesContext cascadesContext = MemoTestUtils.createCascadesContext(filter); + List rules = Lists.newArrayList(new PruneOlapScanPartition().build()); + cascadesContext.topDownRewrite(rules); + Plan resultPlan = cascadesContext.getMemo().copyOut(); + LogicalOlapScan rewrittenOlapScan = (LogicalOlapScan) resultPlan.child(0); + Assertions.assertEquals(0L, rewrittenOlapScan.getSelectedPartitionIds().toArray()[0]); + + Expression lessThan0 = new LessThan(slotRef, new IntegerLiteral(0)); + Expression greaterThan6 = new GreaterThan(slotRef, new IntegerLiteral(6)); + Or lessThan0OrGreaterThan6 = new Or(lessThan0, greaterThan6); + filter = new LogicalFilter<>(lessThan0OrGreaterThan6, scan); + scan = new LogicalOlapScan(olapTable); + cascadesContext = MemoTestUtils.createCascadesContext(filter); + rules = Lists.newArrayList(new PruneOlapScanPartition().build()); + cascadesContext.topDownRewrite(rules); + resultPlan = cascadesContext.getMemo().copyOut(); + rewrittenOlapScan = (LogicalOlapScan) resultPlan.child(0); + Assertions.assertEquals(1L, rewrittenOlapScan.getSelectedPartitionIds().toArray()[0]); + + Expression greaterThanEqual0 = + new GreaterThanEqual( + slotRef, new IntegerLiteral(0)); + Expression lessThanEqual5 = + new LessThanEqual(slotRef, new IntegerLiteral(5)); + And greaterThanEqual0AndLessThanEqual5 = new And(greaterThanEqual0, lessThanEqual5); + scan = new LogicalOlapScan(olapTable); + filter = new LogicalFilter<>(greaterThanEqual0AndLessThanEqual5, scan); + cascadesContext = MemoTestUtils.createCascadesContext(filter); + rules = Lists.newArrayList(new PruneOlapScanPartition().build()); + cascadesContext.topDownRewrite(rules); + resultPlan = cascadesContext.getMemo().copyOut(); + rewrittenOlapScan = (LogicalOlapScan) resultPlan.child(0); + Assertions.assertEquals(0L, rewrittenOlapScan.getSelectedPartitionIds().toArray()[0]); + Assertions.assertEquals(2, rewrittenOlapScan.getSelectedPartitionIds().toArray().length); + } + + @Test + public void testOlapScanPartitionPruneWithMultiColumnCase(@Mocked OlapTable olapTable) throws Exception { + List columnNameList = new ArrayList<>(); + columnNameList.add(new Column("col1", Type.INT.getPrimitiveType())); + columnNameList.add(new Column("col2", Type.INT.getPrimitiveType())); + Map keyItemMap = new HashMap<>(); + PartitionKey k0 = new PartitionKey(); + k0.pushColumn(new IntLiteral(1), Type.INT.getPrimitiveType()); + k0.pushColumn(new IntLiteral(10), Type.INT.getPrimitiveType()); + PartitionKey k1 = new PartitionKey(); + k1.pushColumn(new IntLiteral(4), Type.INT.getPrimitiveType()); + k1.pushColumn(new IntLiteral(5), Type.INT.getPrimitiveType()); + keyItemMap.put(0L, new RangePartitionItem(Range.range(k0, BoundType.CLOSED, k1, BoundType.OPEN))); + RangePartitionInfo rangePartitionInfo = new RangePartitionInfo(columnNameList); + Deencapsulation.setField(rangePartitionInfo, "idToItem", keyItemMap); + new Expectations() {{ + olapTable.getPartitionInfo(); + result = rangePartitionInfo; + olapTable.getPartitionColumnNames(); + result = rangePartitionInfo.getPartitionColumns().stream().map(c -> c.getName().toLowerCase()) + .collect(Collectors.toSet()); + }}; + LogicalOlapScan scan = new LogicalOlapScan(olapTable); + Expression left = new LessThan(new SlotReference("col1", IntegerType.INSTANCE), new IntegerLiteral(4)); + Expression right = new GreaterThan(new SlotReference("col2", IntegerType.INSTANCE), new IntegerLiteral(11)); + CompoundPredicate and = new And(left, right); + LogicalFilter filter = new LogicalFilter<>(and, scan); + CascadesContext cascadesContext = MemoTestUtils.createCascadesContext(filter); + List rules = Lists.newArrayList(new PruneOlapScanPartition().build()); + cascadesContext.topDownRewrite(rules); + Plan resultPlan = cascadesContext.getMemo().copyOut(); + LogicalOlapScan rewrittenOlapScan = (LogicalOlapScan) resultPlan.child(0); + Assertions.assertEquals(0L, rewrittenOlapScan.getSelectedPartitionIds().toArray()[0]); + } + +}