This is an automated email from the ASF dual-hosted git repository.
huajianlan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new d0731c8a309 [opt](nereids) support prune partitions by specified
tablet ids (#50424)
d0731c8a309 is described below
commit d0731c8a309d61c81f8dcf20c063d8ad2f9d4971
Author: zhaorongsheng <[email protected]>
AuthorDate: Mon May 12 22:19:42 2025 +0800
[opt](nereids) support prune partitions by specified tablet ids (#50424)
Issue Number: close #50255
support prune partitions by specified tablet ids
---
.../rules/rewrite/PruneOlapScanPartition.java | 187 +++++++++++++++------
.../rules/rewrite/PruneOlapScanPartitionTest.java | 47 +++++-
2 files changed, 183 insertions(+), 51 deletions(-)
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartition.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartition.java
index debf4e880cc..023a17bec5a 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartition.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartition.java
@@ -19,10 +19,14 @@ package org.apache.doris.nereids.rules.rewrite;
import org.apache.doris.catalog.Column;
import org.apache.doris.catalog.Env;
+import org.apache.doris.catalog.MaterializedIndex;
import org.apache.doris.catalog.OlapTable;
+import org.apache.doris.catalog.Partition;
import org.apache.doris.catalog.PartitionInfo;
import org.apache.doris.catalog.PartitionItem;
+import org.apache.doris.catalog.Tablet;
import org.apache.doris.common.cache.NereidsSortedPartitionsCacheManager;
+import org.apache.doris.nereids.pattern.MatchingContext;
import org.apache.doris.nereids.rules.Rule;
import org.apache.doris.nereids.rules.RuleType;
import org.apache.doris.nereids.rules.expression.rules.PartitionPruner;
@@ -32,12 +36,15 @@ import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.plans.logical.LogicalEmptyRelation;
import org.apache.doris.nereids.trees.plans.logical.LogicalFilter;
import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalRelation;
import org.apache.doris.nereids.util.Utils;
import org.apache.doris.qe.ConnectContext;
import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
import java.util.ArrayList;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
@@ -49,65 +56,145 @@ import java.util.stream.Collectors;
* 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 {
+public class PruneOlapScanPartition implements RewriteRuleFactory {
@Override
- public Rule build() {
- return logicalFilter(logicalOlapScan()).when(p ->
!p.child().isPartitionPruned()).thenApply(ctx -> {
- LogicalFilter<LogicalOlapScan> filter = ctx.root;
- LogicalOlapScan scan = filter.child();
- OlapTable table = scan.getTable();
- Set<String> partitionColumnNameSet =
Utils.execWithReturnVal(table::getPartitionColumnNames);
- if (partitionColumnNameSet.isEmpty()) {
- return null;
- }
- List<Slot> output = scan.getOutput();
- PartitionInfo partitionInfo = table.getPartitionInfo();
- List<Column> partitionColumns =
partitionInfo.getPartitionColumns();
- List<Slot> partitionSlots = new
ArrayList<>(partitionColumns.size());
- for (Column column : partitionColumns) {
- Slot partitionSlot = null;
- // loop search is faster than build a map
- for (Slot slot : output) {
- if (slot.getName().equalsIgnoreCase(column.getName())) {
- partitionSlot = slot;
- break;
- }
- }
- if (partitionSlot == null) {
- return null;
- } else {
- partitionSlots.add(partitionSlot);
+ public List<Rule> buildRules() {
+ return ImmutableList.of(
+ logicalOlapScan()
+ .when(scan -> !scan.isPartitionPruned()
+ && !scan.getManuallySpecifiedTabletIds().isEmpty()
+ && scan.getTable().isPartitionedTable()
+ )
+ .thenApply(ctx -> {
+ // Case1: sql without filter condition, e.g. SELECT *
FROM tbl (${tabletID})
+ LogicalOlapScan scan = ctx.root;
+ OlapTable table = scan.getTable();
+ return prunePartition(scan, table, null, ctx);
+ }).toRule(RuleType.OLAP_SCAN_PARTITION_PRUNE),
+ logicalFilter(logicalOlapScan()
+ .whenNot(LogicalOlapScan::isPartitionPruned))
+ .thenApply(ctx -> {
+ // Case2: sql with filter condition, e.g. SELECT *
FROM tbl (${tabletID}) WHERE part_column='x'
+ LogicalFilter<LogicalOlapScan> filter = ctx.root;
+ LogicalOlapScan scan = filter.child();
+ OlapTable table = scan.getTable();
+ LogicalRelation rewrittenLogicalRelation =
prunePartition(scan, table, filter, ctx);
+ if (rewrittenLogicalRelation == null) {
+ return null;
+ }
+ if (rewrittenLogicalRelation instanceof
LogicalEmptyRelation) {
+ return rewrittenLogicalRelation;
+ } else {
+ LogicalOlapScan rewrittenScan = (LogicalOlapScan)
rewrittenLogicalRelation;
+ return
filter.withChildren(ImmutableList.of(rewrittenScan));
+ }
+ }).toRule(RuleType.OLAP_SCAN_PARTITION_PRUNE)
+ );
+ }
+
+ private LogicalRelation prunePartition(LogicalOlapScan scan,
+ OlapTable table,
+ LogicalFilter filter,
+ MatchingContext ctx) {
+ List<Long> prunedPartitionsByFilters = prunePartitionByFilters(scan,
table, filter, ctx);
+ List<Long> prunedPartitions = prunePartitionByTabletIds(scan, table,
prunedPartitionsByFilters);
+ if (prunedPartitions == null) {
+ return null;
+ }
+ if (prunedPartitions.isEmpty()) {
+ return new LogicalEmptyRelation(
+ ConnectContext.get().getStatementContext().getNextRelationId(),
+ ctx.root.getOutput());
+ }
+ return scan.withSelectedPartitionIds(prunedPartitions);
+ }
+
+ private List<Long> prunePartitionByFilters(LogicalOlapScan scan,
+ OlapTable table,
+ LogicalFilter filter,
+ MatchingContext ctx) {
+ Set<String> partitionColumnNameSet =
Utils.execWithReturnVal(table::getPartitionColumnNames);
+ if (partitionColumnNameSet.isEmpty()) {
+ return null;
+ }
+ List<Slot> output = scan.getOutput();
+ PartitionInfo partitionInfo = table.getPartitionInfo();
+ List<Column> partitionColumns = partitionInfo.getPartitionColumns();
+ List<Slot> partitionSlots = new ArrayList<>(partitionColumns.size());
+ for (Column column : partitionColumns) {
+ Slot partitionSlot = null;
+ // loop search is faster than build a map
+ for (Slot slot : output) {
+ if (slot.getName().equalsIgnoreCase(column.getName())) {
+ partitionSlot = slot;
+ break;
}
}
- NereidsSortedPartitionsCacheManager sortedPartitionsCacheManager =
Env.getCurrentEnv()
- .getSortedPartitionsCacheManager();
- List<Long> manuallySpecifiedPartitions =
scan.getManuallySpecifiedPartitions();
- Map<Long, PartitionItem> idToPartitions;
- Optional<SortedPartitionRanges<Long>> sortedPartitionRanges =
Optional.empty();
- if (manuallySpecifiedPartitions.isEmpty()) {
- Optional<SortedPartitionRanges<?>> sortedPartitionRangesOpt
- = sortedPartitionsCacheManager.get(table, scan);
- if (sortedPartitionRangesOpt.isPresent()) {
- sortedPartitionRanges = (Optional)
sortedPartitionRangesOpt;
- }
- idToPartitions = partitionInfo.getIdToItem(false);
+ if (partitionSlot == null) {
+ return null;
} else {
- Map<Long, PartitionItem> allPartitions =
partitionInfo.getAllPartitions();
- idToPartitions = allPartitions.keySet().stream()
- .filter(manuallySpecifiedPartitions::contains)
- .collect(Collectors.toMap(Function.identity(),
allPartitions::get));
+ partitionSlots.add(partitionSlot);
+ }
+ }
+ NereidsSortedPartitionsCacheManager sortedPartitionsCacheManager =
Env.getCurrentEnv()
+ .getSortedPartitionsCacheManager();
+ List<Long> manuallySpecifiedPartitions =
scan.getManuallySpecifiedPartitions();
+ Map<Long, PartitionItem> idToPartitions;
+ Optional<SortedPartitionRanges<Long>> sortedPartitionRanges =
Optional.empty();
+ if (manuallySpecifiedPartitions.isEmpty()) {
+ Optional<SortedPartitionRanges<?>> sortedPartitionRangesOpt
+ = sortedPartitionsCacheManager.get(table, scan);
+ if (sortedPartitionRangesOpt.isPresent()) {
+ sortedPartitionRanges = (Optional) sortedPartitionRangesOpt;
}
+ idToPartitions = partitionInfo.getIdToItem(false);
+ } else {
+ Map<Long, PartitionItem> allPartitions =
partitionInfo.getAllPartitions();
+ idToPartitions = allPartitions.keySet().stream()
+ .filter(manuallySpecifiedPartitions::contains)
+ .collect(Collectors.toMap(Function.identity(),
allPartitions::get));
+ }
+ if (filter != null) {
List<Long> prunedPartitions = PartitionPruner.prune(
partitionSlots, filter.getPredicate(), idToPartitions,
ctx.cascadesContext,
PartitionTableType.OLAP, sortedPartitionRanges);
- if (prunedPartitions.isEmpty()) {
- return new LogicalEmptyRelation(
-
ConnectContext.get().getStatementContext().getNextRelationId(),
- filter.getOutput());
+ return prunedPartitions;
+ } else if (!manuallySpecifiedPartitions.isEmpty()) {
+ return Utils.fastToImmutableList(idToPartitions.keySet());
+ } else {
+ return null;
+ }
+ }
+
+ private List<Long> prunePartitionByTabletIds(LogicalOlapScan scan,
+ OlapTable table,
+ List<Long>
prunedPartitionsByFilters) {
+ if (scan.getManuallySpecifiedTabletIds().size() == 0
+ || (prunedPartitionsByFilters != null &&
prunedPartitionsByFilters.isEmpty())) {
+ // `prunedPartitionsByFilters is not null and is empty` means
empty partitions after pruner
+ return prunedPartitionsByFilters;
+ }
+
+ Set<Long> selectedPartitions = new LinkedHashSet<>();
+ if (prunedPartitionsByFilters != null) {
+ selectedPartitions.addAll(prunedPartitionsByFilters);
+ }
+
+ Set<Long> manuallySpecifiedTabletIds =
ImmutableSet.copyOf(scan.getManuallySpecifiedTabletIds());
+ List<Long> selectPartitionIds = new ArrayList<>();
+ for (Partition partition : table.getPartitions()) {
+ if (!selectedPartitions.isEmpty() &&
!selectedPartitions.contains(partition.getId())) {
+ continue;
+ }
+ MaterializedIndex baseIndex = partition.getBaseIndex();
+ for (Tablet tablet : baseIndex.getTablets()) {
+ if (manuallySpecifiedTabletIds.contains(tablet.getId())) {
+ selectPartitionIds.add(partition.getId());
+ break;
+ }
}
- LogicalOlapScan rewrittenScan =
scan.withSelectedPartitionIds(prunedPartitions);
- return filter.withChildren(ImmutableList.of(rewrittenScan));
- }).toRule(RuleType.OLAP_SCAN_PARTITION_PRUNE);
+ }
+ return selectPartitionIds;
}
}
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartitionTest.java
b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartitionTest.java
index de54a05bc5e..f87f51253ad 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartitionTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartitionTest.java
@@ -17,6 +17,11 @@
package org.apache.doris.nereids.rules.rewrite;
+import org.apache.doris.catalog.Database;
+import org.apache.doris.catalog.Env;
+import org.apache.doris.catalog.MaterializedIndex;
+import org.apache.doris.catalog.OlapTable;
+import org.apache.doris.catalog.Tablet;
import org.apache.doris.common.FeConstants;
import org.apache.doris.nereids.util.MemoPatternMatchSupported;
import org.apache.doris.nereids.util.PlanChecker;
@@ -152,6 +157,30 @@ class PruneOlapScanPartitionTest extends TestWithFeService
implements MemoPatter
test("testOlapScanPartitionWithSingleColumnCase", "col1 >= 0 and col1
<= 5", 2);
}
+ @Test
+ void testOlapScanPartitionPruneWithTablet() throws Exception {
+ Database db =
Env.getCurrentInternalCatalog().getDbOrMetaException("test");
+ OlapTable tbl = (OlapTable)
db.getTableOrMetaException("single_not_null");
+ Assertions.assertNotNull(tbl);
+ Tablet tablet1 = tbl.getPartition("p20211122")
+
.getMaterializedIndices(MaterializedIndex.IndexExtState.ALL).iterator().next()
+ .getTablets().iterator().next();
+ Tablet tablet2 = tbl.getPartition("p20211125")
+
.getMaterializedIndices(MaterializedIndex.IndexExtState.ALL).iterator().next()
+ .getTablets().iterator().next();
+
+ test("test.single_not_null", "", String.valueOf(tablet1.getId()), 1);
+ test("test.single_not_null", "", tablet1.getId() + ", " +
tablet2.getId(), 2);
+ test("test.single_not_null", "dt in (20211121)", tablet1.getId() + ",
" + tablet2.getId(), 1);
+ test("test.single_not_null", "dt in (20211121, 20211124)",
tablet1.getId() + ", " + tablet2.getId(), 2);
+ test("test.single_not_null", "dt in (20211121)",
String.valueOf(tablet2.getId()), 0);
+ test("test.single_not_null", "", "", 7);
+ test("test.single_not_null", "", String.valueOf(tablet1.getId()),
"p20211122", 1);
+ test("test.single_not_null", "", String.valueOf(tablet2.getId()),
"p20211122", 0);
+ test("test.single_not_null", "dt in (20211124)",
String.valueOf(tablet1.getId()), "p20211122", 0);
+
+ }
+
@Test
void testOlapScanPartitionPruneWithMultiColumnCase() throws Exception {
createTable("create table
testOlapScanPartitionPruneWithMultiColumnCase("
@@ -389,8 +418,24 @@ class PruneOlapScanPartitionTest extends TestWithFeService
implements MemoPatter
}
private void test(String table, String filter, int expectScanPartitionNum)
{
+ test(table, filter, "", expectScanPartitionNum);
+ }
+
+ private void test(String table, String filter, String tabletsFilter, int
expectScanPartitionNum) {
+ test(table, filter, tabletsFilter, "", expectScanPartitionNum);
+ }
+
+ private void test(String table,
+ String filter,
+ String tabletsFilter,
+ String specifiedPartition,
+ int expectScanPartitionNum) {
PlanChecker planChecker = PlanChecker.from(connectContext)
- .analyze("select * from " + table + (filter.isEmpty() ? "" : "
where " + filter))
+ .analyze("select * from " + table
+ + (specifiedPartition.isEmpty() ? "" : String.format("
PARTITION %s", specifiedPartition))
+ + (tabletsFilter.isEmpty() ? "" : String.format("
tablet(%s)", tabletsFilter))
+ + (filter.isEmpty() ? "" : " where " + filter)
+ )
.rewrite()
.printlnTree();
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
