Add WorkerManager extension points for customizing leaf-stage segment assignment#18645
Merged
KKcorps merged 1 commit intoJun 2, 2026
Merged
Conversation
… assignment 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 setSegmentsForReplicatedLeafFragment). A subclass can rewrite DispatchablePlanMetadata's worker/replicated segment maps via the existing setters. Both default to no-ops, so behavior is unchanged.
KKcorps
force-pushed
the
workermanager-leaf-segment-filter-hooks
branch
from
d656769 to
24d4d54
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #18645 +/- ##
=========================================
Coverage 64.39% 64.40%
Complexity 1282 1282
=========================================
Files 3362 3362
Lines 207915 207919 +4
Branches 32463 32464 +1
=========================================
+ Hits 133883 133900 +17
+ Misses 63258 63245 -13
Partials 10774 10774
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
yashmayya
approved these changes
Jun 2, 2026
yashmayya
added
multi-stage
xiangfu0
added a commit
to pinot-contrib/pinot-docs
that referenced
this pull request
Jun 2, 2026
Contributor
|
Opened the docs follow-up PR for this change: pinot-contrib/pinot-docs#848 |
xiangfu0
added a commit
to pinot-contrib/pinot-docs
that referenced
this pull request
Jun 2, 2026
## Summary - document the new post-assignment `WorkerManager` hooks for customizing multi-stage leaf segment routing - explain where subclasses can read query options and rewrite `DispatchablePlanMetadata` segment maps - position the hooks as an upgrade-sensitive code-level extension, not a stable plugin family ## Cross-check - verified the merged Apache Pinot behavior against `apache/pinot#18645`, including the new `filterLeafStageSegments(...)` and `filterReplicatedLeafStageSegments(...)` hooks and the existing candidate-server hooks ## Validation - `git diff --check`
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Subclasses of
WorkerManagercan already customize multi-stage leaf-stage worker selection by overridinggetCandidateServers/getCandidateServersForReplicatedLeaf, but the per-worker segment assignment is finalizedinternally with no extension point to adjust it.
This PR adds two
protectedno-op extension points, invoked once the leaf-stage segment assignment has been built:filterLeafStageSegments(DispatchablePlanContext, DispatchablePlanMetadata)— invoked fromupdateContextForLeafStage, which the non-partitioned, partitioned, and logical-table leaf paths all funnel through.filterReplicatedLeafStageSegments(DispatchablePlanContext, DispatchablePlanMetadata)— invoked fromsetSegmentsForReplicatedLeafFragmentfor the replicated (broadcast) leaf path.A subclass can rewrite
DispatchablePlanMetadata#getWorkerIdToSegmentsMap()/#getReplicatedSegments()through theexisting setters — for example to apply a custom segment-pruning policy — with access to the query options via
DispatchablePlanContext#getPlannerContext().getOptions().Backward compatibility
Purely additive. Both hooks default to no-ops, so behavior is unchanged for the base
WorkerManagerand all existingsubclasses.
Testing
Existing multi-stage planning /
WorkerManagertests continue to pass; the hooks are no-ops by default.Release notes
No behavior change — new
protectedextension points only.