xiangfu0 commented on code in PR #19537:
URL: https://github.com/apache/pinot/pull/19537#discussion_r4221916889
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -1169,9 +1185,17 @@ private void
setSegmentsForReplicatedLeafFragment(DispatchablePlanMetadata metad
}
}
- // TODO: Support unavailable segments and optional segments for replicated
leaf stage
metadata.setReplicatedSegments(segmentsMap);
filterReplicatedLeafStageSegments(context, metadata);
+
+ Map<String, RoutingTable> routingTableMap =
Review Comment:
**Second full routing computation per replicated-leaf query.** This
`getRoutingTable` call runs complete instance selection + pruning solely to
read `getUnavailableSegments()`, on top of the `getSegments()` call above —
every replicated-table query pays it even when nothing is unavailable. It also
snapshots routing at a different time than `getSegments()`, so the unavailable
list can disagree with the segment list under concurrent state changes. A
routing-manager API that exposes only unavailable segments (or reusing a single
routing acquisition for both) would avoid the duplicated work.
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -2075,4 +2108,35 @@ private static Map<String, List<String>>
getSegmentsMap(@Nullable List<String> o
return Map.of(TableType.REALTIME.name(), realtimeSegments);
}
}
+
+ private void
assignWorkersForNonPartitionedLeafSegmentsWhenNoServersHaveSegments(DispatchablePlanMetadata
metadata,
+ String tableName, Map<String, RoutingTable> routingTableMap,
DispatchablePlanContext context,
+ @Nullable TimeBoundaryInfo timeBoundaryInfo) {
+ Set<String> serversRoutingToTable =
_routingManager.getServingInstances(tableName);
+ if (CollectionUtils.isEmpty(serversRoutingToTable)) {
+ LOGGER.error("[RequestId: {}] No routable server for empty routing on
table: {}", context.getRequestId(),
+ tableName);
+ throw new IllegalStateException("No routable server for empty routing on
table: " + tableName);
+ }
+
+ Map<String, ServerInstance> serverInstanceMap =
_routingManager.getRoutableServerInstanceMap();
+ ServerInstance serverInstance =
serverInstanceMap.get(serversRoutingToTable.iterator().next());
+ Map<String, List<String>> emptySegmentsMap = new HashMap<>();
+ for (String tableType : routingTableMap.keySet()) {
+ emptySegmentsMap.put(tableType, List.of());
+ }
+ if (emptySegmentsMap.isEmpty()) {
Review Comment:
**Dead branch + immutable collections handed to an extension point.**
`routingTableMap.isEmpty()` already throws `BROKER_RESOURCE_MISSING` before
this method is reachable, and the hybrid OFFLINE removal only happens when
`size() > 1`, so the loop above always inserts at least one key — this whole
`if (emptySegmentsMap.isEmpty())` block (and the `timeBoundaryInfo` parameter
that only it reads) can never execute. Separately, `updateContextForLeafStage`
then runs `filterLeafStageSegments`, whose contract says the handed-over
assignment "is nevertheless kept safe to edit in place"; `Map.of(0, ...)` /
`List.of()` would throw `UnsupportedOperationException` on a subclass that
edits in place.
##########
pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java:
##########
@@ -3115,4 +3115,133 @@ public boolean isTableDisabled(String
tableNameWithType) {
return false;
}
}
+
+ @Test
+ public void
testNonPartitionedLeafPropagatesOptionalSegmentsToWorkerMetadata() {
+ Schema tableWithOptionalSegmentsSchema =
getSchemaBuilder("tableWithOptionalSegments").build();
+
+ // Create server instance
+ ServerInstance server = getServerInstance("localhost", 1);
+ Map<String, ServerInstance> serverInstanceMap = new HashMap<>();
+ serverInstanceMap.put(server.getInstanceId(), server);
+
+ // Create a routing table with both optional and normal segments
+ RoutingTable routingTable =
+ new RoutingTable(Map.of(server, new SegmentsToQuery(List.of("seg"),
List.of("optSeg"))),
+ List.of("unavailableSeg"), 0);
+
+ // Create mock routing manager
+ CapturingRoutingManager routingManager = new
CapturingRoutingManager(serverInstanceMap,
+ Map.of("tableWithOptionalSegments_OFFLINE", routingTable));
+
+ // Create mock table cache
+ Map<String, String> tableNameMap = new HashMap<>();
+ tableNameMap.put("tableWithOptionalSegments_OFFLINE",
"tableWithOptionalSegments_OFFLINE");
+ tableNameMap.put("tableWithOptionalSegments", "tableWithOptionalSegments");
+
+ TableCache tableCache = mock(TableCache.class);
+ when(tableCache.getTableNameMap()).thenReturn(tableNameMap);
+ when(tableCache.getActualTableName(anyString())).thenAnswer(inv ->
tableNameMap.get(inv.getArgument(0)));
+
when(tableCache.getSchema(anyString())).thenReturn(tableWithOptionalSegmentsSchema);
+
when(tableCache.getTableConfig("tableWithOptionalSegments_OFFLINE")).thenReturn(mock(TableConfig.class));
+
+ WorkerManager workerManager = new WorkerManager("Broker_localhost",
"localhost", 3, routingManager);
+ QueryEnvironment queryEnvironment = new
QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
+ workerManager);
+
+ String query = "SELECT * FROM tableWithOptionalSegments LIMIT 10";
+
+ try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(query)) {
+ DispatchableSubPlan dispatchableSubPlan =
compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan);
+ assertEquals(List.of("unavailableSeg"),
+
dispatchableSubPlan.getTableToUnavailableSegmentsMap().get("tableWithOptionalSegments"));
+ // This stage should handle basic data reads
+ DispatchablePlanFragment dataReadPlanFragment =
dispatchableSubPlan.getQueryStages().getLast();
+ List<WorkerMetadata> workerMetadataList =
dataReadPlanFragment.getWorkerMetadataList();
+
assertNotNull(workerMetadataList.getFirst().getOptionalTableSegmentsMap());
+
assertEquals(workerMetadataList.getFirst().getOptionalTableSegmentsMap().size(),
1);
+
assertTrue(workerMetadataList.getFirst().getOptionalTableSegmentsMap().containsKey("OFFLINE"));
+ assertEquals(List.of("optSeg"),
workerMetadataList.getFirst().getOptionalTableSegmentsMap().get("OFFLINE"));
+ assertNotNull(workerMetadataList.getFirst().getTableSegmentsMap());
+ assertEquals(workerMetadataList.getFirst().getTableSegmentsMap().size(),
1);
+
assertTrue(workerMetadataList.getFirst().getTableSegmentsMap().containsKey("OFFLINE"));
+ assertEquals(List.of("seg"),
workerMetadataList.getFirst().getTableSegmentsMap().get("OFFLINE"));
+ }
+ }
+
+ @Test
+ public void
testWorkerAssignmentForNonPartitionedLeadSegmentsWhenNoServersHaveSegments() {
+ Schema tableWithNoServersHavingSegmentsSchema =
getSchemaBuilder("tableWithNoServersHavingSegments").build();
+ Schema anotherTableSchema = getSchemaBuilder("anotherTable").build();
+
+ // Create server instance to server the table have no segments
+ ServerInstance server = getServerInstance("localhost", 1);
+ Map<String, ServerInstance> serverInstanceMap = new HashMap<>();
+ serverInstanceMap.put(server.getInstanceId(), server);
+
+ // Create server instance to server the other normal table
+ ServerInstance server2 = getServerInstance("localhost", 2);
+ serverInstanceMap.put(server2.getInstanceId(), server);
Review Comment:
Copy-paste bug: this registers `server` (port 1) under `server2`'s instance
id, so both ids resolve to the same `ServerInstance` and the test cannot tell
which server the fallback actually picked — it would keep passing even if the
implementation chose the wrong instance id.
```suggestion
serverInstanceMap.put(server2.getInstanceId(), server2);
```
--
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]
