[#12507] Implement adaptive pool selection for StrictReplicaGroupInstanceSelector

pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BalancedInstanceSelector.java

@@ -21,7 +21,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pinot.broker.routing.adaptiveserverselector.ServerSelectionContext;
 import org.apache.pinot.common.metrics.BrokerMeter;
 import org.apache.pinot.common.metrics.BrokerMetrics;
@@ -45,7 +44,7 @@
 public class BalancedInstanceSelector extends BaseInstanceSelector {
 
   @Override
-  public Pair<Map<String, String>, Map<String, String>> select(List<String> segments, int requestId,
+  public InstanceMapping select(List<String> segments, int requestId,
       SegmentStates segmentStates, Map<String, String> queryOptions) {
     Map<String, String> segmentToSelectedInstanceMap = new HashMap<>(HashUtil.getHashMapCapacity(segments.size()));
     // No need to adjust this map per total segment numbers, as optional segments should be empty most of the time.
@@ -88,6 +87,6 @@ public Pair<Map<String, String>, Map<String, String>> select(List<String> segmen
       _brokerMetrics.addMeteredValue(BrokerMeter.POOL_SEG_QUERIES, entry.getValue(),
           BrokerMetrics.getTagForPreferredPool(queryOptions), String.valueOf(entry.getKey()));
     }
-    return Pair.of(segmentToSelectedInstanceMap, optionalSegmentToInstanceMap);
+    return new InstanceMapping(segmentToSelectedInstanceMap, optionalSegmentToInstanceMap);
   }
 }

pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java

@@ -30,7 +30,6 @@
 import java.util.TreeMap;
 import java.util.TreeSet;
 import javax.annotation.Nullable;
-import org.apache.commons.lang3.tuple.Pair;
 import org.apache.helix.AccessOption;
 import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.IdealState;
@@ -465,20 +464,20 @@ public SelectionResult select(BrokerRequest brokerRequest, List<String> segments
     // Copy the volatile reference so that segmentToInstanceMap and unavailableSegments can have a consistent view of
     // the state.
     SegmentStates segmentStates = _segmentStates;
-    Pair<Map<String, String>, Map<String, String>> segmentToInstanceMap =
-        select(segments, requestIdInt, segmentStates, queryOptions);
+    InstanceMapping mapping = select(segments, requestIdInt, segmentStates, queryOptions);
     Set<String> unavailableSegments = segmentStates.getUnavailableSegments();
+    List<String> mappingUnavailable = mapping.unavailableSegments();
 
-    if (unavailableSegments.isEmpty()) {
-      return new SelectionResult(segmentToInstanceMap, List.of(), 0);
+    if (unavailableSegments.isEmpty() && mappingUnavailable.isEmpty()) {
+      return new SelectionResult(mapping, List.of(), 0);
     } else {
-      List<String> unavailableSegmentsForRequest = new ArrayList<>();
+      List<String> unavailableSegmentsForRequest = new ArrayList<>(mappingUnavailable);
       for (String segment : segments) {
         if (unavailableSegments.contains(segment)) {
           unavailableSegmentsForRequest.add(segment);
         }
       }
-      return new SelectionResult(segmentToInstanceMap, unavailableSegmentsForRequest, 0);
+      return new SelectionResult(mapping, unavailableSegmentsForRequest, 0);
     }
   }
 
@@ -501,11 +500,11 @@ int getPool(String instanceID) {
   }
 
   /**
-   * Selects the server instances for the given segments based on the request id and segment states. Returns two maps
-   * from segment to selected server instance hosting the segment. The 2nd map is for optional segments. The optional
-   * segments are used to get the new segments that are not online yet. Instead of simply skipping them by broker at
-   * routing time, we can send them to servers and let servers decide how to handle them.
+   * Selects the server instances for the given segments based on the request id and segment states. Returns an
+   * {@link InstanceSelector.InstanceMapping} containing two maps from segment to selected server instance. The optional
+   * map covers new segments that are not online yet; instead of simply skipping them at the broker, we can send them to
+   * servers and let servers decide how to handle them.
    */
-  protected abstract Pair<Map<String, String>, Map<String, String>/*optional segments*/> select(List<String> segments,
-      int requestId, SegmentStates segmentStates, Map<String, String> queryOptions);
+  protected abstract InstanceMapping select(List<String> segments, int requestId, SegmentStates segmentStates,
+      Map<String, String> queryOptions);
 }

pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java

@@ -23,7 +23,6 @@
 import java.util.Map;
 import java.util.Set;
 import javax.annotation.Nullable;
-import org.apache.commons.lang3.tuple.Pair;
 import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.IdealState;
 import org.apache.helix.store.zk.ZkHelixPropertyStore;
@@ -83,31 +82,51 @@ void init(TableConfig tableConfig, ZkHelixPropertyStore<ZNRecord> propertyStore,
    */
   Set<String> getServingInstances();
 
+  /**
+   * Holds the result of an instance selection: {@code segmentToInstanceMap} maps each segment to its selected server
+   * instance, {@code optionalSegmentToInstanceMap} maps segments not yet fully online that the server may skip, and
+   * {@code unavailableSegments} lists segments that have candidates but could not be routed.
+   */
+  record InstanceMapping(Map<String, String> segmentToInstanceMap,
+      Map<String, String> optionalSegmentToInstanceMap,
+      List<String> unavailableSegments) {
+    static final InstanceMapping EMPTY =
+        new InstanceMapping(Map.of(), Map.of(), List.of());
+
+    InstanceMapping(Map<String, String> segmentToInstanceMap, Map<String, String> optionalSegmentToInstanceMap) {
+      this(segmentToInstanceMap, optionalSegmentToInstanceMap, List.of());
+    }
+  }
+
   class SelectionResult {
-    private final Pair<Map<String, String>, Map<String, String>/*optional segments*/> _segmentToInstanceMap;
+    private final InstanceMapping _instanceMapping;
     private final List<String> _unavailableSegments;
     private int _numPrunedSegments;
 
-    public SelectionResult(Pair<Map<String, String>, Map<String, String>> segmentToInstanceMap,
+    public SelectionResult(InstanceMapping instanceMapping,
         List<String> unavailableSegments, int numPrunedSegments) {
-      _segmentToInstanceMap = segmentToInstanceMap;
+      _instanceMapping = instanceMapping;
       _unavailableSegments = unavailableSegments;
       _numPrunedSegments = numPrunedSegments;
     }
 
+    public static SelectionResult empty(int numPrunedSegments) {
+      return new SelectionResult(InstanceMapping.EMPTY, List.of(), numPrunedSegments);
+    }
+
     /**
      * Returns the map from segment to selected server instance hosting the segment.
      */
     public Map<String, String> getSegmentToInstanceMap() {
-      return _segmentToInstanceMap.getLeft();
+      return _instanceMapping.segmentToInstanceMap();
     }
 
     /**
      * Returns the map from optional segment to selected server instance hosting the optional segment.
      * Optional segments can be skipped by broker or server upon any issue w/o failing the query.
      */
     public Map<String, String> getOptionalSegmentToInstanceMap() {
-      return _segmentToInstanceMap.getRight();
+      return _instanceMapping.optionalSegmentToInstanceMap();
     }
 
     /**

pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelectorFactory.java

@@ -106,6 +106,14 @@ public static InstanceSelector getInstanceSelector(TableConfig tableConfig,
           }
           case RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE: {
             LOGGER.info("Using StrictReplicaGroupInstanceSelector for table: {}", tableNameWithType);
+            boolean enableStrictReplicaGroupAdaptiveRouting = brokerConfig.getProperty(
+                CommonConstants.Broker.AdaptiveServerSelector.CONFIG_OF_ENABLE_STRICT_REPLICA_GROUP,
+                CommonConstants.Broker.AdaptiveServerSelector.DEFAULT_ENABLE_STRICT_REPLICA_GROUP);
+            if (!enableStrictReplicaGroupAdaptiveRouting && adaptiveServerSelector != null) {
+              LOGGER.info("Adaptive routing disabled for StrictReplicaGroupInstanceSelector table: {}",
+                  tableNameWithType);
+              adaptiveServerSelector = null;
+            }
             instanceSelector = new StrictReplicaGroupInstanceSelector();
             break;
           }

pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/MultiStageReplicaGroupSelector.java

@@ -29,7 +29,6 @@
 import java.util.Set;
 import javax.annotation.Nullable;
 import org.apache.commons.collections4.CollectionUtils;
-import org.apache.commons.lang3.tuple.Pair;
 import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.IdealState;
 import org.apache.helix.store.zk.ZkHelixPropertyStore;
@@ -86,7 +85,7 @@ public void onAssignmentChange(IdealState idealState, ExternalView externalView,
   }
 
   @Override
-  public Pair<Map<String, String>, Map<String, String>> select(List<String> segments, int requestId,
+  public InstanceMapping select(List<String> segments, int requestId,
       SegmentStates segmentStates, Map<String, String> queryOptions) {
     // Create a copy of InstancePartitions to avoid race-condition with event-listeners above.
     InstancePartitions instancePartitions = _instancePartitions;
@@ -110,8 +109,9 @@ public Pair<Map<String, String>, Map<String, String>> select(List<String> segmen
    * Returns a map from the segmentName to the corresponding server. It tries to select all servers from the
    * preferredReplicaGroup, but if it fails, it will try to select the relevant server from other instance partitions.
    *
-   * @return A pair of maps, where the first map contains the segments that are assigned to a server and the second
-   * map contains the segments that are optional (i.e., the server is not online to serve that segment).
+   * @return An {@link InstanceMapping} whose {@code segmentToInstanceMap} contains the segments assigned to a server
+   * and whose {@code optionalSegmentToInstanceMap} contains segments that are optional (i.e., the server is not
+   * online to serve that segment).
    * Example:
    * {
    *   "required_segments": {
@@ -124,7 +124,7 @@ public Pair<Map<String, String>, Map<String, String>> select(List<String> segmen
    *   }
    * }
    */
-  private Pair<Map<String, String>, Map<String, String>> assign(Set<String> segments,
+  private InstanceMapping assign(Set<String> segments,
       SegmentStates segmentStates, InstancePartitions instancePartitions, int preferredReplicaId) {
     Map<String, Integer> instanceToPartitionMap = instancePartitions.getInstanceToPartitionIdMap();
     Map<String, Set<String>> instanceToSegmentsMap = new HashMap<>();
@@ -229,7 +229,7 @@ private void getSelectedInstancesForPartition(Map<String, Set<String>> instanceT
    * Based on whether the selected instance for the segment is online to serve that segment or not,
    * this method computes the segments that are optional and the segments that are not.
    */
-  private Pair<Map<String, String>, Map<String, String>> computeOptionalSegments(
+  private InstanceMapping computeOptionalSegments(
       Map<String, String> segmentToSelectedInstanceMap, SegmentStates segmentStates) {
 
     Map<String, String> segmentsToInstanceMap = new HashMap<>();
@@ -258,7 +258,7 @@ private Pair<Map<String, String>, Map<String, String>> computeOptionalSegments(
       }
     }
 
-    return Pair.of(segmentsToInstanceMap, optionalSegmentToInstanceMap);
+    return new InstanceMapping(segmentsToInstanceMap, optionalSegmentToInstanceMap);
   }
 
   @VisibleForTesting

pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupInstanceSelector.java

@@ -27,7 +27,6 @@
 import java.util.Set;
 import javax.annotation.Nullable;
 import org.apache.commons.collections.MapUtils;
-import org.apache.commons.lang3.tuple.Pair;
 import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.IdealState;
 import org.apache.pinot.broker.routing.adaptiveserverselector.ServerSelectionContext;
@@ -70,7 +69,7 @@ public class ReplicaGroupInstanceSelector extends BaseInstanceSelector {
   private static final Logger LOGGER = LoggerFactory.getLogger(ReplicaGroupInstanceSelector.class);
 
   @Override
-  public Pair<Map<String, String>, Map<String, String>> select(List<String> segments, int requestId,
+  public InstanceMapping select(List<String> segments, int requestId,
       SegmentStates segmentStates, Map<String, String> queryOptions) {
     ServerSelectionContext ctx = new ServerSelectionContext(queryOptions, _config);
     if (_adaptiveServerSelector != null) {
@@ -91,7 +90,7 @@ public Pair<Map<String, String>, Map<String, String>> select(List<String> segmen
     }
   }
 
-  private Pair<Map<String, String>, Map<String, String>> selectServers(List<String> segments, int requestId,
+  protected InstanceMapping selectServers(List<String> segments, int requestId,
       SegmentStates segmentStates, @Nullable Map<String, Integer> serverRankMap, ServerSelectionContext ctx) {
 
     Map<String, String> segmentToSelectedInstanceMap = new HashMap<>(HashUtil.getHashMapCapacity(segments.size()));
@@ -121,7 +120,7 @@ private Pair<Map<String, String>, Map<String, String>> selectServers(List<String
       } else if (MapUtils.isNotEmpty(serverRankMap)) {
         // Adaptive Server Selection is enabled.
         // Use the instance with the best rank if all servers have stats populated, else use the round-robin selected
-        // instance
+        // instance. As of 8 July 2026, this fallback is unreachable, but new implementations could require it.
         selectedInstance = candidates.stream()
             .anyMatch(candidate -> !serverRankMap.containsKey(candidate.getInstance()))
             ? selectedInstance
@@ -147,7 +146,7 @@ private Pair<Map<String, String>, Map<String, String>> selectServers(List<String
       _brokerMetrics.addMeteredValue(BrokerMeter.POOL_SEG_QUERIES, entry.getValue(),
           BrokerMetrics.getTagForPreferredPool(ctx.getQueryOptions()), String.valueOf(entry.getKey()));
     }
-    return Pair.of(segmentToSelectedInstanceMap, optionalSegmentToInstanceMap);
+    return new InstanceMapping(segmentToSelectedInstanceMap, optionalSegmentToInstanceMap);
   }
 
   private List<SegmentInstanceCandidate> fetchCandidateServersForQuery(List<String> segments,