Copilot commented on code in PR #18947: URL: https://github.com/apache/pinot/pull/18947#discussion_r3669268494
########## 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"); Review Comment: Using raw `Collections.EMPTY_MAP` forces the test to suppress unchecked warnings. Prefer `Map.of()` for a typed empty map (also aligns with Pinot's general preference for `Map.of()` literals). ########## 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); + continue; + } Review Comment: Segments that have candidates but no candidate for the chosen replica-group are always added to `unavailableSegments`. This can accidentally treat *new segments* as unavailable, even though `BaseInstanceSelector` explicitly avoids counting new segments as unavailable (new segments are allowed to be partially available). Consider skipping these segments (same as the `candidates == null` case) when they are new, and only reporting old segments as unavailable. ########## 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: `Math.abs(requestId)` can overflow for `Integer.MIN_VALUE`, which would make the modulo negative and can throw `IndexOutOfBoundsException`. Also, `replicaGroupToQueryServers.keySet()` iteration order is insertion-dependent, so the round-robin fallback may not be stable across runs. Use `Math.floorMod` and a deterministic ordering (e.g. sort group ids) for the empty-ranking fallback. -- 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]
