xiangfu0 commented on code in PR #19537:
URL: https://github.com/apache/pinot/pull/19537#discussion_r4240980193
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -1343,22 +1375,38 @@ private static void
assignTableSegmentsToWorkers(LogicalTableRouteInfo logicalTa
workerIdToServerInstanceMap.put(workerId, server);
workerIdToLogicalTableSegmentsMap.put(workerId, segmentsMap);
+
+ if (serverInstanceToLogicalOptionalSegmentsMap.containsKey(server)) {
+ Map<String, List<String>> optionalSegmentsMap =
serverInstanceToLogicalOptionalSegmentsMap.get(server);
+ if (MapUtils.isNotEmpty(optionalSegmentsMap)) {
+ workerIdTOptionalLogicalTableSegmentsMap.put(workerId,
+ optionalSegmentsMap);
+ }
+ }
}
metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap);
metadata.setWorkerIdToTableSegmentsMap(workerIdToLogicalTableSegmentsMap);
+ if (MapUtils.isNotEmpty(workerIdTOptionalLogicalTableSegmentsMap)) {
+
metadata.setWorkerIdToOptionalSegmentsMap(workerIdTOptionalLogicalTableSegmentsMap);
Review Comment:
**[CRITICAL] Logical-table optional segments are stored into the wrong
metadata field and silently dropped.** This stores the logical-path map via
`setWorkerIdToOptionalSegmentsMap`, but
`DispatchablePlanContext.constructDispatchablePlanFragmentMap` reads the
logical branch from `getWorkerIdToOptionalTableSegmentsMap()` — and
`setWorkerIdToOptionalTableSegmentsMap` has zero callers in the entire PR.
Since the logical path never sets `workerIdToSegmentsMap`, the physical branch
that would read `getWorkerIdToOptionalSegmentsMap()` is also skipped. Net
effect: for logical tables, optional segments never reach `WorkerMetadata` at
all (even broker-side), and
`DispatchablePlanMetadata._workerIdToOptionalTableSegmentsMap` plus its
getter/setter are dead code. This should call
`setWorkerIdToOptionalTableSegmentsMap` (or better: collapse the two
near-identical fields into one, mirroring the single
`WorkerMetadata._optionalTableSegmentsMap`).
##########
pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java:
##########
@@ -3115,4 +3121,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() {
Review Comment:
**[MAJOR] This test exercises behavior the PR does not implement — it cannot
pass against this PR's production code.**
`assignWorkersToNonPartitionedLeafFragment` has no fallback when
`serverInstanceToSegmentsMap` is empty: with `RoutingTable(Map.of(),
List.of("seg","seg2"), 0)` it assigns `numWorkers = 0`, so
`dataReadPlanFragment.getServerInstanceToWorkerIdMap().get(queryServerInstance)`
returns null and `workerIds.size()` throws NPE (or planning fails earlier on
an empty leaf stage). Either the production fallback (assigning a worker from
`getServingInstances` when no server has segments) is missing from this PR, or
this test belongs to a different change. Also note the typo in the name:
"NonPartitionedLead" → "NonPartitionedLeaf".
##########
pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java:
##########
@@ -2941,6 +2943,10 @@ public TablePartitionReplicatedServersInfo
getTablePartitionReplicatedServersInf
@Override
public Set<String> getServingInstances(String tableNameWithType) {
+ if(tableNameWithType!=null){
Review Comment:
**[MAJOR] These added lines in the
`CapturingRoutingManager.getServingInstances` mock are a
dereference-and-discard landmine.**
`_routingTableByName.get(tableNameWithType)` returns null for any table not in
the map — e.g. `getCandidateServersPerTables` calls `getServingInstances` with
both `_OFFLINE` and `_REALTIME` names, so an OFFLINE-only test table NPEs on
the REALTIME lookup — and
`rt.getServerInstanceToSegmentsMap().keySet().iterator().next()` throws
`NoSuchElementException` for an empty routing table (exactly the `Map.of()`
routing table the new test below registers). The result of `.next()` is
discarded, so the statements have no purpose; they should be deleted. The block
also fails formatting/checkstyle (`if(tableNameWithType!=null){`, 5-space
indent) — per CLAUDE.md, run `./mvnw spotless:apply` and `./mvnw
checkstyle:check`.
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -1169,9 +1187,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 =
+ getRoutingTable(tableName, context.getRequestId(),
context.getPlannerContext().getOptions());
Review Comment:
**[MAJOR] Full extra routing-table computation per query just to read
unavailable segments, from an inconsistent snapshot.** `getRoutingTable`
compiles a `SELECT *` broker request per table type and computes the complete
server-to-segments assignment, all of which is discarded except
`getUnavailableSegments()` — a nontrivial cost added to every replicated-leaf
query plan. Worse, the replicated segment list above comes from `getSegments()`
(instance-selector state read at a different time), so the two snapshots can
disagree under segment movement. Also, for a hybrid table where the time
boundary was not acquired, `segmentsMap.remove(TableType.OFFLINE.name())` drops
the OFFLINE side above, yet this loop still reports OFFLINE unavailable
segments for a side the query will not scan. Consider exposing unavailable
segments from the same instance-selector lookup `getSegments()` uses rather
than a second routing pass.
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -31,6 +31,7 @@
import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
Review Comment:
**[MAJOR][conventions] Unused import — fails checkstyle.**
`java.util.Optional` is not referenced anywhere in this file;
`config/checkstyle.xml` enables `UnusedImports` (line 209), so `./mvnw
checkstyle:check` fails on this PR. Same problem in `WorkerManagerTest.java`:
the added imports `org.apache.pinot.spi.utils.builder.TableNameBuilder` (line
54) and `org.apache.pinot.sql.parsers.parser.TableNameExtractor` (line 57) are
unused. CLAUDE.md: "Run `./mvnw checkstyle:check` to validate style."
--
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]
