[Enhancement](Short Circuit) short circuit query supports IN

be/src/olap/base_tablet.cpp

@@ -486,7 +486,6 @@ Status BaseTablet::lookup_row_key(const Slice& encoded_key, TabletSchema* latest
         }
         auto& segments = segment_caches[i]->get_segments();
         DCHECK_EQ(segments.size(), num_segments);
-
         for (auto id : picked_segments) {
             Status s = segments[id]->lookup_row_key(encoded_key, schema, with_seq_col, with_rowid,
                                                     &loc);

be/src/service/internal_service.cpp

@@ -912,6 +912,35 @@ void PInternalService::tablet_fetch_data(google::protobuf::RpcController* contro
     }
 }
 
+void PInternalService::tablet_batch_fetch_data(google::protobuf::RpcController* controller,
+                                               const PTabletBatchKeyLookupRequest* batchRequest,
+                                               PTabletBatchKeyLookupResponse* batchResponse,
+                                               google::protobuf::Closure* done) {
+    int request_count = batchRequest->sub_key_lookup_req_size();
+    batchResponse->mutable_sub_key_lookup_res()->Reserve(request_count);
+    [[maybe_unused]] auto* cntl = static_cast<brpc::Controller*>(controller);
+    bool ret =
+            _light_work_pool.try_offer([this, batchRequest, batchResponse, done, request_count]() {
+                Status st = Status::OK();
+                brpc::ClosureGuard guard(done);
+                for (int i = 0; i < request_count; ++i) {
+                    batchResponse->add_sub_key_lookup_res();
+                    const PTabletKeyLookupRequest* request = &batchRequest->sub_key_lookup_req(i);
+                    PTabletKeyLookupResponse* response =
+                            batchResponse->mutable_sub_key_lookup_res(i);
+                    Status status = _tablet_fetch_data(request, response);
+                    status.to_protobuf(response->mutable_status());
+                    if (!status.ok()) {
+                        st = status;
+                    }
+                }
+                st.to_protobuf(batchResponse->mutable_status());
+            });
+    if (!ret) {
+        offer_failed(batchResponse, done, _light_work_pool);
+    }
+}
+
 void PInternalService::test_jdbc_connection(google::protobuf::RpcController* controller,
                                             const PJdbcTestConnectionRequest* request,
                                             PJdbcTestConnectionResult* result,

be/src/service/internal_service.h

@@ -220,6 +220,11 @@ class PInternalService : public PBackendService {
                            PTabletKeyLookupResponse* response,
                            google::protobuf::Closure* done) override;
 
+    void tablet_batch_fetch_data(google::protobuf::RpcController* controller,
+                                 const PTabletBatchKeyLookupRequest* batchRequest,
+                                 PTabletBatchKeyLookupResponse* batchResponse,
+                                 google::protobuf::Closure* done) override;
+
     void test_jdbc_connection(google::protobuf::RpcController* controller,
                               const PJdbcTestConnectionRequest* request,
                               PJdbcTestConnectionResult* result,

be/src/service/point_query_executor.cpp

@@ -260,6 +260,11 @@ Status PointQueryExecutor::init(const PTabletKeyLookupRequest* request,
     auto cache_handle = LookupConnectionCache::instance()->get(uuid);
     _binary_row_format = request->is_binary_row();
     _tablet = DORIS_TRY(ExecEnv::get_tablet(request->tablet_id()));
+
+    if (_tablet->tablet_meta()->replica_id() != request->replica_id()) {
+        return Status::OK();
+    }
+
     if (cache_handle != nullptr) {
         _reusable = cache_handle;
         _profile_metrics.hit_lookup_cache = true;

fe/fe-core/src/main/java/org/apache/doris/catalog/PartitionKey.java

@@ -58,6 +58,7 @@
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
 import java.time.LocalDateTime;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.zip.CRC32;
@@ -79,6 +80,24 @@ public PartitionKey() {
         types = Lists.newArrayList();
     }
 
+    private PartitionKey(PartitionKey other) {
+        this.keys = new ArrayList<>(other.keys.size());
+        for (LiteralExpr expr : other.keys) {
+            try {
+                String value = expr.getStringValue();
+                if ("null".equalsIgnoreCase(value)) {
+                    this.keys.add(NullLiteral.create(expr.getType()));
+                } else {
+                    this.keys.add(LiteralExpr.create(value, expr.getType()));
+                }
+            } catch (Exception e) {
+                throw new RuntimeException("Create partition key failed: " + e.getMessage());
+            }
+        }
+        this.originHiveKeys = new ArrayList<>(other.originHiveKeys);
+        this.types = new ArrayList<>(other.types);
+    }
+
     public void setDefaultListPartition(boolean isDefaultListPartitionKey) {
         this.isDefaultListPartitionKey = isDefaultListPartitionKey;
     }
@@ -205,6 +224,10 @@ public static PartitionKey createListPartitionKey(List<PartitionValue> values, L
         return createListPartitionKeyWithTypes(values, types, false);
     }
 
+    public static PartitionKey clone(PartitionKey other) {
+        return new PartitionKey(other);
+    }
+
     public void pushColumn(LiteralExpr keyValue, PrimitiveType keyType) {
         keys.add(keyValue);
         types.add(keyType);

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

@@ -25,6 +25,7 @@
 import org.apache.doris.nereids.trees.expressions.Cast;
 import org.apache.doris.nereids.trees.expressions.EqualTo;
 import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.InPredicate;
 import org.apache.doris.nereids.trees.expressions.SlotReference;
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalFilter;
@@ -53,7 +54,7 @@ private Expression removeCast(Expression expression) {
     private boolean filterMatchShortCircuitCondition(LogicalFilter<LogicalOlapScan> filter) {
         return filter.getConjuncts().stream().allMatch(
                 // all conjuncts match with pattern `key = ?`
-                expression -> (expression instanceof EqualTo)
+                expression -> ((expression instanceof EqualTo) || expression instanceof InPredicate)
                         && (removeCast(expression.child(0)).isKeyColumnFromTable()
                         || (expression.child(0) instanceof SlotReference
                         && ((SlotReference) expression.child(0)).getName().equals(Column.DELETE_SIGN)))

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

@@ -26,9 +26,8 @@
 import org.apache.doris.common.Config;
 
 import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
 
 import java.util.Collection;
 import java.util.List;
@@ -49,8 +48,6 @@
  * If depth is larger than 'max_distribution_pruner_recursion_depth', all buckets will be return without pruning.
  */
 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;
     // partition columns
@@ -61,6 +58,18 @@ public class HashDistributionPruner implements DistributionPruner {
 
     private boolean isBaseIndexSelected;
 
+    /*
+     * This map maintains a relationship between distribution keys and their corresponding tablet IDs.
+     * For example, if the distribution columns are (k1, k2, k3),
+     * and the tuple (1, 2, 3) is hashed into bucket 1001,
+     * then the `distributionKey2TabletIDs` would map the key (1, 2, 3) to the tablet ID 1001.
+     * (1, 2, 3) -> 1001
+     * Map structure:
+     * - Key: PartitionKey, representing a specific combination of distribution columns (e.g., k1, k2, k3).
+     * - Value: Set<Long>, containing the tablet IDs associated with the corresponding distribution key.
+     */
+    private Map<PartitionKey, Set<Long>> distributionKey2TabletIDs = Maps.newHashMap();
+
     public HashDistributionPruner(List<Long> bucketsList, List<Column> columns,
                            Map<String, PartitionColumnFilter> filters, int hashMod, boolean isBaseIndexSelected) {
         this.bucketsList = bucketsList;
@@ -70,13 +79,21 @@ public HashDistributionPruner(List<Long> bucketsList, List<Column> columns,
         this.isBaseIndexSelected = isBaseIndexSelected;
     }
 
+    public Map<PartitionKey, Set<Long>> getDistributionKeysTabletIDs() {
+        return distributionKey2TabletIDs;
+    }
+
     // columnId: which column to compute
     // hashKey: the key which to compute hash value
     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)));
+            List<Long> result =
+                    Lists.newArrayList(bucketsList.get((int) ((hashValue & 0xffffffff) % hashMod)));
+            distributionKey2TabletIDs.computeIfAbsent(PartitionKey.clone(hashKey),
+                                                      k -> Sets.newHashSet(result)).addAll(result);
+            return result;
         }
         Column keyColumn = distributionColumns.get(columnId);
         String columnName = isBaseIndexSelected ? keyColumn.getName()
@@ -119,9 +136,6 @@ 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()) {
-                break;
-            }
         }
         return resultSet;
     }

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

@@ -27,7 +27,6 @@
 import org.apache.doris.analysis.FunctionCallExpr;
 import org.apache.doris.analysis.InPredicate;
 import org.apache.doris.analysis.IntLiteral;
-import org.apache.doris.analysis.LiteralExpr;
 import org.apache.doris.analysis.PartitionNames;
 import org.apache.doris.analysis.SlotDescriptor;
 import org.apache.doris.analysis.SlotId;
@@ -52,6 +51,7 @@
 import org.apache.doris.catalog.Partition.PartitionState;
 import org.apache.doris.catalog.PartitionInfo;
 import org.apache.doris.catalog.PartitionItem;
+import org.apache.doris.catalog.PartitionKey;
 import org.apache.doris.catalog.PartitionType;
 import org.apache.doris.catalog.Replica;
 import org.apache.doris.catalog.ScalarType;
@@ -98,9 +98,13 @@
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.HashBasedTable;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Range;
+import com.google.common.collect.RangeMap;
 import com.google.common.collect.Sets;
+import com.google.common.collect.Table;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -194,7 +198,19 @@ public class OlapScanNode extends ScanNode {
 
     private HashSet<Long> scanBackendIds = new HashSet<>();
 
+    private PartitionPruner partitionPruner = null;
+
     private Map<Long, Integer> tabletId2BucketSeq = Maps.newHashMap();
+
+    // Maps partition column names to a RangeMap that associates ColumnBound ranges with lists of partition IDs,
+    // similar to the implementation in PartitionPrunerV2Base.
+    private Map<String, RangeMap<ColumnBound, List<Long>>> partitionCol2PartitionID = Maps.newHashMap();
+
+    private Map<PartitionKey, Set<Long>> distributionKeys2TabletID = Maps.newHashMap();
+
+    /// tablet id -> (backend id -> replica)
+    private Table<Long, Long, Replica> scanBackendReplicaTable = HashBasedTable.create();
+
     // a bucket seq may map to many tablets, and each tablet has a
     // TScanRangeLocations.
     public ArrayListMultimap<Integer, TScanRangeLocations> bucketSeq2locations = ArrayListMultimap.create();
@@ -259,6 +275,14 @@ public HashSet<Long> getScanBackendIds() {
         return scanBackendIds;
     }
 
+    public Map<String, RangeMap<ColumnBound, List<Long>>> getPartitionCol2PartitionID() {
+        return partitionCol2PartitionID;
+    }
+
+    public Map<PartitionKey, Set<Long>> getDistributionKeys2TabletID() {
+        return distributionKeys2TabletID;
+    }
+
     public void setSampleTabletIds(List<Long> sampleTablets) {
         if (sampleTablets != null) {
             this.sampleTabletIds.addAll(sampleTablets);
@@ -294,6 +318,10 @@ public ArrayList<Long> getScanTabletIds() {
         return scanTabletIds;
     }
 
+    public Table<Long, Long, Replica> getScanBackendReplicaTable() {
+        return scanBackendReplicaTable;
+    }
+
     public void setForceOpenPreAgg(boolean forceOpenPreAgg) {
         this.forceOpenPreAgg = forceOpenPreAgg;
     }
@@ -658,9 +686,9 @@ private void computeInaccurateCardinality() throws UserException {
         cardinality = (long) statsDeriveResult.getRowCount();
     }
 
+    // get the pruned partition IDs
     private Collection<Long> partitionPrune(PartitionInfo partitionInfo,
             PartitionNames partitionNames) throws AnalysisException {
-        PartitionPruner partitionPruner = null;
         Map<Long, PartitionItem> keyItemMap;
         if (partitionNames != null) {
             keyItemMap = Maps.newHashMap();
@@ -677,13 +705,12 @@ private Collection<Long> partitionPrune(PartitionInfo partitionInfo,
         if (partitionInfo.getType() == PartitionType.RANGE) {
             if (isPointQuery() && partitionInfo.getPartitionColumns().size() == 1) {
                 // short circuit, a quick path to find partition
-                ColumnRange filterRange = columnNameToRange.get(partitionInfo.getPartitionColumns().get(0).getName());
-                LiteralExpr lowerBound = filterRange.getRangeSet().get().asRanges().stream()
-                        .findFirst().get().lowerEndpoint().getValue();
-                LiteralExpr upperBound = filterRange.getRangeSet().get().asRanges().stream()
-                        .findFirst().get().upperEndpoint().getValue();
+                Column col = partitionInfo.getPartitionColumns().get(0);
+                // todo: support range query
+                Set<Range<ColumnBound>> filterRanges =
+                        columnNameToRange.get(col.getName()).getRangeSet().get().asRanges();
                 cachedPartitionPruner.update(keyItemMap);
-                return cachedPartitionPruner.prune(lowerBound, upperBound);
+                return cachedPartitionPruner.prune(filterRanges, col.getName(), partitionCol2PartitionID);
             }
             partitionPruner = new RangePartitionPrunerV2(keyItemMap,
                     partitionInfo.getPartitionColumns(), columnNameToRange);
@@ -701,12 +728,22 @@ private Collection<Long> distributionPrune(
         switch (distributionInfo.getType()) {
             case HASH: {
                 HashDistributionInfo info = (HashDistributionInfo) distributionInfo;
-                distributionPruner = new HashDistributionPruner(table.getTabletIdsInOrder(),
+                distributionPruner =
+                        new HashDistributionPruner(table.getTabletIdsInOrder(),
                         info.getDistributionColumns(),
                         columnFilters,
                         info.getBucketNum(),
                         getSelectedIndexId() == olapTable.getBaseIndexId());
-                return distributionPruner.prune();
+                HashDistributionPruner hashPruner = (HashDistributionPruner) distributionPruner;
+                Collection<Long> resultIDs = hashPruner.prune();
+                Map<PartitionKey, Set<Long>> newPrunedIDs = hashPruner.getDistributionKeysTabletIDs();
+                for (Map.Entry<PartitionKey, Set<Long>> entry : newPrunedIDs.entrySet()) {
+                    distributionKeys2TabletID.merge(entry.getKey(), entry.getValue(), (existingSet, newSet) -> {
+                        existingSet.addAll(newSet);
+                        return existingSet;
+                    });
+                }
+                return resultIDs;
             }
             case RANDOM: {
                 return null;
@@ -946,6 +983,7 @@ private void addScanRangeLocations(Partition partition,
                     collectedStat = true;
                 }
                 scanBackendIds.add(backend.getId());
+                scanBackendReplicaTable.put(tabletId, backend.getId(), replica);
                 // For skipping missing version of tablet, we only select the backend with the highest last
                 // success version replica to save as much data as possible.
                 if (skipMissingVersion) {
@@ -981,10 +1019,20 @@ private void computePartitionInfo() throws AnalysisException {
         // Step1: compute partition ids
         PartitionNames partitionNames = ((BaseTableRef) desc.getRef()).getPartitionNames();
         PartitionInfo partitionInfo = olapTable.getPartitionInfo();
-        if (partitionInfo.getType() == PartitionType.RANGE || partitionInfo.getType() == PartitionType.LIST) {
-            selectedPartitionIds = partitionPrune(partitionInfo, partitionNames);
-        } else {
-            selectedPartitionIds = olapTable.getPartitionIds();
+        switch (partitionInfo.getType()) {
+            case RANGE:
+                selectedPartitionIds = partitionPrune(partitionInfo, partitionNames);
+                if (isPointQuery() && partitionPruner instanceof RangePartitionPrunerV2) {
+                    RangePartitionPrunerV2 rangePartitionPruner = (RangePartitionPrunerV2) partitionPruner;
+                    this.partitionCol2PartitionID = rangePartitionPruner.getPartitionCol2PartitionID();
+                }
+                break;
+            case LIST:
+                selectedPartitionIds = partitionPrune(partitionInfo, partitionNames);
+                break;
+            default:
+                selectedPartitionIds = olapTable.getPartitionIds();
+                break;
         }
         selectedPartitionIds = olapTable.selectNonEmptyPartitionIds(selectedPartitionIds);
         selectedPartitionNum = selectedPartitionIds.size();
@@ -1287,13 +1335,16 @@ public List<TScanRangeLocations> lazyEvaluateRangeLocations() throws UserExcepti
         // Lazy evaluation
         selectedIndexId = olapTable.getBaseIndexId();
         // Only key columns
+        distributionKeys2TabletID.clear();
+        partitionCol2PartitionID.clear();
         computeColumnsFilter(olapTable.getBaseSchemaKeyColumns(), olapTable.getPartitionInfo());
         computePartitionInfo();
         scanBackendIds.clear();
         scanTabletIds.clear();
         bucketSeq2locations.clear();
         scanReplicaIds.clear();
         sampleTabletIds.clear();
+        scanBackendReplicaTable.clear();
         try {
             createScanRangeLocations();
         } catch (AnalysisException e) {