Multi-stage leaf routing: optional segments, unavailable metadata, and empty partitions

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java

@@ -156,6 +156,10 @@ public Map<Integer, DispatchablePlanFragment> constructDispatchablePlanFragmentM
           dispatchablePlanMetadata.getWorkerIdToMailboxesMap();
       Preconditions.checkArgument(workerIdToSegmentsMap == null || workerIdToTableNameSegmentsMap == null,
           "Both workerIdToSegmentsMap and workerIdToTableNameSegmentsMap cannot be set at the same time");
+      Map<Integer, Map<String, List<String>>> workerIdToOptionalSegmentsMap =
+          dispatchablePlanMetadata.getWorkerIdToOptionalSegmentsMap();
+      Map<Integer, Map<String, List<String>>> workerIdToOptionalTableSegmentsMap =
+          dispatchablePlanMetadata.getWorkerIdToOptionalTableSegmentsMap();
       Map<QueryServerInstance, List<Integer>> serverInstanceToWorkerIdsMap = new HashMap<>();
       WorkerMetadata[] workerMetadataArray = new WorkerMetadata[workerIdToServerInstanceMap.size()];
       for (Map.Entry<Integer, QueryServerInstance> serverEntry : workerIdToServerInstanceMap.entrySet()) {
@@ -165,9 +169,21 @@ public Map<Integer, DispatchablePlanFragment> constructDispatchablePlanFragmentM
         WorkerMetadata workerMetadata = new WorkerMetadata(workerId, workerIdToMailboxesMap.get(workerId));
         if (workerIdToSegmentsMap != null) {
           workerMetadata.setTableSegmentsMap(workerIdToSegmentsMap.get(workerId));
+          if (workerIdToOptionalSegmentsMap != null) {
+            Map<String, List<String>> optionalForWorker = workerIdToOptionalSegmentsMap.get(workerId);
+            if (optionalForWorker != null && !optionalForWorker.isEmpty()) {
+              workerMetadata.setOptionalTableSegmentsMap(optionalForWorker);
+            }
+          }
         }
         if (workerIdToTableNameSegmentsMap != null) {
           workerMetadata.setLogicalTableSegmentsMap(workerIdToTableNameSegmentsMap.get(workerId));
+          if (workerIdToOptionalTableSegmentsMap != null) {
+            Map<String, List<String>> optionalLogical = workerIdToOptionalTableSegmentsMap.get(workerId);
+            if (optionalLogical != null && !optionalLogical.isEmpty()) {
+              workerMetadata.setOptionalTableSegmentsMap(optionalLogical);
+            }
+          }
         }
         workerMetadataArray[workerId] = workerMetadata;
       }

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java

@@ -71,6 +71,9 @@ public class DispatchablePlanMetadata implements Serializable {
   // Available for leaf stage only
   // Map from workerId -> {tableType -> segments}
   private Map<Integer, Map<String, List<String>>> _workerIdToSegmentsMap;
+  // Map from workerId -> {tableType -> optional segments}, parallel to _workerIdToSegmentsMap when present
+  @Nullable
+  private Map<Integer, Map<String, List<String>>> _workerIdToOptionalSegmentsMap;
   // Map from tableType -> segments, available when 'is_replicated' hint is set to true
   private Map<String, List<String>> _replicatedSegments;
   private TimeBoundaryInfo _timeBoundaryInfo;
@@ -85,6 +88,8 @@ public class DispatchablePlanMetadata implements Serializable {
    * Map from workerId -> {physicalTableName -> segments} is required for logical tables.
    */
   private Map<Integer, Map<String, List<String>>> _workerIdToTableSegmentsMap;
+  @Nullable
+  private Map<Integer, Map<String, List<String>>> _workerIdToOptionalTableSegmentsMap;
   private LogicalTableRouteInfo _logicalTableRouteInfo;
 
   public List<String> getScannedTables() {
@@ -125,6 +130,16 @@ public void setWorkerIdToSegmentsMap(Map<Integer, Map<String, List<String>>> wor
     _workerIdToSegmentsMap = workerIdToSegmentsMap;
   }
 
+  @Nullable
+  public Map<Integer, Map<String, List<String>>> getWorkerIdToOptionalSegmentsMap() {
+    return _workerIdToOptionalSegmentsMap;
+  }
+
+  public void setWorkerIdToOptionalSegmentsMap(
+      @Nullable Map<Integer, Map<String, List<String>>> workerIdToOptionalSegmentsMap) {
+    _workerIdToOptionalSegmentsMap = workerIdToOptionalSegmentsMap;
+  }
+
   @Nullable
   public Map<String, List<String>> getReplicatedSegments() {
     return _replicatedSegments;
@@ -204,4 +219,14 @@ public void setWorkerIdToTableSegmentsMap(
       Map<Integer, Map<String, List<String>>> workerIdToTableSegmentsMap) {
     _workerIdToTableSegmentsMap = workerIdToTableSegmentsMap;
   }
+
+  @Nullable
+  public Map<Integer, Map<String, List<String>>> getWorkerIdToOptionalTableSegmentsMap() {
+    return _workerIdToOptionalTableSegmentsMap;
+  }
+
+  public void setWorkerIdToOptionalTableSegmentsMap(
+      @Nullable Map<Integer, Map<String, List<String>>> workerIdToOptionalTableSegmentsMap) {
+    _workerIdToOptionalTableSegmentsMap = workerIdToOptionalTableSegmentsMap;
+  }
 }