database/doris
2
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

[feature](Nereids)add range partition prune (#11964)

Browse Source 
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.
This commit is contained in:
Kikyou1997
2022-08-28 23:39:09 +08:00
committed by GitHub
parent 6e6269c682
commit acd7ab379d
15 changed files with 520 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
fe/fe-core/src/main/java/org/apache/doris/nereids/NereidsPlanner.java
View File

@ -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: <? extends LogicalPlanAdapter>");
}
@ -80,8 +79,8 @@ public class NereidsPlanner extends Planner {
// set output exprs
logicalPlanAdapter.setResultExprs(root.getOutputExprs());
ArrayList<String> columnLabelList = physicalPlan.getOutput().stream()
.map(NamedExpression::getName).collect(Collectors.toCollection(ArrayList::new));
ArrayList<String> 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);
}

11
fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/batch/RewriteJob.java
View File

@ -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<Job> jobs = new ImmutableList.Builder<Job>()
.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);
}
}

2
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java
View File

@ -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),

16
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java
View File

@ -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);
}

179
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/logical/PruneOlapScanPartition.java Normal file
View File

@ -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<LogicalOlapScan> filter = ctx.root;
LogicalOlapScan scan = filter.child();
Expression predicate = filter.getPredicates();
OlapTable table = scan.getTable();
Set<String> 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<Expression> expressionList = ExpressionUtils.extractConjunction(predicate);
// TODO: Process all partition column for now, better to process required column only.
Map<String, ColumnRange> columnNameToRange = Maps.newHashMap();
for (String colName : partitionColumnNameSet) {
ColumnRange columnRange = createColumnRange(colName, expressionList);
columnNameToRange.put(colName, columnRange);
}
Map<Long, PartitionItem> keyItemMap = partitionInfo.getIdToItem(false);
PartitionPruner partitionPruner = new RangePartitionPrunerV2(keyItemMap,
partitionInfo.getPartitionColumns(), columnNameToRange);
Collection<Long> 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<Expression> expressionList) {
ColumnRange result = ColumnRange.create();
for (Expression expression : expressionList) {
List<SlotReference> 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<Expression> disjunctiveList = ExpressionUtils.extractDisjunction(expression);
if (disjunctiveList.isEmpty()) {
continue;
}
List<Range<ColumnBound>> disjunctiveRanges = Lists.newArrayList();
Set<Boolean> 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<Range<ColumnBound>> 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);
}
}
}

56
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/logical/SwapFilterAndProject.java Normal file
View File

@ -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<LogicalProject<GroupPlan>> filter = ctx.root;
LogicalProject<GroupPlan> project = filter.child();
List<NamedExpression> namedExpressionList = project.getProjects();
Map<Expression, Expression> 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<LogicalPlan> rewrittenFilter =
new LogicalFilter<LogicalPlan>(rewrittenPredicate, project.child());
return new LogicalProject(project.getProjects(), rewrittenFilter);
}).toRule(RuleType.SWAP_FILTER_AND_PROJECT);
}
}

20
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Expression.java
View File

@ -146,4 +146,24 @@ public abstract class Expression extends AbstractTreeNode<Expression> {
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;
}
}

13
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/InPredicate.java
View File

@ -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<Expression> 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;
}
}

5
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SlotReference.java
View File

@ -154,4 +154,9 @@ public class SlotReference extends Slot {
public Slot withQualifier(List<String> qualifiers) {
return new SlotReference(exprId, name, dataType, nullable, qualifiers);
}
@Override
public boolean boundToColumn(String name) {
return this.name.equals(name);
}
}

62
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java
View File

@ -42,14 +42,20 @@ public class LogicalOlapScan extends LogicalRelation {
private final long selectedIndexId;
private final List<Long> selectedTabletId;
private final List<Long> selectedPartitionId;
private final boolean partitionPruned;
public LogicalOlapScan(OlapTable table) {
this(table, ImmutableList.of());
}
public LogicalOlapScan(OlapTable table, List<String> qualifier) {
this(table, qualifier, Optional.empty(), Optional.empty());
this(table, qualifier, Optional.empty(), Optional.empty(),
table.getPartitionIds(), false);
}
public LogicalOlapScan(Table table, List<String> 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<String> qualifier,
Optional<GroupExpression> groupExpression, Optional<LogicalProperties> logicalProperties) {
super(PlanType.LOGICAL_OLAP_SCAN, table, qualifier, groupExpression, logicalProperties);
public LogicalOlapScan(Table table, List<String> qualifier, Optional<GroupExpression> groupExpression,
Optional<LogicalProperties> logicalProperties, List<Long> 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<Long> getSelectedTabletId() {
return selectedTabletId;
}
public long getSelectedIndexId() {
return selectedIndexId;
}
public List<Long> 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> 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> logicalProperties) {
return new LogicalOlapScan(table, qualifier, Optional.empty(), logicalProperties);
return new LogicalOlapScan(table, qualifier, Optional.empty(), logicalProperties, selectedPartitionIds,
partitionPruned);
}
public LogicalOlapScan withSelectedPartitionId(List<Long> selectedPartitionId) {
return new LogicalOlapScan(table, qualifier, Optional.empty(), Optional.of(logicalProperties),
selectedPartitionId, true);
}
@Override
public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
return visitor.visitLogicalOlapScan(this, context);
}
public boolean isPartitionPruned() {
return partitionPruned;
}
public List<Long> getSelectedTabletId() {
return selectedTabletId;
}
public long getSelectedIndexId() {
return selectedIndexId;
}
}

16
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalRelation.java
View File

@ -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<String> qualifier;
protected List<Long> selectedPartitionIds = Lists.newArrayList();
public LogicalRelation(PlanType type, Table table, List<String> 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<String> qualifier,
Optional<GroupExpression> groupExpression, Optional<LogicalProperties> logicalProperties) {
Optional<GroupExpression> groupExpression,
Optional<LogicalProperties> logicalProperties,
List<Long> 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<Long> getSelectedPartitionIds() {
return selectedPartitionIds;
}
}

2
fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java
View File

@ -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<Long> selectedPartitionIds) {
this.selectedPartitionIds = selectedPartitionIds;
}

8
fe/fe-core/src/main/java/org/apache/doris/planner/ScanNode.java
View File

@ -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<Range<ColumnBound>> ranges;
public final Type type;
public final List<Range<ColumnBound>> ranges;
private ColumnRanges(Type type, List<Range<ColumnBound>> ranges) {
this.type = type;

4
fe/fe-core/src/test/java/org/apache/doris/nereids/jobs/RewriteTopDownJobTest.java
View File

@ -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<String> qualifier, Optional<GroupExpression> groupExpression,
Optional<LogicalProperties> logicalProperties) {
super(PlanType.LOGICAL_BOUND_RELATION, table, qualifier, groupExpression, logicalProperties);
super(PlanType.LOGICAL_BOUND_RELATION, table, qualifier, groupExpression, logicalProperties,
Collections.emptyList());
}
@Override

165
fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/logical/PruneOlapScanPartitionTest.java Normal file
View File

@ -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<Column> columnNameList = new ArrayList<>();
columnNameList.add(new Column("col1", Type.INT.getPrimitiveType()));
columnNameList.add(new Column("col2", Type.INT.getPrimitiveType()));
Map<Long, PartitionItem> 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<LogicalOlapScan> filter = new LogicalFilter<>(expression, scan);
CascadesContext cascadesContext = MemoTestUtils.createCascadesContext(filter);
List<Rule> 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<Column> columnNameList = new ArrayList<>();
columnNameList.add(new Column("col1", Type.INT.getPrimitiveType()));
columnNameList.add(new Column("col2", Type.INT.getPrimitiveType()));
Map<Long, PartitionItem> 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<LogicalOlapScan> filter = new LogicalFilter<>(and, scan);
CascadesContext cascadesContext = MemoTestUtils.createCascadesContext(filter);
List<Rule> 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]);
}
}