[Improvement](planner)support select tablets with nereids optimizer

fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4

@@ -259,7 +259,8 @@ identifierSeq
     ;
 
 relationPrimary
-    : multipartIdentifier specifiedPartition? tableAlias sample? relationHint? lateralView*           #tableName
+    : multipartIdentifier specifiedPartition?
+       tabletList? tableAlias sample? relationHint? lateralView*           #tableName
     | LEFT_PAREN query RIGHT_PAREN tableAlias lateralView*                                    #aliasedQuery
     | tvfName=identifier LEFT_PAREN
       (properties+=property (COMMA properties+=property)*)?
@@ -280,6 +281,10 @@ multipartIdentifier
     : parts+=errorCapturingIdentifier (DOT parts+=errorCapturingIdentifier)*
     ;
 
+tabletList
+    : TABLET LEFT_PAREN tabletIdList+=INTEGER_VALUE (COMMA tabletIdList+=INTEGER_VALUE)*  RIGHT_PAREN
+    ;
+
 // -----------------Expression-----------------
 namedExpression
     : expression (AS? (identifierOrText))?

fe/fe-core/src/main/java/org/apache/doris/nereids/analyzer/UnboundRelation.java

@@ -47,35 +47,37 @@ public class UnboundRelation extends LogicalRelation implements Unbound {
 
     private final List<String> nameParts;
     private final List<String> partNames;
+    private final List<Long> tabletIds;
     private final boolean isTempPart;
     private final List<String> hints;
     private final Optional<TableSample> tableSample;
 
     public UnboundRelation(RelationId id, List<String> nameParts) {
-        this(id, nameParts, Optional.empty(), Optional.empty(), ImmutableList.of(), false,
+        this(id, nameParts, Optional.empty(), Optional.empty(), ImmutableList.of(), false, ImmutableList.of(),
                 ImmutableList.of(), Optional.empty());
     }
 
     public UnboundRelation(RelationId id, List<String> nameParts, List<String> partNames, boolean isTempPart) {
-        this(id, nameParts, Optional.empty(), Optional.empty(), partNames, isTempPart,
+        this(id, nameParts, Optional.empty(), Optional.empty(), partNames, isTempPart, ImmutableList.of(),
                 ImmutableList.of(), Optional.empty());
     }
 
     public UnboundRelation(RelationId id, List<String> nameParts, List<String> partNames, boolean isTempPart,
-            List<String> hints, Optional<TableSample> tableSample) {
+            List<Long> tabletIds, List<String> hints, Optional<TableSample> tableSample) {
         this(id, nameParts, Optional.empty(), Optional.empty(),
-                partNames, isTempPart, hints, tableSample);
+                partNames, isTempPart, tabletIds, hints, tableSample);
     }
 
     /**
-     * Constructor.
+     * constructor of UnboundRelation
      */
     public UnboundRelation(RelationId id, List<String> nameParts, Optional<GroupExpression> groupExpression,
             Optional<LogicalProperties> logicalProperties, List<String> partNames, boolean isTempPart,
-            List<String> hints, Optional<TableSample> tableSample) {
+            List<Long> tabletIds, List<String> hints, Optional<TableSample> tableSample) {
         super(id, PlanType.LOGICAL_UNBOUND_RELATION, groupExpression, logicalProperties);
         this.nameParts = ImmutableList.copyOf(Objects.requireNonNull(nameParts, "nameParts should not null"));
         this.partNames = ImmutableList.copyOf(Objects.requireNonNull(partNames, "partNames should not null"));
+        this.tabletIds = ImmutableList.copyOf(Objects.requireNonNull(tabletIds, "tabletIds should not null"));
         this.isTempPart = isTempPart;
         this.hints = ImmutableList.copyOf(Objects.requireNonNull(hints, "hints should not be null."));
         this.tableSample = tableSample;
@@ -99,15 +101,14 @@ public LogicalProperties computeLogicalProperties() {
     public Plan withGroupExpression(Optional<GroupExpression> groupExpression) {
         return new UnboundRelation(relationId, nameParts,
                 groupExpression, Optional.of(getLogicalProperties()),
-                partNames, isTempPart, hints, tableSample);
-
+                partNames, isTempPart, tabletIds, hints, tableSample);
     }
 
     @Override
     public Plan withGroupExprLogicalPropChildren(Optional<GroupExpression> groupExpression,
             Optional<LogicalProperties> logicalProperties, List<Plan> children) {
         return new UnboundRelation(relationId, nameParts, groupExpression, logicalProperties, partNames,
-                isTempPart, hints, tableSample);
+                isTempPart, tabletIds, hints, tableSample);
     }
 
     @Override
@@ -146,6 +147,10 @@ public boolean isTempPart() {
         return isTempPart;
     }
 
+    public List<Long> getTabletIds() {
+        return tabletIds;
+    }
+
     public List<String> getHints() {
         return hints;
     }

fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java

@@ -630,6 +630,13 @@ public LogicalPlan visitTableName(TableNameContext ctx) {
             }
         }
 
+        List<Long> tabletIdLists = new ArrayList<>();
+        if (ctx.tabletList() != null) {
+            ctx.tabletList().tabletIdList.stream().forEach(tabletToken -> {
+                tabletIdLists.add(Long.parseLong(tabletToken.getText()));
+            });
+        }
+
         final List<String> relationHints;
         if (ctx.relationHint() != null) {
             relationHints = typedVisit(ctx.relationHint());
@@ -640,7 +647,7 @@ public LogicalPlan visitTableName(TableNameContext ctx) {
         TableSample tableSample = ctx.sample() == null ? null : (TableSample) visit(ctx.sample());
         LogicalPlan checkedRelation = withCheckPolicy(
                 new UnboundRelation(StatementScopeIdGenerator.newRelationId(),
-                        tableId, partitionNames, isTempPart, relationHints,
+                        tableId, partitionNames, isTempPart, tabletIdLists, relationHints,
                         Optional.ofNullable(tableSample)));
         LogicalPlan plan = withTableAlias(checkedRelation, ctx.tableAlias());
         for (LateralViewContext lateralViewContext : ctx.lateralView()) {

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindRelation.java

@@ -167,13 +167,14 @@ private LogicalPlan bind(CascadesContext cascadesContext, UnboundRelation unboun
     private LogicalPlan makeOlapScan(TableIf table, UnboundRelation unboundRelation, List<String> tableQualifier) {
         LogicalOlapScan scan;
         List<Long> partIds = getPartitionIds(table, unboundRelation);
+        List<Long> tabletIds = unboundRelation.getTabletIds();
         if (!CollectionUtils.isEmpty(partIds)) {
             scan = new LogicalOlapScan(unboundRelation.getRelationId(),
                     (OlapTable) table, ImmutableList.of(tableQualifier.get(1)), partIds,
-                    unboundRelation.getHints(), unboundRelation.getTableSample());
+                    tabletIds, unboundRelation.getHints(), unboundRelation.getTableSample());
         } else {
             scan = new LogicalOlapScan(unboundRelation.getRelationId(),
-                    (OlapTable) table, ImmutableList.of(tableQualifier.get(1)), unboundRelation.getHints(),
+                    (OlapTable) table, ImmutableList.of(tableQualifier.get(1)), tabletIds, unboundRelation.getHints(),
                     unboundRelation.getTableSample());
         }
         if (!Util.showHiddenColumns() && scan.getTable().hasDeleteSign()

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

@@ -55,14 +55,18 @@ public Rule build() {
                     LogicalOlapScan olapScan = filter.child();
                     OlapTable table = olapScan.getTable();
                     Builder<Long> selectedTabletIdsBuilder = ImmutableList.builder();
-                    for (Long id : olapScan.getSelectedPartitionIds()) {
-                        Partition partition = table.getPartition(id);
-                        MaterializedIndex index = partition.getIndex(olapScan.getSelectedIndexId());
-                        selectedTabletIdsBuilder
-                                .addAll(getSelectedTabletIds(filter.getConjuncts(), index,
-                                        olapScan.getSelectedIndexId() == olapScan.getTable()
-                                                .getBaseIndexId(),
-                                        partition.getDistributionInfo()));
+                    if (olapScan.getSelectedTabletIds().isEmpty()) {
+                        for (Long id : olapScan.getSelectedPartitionIds()) {
+                            Partition partition = table.getPartition(id);
+                            MaterializedIndex index = partition.getIndex(olapScan.getSelectedIndexId());
+                            selectedTabletIdsBuilder
+                                    .addAll(getSelectedTabletIds(filter.getConjuncts(), index,
+                                            olapScan.getSelectedIndexId() == olapScan.getTable()
+                                                    .getBaseIndexId(),
+                                            partition.getDistributionInfo()));
+                        }
+                    } else {
+                        selectedTabletIdsBuilder.addAll(olapScan.getSelectedTabletIds());
                     }
                     List<Long> selectedTabletIds = selectedTabletIdsBuilder.build();
                     if (new HashSet(selectedTabletIds).equals(new HashSet(olapScan.getSelectedTabletIds()))) {

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopier.java

@@ -160,11 +160,13 @@ public Plan visitLogicalOlapScan(LogicalOlapScan olapScan, DeepCopierContext con
         LogicalOlapScan newOlapScan;
         if (olapScan.getManuallySpecifiedPartitions().isEmpty()) {
             newOlapScan = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(),
-                    olapScan.getTable(), olapScan.getQualifier(), olapScan.getHints(), olapScan.getTableSample());
+                    olapScan.getTable(), olapScan.getQualifier(), olapScan.getSelectedTabletIds(),
+                    olapScan.getHints(), olapScan.getTableSample());
         } else {
             newOlapScan = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(),
                     olapScan.getTable(), olapScan.getQualifier(),
-                    olapScan.getManuallySpecifiedPartitions(), olapScan.getHints(), olapScan.getTableSample());
+                    olapScan.getManuallySpecifiedPartitions(), olapScan.getSelectedTabletIds(),
+                    olapScan.getHints(), olapScan.getTableSample());
         }
         newOlapScan.getOutput();
         context.putRelation(olapScan.getRelationId(), newOlapScan);

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

@@ -118,19 +118,19 @@ public LogicalOlapScan(RelationId id, OlapTable table, List<String> qualifier) {
                 Maps.newHashMap(), Optional.empty());
     }
 
-    public LogicalOlapScan(RelationId id, OlapTable table, List<String> qualifier,
+    public LogicalOlapScan(RelationId id, OlapTable table, List<String> qualifier, List<Long> tabletIds,
             List<String> hints, Optional<TableSample> tableSample) {
         this(id, table, qualifier, Optional.empty(), Optional.empty(),
-                table.getPartitionIds(), false, ImmutableList.of(),
+                table.getPartitionIds(), false, tabletIds,
                 -1, false, PreAggStatus.on(), ImmutableList.of(), hints, Maps.newHashMap(),
                 tableSample);
     }
 
     public LogicalOlapScan(RelationId id, OlapTable table, List<String> qualifier, List<Long> specifiedPartitions,
-            List<String> hints, Optional<TableSample> tableSample) {
+            List<Long> tabletIds, List<String> hints, Optional<TableSample> tableSample) {
         this(id, table, qualifier, Optional.empty(), Optional.empty(),
                 // must use specifiedPartitions here for prune partition by sql like 'select * from t partition p1'
-                specifiedPartitions, false, ImmutableList.of(),
+                specifiedPartitions, false, tabletIds,
                 -1, false, PreAggStatus.on(), specifiedPartitions, hints, Maps.newHashMap(),
                 tableSample);
     }

regression-test/data/nereids_p0/select_tablets/select_with_tablets.out

@@ -0,0 +1,58 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !select1 --
+1	doris	19
+2	nereids	18
+
+-- !select2 --
+1	doris	19
+2	nereids	18
+
+-- !select3 --
+1	doris	19
+2	nereids	18
+
+-- !select4 --
+
+-- !select5 --
+1	doris	19
+
+-- !select6 --
+2	nereids	18
+
+-- !select7 --
+1	doris	19
+
+-- !select8 --
+1	doris	19
+
+-- !select9 --
+
+-- !select10 --
+1	doris	19
+
+-- !select11 --
+
+-- !select12 --
+
+-- !select13 --
+
+-- !no_partition_1 --
+1	doris	19
+2	nereids	18
+
+-- !no_partition_2 --
+2	nereids	18
+
+-- !no_partition_3 --
+1	doris	19
+
+-- !no_partition_4 --
+
+-- !no_partition_5 --
+2	nereids	18
+
+-- !no_partition_6 --
+2	nereids	18
+
+-- !no_partition_7 --
+

regression-test/suites/nereids_p0/select_tablets/select_with_tablets.groovy

@@ -0,0 +1,104 @@
+// 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.
+
+suite("select_with_tablets") {
+    sql 'set enable_nereids_planner=true;'
+    sql 'set enable_fallback_to_original_planner=false;'
+
+
+    def table_name1 = "test_table"
+    sql """ DROP TABLE IF EXISTS ${table_name1} """
+    sql """
+    CREATE TABLE IF NOT EXISTS ${table_name1} (
+        `id` int(11) NULL,
+        `name` string NULL,
+        `age` int(11) NULL
+        )
+        PARTITION BY RANGE(id)
+        (
+            PARTITION less_than_20 VALUES LESS THAN ("20"),
+            PARTITION between_20_70 VALUES [("20"),("70")),
+            PARTITION more_than_70 VALUES LESS THAN ("151")
+        )
+        DISTRIBUTED BY HASH(id) BUCKETS 1
+        PROPERTIES("replication_num" = "1");
+    """
+
+    sql """ INSERT INTO ${table_name1} VALUES (1, 'doris', 19), (2, 'nereids', 18) """
+    def insert_res = sql "show last insert;"
+
+    logger.info("insert result: " + insert_res.toString())
+    order_qt_select1 """ SELECT * FROM ${table_name1} """
+
+    def res = sql """ show tablets from ${table_name1} where version = 2 """
+    assertTrue(res.size() == 1)
+    assertTrue(res[0].size() == 21)
+    assertEquals("2", res[0][4])
+
+    order_qt_select2 """ SELECT * FROM ${table_name1} TABLET(${res[0][0]}) """
+    order_qt_select3 """ SELECT * FROM ${table_name1} PARTITION less_than_20 TABLET(${res[0][0]}) """
+    // result should be empty because TABLET(${res[0][0]}) is not belonged to partition between_20_70.
+    order_qt_select4 """ SELECT * FROM ${table_name1} PARTITION between_20_70 TABLET(${res[0][0]}) """
+
+    order_qt_select5 """ SELECT * FROM ${table_name1} where id < 2 """
+    order_qt_select6 """ SELECT * FROM ${table_name1} TABLET(${res[0][0]}) where id = 2 """
+    order_qt_select7 """ SELECT * FROM ${table_name1} TABLET(${res[0][0]}) where id < 2 """
+    order_qt_select8 """ SELECT * FROM ${table_name1} PARTITION less_than_20 TABLET(${res[0][0]}) where id < 2 """
+    // result of order_qt_select9 should be empty
+    order_qt_select9 """ SELECT * FROM ${table_name1} PARTITION between_20_70 TABLET(${res[0][0]}) where id < 2"""
+    order_qt_select10 """ SELECT * FROM ${table_name1} PARTITION less_than_20 where id < 2"""
+    // result of order_qt_select11 should be empty
+    order_qt_select11 """ SELECT * FROM ${table_name1} PARTITION between_20_70 where id < 2"""
+
+    res = sql """ show tablets from ${table_name1} where version = 1 """
+    assertTrue(res.size() == 2)
+    assertEquals("1", res[0][4])
+    assertEquals("1", res[1][4])
+    // result should be empty because TABLET(${res[0][0]}) does not have data.
+    order_qt_select12 """ SELECT * FROM ${table_name1} TABLET(${res[0][0]}) """
+    // result should be empty because TABLET(${res[1][0]}) does not have data.
+    order_qt_select13 """ SELECT * FROM ${table_name1} TABLET(${res[1][0]}) """
+
+    // Test non-partitioned table
+    def table_no_partition = "table_no_partition"
+    sql """ DROP TABLE IF EXISTS ${table_no_partition} """
+    sql """
+    CREATE TABLE IF NOT EXISTS ${table_no_partition} (
+        `id` int(11) NULL,
+        `name` string NULL,
+        `age` int(11) NULL
+        )
+        DISTRIBUTED BY HASH(id) BUCKETS 3
+        PROPERTIES("replication_num" = "1");
+    """
+
+    sql """ INSERT INTO ${table_no_partition} VALUES (1, 'doris', 19), (2, 'nereids', 18) """
+    insert_res = sql "show last insert;"
+
+    logger.info("insert result: " + insert_res.toString())
+    order_qt_no_partition_1 """ SELECT * FROM ${table_no_partition} """
+
+    res = sql """ show tablets from ${table_no_partition} where version = 2 """
+
+    order_qt_no_partition_2 """ SELECT * FROM ${table_no_partition} TABLET(${res[0][0]}) """
+    order_qt_no_partition_3 """ SELECT * FROM ${table_no_partition} TABLET(${res[1][0]}) """
+    order_qt_no_partition_4 """ SELECT * FROM ${table_no_partition} TABLET(${res[2][0]}) """
+
+    order_qt_no_partition_5 """ SELECT * FROM ${table_no_partition} where id = 2 """
+    order_qt_no_partition_6 """ SELECT * FROM ${table_no_partition} TABLET(${res[0][0]}) where id = 2 """
+    order_qt_no_partition_7 """ SELECT * FROM ${table_no_partition} TABLET(${res[0][0]}) where id > 2 """
+}