[Enhancement] Align range-colocate scan-range dispatch by ColocateRange index (backport #72278) (#72804)

Signed-off-by: xiangguangyxg <xiangguangyxg@gmail.com>
Co-authored-by: xiangguangyxg <110401425+xiangguangyxg@users.noreply.github.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 3 people

fe/fe-core/src/main/java/com/starrocks/catalog/ColocateRangeMgr.java

@@ -66,13 +66,39 @@ public List<ColocateRange> getColocateRanges(long colocateGroupId) {
      * @return the ColocateRange containing the value, or null if not found
      */
     public ColocateRange getColocateRange(long colocateGroupId, Tuple colocateValue) {
+        int index = getColocateRangeIndex(colocateGroupId, colocateValue);
+        if (index < 0) {
+            return null;
+        }
+        return colocateGroupToRanges.get(colocateGroupId).get(index);
+    }
+
+    /**
+     * Returns the index of the ColocateRange that contains the given colocate value.
+     *
+     * <p>The index is stable across all tables sharing this colocate group: the i-th
+     * range in {@link #getColocateRanges(long)} corresponds to the same PACK shard group
+     * across every partition/table/database in the group. Coordinator-side scan-range
+     * dispatch uses this index as the bucket sequence so that scan ranges from joined
+     * tables that share a colocate range are routed to the same fragment instance.
+     *
+     * @param colocateGroupId the colocate group id
+     * @param colocateValue the colocate prefix to look up; null means the global lower
+     *                      bound (-inf), which always lands in the first range
+     * @return the index in {@code [0, getColocateRanges(grpId).size())}, or -1 if the
+     *         group does not exist or the value is not covered
+     */
+    public int getColocateRangeIndex(long colocateGroupId, Tuple colocateValue) {
         List<ColocateRange> colocateRanges = colocateGroupToRanges.get(colocateGroupId);
         if (colocateRanges == null || colocateRanges.isEmpty()) {
-            return null;
+            return -1;
+        }
+        if (colocateValue == null) {
+            return 0;
         }
         Range<Tuple> pointRange = Range.gele(colocateValue, colocateValue);
         int index = Collections.binarySearch(colocateRanges, pointRange);
-        return index >= 0 ? colocateRanges.get(index) : null;
+        return index >= 0 ? index : -1;
     }
 
     /**

fe/fe-core/src/main/java/com/starrocks/catalog/ColocateRangeUtils.java

@@ -83,4 +83,43 @@ private static Tuple extendTupleWithNull(Tuple tuple, List<Column> sortKeyColumn
         }
         return new Tuple(values);
     }
+
+    /**
+     * Extracts the colocate column prefix from a tablet's full sort-key range.
+     *
+     * <p>Inverse of {@link #expandToFullSortKey}: a tablet whose range was produced by
+     * expansion stores a full sort-key tuple in its lower bound, but the colocate-range
+     * lookup keys on the colocate prefix only.
+     *
+     * <p>If the tablet range is unbounded below (lower bound is -inf), this returns
+     * {@code null}, signaling the caller to fall back to the first colocate range
+     * (which always begins at -inf by the {@link ColocateRangeMgr} invariant).
+     *
+     * <p>Unlike {@link #expandToFullSortKey}, this method requires
+     * {@code colocateColumnCount > 0}: a colocate group with zero colocate columns
+     * is not a meaningful concept on the lookup side (every value would map to the
+     * same range, which is just the no-colocate case).
+     *
+     * @param tabletRange the tablet's full sort-key range
+     * @param colocateColumnCount the number of colocate columns (sort key prefix length),
+     *                            must be positive
+     * @return the colocate prefix Tuple, or {@code null} if the range is unbounded below
+     */
+    public static Tuple extractColocatePrefix(Range<Tuple> tabletRange, int colocateColumnCount) {
+        Preconditions.checkArgument(colocateColumnCount > 0,
+                "colocateColumnCount must be positive, got %s", colocateColumnCount);
+        if (tabletRange.isMinimum()) {
+            return null;
+        }
+        List<Variant> values = tabletRange.getLowerBound().getValues();
+        Preconditions.checkArgument(values.size() >= colocateColumnCount,
+                "tablet lower bound has %s values, fewer than colocateColumnCount %s",
+                values.size(), colocateColumnCount);
+        if (values.size() == colocateColumnCount) {
+            return tabletRange.getLowerBound();
+        }
+        // subList returns a view backed by the original list; copy so that the
+        // returned Tuple does not retain a reference to the tablet's bound.
+        return new Tuple(new ArrayList<>(values.subList(0, colocateColumnCount)));
+    }
 }

fe/fe-core/src/main/java/com/starrocks/common/Range.java

@@ -346,6 +346,53 @@ public boolean isOverlapping(Range<T> other) {
         return !(isLessThan(other) || isGreaterThan(other));
     }
 
+    /**
+     * Checks if this range contains every value of another range.
+     *
+     * <p>Handles {@code -∞} / {@code +∞} bounds and inclusive/exclusive edges
+     * explicitly. Unlike {@link #compareTo(Range)}, which returns 0 for any
+     * overlapping pair, this returns true only when {@code other} is a subset
+     * of {@code this}.
+     *
+     * @param other the candidate sub-range, must not be null
+     * @return true if every value in {@code other} is also in this range
+     * @throws NullPointerException if other is null
+     */
+    public boolean contains(Range<T> other) {
+        Objects.requireNonNull(other, "Range must not be null");
+        return lowerBoundContainsLowerOf(other) && upperBoundContainsUpperOf(other);
+    }
+
+    private boolean lowerBoundContainsLowerOf(Range<T> other) {
+        if (isMinimum()) {
+            return true;
+        }
+        if (other.isMinimum()) {
+            return false;
+        }
+        int cmp = getLowerBound().compareTo(other.getLowerBound());
+        if (cmp != 0) {
+            return cmp < 0;
+        }
+        // Equal lower bound values: this range must include the bound, or
+        // other must also exclude it.
+        return isLowerBoundIncluded() || !other.isLowerBoundIncluded();
+    }
+
+    private boolean upperBoundContainsUpperOf(Range<T> other) {
+        if (isMaximum()) {
+            return true;
+        }
+        if (other.isMaximum()) {
+            return false;
+        }
+        int cmp = other.getUpperBound().compareTo(getUpperBound());
+        if (cmp != 0) {
+            return cmp < 0;
+        }
+        return isUpperBoundIncluded() || !other.isUpperBoundIncluded();
+    }
+
     /**
      * Compares this range with another for ordering.
      *

fe/fe-core/src/main/java/com/starrocks/planner/OlapScanNode.java

@@ -131,6 +131,8 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import javax.annotation.Nullable;
+
 public class OlapScanNode extends AbstractOlapTableScanNode {
     private static final Logger LOG = LogManager.getLogger(OlapScanNode.class);
 
@@ -316,7 +318,12 @@ public List<Expr> getBucketExprs() {
 
     @Override
     public int getBucketNums() {
-        int bucketNum = olapTable.getDefaultDistributionInfo().getBucketNum();
+        DistributionInfo distInfo = olapTable.getDefaultDistributionInfo();
+        if (distInfo.getType() == DistributionInfo.DistributionInfoType.RANGE) {
+            return getRangeDistributionBucketNums(distInfo);
+        }
+        // HASH path.
+        int bucketNum = distInfo.getBucketNum();
         if (getSelectedPartitionIds().size() <= 1) {
             for (Long pid : getSelectedPartitionIds()) {
                 bucketNum = olapTable.getPartition(pid).getDistributionInfo().getBucketNum();
@@ -325,6 +332,54 @@ public int getBucketNums() {
         return bucketNum;
     }
 
+    private int getRangeDistributionBucketNums(DistributionInfo distInfo) {
+        RangeColocateScanDispatch dispatch = RangeColocateScanDispatch.forTable(olapTable);
+        if (dispatch != null) {
+            // getBucketNums() is invoked from ExecutionFragment.getOrCreateColocatedAssignment
+            // only when BackendSelectorFactory has chosen a colocate-dispatch path. Verify
+            // alignment HERE, after the dispatch decision: a misaligned ColocateRangeMgr
+            // would silently produce wrong join results under colocate dispatch.
+            // Non-colocate scans go through NormalBackendSelector and never reach this point.
+            dispatch.requireAligned(getSelectedPhysicalPartitions(), index.indexMetaId);
+            return dispatch.bucketCount();
+        }
+        // Range distribution without a colocate group: RangeDistributionInfo always
+        // reports 1 (one tablet per partition by design).
+        return distInfo.getBucketNum();
+    }
+
+    /**
+     * Returns the physical partitions that actually contribute scan tablets.
+     * Optimizer path: derived from {@link #partitionToScanTabletMap} (skip
+     * empty entries — those were pruned out). Legacy path: every sub-partition
+     * of every selected logical partition.
+     */
+    private List<PhysicalPartition> getSelectedPhysicalPartitions() {
+        List<PhysicalPartition> result = new ArrayList<>();
+        if (partitionToScanTabletMap != null) {
+            for (Map.Entry<Long, List<Long>> entry : partitionToScanTabletMap.entrySet()) {
+                if (entry.getValue() == null || entry.getValue().isEmpty()) {
+                    continue;
+                }
+                PhysicalPartition physicalPartition = olapTable.getPhysicalPartition(entry.getKey());
+                if (physicalPartition != null) {
+                    result.add(physicalPartition);
+                }
+            }
+            return result;
+        }
+        for (Long partitionId : selectedPartitionIds) {
+            Partition partition = olapTable.getPartition(partitionId);
+            if (partition == null) {
+                continue;
+            }
+            for (PhysicalPartition physicalPartition : partition.getSubPartitions()) {
+                result.add(physicalPartition);
+            }
+        }
+        return result;
+    }
+
     @Override
     public Optional<List<BucketProperty>> getBucketProperties() throws StarRocksException {
         return Optional.empty();
@@ -763,6 +818,11 @@ private void computeTabletInfo() throws StarRocksException {
          */
         Preconditions.checkState(scanBackendIds.size() == 0);
         Preconditions.checkState(scanTabletIds.size() == 0);
+        DistributionInfo distInfo = olapTable.getDefaultDistributionInfo();
+        RangeColocateScanDispatch dispatch = null;
+        if (distInfo.getType() == DistributionInfo.DistributionInfoType.RANGE) {
+            dispatch = RangeColocateScanDispatch.forTable(olapTable);
+        }
         for (Long partitionId : selectedPartitionIds) {
             final Partition partition = olapTable.getPartition(partitionId);
 
@@ -783,16 +843,37 @@ private void computeTabletInfo() throws StarRocksException {
                     scanTabletIds.addAll(allTabletIds);
                 }
 
-                for (int i = 0; i < allTabletIds.size(); i++) {
-                    tabletId2BucketSeq.put(allTabletIds.get(i), i);
-                }
+                fillTabletId2BucketSeq(dispatch, selectedIndex, allTabletIds);
                 totalTabletsNum += selectedIndex.getTablets().size();
                 selectedTabletsNum += tablets.size();
                 addScanRangeLocations(partition, physicalPartition, selectedIndex, tablets, localBeId);
             }
         }
     }
 
+    /**
+     * Populates {@link #tabletId2BucketSeq} for one {@link MaterializedIndex}.
+     * Range-colocate scans use the bucket sequence supplied by the dispatch
+     * facade when alignment holds; everything else (HASH, range non-colocate,
+     * or transiently unaligned range colocate) falls back to position-based
+     * bucketSeq. The dispatch-time alignment guard fires later in
+     * {@link #getBucketNums()} for the colocate-dispatch path.
+     */
+    private void fillTabletId2BucketSeq(@Nullable RangeColocateScanDispatch dispatch,
+                                          MaterializedIndex selectedIndex,
+                                          List<Long> allTabletIds) {
+        if (dispatch != null) {
+            Map<Long, Integer> rangeColocateMap = dispatch.computeBucketSeq(selectedIndex);
+            if (rangeColocateMap != null) {
+                tabletId2BucketSeq.putAll(rangeColocateMap);
+                return;
+            }
+        }
+        for (int i = 0; i < allTabletIds.size(); i++) {
+            tabletId2BucketSeq.put(allTabletIds.get(i), i);
+        }
+    }
+
     public void checkIfScanRangeNumSafe(long scanRangeSize) {
         long totalPartitionNum = 0;
         long totalTabletsNum = 0;