[opt](bucket pruning) Avoid materializing tablet IDs during pruning

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanTablet.java

@@ -106,7 +106,7 @@ private Collection<Long> getSelectedTabletIds(List<Column> schema, Map<String, P
             return index.getTabletIdsInOrder();
         }
         HashDistributionInfo hashInfo = (HashDistributionInfo) info;
-        return new HashDistributionPruner(schema, index.getTabletIdsInOrder(),
+        return new HashDistributionPruner(schema, index,
                 hashInfo.getDistributionColumns(),
                 filterMap,
                 hashInfo.getBucketNum(),

fe/fe-core/src/main/java/org/apache/doris/planner/HashDistributionPruner.java

@@ -21,7 +21,9 @@
 import org.apache.doris.analysis.LiteralExpr;
 import org.apache.doris.analysis.SlotRef;
 import org.apache.doris.catalog.Column;
+import org.apache.doris.catalog.MaterializedIndex;
 import org.apache.doris.catalog.PartitionKey;
+import org.apache.doris.catalog.Tablet;
 import org.apache.doris.common.Config;
 
 import com.google.common.collect.Lists;
@@ -30,6 +32,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
@@ -52,22 +55,21 @@
 public class HashDistributionPruner implements DistributionPruner {
     private static final Logger LOG = LogManager.getLogger(HashDistributionPruner.class);
 
-    // partition list, sort by the hash code
-    private List<Long> bucketsList;
+    // Tablet snapshot in hash bucket order.
+    private final List<Tablet> tablets;
+    private final int bucketNum;
     // partition columns
-    private List<Column>                       distributionColumns;
+    private final List<Column> distributionColumns;
     // partition column filters
-    private Map<String, PartitionColumnFilter> distributionColumnFilters;
-    private int                                hashMod;
+    private final Map<String, PartitionColumnFilter> distributionColumnFilters;
+    private final int hashMod;
 
-    private boolean isBaseIndexSelected;
-
-    public HashDistributionPruner(List<Column> schema, List<Long> bucketsList, List<Column> columns,
+    public HashDistributionPruner(List<Column> schema, MaterializedIndex materializedIndex, List<Column> columns,
             Map<String, PartitionColumnFilter> 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<Long> 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<Long> 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<Long> resultSet = Sets.newHashSet();
         int inElementNum = inPredicate.getInElementNum();
@@ -130,13 +132,25 @@ public Collection<Long> prune(int columnId, PartitionKey hashKey, int complex) {
             Collection<Long> 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<Long> getAllTabletIds() {
+        List<Long> tabletIds = new ArrayList<>(bucketNum);
+        for (Tablet tablet : tablets) {
+            tabletIds.add(tablet.getId());
+        }
+        return tabletIds;
+    }
+
     public Collection<Long> prune() {
         PartitionKey hashKey = new PartitionKey();
         return prune(0, hashKey, 1);

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

@@ -422,7 +422,7 @@ private Collection<Long> partitionPrune(PartitionInfo partitionInfo) throws Anal
 
     private Collection<Long> distributionPrune(
             List<Column> schema,
-            List<Long> tabletIdsInOrder,
+            MaterializedIndex index,
             DistributionInfo distributionInfo,
             boolean pruneTablesByNereids) throws AnalysisException {
         if (pruneTablesByNereids) {
@@ -434,7 +434,8 @@ private Collection<Long> distributionPrune(
             // getTablet hash lookups (most returning null), which dominates plan time
             // when both partition count and pruned tablet count are large.
             List<Long> 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<Long> 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<Tablet> tablets = Lists.newArrayList();
-            List<Long> allTabletIds = selectedTable.getTabletIdsInOrder();
             // point query need prune tablets at this place
             Collection<Long> 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<Tablet> 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));
                 }
             }
 

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<Long> tabletIds = Lists.newArrayListWithExpectedSize(300);
+        List<Tablet> tablets = Lists.newArrayListWithExpectedSize(300);
         for (long i = 0; i < 300; i++) {
-            tabletIds.add(i);
+            tablets.add(new LocalTablet(i));
         }
+        index.appendTablets(tablets);
 
         List<Column> 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);
 

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<Long> tabletIds = Lists.newArrayListWithExpectedSize(300);
+        List<Tablet> 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<Long> results = pruner.prune();
@@ -139,4 +146,43 @@ public void test() {
         Assert.assertEquals(39, tablets.size());
     }
 
+    @Test
+    public void testPruneWithMaterializedIndex() {
+        List<Long> 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<Column> columns = Lists.newArrayList(column);
+
+        List<Expr> 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<String, PartitionColumnFilter> filters = new CaseInsensitiveMap();
+        filters.put("K1", filter);
+
+        Collection<Long> indexResult = new HashDistributionPruner(null, index, columns, filters,
+                tabletIds.size(), true).prune();
+        Set<Long> 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<String, PartitionColumnFilter> emptyFilters = new CaseInsensitiveMap();
+        Collection<Long> allIndexTablets = new HashDistributionPruner(null, index, columns, emptyFilters,
+                tabletIds.size(), true).prune();
+        Assert.assertEquals(tabletIds, Lists.newArrayList(allIndexTablets));
+    }
+
 }

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<Long> tabletIds) {
+        MaterializedIndex index = new MaterializedIndex();
+        List<Tablet> 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<Long> partitions = new ArrayList<>();
-        partitions.add(new Long(0));
-        partitions.add(new Long(1));
-        partitions.add(new Long(2));
+        List<Long> tabletIds = Lists.newArrayList(0L, 1L, 2L);
 
 
         List<Column> 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<Long> partitions = new ArrayList<>();
-        partitions.add(new Long(0));
-        partitions.add(new Long(1));
-        partitions.add(new Long(2));
+        List<Long> tabletIds = Lists.newArrayList(0L, 1L, 2L);
 
         List<Column> 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,