diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanTablet.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanTablet.java index e336bb8a8badbd..20f8c1d6e48d20 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanTablet.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanTablet.java @@ -106,7 +106,7 @@ private Collection getSelectedTabletIds(List schema, Map bucketsList; + // Tablet snapshot in hash bucket order. + private final List tablets; + private final int bucketNum; // partition columns - private List distributionColumns; + private final List distributionColumns; // partition column filters - private Map distributionColumnFilters; - private int hashMod; + private final Map distributionColumnFilters; + private final int hashMod; - private boolean isBaseIndexSelected; - - public HashDistributionPruner(List schema, List bucketsList, List columns, + public HashDistributionPruner(List schema, MaterializedIndex materializedIndex, List columns, Map filters, int hashMod, boolean isBaseIndexSelected) { - this.bucketsList = bucketsList; + this.tablets = materializedIndex.getTablets(); + this.bucketNum = tablets.size(); this.distributionColumns = columns; this.hashMod = hashMod; - this.isBaseIndexSelected = isBaseIndexSelected; if (isBaseIndexSelected) { this.distributionColumnFilters = filters; } else { @@ -91,14 +93,14 @@ public Collection prune(int columnId, PartitionKey hashKey, int complex) { if (columnId == distributionColumns.size()) { // compute Hash Key long hashValue = hashKey.getHashValue(); - return Lists.newArrayList(bucketsList.get((int) ((hashValue & 0xffffffff) % hashMod))); + return Lists.newArrayList(getTabletId((int) ((hashValue & 0xffffffff) % hashMod))); } Column keyColumn = distributionColumns.get(columnId); PartitionColumnFilter filter = distributionColumnFilters.get(keyColumn.getName()); if (null == filter) { // no filter in this column, no partition Key // return all subPartition - return Lists.newArrayList(bucketsList); + return getAllTabletIds(); } InPredicate inPredicate = filter.getInPredicate(); if (null == inPredicate @@ -113,12 +115,12 @@ public Collection prune(int columnId, PartitionKey hashKey, int complex) { return result; } // return all SubPartition - return Lists.newArrayList(bucketsList); + return getAllTabletIds(); } if (!(inPredicate.getChild(0) instanceof SlotRef)) { // return all SubPartition - return Lists.newArrayList(bucketsList); + return getAllTabletIds(); } Set resultSet = Sets.newHashSet(); int inElementNum = inPredicate.getInElementNum(); @@ -130,13 +132,25 @@ public Collection prune(int columnId, PartitionKey hashKey, int complex) { Collection subList = prune(columnId + 1, hashKey, newComplex); resultSet.addAll(subList); hashKey.popColumn(); - if (resultSet.size() >= bucketsList.size()) { + if (resultSet.size() >= bucketNum) { break; } } return resultSet; } + private long getTabletId(int bucket) { + return tablets.get(bucket).getId(); + } + + private List getAllTabletIds() { + List tabletIds = new ArrayList<>(bucketNum); + for (Tablet tablet : tablets) { + tabletIds.add(tablet.getId()); + } + return tabletIds; + } + public Collection prune() { PartitionKey hashKey = new PartitionKey(); return prune(0, hashKey, 1); 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 fd0d219cf9acd0..3418137d57c734 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 @@ -422,7 +422,7 @@ private Collection partitionPrune(PartitionInfo partitionInfo) throws Anal private Collection distributionPrune( List schema, - List tabletIdsInOrder, + MaterializedIndex index, DistributionInfo distributionInfo, boolean pruneTablesByNereids) throws AnalysisException { if (pruneTablesByNereids) { @@ -434,7 +434,8 @@ private Collection distributionPrune( // getTablet hash lookups (most returning null), which dominates plan time // when both partition count and pruned tablet count are large. List result = new ArrayList<>(); - for (Long id : tabletIdsInOrder) { + for (Tablet tablet : index.getTablets()) { + long id = tablet.getId(); if (nereidsPrunedTabletIds.contains(id)) { result.add(id); } @@ -445,7 +446,7 @@ private Collection distributionPrune( switch (distributionInfo.getType()) { case HASH: { HashDistributionInfo info = (HashDistributionInfo) distributionInfo; - distributionPruner = new HashDistributionPruner(schema, tabletIdsInOrder, + distributionPruner = new HashDistributionPruner(schema, index, info.getDistributionColumns(), columnFilters, info.getBucketNum(), @@ -1019,10 +1020,9 @@ private void computeTabletInfo() throws UserException { final Partition partition = olapTable.getPartition(partitionId); final MaterializedIndex selectedTable = olapTable.getPartitionIndex(partition, selectedIndexId); final List tablets = Lists.newArrayList(); - List allTabletIds = selectedTable.getTabletIdsInOrder(); // point query need prune tablets at this place Collection prunedTabletIds = distributionPrune(olapTable.getSchemaByIndexId(selectedIndexId), - allTabletIds, partition.getDistributionInfo(), isNereids && !isPointQuery); + selectedTable, partition.getDistributionInfo(), isNereids && !isPointQuery); if (LOG.isDebugEnabled()) { LOG.debug("distribution prune tablets: {}", prunedTabletIds); } @@ -1058,14 +1058,17 @@ private void computeTabletInfo() throws UserException { } } } else { - tablets.addAll(selectedTable.getTablets()); - scanTabletIds.addAll(allTabletIds); + for (Tablet tablet : selectedTable.getTablets()) { + tablets.add(tablet); + scanTabletIds.add(tablet.getId()); + } } if (!isPointQuery) { + List allTablets = selectedTable.getTablets(); int bucketNum = partition.getDistributionInfo().getBucketNum(); - for (int i = 0; i < allTabletIds.size(); i++) { - tabletId2BucketInfo.put(allTabletIds.get(i), encodeBucketInfo(i, bucketNum)); + for (int i = 0; i < allTablets.size(); i++) { + tabletId2BucketInfo.put(allTablets.get(i).getId(), encodeBucketInfo(i, bucketNum)); } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/RewriteRuleSuiteTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/RewriteRuleSuiteTest.java index 33dd172385ac98..c1024026e027d7 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/RewriteRuleSuiteTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/RewriteRuleSuiteTest.java @@ -24,10 +24,12 @@ import org.apache.doris.catalog.Column; import org.apache.doris.catalog.DistributionInfo; import org.apache.doris.catalog.HashDistributionInfo; +import org.apache.doris.catalog.LocalTablet; import org.apache.doris.catalog.MaterializedIndex; import org.apache.doris.catalog.OlapTable; import org.apache.doris.catalog.Partition; import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.catalog.Tablet; import org.apache.doris.nereids.CascadesContext; import org.apache.doris.nereids.sqltest.SqlTestBase; import org.apache.doris.nereids.trees.expressions.EqualTo; @@ -75,13 +77,14 @@ public class RewriteRuleSuiteTest extends SqlTestBase { void testPruneOlapScanTablet() { OlapTable olapTable = Mockito.mock(OlapTable.class); Partition partition = Mockito.mock(Partition.class); - MaterializedIndex index = Mockito.mock(MaterializedIndex.class); + MaterializedIndex index = new MaterializedIndex(); HashDistributionInfo distributionInfo = Mockito.mock(HashDistributionInfo.class); - List tabletIds = Lists.newArrayListWithExpectedSize(300); + List tablets = Lists.newArrayListWithExpectedSize(300); for (long i = 0; i < 300; i++) { - tabletIds.add(i); + tablets.add(new LocalTablet(i)); } + index.appendTablets(tablets); List columns = Lists.newArrayList( new Column("k0", PrimitiveType.DATE, false), @@ -128,10 +131,9 @@ void testPruneOlapScanTablet() { Mockito.when(partition.getIndex(Mockito.anyLong())).thenReturn(index); Mockito.when(olapTable.getPartitionIndex(Mockito.eq(partition), Mockito.anyLong())).thenReturn(index); Mockito.when(partition.getDistributionInfo()).thenReturn(distributionInfo); - Mockito.when(index.getTabletIdsInOrder()).thenReturn(tabletIds); Mockito.when(distributionInfo.getDistributionColumns()).thenReturn(columns); Mockito.when(distributionInfo.getType()).thenReturn(DistributionInfo.DistributionInfoType.HASH); - Mockito.when(distributionInfo.getBucketNum()).thenReturn(tabletIds.size()); + Mockito.when(distributionInfo.getBucketNum()).thenReturn(tablets.size()); LogicalOlapScan scan = new LogicalOlapScan(RelationId.createGenerator().getNextId(), olapTable); diff --git a/fe/fe-core/src/test/java/org/apache/doris/planner/HashDistributionPrunerTest.java b/fe/fe-core/src/test/java/org/apache/doris/planner/HashDistributionPrunerTest.java index 185a1ac1a7aa73..6cc4194dfe7cc7 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/planner/HashDistributionPrunerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/planner/HashDistributionPrunerTest.java @@ -22,8 +22,11 @@ import org.apache.doris.analysis.SlotRef; import org.apache.doris.analysis.StringLiteral; import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.LocalTablet; +import org.apache.doris.catalog.MaterializedIndex; import org.apache.doris.catalog.PartitionKey; import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.catalog.Tablet; import com.google.common.collect.Lists; import com.google.common.collect.Sets; @@ -41,9 +44,13 @@ public class HashDistributionPrunerTest { @Test public void test() { List tabletIds = Lists.newArrayListWithExpectedSize(300); + List indexTablets = Lists.newArrayListWithExpectedSize(300); for (long i = 0; i < 300; i++) { tabletIds.add(i); + indexTablets.add(new LocalTablet(i)); } + MaterializedIndex index = new MaterializedIndex(); + index.appendTablets(indexTablets); // distribution columns Column dealDate = new Column("dealDate", PrimitiveType.DATE, false); @@ -91,7 +98,7 @@ public void test() { filters.put("CHANNEL", channelFilter); filters.put("SHOP_TYPE", shopTypeFilter); - HashDistributionPruner pruner = new HashDistributionPruner(null, tabletIds, columns, filters, tabletIds.size(), + HashDistributionPruner pruner = new HashDistributionPruner(null, index, columns, filters, tabletIds.size(), true); Collection results = pruner.prune(); @@ -139,4 +146,43 @@ public void test() { Assert.assertEquals(39, tablets.size()); } + @Test + public void testPruneWithMaterializedIndex() { + List tabletIds = Lists.newArrayListWithExpectedSize(8); + MaterializedIndex index = new MaterializedIndex(); + for (long i = 0; i < 8; i++) { + long tabletId = 100 + i; + tabletIds.add(tabletId); + index.addTablet(new LocalTablet(tabletId), null, true); + } + + Column column = new Column("k1", PrimitiveType.CHAR, false); + List columns = Lists.newArrayList(column); + + List inList = Lists.newArrayList(); + inList.add(new StringLiteral("a")); + inList.add(new StringLiteral("b")); + PartitionColumnFilter filter = new PartitionColumnFilter(); + filter.setInPredicate(new InPredicate(new SlotRef(null, "k1"), inList, false)); + + Map filters = new CaseInsensitiveMap(); + filters.put("K1", filter); + + Collection indexResult = new HashDistributionPruner(null, index, columns, filters, + tabletIds.size(), true).prune(); + Set expectedTabletIds = Sets.newHashSet(); + for (Expr literal : inList) { + PartitionKey hashKey = new PartitionKey(); + hashKey.pushColumn((StringLiteral) literal, PrimitiveType.CHAR); + long hashValue = hashKey.getHashValue(); + expectedTabletIds.add(tabletIds.get((int) ((hashValue & 0xffffffff) % tabletIds.size()))); + } + Assert.assertEquals(expectedTabletIds, Sets.newHashSet(indexResult)); + + Map emptyFilters = new CaseInsensitiveMap(); + Collection allIndexTablets = new HashDistributionPruner(null, index, columns, emptyFilters, + tabletIds.size(), true).prune(); + Assert.assertEquals(tabletIds, Lists.newArrayList(allIndexTablets)); + } + } diff --git a/fe/fe-core/src/test/java/org/apache/doris/planner/OlapScanNodeTest.java b/fe/fe-core/src/test/java/org/apache/doris/planner/OlapScanNodeTest.java index 30a0d097d4a8ca..eaba851f69aba1 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/planner/OlapScanNodeTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/planner/OlapScanNodeTest.java @@ -31,6 +31,7 @@ import org.apache.doris.catalog.DiskInfo; import org.apache.doris.catalog.LocalReplica; import org.apache.doris.catalog.LocalTablet; +import org.apache.doris.catalog.MaterializedIndex; import org.apache.doris.catalog.OlapTable; import org.apache.doris.catalog.Partition; import org.apache.doris.catalog.PartitionKey; @@ -60,7 +61,6 @@ import org.junit.Test; import org.mockito.Mockito; -import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.HashMap; @@ -70,14 +70,21 @@ import java.util.stream.Collectors; public class OlapScanNodeTest { + private MaterializedIndex createMaterializedIndex(List tabletIds) { + MaterializedIndex index = new MaterializedIndex(); + List tablets = Lists.newArrayListWithExpectedSize(tabletIds.size()); + for (Long tabletId : tabletIds) { + tablets.add(new LocalTablet(tabletId)); + } + index.appendTablets(tablets); + return index; + } + // columnA in (1) hashmode=3 @Test public void testHashDistributionOneUser() throws AnalysisException { - List partitions = new ArrayList<>(); - partitions.add(new Long(0)); - partitions.add(new Long(1)); - partitions.add(new Long(2)); + List tabletIds = Lists.newArrayList(0L, 1L, 2L); List columns = Lists.newArrayList(); @@ -97,7 +104,7 @@ public void testHashDistributionOneUser() throws AnalysisException { DistributionPruner partitionPruner = new HashDistributionPruner( null, - partitions, + createMaterializedIndex(tabletIds), columns, filterMap, 3, @@ -115,10 +122,7 @@ public void testHashDistributionOneUser() throws AnalysisException { @Test public void testHashPartitionManyUser() throws AnalysisException { - List partitions = new ArrayList<>(); - partitions.add(new Long(0)); - partitions.add(new Long(1)); - partitions.add(new Long(2)); + List tabletIds = Lists.newArrayList(0L, 1L, 2L); List columns = Lists.newArrayList(); columns.add(new Column("columnA", PrimitiveType.BIGINT)); @@ -142,7 +146,7 @@ public void testHashPartitionManyUser() throws AnalysisException { DistributionPruner partitionPruner = new HashDistributionPruner( null, - partitions, + createMaterializedIndex(tabletIds), columns, filterMap, 3,