gortiz commented on code in PR #18947:
URL: https://github.com/apache/pinot/pull/18947#discussion_r3691248694
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java:
##########
@@ -50,14 +64,175 @@
* segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that
we always route the segments to the same
* replica-group.
*
+ * When adaptive server selection is enabled, this selector uses
replica-group-level adaptive routing: it picks the best
+ * replica group for the entire query (using worst-case server rank within
each group) and routes all segments to
+ * that group. This preserves the same-replica-group guarantee while
benefiting from adaptive routing intelligence.
+ *
* Note that new segments won't be used to exclude instances from serving when
the segment is unavailable.
* </pre>
*/
public class StrictReplicaGroupInstanceSelector extends
ReplicaGroupInstanceSelector {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StrictReplicaGroupInstanceSelector.class);
@Override
void updateSegmentMaps(IdealState idealState, ExternalView externalView,
Set<String> onlineSegments,
Map<String, Long> newSegmentCreationTimeMap) {
super.updateSegmentMapsForUpsertTable(idealState, externalView,
onlineSegments, newSegmentCreationTimeMap);
}
+
+ @Override
+ public InstanceMapping select(List<String> segments, int requestId,
+ SegmentStates segmentStates, Map<String, String> queryOptions) {
+
+ if (_adaptiveServerSelector == null || _priorityPoolInstanceSelector ==
null) {
Review Comment:
This short-circuit misses `useFixedReplica`, so with adaptive routing
enabled, fixed-replica routing is now silently ignored for strict-RG tables.
Before this PR, `StrictReplicaGroupInstanceSelector` had no `select()`
override, so it inherited `ReplicaGroupInstanceSelector.select()` →
`selectServers()`, which checks `ctx.isUseFixedReplica()` **before** consulting
the rank map:
```java
if (useFixedReplica) {
// Adaptive Server Selection cannot be used with fixed replica routing.
selectedInstance = candidates.get((_tableNameHashForFixedReplicaRouting +
replicaOffset) % numCandidates);
} else if (MapUtils.isNotEmpty(serverRankMap)) {
```
`useFixedReplica` is not only a query option — `ServerSelectionContext`
falls back to `InstanceSelectorConfig`, which `InstanceSelectorFactory`
populates from `pinot.broker.use.fixed.replica` /
`routingConfig.getUseFixedReplica()`. So a cluster with adaptive routing on and
an upsert table with `useFixedReplica: true` used to get deterministic replica
pinning and now gets adaptive pool selection instead, with nothing in the logs.
Suggestion:
```java
ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
if (_adaptiveServerSelector == null || _priorityPoolInstanceSelector == null
|| ctx.isUseFixedReplica()) {
return selectServers(segments, requestId, segmentStates, null, ctx);
}
```
Related: `numReplicaGroupsToQuery` is also no longer honored on this path.
That's probably the right call for strict RG (it would break the invariant
anyway), but worth stating explicitly in the description since the parent's
javadoc still mentions it.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java:
##########
@@ -50,14 +64,175 @@
* segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that
we always route the segments to the same
* replica-group.
*
+ * When adaptive server selection is enabled, this selector uses
replica-group-level adaptive routing: it picks the best
+ * replica group for the entire query (using worst-case server rank within
each group) and routes all segments to
+ * that group. This preserves the same-replica-group guarantee while
benefiting from adaptive routing intelligence.
+ *
* Note that new segments won't be used to exclude instances from serving when
the segment is unavailable.
* </pre>
*/
public class StrictReplicaGroupInstanceSelector extends
ReplicaGroupInstanceSelector {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StrictReplicaGroupInstanceSelector.class);
@Override
void updateSegmentMaps(IdealState idealState, ExternalView externalView,
Set<String> onlineSegments,
Map<String, Long> newSegmentCreationTimeMap) {
super.updateSegmentMapsForUpsertTable(idealState, externalView,
onlineSegments, newSegmentCreationTimeMap);
}
+
+ @Override
+ public InstanceMapping select(List<String> segments, int requestId,
+ SegmentStates segmentStates, Map<String, String> queryOptions) {
+
+ if (_adaptiveServerSelector == null || _priorityPoolInstanceSelector ==
null) {
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ return selectServers(segments, requestId, segmentStates, null, ctx);
+ }
+
+ // Build a map: idealStateReplicaId (replica group) -> distinct servers in
that group that are candidates for this
Review Comment:
Design question, and the one I'm least confident about: does the
same-replica-group invariant need to hold across the whole query, or only
within each ideal-state instance set?
`updateSegmentMapsForUpsertTable()` computes `unavailableInstances` keyed by
`instancesInIdealState` — i.e. per mirror instance set — which reads to me like
the guarantee is per-partition: all segments of an upsert partition live on one
mirror set, so two segments on *different* instance sets can safely be served
from different replica groups. That's also effectively what today's round-robin
does whenever candidate-list lengths differ between instance sets.
If that's right, picking a single global `replicaId` for the entire query is
stricter than necessary, and it's what produces the dropped-segment behavior
further down: one degraded server in instance set A disqualifies replica group
*i* for every segment in sets B, C, …, and when no group covers everything,
segments get dropped.
An alternative would be to group candidates by their ideal-state instance
set and choose the best `replicaId` independently within each set — same
correctness guarantee, no coverage cliff, and adaptive routing still gets to
avoid the slow server. Happy to be told I'm misreading the partitioning
assumption.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java:
##########
@@ -50,14 +64,175 @@
* segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that
we always route the segments to the same
* replica-group.
*
+ * When adaptive server selection is enabled, this selector uses
replica-group-level adaptive routing: it picks the best
+ * replica group for the entire query (using worst-case server rank within
each group) and routes all segments to
+ * that group. This preserves the same-replica-group guarantee while
benefiting from adaptive routing intelligence.
+ *
* Note that new segments won't be used to exclude instances from serving when
the segment is unavailable.
* </pre>
*/
public class StrictReplicaGroupInstanceSelector extends
ReplicaGroupInstanceSelector {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StrictReplicaGroupInstanceSelector.class);
@Override
void updateSegmentMaps(IdealState idealState, ExternalView externalView,
Set<String> onlineSegments,
Map<String, Long> newSegmentCreationTimeMap) {
super.updateSegmentMapsForUpsertTable(idealState, externalView,
onlineSegments, newSegmentCreationTimeMap);
}
+
+ @Override
+ public InstanceMapping select(List<String> segments, int requestId,
+ SegmentStates segmentStates, Map<String, String> queryOptions) {
+
+ if (_adaptiveServerSelector == null || _priorityPoolInstanceSelector ==
null) {
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ return selectServers(segments, requestId, segmentStates, null, ctx);
+ }
+
+ // Build a map: idealStateReplicaId (replica group) -> distinct servers in
that group that are candidates for this
+ // query. Keyed by instance name to deduplicate (a server hosting multiple
segments appears once per group).
+ // Also track how many query segments each replica group can serve
(segment coverage) in a single pass.
+ // The guard defends against multiple candidates per replica group per
segment (e.g., during rebalancing).
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers = new LinkedHashMap<>();
+ Map<Integer, Integer> replicaGroupSegmentCount = new HashMap<>();
+ int totalSegmentsWithCandidates = 0;
+
+ Set<Integer> seenReplicaIds = new HashSet<>(); // allocate once and clear
per segment to avoid O(n) allocations
+ for (String segment : segments) {
+ List<SegmentInstanceCandidate> candidates =
segmentStates.getCandidates(segment);
+ if (candidates == null) {
+ continue;
+ }
+
+ seenReplicaIds.clear();
+
+ totalSegmentsWithCandidates++;
+ for (SegmentInstanceCandidate candidate : candidates) {
+ int replicaId = candidate.getIdealStateReplicaId();
+ replicaGroupToQueryServers
+ .computeIfAbsent(replicaId, k -> new LinkedHashMap<>())
+ .putIfAbsent(candidate.getInstance(), candidate);
+ if (seenReplicaIds.add(replicaId)) {
+ replicaGroupSegmentCount.merge(replicaId, 1, Integer::sum);
+ }
+ }
+ }
+
+ if (replicaGroupToQueryServers.isEmpty()) {
+ return new InstanceMapping(Map.of(), Map.of());
+ }
+
+ // Collect all distinct query-relevant candidates from
replicaGroupToQueryServers, which already
+ // deduplicates by instance name. This avoids a second traversal of
segments.
+ List<SegmentInstanceCandidate> allQueryCandidates =
replicaGroupToQueryServers.values().stream()
+ .flatMap(m -> m.values().stream())
+ .collect(Collectors.toList());
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ int bestReplicaGroupId = chooseBestReplicaGroup(
+ replicaGroupToQueryServers, replicaGroupSegmentCount,
totalSegmentsWithCandidates,
+ allQueryCandidates, ctx, requestId);
+ return selectServersForReplicaGroup(segments, bestReplicaGroupId,
segmentStates, queryOptions);
+ }
+
+ /**
+ * Scores and ranks replica groups using adaptive server selection stats.
+ *
+ * Each replica group is scored by the worst (maximum) rank among its
query-relevant servers.
+ * Since scatter-gather query latency is bounded by the slowest server, we
pick the group
+ * whose bottleneck server is the best (lowest worst-case rank).
+ *
+ * Groups that can serve ALL query segments (full coverage) are preferred
over groups that
+ * would drop segments. Among full-coverage groups, the one with the best
worst-case rank wins.
+ * If no group has full coverage, the best-ranked group is chosen regardless
of coverage.
+ */
+ private int chooseBestReplicaGroup(
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers,
+ Map<Integer, Integer> replicaGroupSegmentCount,
+ int totalSegmentsWithCandidates,
+ List<SegmentInstanceCandidate> allQueryCandidates,
+ ServerSelectionContext ctx,
+ int requestId) {
+
+ List<String> rankedServers = _priorityPoolInstanceSelector.rank(ctx,
allQueryCandidates);
+
+ // If no stats are available yet (empty ranking), fall back to pool based
round-robin.
+ if (rankedServers.isEmpty()) {
+ List<Integer> groupIds = new
ArrayList<>(replicaGroupToQueryServers.keySet());
+ return groupIds.get(Math.abs(requestId) % groupIds.size());
+ }
+
+ Map<String, Integer> serverRankMap = new
HashMap<>(HashUtil.getHashMapCapacity(rankedServers.size()));
+ for (int i = 0; i < rankedServers.size(); i++) {
+ serverRankMap.put(rankedServers.get(i), i);
+ }
+
+ // Pick the group with full coverage first, then best worst-case rank.
+ // getOrDefault guards against AdaptiveServerSelector implementations that
may not return every
+ // submitted server — unranked servers get rank -1 (best), matching
HybridSelector's convention.
+ // As of 8 July 2026, this fallback is unreachable, but new
implementations could require it.
+ return replicaGroupToQueryServers.entrySet().stream()
+ .min(Comparator.<Map.Entry<Integer, Map<String,
SegmentInstanceCandidate>>>comparingInt(
+ e -> hasFullCoverage(e.getKey(), replicaGroupSegmentCount,
totalSegmentsWithCandidates) ? 0 : 1)
Review Comment:
The coverage key is binary, so among partial-coverage groups the comparator
ignores *how much* each covers — a group covering 1 of 1000 queried segments
can beat one covering 999, purely on worst-case rank. Every uncovered segment
then becomes an unavailable segment.
Adding a coverage-count tiebreak between the two existing keys would make
the degraded path minimize dropped segments first and optimize latency second:
```java
.thenComparingInt(e -> -replicaGroupSegmentCount.getOrDefault(e.getKey(), 0))
```
That also makes the binary `hasFullCoverage` predicate redundant, since full
coverage is just the maximum of that count.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java:
##########
@@ -50,14 +64,175 @@
* segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that
we always route the segments to the same
* replica-group.
*
+ * When adaptive server selection is enabled, this selector uses
replica-group-level adaptive routing: it picks the best
+ * replica group for the entire query (using worst-case server rank within
each group) and routes all segments to
+ * that group. This preserves the same-replica-group guarantee while
benefiting from adaptive routing intelligence.
+ *
* Note that new segments won't be used to exclude instances from serving when
the segment is unavailable.
* </pre>
*/
public class StrictReplicaGroupInstanceSelector extends
ReplicaGroupInstanceSelector {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StrictReplicaGroupInstanceSelector.class);
@Override
void updateSegmentMaps(IdealState idealState, ExternalView externalView,
Set<String> onlineSegments,
Map<String, Long> newSegmentCreationTimeMap) {
super.updateSegmentMapsForUpsertTable(idealState, externalView,
onlineSegments, newSegmentCreationTimeMap);
}
+
+ @Override
+ public InstanceMapping select(List<String> segments, int requestId,
+ SegmentStates segmentStates, Map<String, String> queryOptions) {
+
+ if (_adaptiveServerSelector == null || _priorityPoolInstanceSelector ==
null) {
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ return selectServers(segments, requestId, segmentStates, null, ctx);
+ }
+
+ // Build a map: idealStateReplicaId (replica group) -> distinct servers in
that group that are candidates for this
+ // query. Keyed by instance name to deduplicate (a server hosting multiple
segments appears once per group).
+ // Also track how many query segments each replica group can serve
(segment coverage) in a single pass.
+ // The guard defends against multiple candidates per replica group per
segment (e.g., during rebalancing).
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers = new LinkedHashMap<>();
+ Map<Integer, Integer> replicaGroupSegmentCount = new HashMap<>();
+ int totalSegmentsWithCandidates = 0;
+
+ Set<Integer> seenReplicaIds = new HashSet<>(); // allocate once and clear
per segment to avoid O(n) allocations
+ for (String segment : segments) {
+ List<SegmentInstanceCandidate> candidates =
segmentStates.getCandidates(segment);
+ if (candidates == null) {
+ continue;
+ }
+
+ seenReplicaIds.clear();
+
+ totalSegmentsWithCandidates++;
+ for (SegmentInstanceCandidate candidate : candidates) {
+ int replicaId = candidate.getIdealStateReplicaId();
+ replicaGroupToQueryServers
+ .computeIfAbsent(replicaId, k -> new LinkedHashMap<>())
+ .putIfAbsent(candidate.getInstance(), candidate);
+ if (seenReplicaIds.add(replicaId)) {
+ replicaGroupSegmentCount.merge(replicaId, 1, Integer::sum);
+ }
+ }
+ }
+
+ if (replicaGroupToQueryServers.isEmpty()) {
+ return new InstanceMapping(Map.of(), Map.of());
+ }
+
+ // Collect all distinct query-relevant candidates from
replicaGroupToQueryServers, which already
+ // deduplicates by instance name. This avoids a second traversal of
segments.
+ List<SegmentInstanceCandidate> allQueryCandidates =
replicaGroupToQueryServers.values().stream()
+ .flatMap(m -> m.values().stream())
+ .collect(Collectors.toList());
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ int bestReplicaGroupId = chooseBestReplicaGroup(
+ replicaGroupToQueryServers, replicaGroupSegmentCount,
totalSegmentsWithCandidates,
+ allQueryCandidates, ctx, requestId);
+ return selectServersForReplicaGroup(segments, bestReplicaGroupId,
segmentStates, queryOptions);
+ }
+
+ /**
+ * Scores and ranks replica groups using adaptive server selection stats.
+ *
+ * Each replica group is scored by the worst (maximum) rank among its
query-relevant servers.
+ * Since scatter-gather query latency is bounded by the slowest server, we
pick the group
+ * whose bottleneck server is the best (lowest worst-case rank).
+ *
+ * Groups that can serve ALL query segments (full coverage) are preferred
over groups that
+ * would drop segments. Among full-coverage groups, the one with the best
worst-case rank wins.
+ * If no group has full coverage, the best-ranked group is chosen regardless
of coverage.
+ */
+ private int chooseBestReplicaGroup(
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers,
+ Map<Integer, Integer> replicaGroupSegmentCount,
+ int totalSegmentsWithCandidates,
+ List<SegmentInstanceCandidate> allQueryCandidates,
+ ServerSelectionContext ctx,
+ int requestId) {
+
+ List<String> rankedServers = _priorityPoolInstanceSelector.rank(ctx,
allQueryCandidates);
+
+ // If no stats are available yet (empty ranking), fall back to pool based
round-robin.
+ if (rankedServers.isEmpty()) {
+ List<Integer> groupIds = new
ArrayList<>(replicaGroupToQueryServers.keySet());
+ return groupIds.get(Math.abs(requestId) % groupIds.size());
+ }
+
+ Map<String, Integer> serverRankMap = new
HashMap<>(HashUtil.getHashMapCapacity(rankedServers.size()));
+ for (int i = 0; i < rankedServers.size(); i++) {
+ serverRankMap.put(rankedServers.get(i), i);
+ }
+
+ // Pick the group with full coverage first, then best worst-case rank.
+ // getOrDefault guards against AdaptiveServerSelector implementations that
may not return every
+ // submitted server — unranked servers get rank -1 (best), matching
HybridSelector's convention.
+ // As of 8 July 2026, this fallback is unreachable, but new
implementations could require it.
+ return replicaGroupToQueryServers.entrySet().stream()
+ .min(Comparator.<Map.Entry<Integer, Map<String,
SegmentInstanceCandidate>>>comparingInt(
+ e -> hasFullCoverage(e.getKey(), replicaGroupSegmentCount,
totalSegmentsWithCandidates) ? 0 : 1)
+ .thenComparingInt(e -> e.getValue().values().stream()
+ .mapToInt(c -> serverRankMap.getOrDefault(c.getInstance(), -1))
+ .max().orElse(Integer.MAX_VALUE)))
+ .map(Map.Entry::getKey)
+ .orElse(-1);
+ }
+
+ private boolean hasFullCoverage(int pool, Map<Integer, Integer>
replicaGroupSegmentCount,
+ int totalSegmentsWithCandidates) {
+ return replicaGroupSegmentCount.getOrDefault(pool, 0) >=
totalSegmentsWithCandidates;
+ }
+
+ /**
+ * Routes all segments to the chosen replica group. For each segment, picks
the candidate whose
+ * idealStateReplicaId matches the selected replica group ID. If no
candidate matches (transitional state during
+ * rebalancing), the segment is reported as unavailable.
+ */
+ private InstanceMapping selectServersForReplicaGroup(
+ List<String> segments, int replicaGroupId, SegmentStates segmentStates,
Map<String, String> queryOptions) {
+
+ Map<String, String> segmentToInstance = new
HashMap<>(HashUtil.getHashMapCapacity(segments.size()));
+ Map<String, String> optionalSegmentToInstance = new HashMap<>();
+ List<String> unavailableSegments = new ArrayList<>();
+ Map<Integer, Integer> poolToSegmentCount = new HashMap<>();
+
+ for (String segment : segments) {
+ List<SegmentInstanceCandidate> candidates =
segmentStates.getCandidates(segment);
+ if (candidates == null) {
+ continue;
+ }
+
+ SegmentInstanceCandidate selected = candidates.stream()
Review Comment:
`select()` runs for every query, so I'd avoid the stream pipeline here: for
a 5k-segment table this is ~5k `stream`/`filter`/`findFirst`/`Optional`
pipelines, on top of ~15k `computeIfAbsent`/`putIfAbsent`/`HashSet` operations
in the first pass. Candidate lists are short and already sorted, so a plain
`for` loop with a `break` removes all of it.
Separately, the `seenReplicaIds` set above guards against a duplicate
`idealStateReplicaId` within one segment, but that can't happen by
construction: `updateSegmentMapsForUpsertTable()` increments
`idealStateReplicaId` once per instance in the sorted ideal-state map, and
`refreshSegmentStates()` only filters that list. If it ever did happen,
`replicaGroupSegmentCount` would already be wrong — so an `assert` documents
the invariant better than a per-candidate hash lookup on the query path.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/BaseInstanceSelector.java:
##########
@@ -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);
Review Comment:
Merging the two sources is duplicate-safe, but only by an invariant that
lives in another method: `refreshSegmentStates()` puts a segment into
`instanceCandidatesMap` only when `enabledCandidates` is non-empty, so
`getCandidates()` returns `null` for everything in
`segmentStates.getUnavailableSegments()`, while the mapping-level list can only
contain segments that *had* candidates. The two lists are therefore disjoint
today.
Worth a one-line comment stating that (or a `LinkedHashSet` to be safe), so
a future selector can't make the error message list the same segment twice.
##########
pom.xml:
##########
@@ -2071,7 +2071,9 @@
<importOrder>
<order>,\#</order>
</importOrder>
- <removeUnusedImports />
+ <removeUnusedImports>
+ <engine>cleanthat-javaparser-unnecessaryimport</engine>
Review Comment:
I don't think this change belongs in this PR — could you pull it out?
It switches the unused-import engine for the entire reactor (~100 modules),
not just `pinot-broker`. `cleanthat-javaparser-unnecessaryimport` has different
semantics from google-java-format's `RemoveUnusedImports` (notably around
javadoc-only references), so it changes what `spotless:apply` does everywhere,
and that's invisible until someone's unrelated import silently stops being
cleaned up.
If the root cause is google-java-format's parser rejecting the record, the
narrower fix is to pin a newer `googleJavaFormat` version inside
`removeUnusedImports`. If the engine switch really is the way to go, it
deserves its own PR where it can be evaluated on its own merits — at which
point this PR just depends on it.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java:
##########
@@ -50,14 +64,175 @@
* segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that
we always route the segments to the same
* replica-group.
*
+ * When adaptive server selection is enabled, this selector uses
replica-group-level adaptive routing: it picks the best
+ * replica group for the entire query (using worst-case server rank within
each group) and routes all segments to
+ * that group. This preserves the same-replica-group guarantee while
benefiting from adaptive routing intelligence.
+ *
* Note that new segments won't be used to exclude instances from serving when
the segment is unavailable.
* </pre>
*/
public class StrictReplicaGroupInstanceSelector extends
ReplicaGroupInstanceSelector {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StrictReplicaGroupInstanceSelector.class);
@Override
void updateSegmentMaps(IdealState idealState, ExternalView externalView,
Set<String> onlineSegments,
Map<String, Long> newSegmentCreationTimeMap) {
super.updateSegmentMapsForUpsertTable(idealState, externalView,
onlineSegments, newSegmentCreationTimeMap);
}
+
+ @Override
+ public InstanceMapping select(List<String> segments, int requestId,
+ SegmentStates segmentStates, Map<String, String> queryOptions) {
+
+ if (_adaptiveServerSelector == null || _priorityPoolInstanceSelector ==
null) {
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ return selectServers(segments, requestId, segmentStates, null, ctx);
+ }
+
+ // Build a map: idealStateReplicaId (replica group) -> distinct servers in
that group that are candidates for this
+ // query. Keyed by instance name to deduplicate (a server hosting multiple
segments appears once per group).
+ // Also track how many query segments each replica group can serve
(segment coverage) in a single pass.
+ // The guard defends against multiple candidates per replica group per
segment (e.g., during rebalancing).
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers = new LinkedHashMap<>();
+ Map<Integer, Integer> replicaGroupSegmentCount = new HashMap<>();
+ int totalSegmentsWithCandidates = 0;
+
+ Set<Integer> seenReplicaIds = new HashSet<>(); // allocate once and clear
per segment to avoid O(n) allocations
+ for (String segment : segments) {
+ List<SegmentInstanceCandidate> candidates =
segmentStates.getCandidates(segment);
+ if (candidates == null) {
+ continue;
+ }
+
+ seenReplicaIds.clear();
+
+ totalSegmentsWithCandidates++;
+ for (SegmentInstanceCandidate candidate : candidates) {
+ int replicaId = candidate.getIdealStateReplicaId();
+ replicaGroupToQueryServers
+ .computeIfAbsent(replicaId, k -> new LinkedHashMap<>())
+ .putIfAbsent(candidate.getInstance(), candidate);
+ if (seenReplicaIds.add(replicaId)) {
+ replicaGroupSegmentCount.merge(replicaId, 1, Integer::sum);
+ }
+ }
+ }
+
+ if (replicaGroupToQueryServers.isEmpty()) {
+ return new InstanceMapping(Map.of(), Map.of());
+ }
+
+ // Collect all distinct query-relevant candidates from
replicaGroupToQueryServers, which already
+ // deduplicates by instance name. This avoids a second traversal of
segments.
+ List<SegmentInstanceCandidate> allQueryCandidates =
replicaGroupToQueryServers.values().stream()
+ .flatMap(m -> m.values().stream())
+ .collect(Collectors.toList());
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ int bestReplicaGroupId = chooseBestReplicaGroup(
+ replicaGroupToQueryServers, replicaGroupSegmentCount,
totalSegmentsWithCandidates,
+ allQueryCandidates, ctx, requestId);
+ return selectServersForReplicaGroup(segments, bestReplicaGroupId,
segmentStates, queryOptions);
+ }
+
+ /**
+ * Scores and ranks replica groups using adaptive server selection stats.
+ *
+ * Each replica group is scored by the worst (maximum) rank among its
query-relevant servers.
+ * Since scatter-gather query latency is bounded by the slowest server, we
pick the group
+ * whose bottleneck server is the best (lowest worst-case rank).
+ *
+ * Groups that can serve ALL query segments (full coverage) are preferred
over groups that
+ * would drop segments. Among full-coverage groups, the one with the best
worst-case rank wins.
+ * If no group has full coverage, the best-ranked group is chosen regardless
of coverage.
+ */
+ private int chooseBestReplicaGroup(
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers,
+ Map<Integer, Integer> replicaGroupSegmentCount,
+ int totalSegmentsWithCandidates,
+ List<SegmentInstanceCandidate> allQueryCandidates,
+ ServerSelectionContext ctx,
+ int requestId) {
+
+ List<String> rankedServers = _priorityPoolInstanceSelector.rank(ctx,
allQueryCandidates);
+
+ // If no stats are available yet (empty ranking), fall back to pool based
round-robin.
+ if (rankedServers.isEmpty()) {
+ List<Integer> groupIds = new
ArrayList<>(replicaGroupToQueryServers.keySet());
+ return groupIds.get(Math.abs(requestId) % groupIds.size());
Review Comment:
Two things about this fallback:
1. `replicaGroupToQueryServers` is a `LinkedHashMap` in first-seen order,
and which replica id is seen first depends on which segments this particular
query touched. So `requestId % size` doesn't map consistently to the same
replica group across queries or across brokers — round-robin over a sorted id
list would be both stable and evenly distributed.
2. Unlike the main path, this branch can pick a partial-coverage group and
silently drop segments. Worth reusing the coverage-preferring comparison here.
Also, the comment says "pool based round-robin" but no pool is involved. And
as far as I can tell this branch is unreachable with `HybridSelector`:
`fetchServerRankingsWithScores()` substitutes `-1.0` for a missing score rather
than omitting the server, so a non-empty candidate list always yields a
non-empty ranking. That's worth stating, because the description says "replica
groups with no stats are preferred" — that preference actually comes from the
`-1.0` substitution sorting first, not from this fallback. Same for the
`getOrDefault(..., -1)` below.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java:
##########
@@ -50,14 +64,175 @@
* segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that
we always route the segments to the same
* replica-group.
*
+ * When adaptive server selection is enabled, this selector uses
replica-group-level adaptive routing: it picks the best
+ * replica group for the entire query (using worst-case server rank within
each group) and routes all segments to
+ * that group. This preserves the same-replica-group guarantee while
benefiting from adaptive routing intelligence.
+ *
* Note that new segments won't be used to exclude instances from serving when
the segment is unavailable.
* </pre>
*/
public class StrictReplicaGroupInstanceSelector extends
ReplicaGroupInstanceSelector {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StrictReplicaGroupInstanceSelector.class);
@Override
void updateSegmentMaps(IdealState idealState, ExternalView externalView,
Set<String> onlineSegments,
Map<String, Long> newSegmentCreationTimeMap) {
super.updateSegmentMapsForUpsertTable(idealState, externalView,
onlineSegments, newSegmentCreationTimeMap);
}
+
+ @Override
+ public InstanceMapping select(List<String> segments, int requestId,
+ SegmentStates segmentStates, Map<String, String> queryOptions) {
+
+ if (_adaptiveServerSelector == null || _priorityPoolInstanceSelector ==
null) {
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ return selectServers(segments, requestId, segmentStates, null, ctx);
+ }
+
+ // Build a map: idealStateReplicaId (replica group) -> distinct servers in
that group that are candidates for this
+ // query. Keyed by instance name to deduplicate (a server hosting multiple
segments appears once per group).
+ // Also track how many query segments each replica group can serve
(segment coverage) in a single pass.
+ // The guard defends against multiple candidates per replica group per
segment (e.g., during rebalancing).
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers = new LinkedHashMap<>();
+ Map<Integer, Integer> replicaGroupSegmentCount = new HashMap<>();
+ int totalSegmentsWithCandidates = 0;
+
+ Set<Integer> seenReplicaIds = new HashSet<>(); // allocate once and clear
per segment to avoid O(n) allocations
+ for (String segment : segments) {
+ List<SegmentInstanceCandidate> candidates =
segmentStates.getCandidates(segment);
+ if (candidates == null) {
+ continue;
+ }
+
+ seenReplicaIds.clear();
+
+ totalSegmentsWithCandidates++;
+ for (SegmentInstanceCandidate candidate : candidates) {
+ int replicaId = candidate.getIdealStateReplicaId();
+ replicaGroupToQueryServers
+ .computeIfAbsent(replicaId, k -> new LinkedHashMap<>())
+ .putIfAbsent(candidate.getInstance(), candidate);
+ if (seenReplicaIds.add(replicaId)) {
+ replicaGroupSegmentCount.merge(replicaId, 1, Integer::sum);
+ }
+ }
+ }
+
+ if (replicaGroupToQueryServers.isEmpty()) {
+ return new InstanceMapping(Map.of(), Map.of());
+ }
+
+ // Collect all distinct query-relevant candidates from
replicaGroupToQueryServers, which already
+ // deduplicates by instance name. This avoids a second traversal of
segments.
+ List<SegmentInstanceCandidate> allQueryCandidates =
replicaGroupToQueryServers.values().stream()
+ .flatMap(m -> m.values().stream())
+ .collect(Collectors.toList());
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ int bestReplicaGroupId = chooseBestReplicaGroup(
+ replicaGroupToQueryServers, replicaGroupSegmentCount,
totalSegmentsWithCandidates,
+ allQueryCandidates, ctx, requestId);
+ return selectServersForReplicaGroup(segments, bestReplicaGroupId,
segmentStates, queryOptions);
+ }
+
+ /**
+ * Scores and ranks replica groups using adaptive server selection stats.
+ *
+ * Each replica group is scored by the worst (maximum) rank among its
query-relevant servers.
+ * Since scatter-gather query latency is bounded by the slowest server, we
pick the group
+ * whose bottleneck server is the best (lowest worst-case rank).
+ *
+ * Groups that can serve ALL query segments (full coverage) are preferred
over groups that
+ * would drop segments. Among full-coverage groups, the one with the best
worst-case rank wins.
+ * If no group has full coverage, the best-ranked group is chosen regardless
of coverage.
+ */
+ private int chooseBestReplicaGroup(
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers,
+ Map<Integer, Integer> replicaGroupSegmentCount,
+ int totalSegmentsWithCandidates,
+ List<SegmentInstanceCandidate> allQueryCandidates,
+ ServerSelectionContext ctx,
+ int requestId) {
+
+ List<String> rankedServers = _priorityPoolInstanceSelector.rank(ctx,
allQueryCandidates);
Review Comment:
`PriorityPoolInstanceSelector.rank()` returns a list that has already been
reordered by pool preference (all preferred-pool servers first, then ranked by
score within each pool), so `max(rank)` per group can invert the caller's
`orderedPreferredPools`.
Concrete case with `orderedPreferredPools=[0]`, group X = `{a(pool 0),
b(pool 1)}`, group Y = `{c(pool 1), d(pool 1)}`: the ranked list is `[a, c, d,
b]`, so `worst(X) = rank(b) = 3` and `worst(Y) = 2` — Y wins, even though X
holds the only preferred-pool server.
The test fixtures deliberately break pool ↔ replica-group alignment
(`FALLBACK_POOL_ID` everywhere, and the comment "Pools intentionally do not
match replica groups"), so the code can't assume they line up. If preferred
pools are meant to dominate, I'd score groups on preferred-pool membership
first and use the adaptive rank only as a tiebreak. Either way there's no test
covering `orderedPreferredPools` on the strict path.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java:
##########
@@ -50,14 +64,175 @@
* segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that
we always route the segments to the same
* replica-group.
*
+ * When adaptive server selection is enabled, this selector uses
replica-group-level adaptive routing: it picks the best
+ * replica group for the entire query (using worst-case server rank within
each group) and routes all segments to
+ * that group. This preserves the same-replica-group guarantee while
benefiting from adaptive routing intelligence.
+ *
* Note that new segments won't be used to exclude instances from serving when
the segment is unavailable.
* </pre>
*/
public class StrictReplicaGroupInstanceSelector extends
ReplicaGroupInstanceSelector {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StrictReplicaGroupInstanceSelector.class);
@Override
void updateSegmentMaps(IdealState idealState, ExternalView externalView,
Set<String> onlineSegments,
Map<String, Long> newSegmentCreationTimeMap) {
super.updateSegmentMapsForUpsertTable(idealState, externalView,
onlineSegments, newSegmentCreationTimeMap);
}
+
+ @Override
+ public InstanceMapping select(List<String> segments, int requestId,
+ SegmentStates segmentStates, Map<String, String> queryOptions) {
+
+ if (_adaptiveServerSelector == null || _priorityPoolInstanceSelector ==
null) {
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ return selectServers(segments, requestId, segmentStates, null, ctx);
+ }
+
+ // Build a map: idealStateReplicaId (replica group) -> distinct servers in
that group that are candidates for this
+ // query. Keyed by instance name to deduplicate (a server hosting multiple
segments appears once per group).
+ // Also track how many query segments each replica group can serve
(segment coverage) in a single pass.
+ // The guard defends against multiple candidates per replica group per
segment (e.g., during rebalancing).
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers = new LinkedHashMap<>();
+ Map<Integer, Integer> replicaGroupSegmentCount = new HashMap<>();
+ int totalSegmentsWithCandidates = 0;
+
+ Set<Integer> seenReplicaIds = new HashSet<>(); // allocate once and clear
per segment to avoid O(n) allocations
+ for (String segment : segments) {
+ List<SegmentInstanceCandidate> candidates =
segmentStates.getCandidates(segment);
+ if (candidates == null) {
+ continue;
+ }
+
+ seenReplicaIds.clear();
+
+ totalSegmentsWithCandidates++;
+ for (SegmentInstanceCandidate candidate : candidates) {
+ int replicaId = candidate.getIdealStateReplicaId();
+ replicaGroupToQueryServers
+ .computeIfAbsent(replicaId, k -> new LinkedHashMap<>())
+ .putIfAbsent(candidate.getInstance(), candidate);
+ if (seenReplicaIds.add(replicaId)) {
+ replicaGroupSegmentCount.merge(replicaId, 1, Integer::sum);
+ }
+ }
+ }
+
+ if (replicaGroupToQueryServers.isEmpty()) {
+ return new InstanceMapping(Map.of(), Map.of());
+ }
+
+ // Collect all distinct query-relevant candidates from
replicaGroupToQueryServers, which already
+ // deduplicates by instance name. This avoids a second traversal of
segments.
+ List<SegmentInstanceCandidate> allQueryCandidates =
replicaGroupToQueryServers.values().stream()
+ .flatMap(m -> m.values().stream())
+ .collect(Collectors.toList());
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ int bestReplicaGroupId = chooseBestReplicaGroup(
+ replicaGroupToQueryServers, replicaGroupSegmentCount,
totalSegmentsWithCandidates,
+ allQueryCandidates, ctx, requestId);
+ return selectServersForReplicaGroup(segments, bestReplicaGroupId,
segmentStates, queryOptions);
+ }
+
+ /**
+ * Scores and ranks replica groups using adaptive server selection stats.
+ *
+ * Each replica group is scored by the worst (maximum) rank among its
query-relevant servers.
+ * Since scatter-gather query latency is bounded by the slowest server, we
pick the group
+ * whose bottleneck server is the best (lowest worst-case rank).
+ *
+ * Groups that can serve ALL query segments (full coverage) are preferred
over groups that
+ * would drop segments. Among full-coverage groups, the one with the best
worst-case rank wins.
+ * If no group has full coverage, the best-ranked group is chosen regardless
of coverage.
+ */
+ private int chooseBestReplicaGroup(
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers,
+ Map<Integer, Integer> replicaGroupSegmentCount,
+ int totalSegmentsWithCandidates,
+ List<SegmentInstanceCandidate> allQueryCandidates,
+ ServerSelectionContext ctx,
+ int requestId) {
+
+ List<String> rankedServers = _priorityPoolInstanceSelector.rank(ctx,
allQueryCandidates);
+
+ // If no stats are available yet (empty ranking), fall back to pool based
round-robin.
+ if (rankedServers.isEmpty()) {
+ List<Integer> groupIds = new
ArrayList<>(replicaGroupToQueryServers.keySet());
+ return groupIds.get(Math.abs(requestId) % groupIds.size());
+ }
+
+ Map<String, Integer> serverRankMap = new
HashMap<>(HashUtil.getHashMapCapacity(rankedServers.size()));
+ for (int i = 0; i < rankedServers.size(); i++) {
+ serverRankMap.put(rankedServers.get(i), i);
+ }
+
+ // Pick the group with full coverage first, then best worst-case rank.
+ // getOrDefault guards against AdaptiveServerSelector implementations that
may not return every
+ // submitted server — unranked servers get rank -1 (best), matching
HybridSelector's convention.
+ // As of 8 July 2026, this fallback is unreachable, but new
implementations could require it.
+ return replicaGroupToQueryServers.entrySet().stream()
+ .min(Comparator.<Map.Entry<Integer, Map<String,
SegmentInstanceCandidate>>>comparingInt(
+ e -> hasFullCoverage(e.getKey(), replicaGroupSegmentCount,
totalSegmentsWithCandidates) ? 0 : 1)
+ .thenComparingInt(e -> e.getValue().values().stream()
+ .mapToInt(c -> serverRankMap.getOrDefault(c.getInstance(), -1))
+ .max().orElse(Integer.MAX_VALUE)))
+ .map(Map.Entry::getKey)
+ .orElse(-1);
+ }
+
+ private boolean hasFullCoverage(int pool, Map<Integer, Integer>
replicaGroupSegmentCount,
+ int totalSegmentsWithCandidates) {
+ return replicaGroupSegmentCount.getOrDefault(pool, 0) >=
totalSegmentsWithCandidates;
+ }
+
+ /**
+ * Routes all segments to the chosen replica group. For each segment, picks
the candidate whose
+ * idealStateReplicaId matches the selected replica group ID. If no
candidate matches (transitional state during
+ * rebalancing), the segment is reported as unavailable.
+ */
+ private InstanceMapping selectServersForReplicaGroup(
+ List<String> segments, int replicaGroupId, SegmentStates segmentStates,
Map<String, String> queryOptions) {
+
+ Map<String, String> segmentToInstance = new
HashMap<>(HashUtil.getHashMapCapacity(segments.size()));
+ Map<String, String> optionalSegmentToInstance = new HashMap<>();
+ List<String> unavailableSegments = new ArrayList<>();
+ Map<Integer, Integer> poolToSegmentCount = new HashMap<>();
+
+ for (String segment : segments) {
+ List<SegmentInstanceCandidate> candidates =
segmentStates.getCandidates(segment);
+ if (candidates == null) {
+ continue;
+ }
+
+ SegmentInstanceCandidate selected = candidates.stream()
+ .filter(c -> c.getIdealStateReplicaId() == replicaGroupId)
+ .findFirst()
+ .orElse(null);
+
+ // Skip segments with no candidate in the chosen replica group to
preserve the strict same-replica-group
+ // invariant. We prefer to use groups that cover all queried segments,
so this warning only occurs when we've
+ // selected a partial group.
+ if (selected == null) {
+ LOGGER.debug("No candidate found in replica group {} for segment {};
reporting as unavailable",
+ replicaGroupId, segment);
+ unavailableSegments.add(segment);
Review Comment:
This introduces a failure mode that I don't think any instance selector had
before: a segment with a healthy, online server gets reported as unavailable.
`BaseSingleStageBrokerRequestHandler` turns any unavailable segment into a
`QueryProcessingException(BROKER_SEGMENT_UNAVAILABLE)` attached to the response
and increments `BROKER_RESPONSES_WITH_UNAVAILABLE_SEGMENTS` — which is what a
lot of dashboards and alerts are built on, and many clients treat any exception
in the response as a failed query. Pre-PR, every segment with a non-null
candidate list got a server, so this path could not trigger.
Two asks:
1. Can you characterize how reachable the "no group has full coverage" state
is? Two servers lagging in different replica groups (rebalance, overlapping
restarts, EV convergence lag) seems sufficient, and 2-replica upsert tables are
common.
2. New segments should never be reported unavailable — that invariant is
documented in `BaseInstanceSelector`'s class javadoc ("We don't report new
segment as unavailable segments") and enforced in `refreshSegmentStates()`.
This loop can't distinguish a new segment from an old one, so a new segment
whose ideal state has fewer replicas than the old segments in its instance set
would land here.
`testStrictReplicaGroupAdaptiveNewSegmentNotFalselyUnavailable` only covers the
null-candidates case, not this one.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/StrictReplicaGroupInstanceSelector.java:
##########
@@ -50,14 +64,175 @@
* segments with the same assignment ([S1, S2, S3]) down on S1 to ensure that
we always route the segments to the same
* replica-group.
*
+ * When adaptive server selection is enabled, this selector uses
replica-group-level adaptive routing: it picks the best
+ * replica group for the entire query (using worst-case server rank within
each group) and routes all segments to
+ * that group. This preserves the same-replica-group guarantee while
benefiting from adaptive routing intelligence.
+ *
* Note that new segments won't be used to exclude instances from serving when
the segment is unavailable.
* </pre>
*/
public class StrictReplicaGroupInstanceSelector extends
ReplicaGroupInstanceSelector {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StrictReplicaGroupInstanceSelector.class);
@Override
void updateSegmentMaps(IdealState idealState, ExternalView externalView,
Set<String> onlineSegments,
Map<String, Long> newSegmentCreationTimeMap) {
super.updateSegmentMapsForUpsertTable(idealState, externalView,
onlineSegments, newSegmentCreationTimeMap);
}
+
+ @Override
+ public InstanceMapping select(List<String> segments, int requestId,
+ SegmentStates segmentStates, Map<String, String> queryOptions) {
+
+ if (_adaptiveServerSelector == null || _priorityPoolInstanceSelector ==
null) {
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ return selectServers(segments, requestId, segmentStates, null, ctx);
+ }
+
+ // Build a map: idealStateReplicaId (replica group) -> distinct servers in
that group that are candidates for this
+ // query. Keyed by instance name to deduplicate (a server hosting multiple
segments appears once per group).
+ // Also track how many query segments each replica group can serve
(segment coverage) in a single pass.
+ // The guard defends against multiple candidates per replica group per
segment (e.g., during rebalancing).
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers = new LinkedHashMap<>();
+ Map<Integer, Integer> replicaGroupSegmentCount = new HashMap<>();
+ int totalSegmentsWithCandidates = 0;
+
+ Set<Integer> seenReplicaIds = new HashSet<>(); // allocate once and clear
per segment to avoid O(n) allocations
+ for (String segment : segments) {
+ List<SegmentInstanceCandidate> candidates =
segmentStates.getCandidates(segment);
+ if (candidates == null) {
+ continue;
+ }
+
+ seenReplicaIds.clear();
+
+ totalSegmentsWithCandidates++;
+ for (SegmentInstanceCandidate candidate : candidates) {
+ int replicaId = candidate.getIdealStateReplicaId();
+ replicaGroupToQueryServers
+ .computeIfAbsent(replicaId, k -> new LinkedHashMap<>())
+ .putIfAbsent(candidate.getInstance(), candidate);
+ if (seenReplicaIds.add(replicaId)) {
+ replicaGroupSegmentCount.merge(replicaId, 1, Integer::sum);
+ }
+ }
+ }
+
+ if (replicaGroupToQueryServers.isEmpty()) {
+ return new InstanceMapping(Map.of(), Map.of());
+ }
+
+ // Collect all distinct query-relevant candidates from
replicaGroupToQueryServers, which already
+ // deduplicates by instance name. This avoids a second traversal of
segments.
+ List<SegmentInstanceCandidate> allQueryCandidates =
replicaGroupToQueryServers.values().stream()
+ .flatMap(m -> m.values().stream())
+ .collect(Collectors.toList());
+ ServerSelectionContext ctx = new ServerSelectionContext(queryOptions,
_config);
+ int bestReplicaGroupId = chooseBestReplicaGroup(
+ replicaGroupToQueryServers, replicaGroupSegmentCount,
totalSegmentsWithCandidates,
+ allQueryCandidates, ctx, requestId);
+ return selectServersForReplicaGroup(segments, bestReplicaGroupId,
segmentStates, queryOptions);
+ }
+
+ /**
+ * Scores and ranks replica groups using adaptive server selection stats.
+ *
+ * Each replica group is scored by the worst (maximum) rank among its
query-relevant servers.
+ * Since scatter-gather query latency is bounded by the slowest server, we
pick the group
+ * whose bottleneck server is the best (lowest worst-case rank).
+ *
+ * Groups that can serve ALL query segments (full coverage) are preferred
over groups that
+ * would drop segments. Among full-coverage groups, the one with the best
worst-case rank wins.
+ * If no group has full coverage, the best-ranked group is chosen regardless
of coverage.
+ */
+ private int chooseBestReplicaGroup(
+ Map<Integer, Map<String, SegmentInstanceCandidate>>
replicaGroupToQueryServers,
+ Map<Integer, Integer> replicaGroupSegmentCount,
+ int totalSegmentsWithCandidates,
+ List<SegmentInstanceCandidate> allQueryCandidates,
+ ServerSelectionContext ctx,
+ int requestId) {
+
+ List<String> rankedServers = _priorityPoolInstanceSelector.rank(ctx,
allQueryCandidates);
+
+ // If no stats are available yet (empty ranking), fall back to pool based
round-robin.
+ if (rankedServers.isEmpty()) {
+ List<Integer> groupIds = new
ArrayList<>(replicaGroupToQueryServers.keySet());
+ return groupIds.get(Math.abs(requestId) % groupIds.size());
+ }
+
+ Map<String, Integer> serverRankMap = new
HashMap<>(HashUtil.getHashMapCapacity(rankedServers.size()));
+ for (int i = 0; i < rankedServers.size(); i++) {
+ serverRankMap.put(rankedServers.get(i), i);
+ }
+
+ // Pick the group with full coverage first, then best worst-case rank.
+ // getOrDefault guards against AdaptiveServerSelector implementations that
may not return every
+ // submitted server — unranked servers get rank -1 (best), matching
HybridSelector's convention.
+ // As of 8 July 2026, this fallback is unreachable, but new
implementations could require it.
+ return replicaGroupToQueryServers.entrySet().stream()
+ .min(Comparator.<Map.Entry<Integer, Map<String,
SegmentInstanceCandidate>>>comparingInt(
+ e -> hasFullCoverage(e.getKey(), replicaGroupSegmentCount,
totalSegmentsWithCandidates) ? 0 : 1)
+ .thenComparingInt(e -> e.getValue().values().stream()
+ .mapToInt(c -> serverRankMap.getOrDefault(c.getInstance(), -1))
+ .max().orElse(Integer.MAX_VALUE)))
+ .map(Map.Entry::getKey)
+ .orElse(-1);
+ }
+
+ private boolean hasFullCoverage(int pool, Map<Integer, Integer>
replicaGroupSegmentCount,
Review Comment:
Nit: the parameter is named `pool` but it carries an `idealStateReplicaId` —
which is exactly the distinction the "Group by idealStateReplicaId not pool"
commit was drawing. `replicaGroupId` would be clearer, especially since
`poolToSegmentCount` a few lines below genuinely is keyed by pool.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/instanceselector/InstanceSelector.java:
##########
@@ -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,
Review Comment:
Good change — the positional `Pair` was easy to get backwards. One
compatibility note: `routingConfig.instanceSelectorType` falls through to
`PluginManager.createInstance()`, so instance selectors are a supported
extension point, and any out-of-tree selector extending `BaseInstanceSelector`
will fail to compile against the new `protected abstract select()` signature. I
think that's an acceptable trade, but it warrants the `backward-incompat` label
(adding it) and a line in the release notes.
Two smaller things:
- The 2-arg convenience constructor and `EMPTY` are package-private, so an
out-of-package subclass can only use the 3-arg canonical form. Consider making
the 2-arg constructor public.
- The javadoc's "segments that have candidates but could not be routed" is
new behavior at the interface level — no selector could previously return such
segments. Worth calling out here that these are merged into the query's
unavailable-segment list and become a `BROKER_SEGMENT_UNAVAILABLE` error.
##########
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(
Review Comment:
Two things here.
**Naming.** Config keys are effectively permanent, and
`pinot.broker.adaptive.server.selector.enable.strict.replica.group` reads like
"enable strict replica group" rather than "enable adaptive routing *for*
strict-RG tables". Something like
`...adaptive.server.selector.strict.replica.group.enabled` or
`...enable.for.strict.replica.group` would age better.
**Scope.** Upsert/dedup validation allows `strictReplicaGroup` or
`multiStageReplicaGroup`, and `MultiStageReplicaGroupSelector` ignores the
adaptive selector, so new tables look covered. But
`ReplicaGroupInstanceSelector.updateSegmentMaps()` still branches on
`isUpsertEnabled() || isDedupEnabled()`, which suggests grandfathered upsert
tables on plain `replicaGroup` routing exist somewhere. Those would keep the
per-segment adaptive routing this PR is fixing, and neither the flag nor the
new logic touches them. Is that combination actually reachable?
Minor: reassigning the `adaptiveServerSelector` parameter inside the case
block works, but a local would read better given it's consumed much further
down at `init()`.
##########
pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/ReplicaGroupSelectorTest.java:
##########
@@ -0,0 +1,913 @@
+/**
+ * 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.
+ */
+package org.apache.pinot.broker.routing.instanceselector;
+
+import java.time.Clock;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeMap;
+import org.apache.commons.lang3.tuple.ImmutablePair;
+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;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.broker.routing.adaptiveserverselector.HybridSelector;
+import org.apache.pinot.common.metrics.BrokerMeter;
+import org.apache.pinot.common.metrics.BrokerMetrics;
+import org.apache.pinot.common.request.BrokerRequest;
+import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.core.transport.ServerInstance;
+import org.apache.pinot.spi.config.table.RoutingConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import static
org.apache.pinot.spi.config.table.RoutingConfig.REPLICA_GROUP_INSTANCE_SELECTOR_TYPE;
+import static
org.apache.pinot.spi.config.table.RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE;
+import static
org.apache.pinot.spi.utils.CommonConstants.Broker.FALLBACK_POOL_ID;
+import static
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel.ONLINE;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+
+
+/**
+ * Tests for {@link ReplicaGroupInstanceSelector} and {@link
StrictReplicaGroupInstanceSelector},
+ * including adaptive server selection (replica-group-level routing for strict
replica groups).
+ */
+@SuppressWarnings("unchecked")
+public class ReplicaGroupSelectorTest {
+ private AutoCloseable _mocks;
+
+ @Mock
+ private TableConfig _tableConfig;
+
+ private static final String TABLE_NAME = "testTable_OFFLINE";
+ private static final Map<String, ServerInstance> EMPTY_SERVER_MAP =
Collections.EMPTY_MAP;
+ private static final InstanceSelectorConfig INSTANCE_SELECTOR_CONFIG = new
InstanceSelectorConfig(false, 300, false);
+ private static final List<String> SEGMENTS =
+ Arrays.asList("segment0", "segment1", "segment2", "segment3",
"segment4", "segment5", "segment6", "segment7",
+ "segment8", "segment9", "segment10", "segment11");
+
+ @BeforeMethod
+ public void setUp() {
+ _mocks = MockitoAnnotations.openMocks(this);
+ when(_tableConfig.getTableName()).thenReturn(TABLE_NAME);
+ }
+
+ @AfterMethod
+ public void tearDown()
+ throws Exception {
+ _mocks.close();
+ }
+
+ // --- Shared helpers ---
+
+ static IdealState createIdealState(Map<String, List<Pair<String, String>>>
segmentState) {
+ IdealState idealState = new IdealState(TABLE_NAME);
+ Map<String, Map<String, String>> idealStateSegmentAssignment =
idealState.getRecord().getMapFields();
+ for (Map.Entry<String, List<Pair<String, String>>> entry :
segmentState.entrySet()) {
+ Map<String, String> instanceStateMap = new TreeMap<>();
+ for (Pair<String, String> instanceState : entry.getValue()) {
+ instanceStateMap.put(instanceState.getLeft(),
instanceState.getRight());
+ }
+ idealStateSegmentAssignment.put(entry.getKey(), instanceStateMap);
+ }
+ return idealState;
+ }
+
+ static ExternalView createExternalView(Map<String, List<Pair<String,
String>>> segmentState) {
+ ExternalView externalView = new ExternalView(TABLE_NAME);
+ Map<String, Map<String, String>> externalViewSegmentAssignment =
externalView.getRecord().getMapFields();
+ for (Map.Entry<String, List<Pair<String, String>>> entry :
segmentState.entrySet()) {
+ Map<String, String> instanceStateMap = new TreeMap<>();
+ for (Pair<String, String> instanceState : entry.getValue()) {
+ instanceStateMap.put(instanceState.getLeft(),
instanceState.getRight());
+ }
+ externalViewSegmentAssignment.put(entry.getKey(), instanceStateMap);
+ }
+ return externalView;
+ }
+
+ // --- ReplicaGroupInstanceSelector: numReplicaGroupsToQuery tests ---
+
+ @Test
+ public void testReplicaGroupInstanceSelectorNumReplicaGroupsToQuery() {
+ String offlineTableName = "testTable_OFFLINE";
+ ZkHelixPropertyStore<ZNRecord> propertyStore =
mock(ZkHelixPropertyStore.class);
+ BrokerMetrics brokerMetrics = mock(BrokerMetrics.class);
+ BrokerRequest brokerRequest = mock(BrokerRequest.class);
+ PinotQuery pinotQuery = mock(PinotQuery.class);
+ Map<String, String> queryOptions = new HashMap<>();
+ // numReplicas = 3, fanning the query to 2 replica groups
+ queryOptions.put("numReplicaGroupsToQuery", "2");
+ when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery);
+ when(pinotQuery.getQueryOptions()).thenReturn(queryOptions);
+
+ ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new
ReplicaGroupInstanceSelector();
+
+ Set<String> enabledInstances = new HashSet<>();
+ IdealState idealState = new IdealState(offlineTableName);
+ Map<String, Map<String, String>> idealStateSegmentAssignment =
idealState.getRecord().getMapFields();
+ ExternalView externalView = new ExternalView(offlineTableName);
+ Map<String, Map<String, String>> externalViewSegmentAssignment =
externalView.getRecord().getMapFields();
+ Set<String> onlineSegments = new HashSet<>();
+
+ // 12 online segments with each segment having all 3 instances as online
+ // replicas are 3
+ String instance0 = "instance0";
+ String instance1 = "instance1";
+ String instance2 = "instance2";
+ enabledInstances.add(instance0);
+ enabledInstances.add(instance1);
+ enabledInstances.add(instance2);
+
+ Map<String, String> idealStateInstanceStateMap0 = new TreeMap<>();
+ Map<String, String> externalViewInstanceStateMap0 = new TreeMap<>();
+
+ for (String instance : enabledInstances) {
+ idealStateInstanceStateMap0.put(instance, ONLINE);
+ externalViewInstanceStateMap0.put(instance, ONLINE);
+ }
+
+ List<String> segments = SEGMENTS;
+ // add all segments to both idealStateSegmentAssignment and
externalViewSegmentAssignment maps and also to online
+ // segments
+ for (String segment : segments) {
+ idealStateSegmentAssignment.put(segment, idealStateInstanceStateMap0);
+ externalViewSegmentAssignment.put(segment,
externalViewInstanceStateMap0);
+ onlineSegments.add(segment);
+ }
+
+ replicaGroupInstanceSelector.init(_tableConfig, propertyStore,
brokerMetrics, null, Clock.systemUTC(),
+ INSTANCE_SELECTOR_CONFIG, enabledInstances, EMPTY_SERVER_MAP,
idealState, externalView, onlineSegments);
+
+ Map<String, String> expectedReplicaGroupInstanceSelectorResult = new
HashMap<>();
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(0), instance0);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(1), instance1);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(2), instance0);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(3), instance1);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(4), instance0);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(5), instance1);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(6), instance0);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(7), instance1);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(8), instance0);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(9), instance1);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(10),
instance0);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(11),
instance1);
+ InstanceSelector.SelectionResult selectionResult =
replicaGroupInstanceSelector.select(brokerRequest, segments, 0);
+ assertEquals(selectionResult.getSegmentToInstanceMap(),
expectedReplicaGroupInstanceSelectorResult);
+ assertTrue(selectionResult.getUnavailableSegments().isEmpty());
+ }
+
+ @Test
+ public void
testReplicaGroupInstanceSelectorNumReplicaGroupsToQueryGreaterThanReplicas() {
+ String offlineTableName = "testTable_OFFLINE";
+ ZkHelixPropertyStore<ZNRecord> propertyStore =
mock(ZkHelixPropertyStore.class);
+ BrokerMetrics brokerMetrics = mock(BrokerMetrics.class);
+ BrokerRequest brokerRequest = mock(BrokerRequest.class);
+ PinotQuery pinotQuery = mock(PinotQuery.class);
+ Map<String, String> queryOptions = new HashMap<>();
+ queryOptions.put("numReplicaGroupsToQuery", "4");
+
+ when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery);
+ when(pinotQuery.getQueryOptions()).thenReturn(queryOptions);
+
+ ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new
ReplicaGroupInstanceSelector();
+
+ Set<String> enabledInstances = new HashSet<>();
+ IdealState idealState = new IdealState(offlineTableName);
+ Map<String, Map<String, String>> idealStateSegmentAssignment =
idealState.getRecord().getMapFields();
+ ExternalView externalView = new ExternalView(offlineTableName);
+ Map<String, Map<String, String>> externalViewSegmentAssignment =
externalView.getRecord().getMapFields();
+ Set<String> onlineSegments = new HashSet<>();
+
+ // 12 online segments with each segment having all 3 instances as online
+ // replicas are 3
+ String instance0 = "instance0";
+ String instance1 = "instance1";
+ String instance2 = "instance2";
+ enabledInstances.add(instance0);
+ enabledInstances.add(instance1);
+ enabledInstances.add(instance2);
+
+ List<String> segments = SEGMENTS;
+
+ Map<String, String> idealStateInstanceStateMap0 = new TreeMap<>();
+ Map<String, String> externalViewInstanceStateMap0 = new TreeMap<>();
+
+ for (String instance : enabledInstances) {
+ idealStateInstanceStateMap0.put(instance, ONLINE);
+ externalViewInstanceStateMap0.put(instance, ONLINE);
+ }
+
+ // add all segments to both idealStateSegmentAssignment and
externalViewSegmentAssignment maps and also to online
+ // segments
+ for (String segment : segments) {
+ idealStateSegmentAssignment.put(segment, idealStateInstanceStateMap0);
+ externalViewSegmentAssignment.put(segment,
externalViewInstanceStateMap0);
+ onlineSegments.add(segment);
+ }
+
+ replicaGroupInstanceSelector.init(_tableConfig, propertyStore,
brokerMetrics, null, Clock.systemUTC(),
+ INSTANCE_SELECTOR_CONFIG, enabledInstances, EMPTY_SERVER_MAP,
idealState, externalView, onlineSegments);
+
+ Map<String, String> expectedReplicaGroupInstanceSelectorResult = new
HashMap<>();
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(0), instance0);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(1), instance1);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(2), instance2);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(3), instance0);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(4), instance1);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(5), instance2);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(6), instance0);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(7), instance1);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(8), instance2);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(9), instance0);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(10),
instance1);
+ expectedReplicaGroupInstanceSelectorResult.put(segments.get(11),
instance2);
+ InstanceSelector.SelectionResult selectionResult =
replicaGroupInstanceSelector.select(brokerRequest, segments, 0);
+ assertEquals(selectionResult.getSegmentToInstanceMap(),
expectedReplicaGroupInstanceSelectorResult);
+ assertTrue(selectionResult.getUnavailableSegments().isEmpty());
+ }
+
+ @Test
+ public void testReplicaGroupInstanceSelectorNumReplicaGroupsNotSet() {
+ String offlineTableName = "testTable_OFFLINE";
+ ZkHelixPropertyStore<ZNRecord> propertyStore =
mock(ZkHelixPropertyStore.class);
+ BrokerMetrics brokerMetrics = mock(BrokerMetrics.class);
+ BrokerRequest brokerRequest = mock(BrokerRequest.class);
+ PinotQuery pinotQuery = mock(PinotQuery.class);
+ Map<String, String> queryOptions = new HashMap<>();
+
+ when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery);
+ when(pinotQuery.getQueryOptions()).thenReturn(queryOptions);
+
+ ReplicaGroupInstanceSelector replicaGroupInstanceSelector = new
ReplicaGroupInstanceSelector();
+
+ Set<String> enabledInstances = new HashSet<>();
+ IdealState idealState = new IdealState(offlineTableName);
+ Map<String, Map<String, String>> idealStateSegmentAssignment =
idealState.getRecord().getMapFields();
+ ExternalView externalView = new ExternalView(offlineTableName);
+ Map<String, Map<String, String>> externalViewSegmentAssignment =
externalView.getRecord().getMapFields();
+ Set<String> onlineSegments = new HashSet<>();
+
+ // 12 online segments with each segment having all 3 instances as online
+ // replicas are 3
+ String instance0 = "instance0";
+ String instance1 = "instance1";
+ String instance2 = "instance2";
+ enabledInstances.add(instance0);
+ enabledInstances.add(instance1);
+ enabledInstances.add(instance2);
+
+ List<String> segments = SEGMENTS;
+
+ Map<String, String> idealStateInstanceStateMap0 = new TreeMap<>();
+ Map<String, String> externalViewInstanceStateMap0 = new TreeMap<>();
+
+ for (String instance : enabledInstances) {
+ idealStateInstanceStateMap0.put(instance, ONLINE);
+ externalViewInstanceStateMap0.put(instance, ONLINE);
+ }
+
+ // add all segments to both idealStateSegmentAssignment and
externalViewSegmentAssignment maps and also to online
+ // segments
+ for (String segment : segments) {
+ idealStateSegmentAssignment.put(segment, idealStateInstanceStateMap0);
+ externalViewSegmentAssignment.put(segment,
externalViewInstanceStateMap0);
+ onlineSegments.add(segment);
+ }
+
+ replicaGroupInstanceSelector.init(_tableConfig, propertyStore,
brokerMetrics, null, Clock.systemUTC(),
+ INSTANCE_SELECTOR_CONFIG, enabledInstances, EMPTY_SERVER_MAP,
idealState, externalView, onlineSegments);
+ // since numReplicaGroupsToQuery is not set, first query should go to
first replica group,
+ // 2nd query should go to next replica group
+
+ Map<String, String> expectedReplicaGroupInstanceSelectorResult = new
HashMap<>();
+ for (String segment : segments) {
+ expectedReplicaGroupInstanceSelectorResult.put(segment, instance0);
+ }
+ InstanceSelector.SelectionResult selectionResult =
replicaGroupInstanceSelector.select(brokerRequest, segments, 0);
+ assertEquals(selectionResult.getSegmentToInstanceMap(),
expectedReplicaGroupInstanceSelectorResult);
+
+ for (String segment : segments) {
+ expectedReplicaGroupInstanceSelectorResult.put(segment, instance1);
+ }
+ selectionResult = replicaGroupInstanceSelector.select(brokerRequest,
segments, 1);
+ assertEquals(selectionResult.getSegmentToInstanceMap(),
expectedReplicaGroupInstanceSelectorResult);
+ }
+
+ // --- Adaptive server selection tests ---
+
+ // Shared topology for AR tests: 3 segments across 5 instances in two pools
/ replica groups.
+ // segment2 intentionally has AR_P0_RG0_SERVER_E (unranked) so AR falls back
to round-robin for that segment.
+ private static final String AR_P0_RG0_SERVER_A = "ar_p0_rg0_server_a";
+ private static final String AR_P1_RG1_SERVER_B = "ar_p1_rg1_server_b";
+ private static final String AR_P0_RG0_SERVER_C = "ar_p0_rg0_server_c";
+ private static final String AR_P1_RG1_SERVER_D = "ar_p1_rg1_server_d";
+ private static final String AR_P0_RG0_SERVER_E = "ar_p0_rg0_server_e";
+ private static final String AR_SEGMENT0 = "segment0";
+ private static final String AR_SEGMENT1 = "segment1";
+ private static final String AR_SEGMENT2 = "segment2";
+ private static final List<String> AR_SEGMENTS =
+ Arrays.asList(AR_SEGMENT0, AR_SEGMENT1, AR_SEGMENT2);
+ // Rankings: D best → C → B → A worst; E absent (triggers AR fallback)
+ private static final List<Pair<String, Double>> AR_SERVER_RANKS =
Arrays.asList(
+ new ImmutablePair<>(AR_P1_RG1_SERVER_D, 1.0),
+ new ImmutablePair<>(AR_P0_RG0_SERVER_C, 2.0),
+ new ImmutablePair<>(AR_P1_RG1_SERVER_B, 3.0),
+ new ImmutablePair<>(AR_P0_RG0_SERVER_A, 4.0)
+ );
+
+ private SegmentStates buildArSegmentStates() {
+ Map<String, List<SegmentInstanceCandidate>> candidatesMap = new
HashMap<>();
+ // segment0 → pool 0 / replica group 0 server A, pool 1 / replica group 1
server B
+ candidatesMap.put(AR_SEGMENT0, Arrays.asList(
+ new SegmentInstanceCandidate(AR_P0_RG0_SERVER_A, true, 0, 0),
+ new SegmentInstanceCandidate(AR_P1_RG1_SERVER_B, true, 1, 1)));
+ // segment1 → pool 0 / replica group 0 server C, pool 1 / replica group 1
server D
+ candidatesMap.put(AR_SEGMENT1, Arrays.asList(
+ new SegmentInstanceCandidate(AR_P0_RG0_SERVER_C, true, 0, 0),
+ new SegmentInstanceCandidate(AR_P1_RG1_SERVER_D, true, 1, 1)));
+ // segment2 → pool 0 / replica group 0 server E (unranked), pool 1 /
replica group 1 server D
+ candidatesMap.put(AR_SEGMENT2, Arrays.asList(
+ new SegmentInstanceCandidate(AR_P0_RG0_SERVER_E, true, 0, 0),
+ new SegmentInstanceCandidate(AR_P1_RG1_SERVER_D, true, 1, 1)));
+ return new SegmentStates(candidatesMap, new HashSet<>(AR_SEGMENTS), null);
+ }
+
+ private ReplicaGroupInstanceSelector buildArSelector(String selectorType,
HybridSelector hybridSelector) {
+ return buildArSelector(selectorType, hybridSelector, new
PinotConfiguration(Map.of()));
+ }
+
+ private ReplicaGroupInstanceSelector buildArSelector(String selectorType,
HybridSelector hybridSelector,
+ PinotConfiguration brokerConfig) {
+ RoutingConfig routingConfig = new RoutingConfig(null, null, selectorType,
false);
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.REALTIME).setTableName("testTable")
+ .setRoutingConfig(routingConfig).build();
+ IdealState idealState = createIdealState(Map.of(
+ AR_SEGMENT0, List.of(Pair.of(AR_P0_RG0_SERVER_A, ONLINE),
Pair.of(AR_P1_RG1_SERVER_B, ONLINE)),
+ AR_SEGMENT1, List.of(Pair.of(AR_P0_RG0_SERVER_C, ONLINE),
Pair.of(AR_P1_RG1_SERVER_D, ONLINE)),
+ AR_SEGMENT2, List.of(Pair.of(AR_P1_RG1_SERVER_D, ONLINE),
Pair.of(AR_P0_RG0_SERVER_E, ONLINE))));
+ ExternalView externalView = createExternalView(Map.of(
+ AR_SEGMENT0, List.of(Pair.of(AR_P0_RG0_SERVER_A, ONLINE),
Pair.of(AR_P1_RG1_SERVER_B, ONLINE)),
+ AR_SEGMENT1, List.of(Pair.of(AR_P0_RG0_SERVER_C, ONLINE),
Pair.of(AR_P1_RG1_SERVER_D, ONLINE)),
+ AR_SEGMENT2, List.of(Pair.of(AR_P1_RG1_SERVER_D, ONLINE),
Pair.of(AR_P0_RG0_SERVER_E, ONLINE))));
+ ServerInstance serverA = mock(ServerInstance.class);
+ when(serverA.getPool()).thenReturn(0);
+ ServerInstance serverB = mock(ServerInstance.class);
+ when(serverB.getPool()).thenReturn(1);
+ ServerInstance serverC = mock(ServerInstance.class);
+ when(serverC.getPool()).thenReturn(0);
+ ServerInstance serverD = mock(ServerInstance.class);
+ when(serverD.getPool()).thenReturn(1);
+ ServerInstance serverE = mock(ServerInstance.class);
+ when(serverE.getPool()).thenReturn(0);
+ Map<String, ServerInstance> serverMap = Map.of(
+ AR_P0_RG0_SERVER_A, serverA,
+ AR_P1_RG1_SERVER_B, serverB,
+ AR_P0_RG0_SERVER_C, serverC,
+ AR_P1_RG1_SERVER_D, serverD,
+ AR_P0_RG0_SERVER_E, serverE);
+ return (ReplicaGroupInstanceSelector)
InstanceSelectorFactory.getInstanceSelector(tableConfig,
+ mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class),
hybridSelector,
+ brokerConfig,
+ Set.of(AR_P0_RG0_SERVER_A, AR_P1_RG1_SERVER_B, AR_P0_RG0_SERVER_C,
AR_P1_RG1_SERVER_D, AR_P0_RG0_SERVER_E),
+ serverMap, idealState, externalView, new HashSet<>(AR_SEGMENTS));
+ }
+
+ @Test
+ public void testReplicaGroupAdaptiveServerSelector() {
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ ReplicaGroupInstanceSelector instanceSelector =
+ buildArSelector(REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, hybridSelector);
+
+ assertTrue(instanceSelector instanceof ReplicaGroupInstanceSelector);
+ assertFalse(instanceSelector instanceof
StrictReplicaGroupInstanceSelector);
+ assertNotNull(instanceSelector._adaptiveServerSelector);
+ assertNotNull(instanceSelector._priorityPoolInstanceSelector);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(AR_SERVER_RANKS);
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(AR_SEGMENTS, 0, buildArSegmentStates(), null);
+
+ // AR prefers the better-ranked server when all candidates are ranked:
+ // segment0: B (rank 3) over A (rank 4)
+ // segment1: D (rank 1) over C (rank 2)
+ // segment2: E unranked → AR falls back to round-robin → index 0 = E
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ AR_SEGMENT0, AR_P1_RG1_SERVER_B,
+ AR_SEGMENT1, AR_P1_RG1_SERVER_D,
+ AR_SEGMENT2, AR_P0_RG0_SERVER_E));
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptiveDisabledByFeatureFlag() {
+ // When the feature flag is disabled, the factory nulls out the adaptive
selector so the selector falls back to the
+ // strict selector's non-adaptive, round-robin-by-replica-index path.
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ BrokerMetrics brokerMetrics = mock(BrokerMetrics.class);
+ PinotConfiguration brokerConfig = new PinotConfiguration(Map.of(
+
CommonConstants.Broker.AdaptiveServerSelector.CONFIG_OF_ENABLE_STRICT_REPLICA_GROUP,
"false"));
+ StrictReplicaGroupInstanceSelector instanceSelector =
+ buildStrictReplicaGroupArSelector(hybridSelector, brokerMetrics,
brokerConfig);
+
+ assertNull(instanceSelector._adaptiveServerSelector);
+ assertNull(instanceSelector._priorityPoolInstanceSelector);
+
+ BaseInstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), null);
+
+ verify(hybridSelector, never()).fetchServerRankingsWithScores(any());
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT1, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT2, STRICT_RG0_SERVER_B));
+
+ result = instanceSelector.select(STRICT_SEGMENTS, 1,
buildStrictReplicaGroupSegmentStates(), null);
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT1, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT2, STRICT_RG1_SERVER_D));
+ }
+
+ @Test
+ public void testStrictReplicaGroupFactoryEnablesAdaptiveRouting() {
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ ReplicaGroupInstanceSelector instanceSelector =
+ buildArSelector(STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE,
hybridSelector);
+
+ assertTrue(instanceSelector instanceof StrictReplicaGroupInstanceSelector);
+ assertNotNull(instanceSelector._adaptiveServerSelector);
+ assertNotNull(instanceSelector._priorityPoolInstanceSelector);
+ }
+
+ // --- Strict replica group adaptive routing tests ---
+
+ // Topology: 2 replica groups, 3 segments
+ // Replica group 0: server_a (segment0, segment1), server_b (segment2)
+ // Replica group 1: server_c (segment0, segment1), server_d (segment2)
+ // If a strict-fixture server name omits `P#`, its pool is FALLBACK_POOL_ID.
+ private static final String STRICT_RG0_SERVER_A = "strict_rg0_server_a";
+ private static final String STRICT_RG0_SERVER_B = "strict_rg0_server_b";
+ private static final String STRICT_RG1_SERVER_C = "strict_rg1_server_c";
+ private static final String STRICT_RG1_SERVER_D = "strict_rg1_server_d";
+ private static final String STRICT_SEGMENT0 = "seg0";
+ private static final String STRICT_SEGMENT1 = "seg1";
+ private static final String STRICT_SEGMENT2 = "seg2";
+ private static final List<String> STRICT_SEGMENTS =
Arrays.asList(STRICT_SEGMENT0, STRICT_SEGMENT1, STRICT_SEGMENT2);
+
+ private List<SegmentInstanceCandidate>
buildStrictReplicaGroupCandidates(String replicaGroup0Instance,
+ String replicaGroup1Instance) {
+ return Arrays.asList(
+ new SegmentInstanceCandidate(replicaGroup0Instance, true,
FALLBACK_POOL_ID, 0),
+ new SegmentInstanceCandidate(replicaGroup1Instance, true,
FALLBACK_POOL_ID, 1));
+ }
+
+ private SegmentStates buildStrictReplicaGroupSegmentStates() {
+ Map<String, List<SegmentInstanceCandidate>> candidatesMap = new
HashMap<>();
+ // All strict-fixture servers intentionally share FALLBACK_POOL_ID so
these tests fail if strict routing regresses
+ // back to grouping or filtering by pool instead of idealStateReplicaId.
+ candidatesMap.put(STRICT_SEGMENT0,
+ buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_A,
STRICT_RG1_SERVER_C));
+ candidatesMap.put(STRICT_SEGMENT1,
+ buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_A,
STRICT_RG1_SERVER_C));
+ candidatesMap.put(STRICT_SEGMENT2,
+ buildStrictReplicaGroupCandidates(STRICT_RG0_SERVER_B,
STRICT_RG1_SERVER_D));
+ return new SegmentStates(candidatesMap, new HashSet<>(STRICT_SEGMENTS),
null);
+ }
+
+ private StrictReplicaGroupInstanceSelector
buildStrictReplicaGroupArSelector(HybridSelector hybridSelector) {
+ return buildStrictReplicaGroupArSelector(hybridSelector,
mock(BrokerMetrics.class));
+ }
+
+ private StrictReplicaGroupInstanceSelector
buildStrictReplicaGroupArSelector(HybridSelector hybridSelector,
+ BrokerMetrics brokerMetrics) {
+ return buildStrictReplicaGroupArSelector(hybridSelector, brokerMetrics,
+ new PinotConfiguration(Map.of()));
+ }
+
+ private StrictReplicaGroupInstanceSelector
buildStrictReplicaGroupArSelector(HybridSelector hybridSelector,
+ BrokerMetrics brokerMetrics, PinotConfiguration brokerConfig) {
+ RoutingConfig routingConfig = new RoutingConfig(null, null,
STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false);
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.REALTIME).setTableName("testUpsertTable")
+ .setRoutingConfig(routingConfig).build();
+ // Ideal state: mirrors the topology above
+ IdealState idealState = createIdealState(Map.of(
+ STRICT_SEGMENT0, List.of(
+ Pair.of(STRICT_RG0_SERVER_A, ONLINE),
+ Pair.of(STRICT_RG1_SERVER_C, ONLINE)),
+ STRICT_SEGMENT1, List.of(
+ Pair.of(STRICT_RG0_SERVER_A, ONLINE),
+ Pair.of(STRICT_RG1_SERVER_C, ONLINE)),
+ STRICT_SEGMENT2, List.of(
+ Pair.of(STRICT_RG0_SERVER_B, ONLINE),
+ Pair.of(STRICT_RG1_SERVER_D, ONLINE))));
+ ExternalView externalView = createExternalView(Map.of(
+ STRICT_SEGMENT0, List.of(
+ Pair.of(STRICT_RG0_SERVER_A, ONLINE),
+ Pair.of(STRICT_RG1_SERVER_C, ONLINE)),
+ STRICT_SEGMENT1, List.of(
+ Pair.of(STRICT_RG0_SERVER_A, ONLINE),
+ Pair.of(STRICT_RG1_SERVER_C, ONLINE)),
+ STRICT_SEGMENT2, List.of(
+ Pair.of(STRICT_RG0_SERVER_B, ONLINE),
+ Pair.of(STRICT_RG1_SERVER_D, ONLINE))));
+ // All servers use the fallback pool to prove strict routing keys off
idealStateReplicaId rather than pool.
+ ServerInstance serverA = mock(ServerInstance.class);
+ when(serverA.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance serverB = mock(ServerInstance.class);
+ when(serverB.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance serverC = mock(ServerInstance.class);
+ when(serverC.getPool()).thenReturn(FALLBACK_POOL_ID);
+ ServerInstance serverD = mock(ServerInstance.class);
+ when(serverD.getPool()).thenReturn(FALLBACK_POOL_ID);
+ Map<String, ServerInstance> serverMap = Map.of(
+ STRICT_RG0_SERVER_A, serverA, STRICT_RG0_SERVER_B, serverB,
+ STRICT_RG1_SERVER_C, serverC, STRICT_RG1_SERVER_D, serverD);
+ return (StrictReplicaGroupInstanceSelector)
InstanceSelectorFactory.getInstanceSelector(tableConfig,
+ mock(ZkHelixPropertyStore.class), brokerMetrics, hybridSelector,
+ brokerConfig,
+ Set.of(STRICT_RG0_SERVER_A, STRICT_RG0_SERVER_B,
+ STRICT_RG1_SERVER_C, STRICT_RG1_SERVER_D),
+ serverMap, idealState, externalView, new HashSet<>(STRICT_SEGMENTS));
+ }
+
+ @Test
+ public void
testStrictReplicaGroupAdaptivePicksBestReplicaGroupByWorstCaseServer() {
+ // Replica group 0 servers: server_a (rank 2), server_b (rank 3) → worst =
3
+ // Replica group 1 servers: server_c (rank 0), server_d (rank 1) → worst =
1
+ // Replica group 1 is better (lower worst-case rank), so all segments
should route there
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ BrokerMetrics brokerMetrics = mock(BrokerMetrics.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector,
+ brokerMetrics);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ new ImmutablePair<>(STRICT_RG1_SERVER_C, 1.0), // rank 0 (best)
+ new ImmutablePair<>(STRICT_RG1_SERVER_D, 2.0), // rank 1
+ new ImmutablePair<>(STRICT_RG0_SERVER_A, 3.0), // rank 2
+ new ImmutablePair<>(STRICT_RG0_SERVER_B, 4.0) // rank 3 (worst)
+ ));
+
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), null);
+
+ // All segments routed to replica group 1
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT1, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT2, STRICT_RG1_SERVER_D));
+ // Metrics are still reported per actual pool, even though routing is
chosen by replica group.
+ verify(brokerMetrics).addMeteredValue(eq(BrokerMeter.POOL_SEG_QUERIES),
eq(3L),
+ eq(BrokerMetrics.getTagForPreferredPool(null)),
eq(String.valueOf(FALLBACK_POOL_ID)));
+ }
+
+ @Test
+ public void
testStrictReplicaGroupAdaptivePicksGroup0WhenItHasBetterWorstCase() {
+ // Replica group 0 servers: server_a (rank 0), server_b (rank 1) → worst =
1
+ // Replica group 1 servers: server_c (rank 2), server_d (rank 3) → worst =
3
+ // Replica group 0 is better
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ new ImmutablePair<>(STRICT_RG0_SERVER_A, 1.0),
+ new ImmutablePair<>(STRICT_RG0_SERVER_B, 2.0),
+ new ImmutablePair<>(STRICT_RG1_SERVER_C, 3.0),
+ new ImmutablePair<>(STRICT_RG1_SERVER_D, 4.0)
+ ));
+
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), null);
+
+ // All segments routed to replica group 0
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT1, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT2, STRICT_RG0_SERVER_B));
+ }
+
+ @Test
+ public void
testStrictReplicaGroupAdaptiveEmptyRankingsFallBackToRoundRobin() {
+ // When adaptive ranking returns no servers, fall back to round-robin by
replica-group index.
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(List.of());
+
+ // requestId=0 → picks group at index 0 (replica group 0)
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0,
buildStrictReplicaGroupSegmentStates(), null);
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT1, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT2, STRICT_RG0_SERVER_B));
+
+ // requestId=1 → picks group at index 1 (replica group 1)
+ result = instanceSelector.select(STRICT_SEGMENTS, 1,
buildStrictReplicaGroupSegmentStates(), null);
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT1, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT2, STRICT_RG1_SERVER_D));
+ }
+
+ @Test
+ public void
testStrictReplicaGroupAdaptiveConsistencyAllSegmentsSameReplicaGroup() {
+ // When adaptive scores favor replica group 1, all segments must route to
that same replica group.
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+ StrictReplicaGroupInstanceSelector instanceSelector =
buildStrictReplicaGroupArSelector(hybridSelector);
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ new ImmutablePair<>(STRICT_RG1_SERVER_C, 1.0),
+ new ImmutablePair<>(STRICT_RG1_SERVER_D, 2.0),
+ new ImmutablePair<>(STRICT_RG0_SERVER_A, 3.0),
+ new ImmutablePair<>(STRICT_RG0_SERVER_B, 4.0)
+ ));
+
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 42,
buildStrictReplicaGroupSegmentStates(), null);
+
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT1, STRICT_RG1_SERVER_C,
+ STRICT_SEGMENT2, STRICT_RG1_SERVER_D));
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptivePrefersFullCoverageGroup() {
+ // Topology: replica group 0 has all 3 segments, replica group 1 only has
segment0 and segment1.
+ // Even though replica group 1 has better adaptive ranks, replica group 0
should be chosen (full coverage).
+ // Replica group 0: server_a (seg0, seg1), server_b (seg2) — full coverage
+ // Replica group 1: server_c (seg0, seg1) — missing seg2, partial coverage
+ // Pools intentionally do not match replica groups: server_a and server_c
share pool 0, while server_b is pool 1.
+ String serverC = "server_c";
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+
+ RoutingConfig routingConfig = new RoutingConfig(null, null,
STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false);
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.REALTIME).setTableName("testUpsertTable")
+ .setRoutingConfig(routingConfig).build();
+ // Replica group 1 only has seg0 and seg1 (no seg2)
+ IdealState idealState = createIdealState(Map.of(
+ STRICT_SEGMENT0, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE),
Pair.of(serverC, ONLINE)),
+ STRICT_SEGMENT1, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE),
Pair.of(serverC, ONLINE)),
+ STRICT_SEGMENT2, List.of(Pair.of(STRICT_RG0_SERVER_B, ONLINE))));
+ ExternalView externalView = createExternalView(Map.of(
+ STRICT_SEGMENT0, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE),
Pair.of(serverC, ONLINE)),
+ STRICT_SEGMENT1, List.of(Pair.of(STRICT_RG0_SERVER_A, ONLINE),
Pair.of(serverC, ONLINE)),
+ STRICT_SEGMENT2, List.of(Pair.of(STRICT_RG0_SERVER_B, ONLINE))));
+ ServerInstance serverA = mock(ServerInstance.class);
+ when(serverA.getPool()).thenReturn(0);
+ ServerInstance serverB = mock(ServerInstance.class);
+ when(serverB.getPool()).thenReturn(1);
+ ServerInstance serverCInstance = mock(ServerInstance.class);
+ when(serverCInstance.getPool()).thenReturn(0);
+ Map<String, ServerInstance> serverMap = Map.of(
+ STRICT_RG0_SERVER_A, serverA, STRICT_RG0_SERVER_B, serverB, serverC,
serverCInstance);
+ StrictReplicaGroupInstanceSelector instanceSelector =
+ (StrictReplicaGroupInstanceSelector)
InstanceSelectorFactory.getInstanceSelector(tableConfig,
+ mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class),
hybridSelector,
+ new PinotConfiguration(Map.of()),
+ Set.of(STRICT_RG0_SERVER_A, STRICT_RG0_SERVER_B, serverC),
+ serverMap, idealState, externalView, new
HashSet<>(STRICT_SEGMENTS));
+
+ // Replica group 1 (server_c) has rank 0 — best rank, but only covers 2 of
3 segments
+ // Replica group 0 (server_a rank 1, server_b rank 2) — worse ranks, but
full coverage
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ new ImmutablePair<>(serverC, 1.0), // rank 0 (best)
+ new ImmutablePair<>(STRICT_RG0_SERVER_A, 3.0), // rank 1
+ new ImmutablePair<>(STRICT_RG0_SERVER_B, 4.0) // rank 2
+ ));
+
+ // Build segment states where replica group 1 is missing seg2
+ Map<String, List<SegmentInstanceCandidate>> candidatesMap = new
HashMap<>();
+ candidatesMap.put(STRICT_SEGMENT0, Arrays.asList(
+ new SegmentInstanceCandidate(STRICT_RG0_SERVER_A, true, 0, 0),
+ new SegmentInstanceCandidate(serverC, true, 0, 1)));
+ candidatesMap.put(STRICT_SEGMENT1, Arrays.asList(
+ new SegmentInstanceCandidate(STRICT_RG0_SERVER_A, true, 0, 0),
+ new SegmentInstanceCandidate(serverC, true, 0, 1)));
+ candidatesMap.put(STRICT_SEGMENT2, List.of(
+ new SegmentInstanceCandidate(STRICT_RG0_SERVER_B, true, 1, 0)));
+ SegmentStates segmentStates = new SegmentStates(candidatesMap, new
HashSet<>(STRICT_SEGMENTS), null);
+
+ InstanceSelector.InstanceMapping result =
+ instanceSelector.select(STRICT_SEGMENTS, 0, segmentStates, null);
+
+ // Replica group 0 should be chosen (full coverage) even though replica
group 1 has better ranks
+ assertEquals(result.segmentToInstanceMap(), Map.of(
+ STRICT_SEGMENT0, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT1, STRICT_RG0_SERVER_A,
+ STRICT_SEGMENT2, STRICT_RG0_SERVER_B));
+ }
+
+ @Test
+ public void testStrictReplicaGroupAdaptiveNewSegmentNotFalselyUnavailable() {
+ // Topology: server_a (replica group 0) hosts seg0 only. seg1 has no
online instance (simulating a
+ // new/unassigned segment with null candidates in segmentStates). Replica
group 0 is chosen (only group).
+ // seg1 should NOT appear in unavailable — it is a new segment, not a
dropped one.
+ String serverA = "new_server_a";
+ String seg0 = "new_seg0";
+ String seg1 = "new_seg1";
+ List<String> segments = Arrays.asList(seg0, seg1);
+ HybridSelector hybridSelector = mock(HybridSelector.class);
+
+ RoutingConfig routingConfig = new RoutingConfig(null, null,
STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false);
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.REALTIME).setTableName("testNewSegTable")
+ .setRoutingConfig(routingConfig).build();
+ // seg1 intentionally absent from idealState/externalView → null
candidates (new segment)
+ IdealState idealState = createIdealState(Map.of(
+ seg0, List.of(Pair.of(serverA, ONLINE))));
+ ExternalView externalView = createExternalView(Map.of(
+ seg0, List.of(Pair.of(serverA, ONLINE))));
+ ServerInstance serverAInstance = mock(ServerInstance.class);
+ when(serverAInstance.getPool()).thenReturn(FALLBACK_POOL_ID);
+ StrictReplicaGroupInstanceSelector instanceSelector =
+ (StrictReplicaGroupInstanceSelector)
InstanceSelectorFactory.getInstanceSelector(tableConfig,
+ mock(ZkHelixPropertyStore.class), mock(BrokerMetrics.class),
hybridSelector,
+ new PinotConfiguration(Map.of()),
+ Set.of(serverA),
+ Map.of(serverA, serverAInstance),
+ idealState, externalView, new HashSet<>(List.of(seg0)));
+
+
when(hybridSelector.fetchServerRankingsWithScores(any())).thenReturn(Arrays.asList(
+ new ImmutablePair<>(serverA, 1.0)));
+
+ BrokerRequest brokerRequest = mock(BrokerRequest.class);
+ PinotQuery pinotQuery = mock(PinotQuery.class);
+ when(brokerRequest.getPinotQuery()).thenReturn(pinotQuery);
+ when(pinotQuery.getQueryOptions()).thenReturn(null);
+
+ InstanceSelector.SelectionResult result =
instanceSelector.select(brokerRequest, segments, 0);
+ assertEquals(result.getSegmentToInstanceMap(), Map.of(seg0, serverA));
+ // seg1 has null candidates (not yet in the selector's state) → must NOT
appear as unavailable
+ assertTrue(result.getUnavailableSegments().isEmpty());
+ }
+
+ @Test
+ public void
testStrictReplicaGroupAdaptiveDroppedSegmentReportedAsUnavailable() {
Review Comment:
This test locks in the availability change as intended behavior — worth
confirming that's the deliberate decision: `seg0` has a live, online
`server_a`, and the assertion requires it to come back as unavailable, i.e. a
`BROKER_SEGMENT_UNAVAILABLE` error on the response. Pre-PR all three segments
would have been served.
If the per-instance-set idea in my comment on
`StrictReplicaGroupInstanceSelector` pans out, this expectation flips, so it's
worth settling that first.
The move into this file is otherwise clean — I diffed the test lists and all
four relocated tests are accounted for. Gaps in the new strict-path suite that
stood out: nothing with `orderedPreferredPools`, nothing with `useFixedReplica:
true`, no test asserting that traffic actually spreads across replica groups
over a range of `requestId`s (the steady-state graphs are manual), and no case
for a new segment (offline candidate) that's missing from the chosen group.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
