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 c658062f1ed2..d7b6383a3bf9 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 @@ private void assignWorkersToLeafFragment(PlanFragment fragment, DispatchablePlan
   }
 
   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 @@ private void setSegmentsForReplicatedLeafFragment(DispatchablePlanMetadata metad
 
     // 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) {
   }
 
   /**