This is an automated email from the ASF dual-hosted git repository.
KKcorps pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 36934936ba4 Add WorkerManager extension points for customizing
leaf-stage segment assignment (#18645)
36934936ba4 is described below
commit 36934936ba425b7fe8b83893aedf065fe29dc848
Author: Kartik Khare <[email protected]>
AuthorDate: Tue Jun 2 11:31:13 2026 +0530
Add WorkerManager extension points for customizing leaf-stage segment
assignment (#18645)
Subclasses can already customize multi-stage worker selection by overriding
getCandidateServers / getCandidateServersForReplicatedLeaf, but the per-worker
segment assignment is finalized internally with no extension point. This adds
two protected no-op hooks invoked once the assignment is built:
filterLeafStageSegments (from updateContextForLeafStage, covering the
non-partitioned, partitioned and logical-table paths) and
filterReplicatedLeafStageSegments (from setSegmentsForReplicated [...]
---
.../java/org/apache/pinot/query/routing/WorkerManager.java | 10 ++++++++++
1 file changed, 10 insertions(+)
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java
index c658062f1ed..d7b6383a3bf 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java
@@ -508,6 +508,7 @@ public class WorkerManager {
}
private void updateContextForLeafStage(DispatchablePlanMetadata metadata,
DispatchablePlanContext context) {
+ filterLeafStageSegments(context, metadata);
if (context.isUseLeafServerForIntermediateStage()) {
Map<Integer, QueryServerInstance> workerIdToServerInstanceMap =
metadata.getWorkerIdToServerInstanceMap();
assert workerIdToServerInstanceMap != null;
@@ -735,6 +736,15 @@ public class WorkerManager {
// TODO: Support unavailable segments and optional segments for replicated
leaf stage
metadata.setReplicatedSegments(segmentsMap);
+ filterReplicatedLeafStageSegments(context, metadata);
+ }
+
+ /** Extension point to filter the non-replicated leaf-stage per-worker
segment assignment; no-op by default. */
+ protected void filterLeafStageSegments(DispatchablePlanContext context,
DispatchablePlanMetadata metadata) {
+ }
+
+ /** Extension point to filter the replicated leaf-stage segments; no-op by
default. */
+ protected void filterReplicatedLeafStageSegments(DispatchablePlanContext
context, DispatchablePlanMetadata metadata) {
}
/**
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
