gortiz commented on code in PR #19166:
URL: https://github.com/apache/pinot/pull/19166#discussion_r3804340716
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -1069,63 +1300,100 @@ private static long
countPrunedSegments(PartitionInfo[] partitionInfoMap, Set<In
return numPrunedSegments;
}
- /// Pick one worker per partition for partitioned leaf stage. When
`partitionsToKeep` is non-null (broker
- /// pruning is active), partitions absent from the set are pruned by the
query filter and skipped.
+ /// Pick one worker per partition for partitioned leaf stage. There is one
partition per class here, so a class id is
+ /// a partition id; which of them get a worker is decided by
[#selectPartitionsToAssign].
private void assignOnePartitionPerWorker(String tableName, long requestId,
PartitionInfo[] partitionInfoMap,
- @Nullable Set<Integer> partitionsToKeep, Map<String, ServerInstance>
enabledServerInstanceMap,
+ @Nullable int[] partitionClassIds, @Nullable Set<Integer>
partitionsToKeep, @Nullable PaddingInfo paddingInfo,
+ Map<String, ServerInstance> enabledServerInstanceMap,
Map<Integer, QueryServerInstance> workerIdToServerInstanceMap,
Map<Integer, Map<String, List<String>>> workerIdToSegmentsMap) {
- int numPartitions = partitionInfoMap.length;
- int workerId = 0;
- for (int i = 0; i < numPartitions; i++) {
- // Skip partitions pruned by the broker filter. Empty partitions are
never in partitionsToKeep, so under pruning
- // they are skipped here too; the precondition below only fires when
pruning is inactive (partitionsToKeep null).
- if (partitionsToKeep != null && !partitionsToKeep.contains(i)) {
+ int[] partitionIds = selectPartitionsToAssign(partitionInfoMap.length,
partitionClassIds, partitionsToKeep);
+ for (int workerId = 0; workerId < partitionIds.length; workerId++) {
+ int partitionId = partitionIds[workerId];
+ PartitionInfo partitionInfo = partitionInfoMap[partitionId];
+ if (partitionInfo == null) {
+ // Pad a class the colocated group keeps but this table has no data
in, see assignPaddedWorker.
+ // TODO: Currently we don't support the case when a partition doesn't
contain any segment outside of a colocated
+ // group, where there is nothing to keep the worker ids aligned
with. The reason is that the leaf stage
+ // won't be able to directly return empty response.
+ Preconditions.checkState(paddingInfo != null &&
paddingInfo._classCandidates.containsKey(partitionId),
+ "Failed to find any segment for table: %s, partition: %s",
tableName, partitionId);
+ assignPaddedWorker(tableName, requestId, partitionId, paddingInfo,
enabledServerInstanceMap, workerId,
+ workerIdToServerInstanceMap, workerIdToSegmentsMap);
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);
// NOTE: Pick worker based on the request id plus the partition id (not
a running counter) so that the same worker
// is picked across different table scans when the segments for
the same partition are colocated, and so
- // that skipping pruned partitions does not shift the server
assignment of the surviving ones.
+ // that skipping pruned or empty partitions does not shift the
server assignment of the surviving ones.
ServerInstance serverInstance =
- pickEnabledServer(partitionInfo._fullyReplicatedServers,
enabledServerInstanceMap, requestId + i);
+ pickEnabledServer(partitionInfo._fullyReplicatedServers,
enabledServerInstanceMap, requestId + partitionId);
Preconditions.checkState(serverInstance != null,
- "Failed to find enabled fully replicated server for table: %s,
partition: %s", tableName, i);
+ "Failed to find enabled fully replicated server for table: %s,
partition: %s", tableName, partitionId);
workerIdToServerInstanceMap.put(workerId, new
QueryServerInstance(serverInstance));
+ // NOTE: Copy the segment lists. Unlike the
multiple-partitions-per-worker path (which merges into fresh lists),
+ // these are the broker's published metadata, shared across
queries and never to be mutated (see
+ // filterLeafStageSegments).
Review Comment:
I don't get the reference to `filterLeafStageSegments`. Its javadoc says the
lists cannot be mutated. Why do you use it as an example here?
--
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]
