Copilot commented on code in PR #19144:
URL: https://github.com/apache/pinot/pull/19144#discussion_r3715614720
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -918,9 +933,17 @@ private static void
transferToServerInstanceLogicalSegmentsMap(String physicalTa
for (Map.Entry<ServerInstance, SegmentsToQuery> serverEntry :
segmentsMap.entrySet()) {
Map<String, List<String>> tableNameToSegmentsMap =
serverInstanceToLogicalSegmentsMap.computeIfAbsent(serverEntry.getKey(), k ->
new HashMap<>());
- // TODO: support optional segments for multi-stage engine.
+ // Merge required and optional segments for multi-stage engine.
+ List<String> segments = serverEntry.getValue().getSegments();
+ List<String> optionalSegments =
serverEntry.getValue().getOptionalSegments();
+ if (optionalSegments != null && !optionalSegments.isEmpty()) {
+ List<String> combinedSegments = new ArrayList<>(segments.size() +
optionalSegments.size());
+ combinedSegments.addAll(segments);
+ combinedSegments.addAll(optionalSegments);
+ segments = Collections.unmodifiableList(combinedSegments);
+ }
Review Comment:
The required+optional segment merge logic is duplicated in multiple places.
To reduce drift and future fixes needing to be applied twice, extract this into
a small helper (e.g., `mergeSegments(SegmentsToQuery)` or similar) and reuse it
at both call sites.
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -737,7 +746,13 @@ private void
setSegmentsForReplicatedLeafFragment(DispatchablePlanMetadata metad
}
}
- // TODO: Support unavailable segments and optional segments for replicated
leaf stage
+ // Attach unavailable segments to metadata for replicated leaf stage
+ Map<String, RoutingTable> routingTableMap = getRoutingTable(tableName,
context.getRequestId());
+ for (Map.Entry<String, RoutingTable> entry : routingTableMap.entrySet()) {
+ if (!entry.getValue().getUnavailableSegments().isEmpty()) {
+ metadata.addUnavailableSegments(tableName,
entry.getValue().getUnavailableSegments());
+ }
+ }
Review Comment:
`routingTableMap` is iterated by entry, but the key (`entry.getKey()`,
likely table type or similar discriminator) is ignored and `tableName` is
always used when calling `metadata.addUnavailableSegments(...)`. If
`routingTableMap` contains multiple routing tables (e.g., OFFLINE/REALTIME),
this can merge or overwrite unavailable-segment info under a single key and
lose table-type separation. Consider passing the correct discriminator (e.g.,
`entry.getKey()` or a derived physical table name) into
`addUnavailableSegments`, or explicitly merging while preserving type.
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -1084,10 +1107,12 @@ private void assignOnePartitionPerWorker(String
tableName, long requestId, Parti
continue;
}
PartitionInfo partitionInfo = partitionInfoMap[i];
- // TODO: Currently we don't support the case when a partition doesn't
contain any segment. The reason is that
- // the leaf stage won't be able to directly return empty response.
- Preconditions.checkState(partitionInfo != null, "Failed to find any
segment for table: %s, partition: %s",
- tableName, i);
+ // Skip partitions that don't contain any segment. The leaf stage can
handle empty partitions by
+ // returning an empty response, which is equivalent to a pruned
partition.
+ if (partitionInfo == null) {
+ LOGGER.warn("No segment found for table: {}, partition: {}, skipping",
tableName, i);
+ continue;
+ }
Review Comment:
This warning is inside a per-partition loop and can become very noisy for
large `numPartitions` (or in partially-loaded clusters), potentially flooding
logs. Consider lowering this to `INFO`/`DEBUG`, or logging an aggregated
summary (e.g., count of empty partitions) after the loop.
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -568,8 +569,16 @@ private void
assignWorkersToNonPartitionedLeafFragment(PlanFragment fragment, Di
for (Map.Entry<ServerInstance, SegmentsToQuery> serverEntry :
segmentsMap.entrySet()) {
Map<String, List<String>> tableTypeToSegmentListMap =
serverInstanceToSegmentsMap.computeIfAbsent(serverEntry.getKey(),
k -> new HashMap<>());
- // TODO: support optional segments for multi-stage engine.
- Preconditions.checkState(tableTypeToSegmentListMap.put(tableType,
serverEntry.getValue().getSegments()) == null,
+ // Merge required and optional segments for multi-stage engine.
+ List<String> segments = serverEntry.getValue().getSegments();
+ List<String> optionalSegments =
serverEntry.getValue().getOptionalSegments();
+ if (optionalSegments != null && !optionalSegments.isEmpty()) {
+ List<String> combinedSegments = new ArrayList<>(segments.size() +
optionalSegments.size());
+ combinedSegments.addAll(segments);
+ combinedSegments.addAll(optionalSegments);
+ segments = Collections.unmodifiableList(combinedSegments);
+ }
Review Comment:
The required+optional segment merge logic is duplicated in multiple places.
To reduce drift and future fixes needing to be applied twice, extract this into
a small helper (e.g., `mergeSegments(SegmentsToQuery)` or similar) and reuse it
at both call sites.
--
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]
