Support broker segment pruning for colocated joins

pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java

@@ -20,6 +20,7 @@
 
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -1162,6 +1163,16 @@ public List<String> getSegments(BrokerRequest brokerRequest, @Nullable String sa
     return routingEntry.getSegments(brokerRequest, samplerName);
   }
 
+  @Nullable
+  @Override
+  public Set<String> getPrunedSegments(BrokerRequest brokerRequest) {
+    RoutingEntry routingEntry = _routingEntryMap.get(brokerRequest.getQuerySource().getTableName());
+    if (routingEntry == null) {
+      return null;
+    }
+    return routingEntry.getPrunedSegments(brokerRequest, extractSamplerName(brokerRequest));
+  }
+
   private static String normalizeSamplerName(String samplerName) {
     return samplerName.trim().toLowerCase(Locale.ROOT);
   }
@@ -1434,22 +1445,33 @@ void refreshSegment(String segment) {
       }
     }
 
-    InstanceSelector.SelectionResult calculateRouting(BrokerRequest brokerRequest, long requestId,
-        @Nullable String samplerName) {
-      SamplerInfo samplerInfo = getSamplerInfo(samplerName);
+    /// Runs selection and then the pruner chain, which is the one place that decides what a query sees. Every caller
+    /// goes through here on purpose: the routing table, the plain segment list and the planner's emptiness proof must
+    /// all be judged by the same selector and the same pruners. A second copy of this sequence that drifted would let
+    /// the planner prove a partition empty that a real query would still have scanned, and that loses rows with no
+    /// error anywhere.
+    private SelectedSegments selectThenPrune(BrokerRequest brokerRequest, @Nullable SamplerInfo samplerInfo) {
       SegmentSelector segmentSelector = samplerInfo != null ? samplerInfo._segmentSelector : _segmentSelector;
-      InstanceSelector instanceSelector = samplerInfo != null ? samplerInfo._instanceSelector : _instanceSelector;
       Set<String> selectedSegments = segmentSelector.select(brokerRequest);
-      int numTotalSelectedSegments = selectedSegments.size();
+      Set<String> survivingSegments = selectedSegments;
       if (!selectedSegments.isEmpty()) {
         for (SegmentPruner segmentPruner : _segmentPruners) {
-          selectedSegments = segmentPruner.prune(brokerRequest, selectedSegments);
+          survivingSegments = segmentPruner.prune(brokerRequest, survivingSegments);
         }
       }
-      int numPrunedSegments = numTotalSelectedSegments - selectedSegments.size();
-      if (!selectedSegments.isEmpty()) {
+      return new SelectedSegments(selectedSegments, survivingSegments);
+    }
+
+    InstanceSelector.SelectionResult calculateRouting(BrokerRequest brokerRequest, long requestId,
+        @Nullable String samplerName) {
+      SamplerInfo samplerInfo = getSamplerInfo(samplerName);
+      InstanceSelector instanceSelector = samplerInfo != null ? samplerInfo._instanceSelector : _instanceSelector;
+      SelectedSegments selectedSegments = selectThenPrune(brokerRequest, samplerInfo);
+      Set<String> survivingSegments = selectedSegments._surviving;
+      int numPrunedSegments = selectedSegments.getNumPruned();
+      if (!survivingSegments.isEmpty()) {
         InstanceSelector.SelectionResult selectionResult =
-            instanceSelector.select(brokerRequest, new ArrayList<>(selectedSegments), requestId);
+            instanceSelector.select(brokerRequest, new ArrayList<>(survivingSegments), requestId);
         selectionResult.setNumPrunedSegments(numPrunedSegments);
         return selectionResult;
       } else {
@@ -1459,15 +1481,47 @@ InstanceSelector.SelectionResult calculateRouting(BrokerRequest brokerRequest, l
     }
 
     List<String> getSegments(BrokerRequest brokerRequest, @Nullable String samplerName) {
-      SamplerInfo samplerInfo = getSamplerInfo(samplerName);
-      SegmentSelector segmentSelector = samplerInfo != null ? samplerInfo._segmentSelector : _segmentSelector;
-      Set<String> selectedSegments = segmentSelector.select(brokerRequest);
-      if (!selectedSegments.isEmpty()) {
-        for (SegmentPruner segmentPruner : _segmentPruners) {
-          selectedSegments = segmentPruner.prune(brokerRequest, selectedSegments);
+      return new ArrayList<>(selectThenPrune(brokerRequest, getSamplerInfo(samplerName))._surviving);
+    }
+
+    /// See [RoutingManager#getPrunedSegments]. The sampler is honoured for the same reason the query path honours it:
+    /// a narrower selection only ever shrinks what this can prove, never widens it.
+    ///
+    /// The pruners return a new set rather than editing the one they are handed, so taking the difference costs
+    /// nothing unless something was actually pruned. If one ever did edit in place the two sets would be the same
+    /// object, the difference would come out empty, and this would fall back to proving nothing -- the safe direction.
+    Set<String> getPrunedSegments(BrokerRequest brokerRequest, @Nullable String samplerName) {
+      SelectedSegments selectedSegments = selectThenPrune(brokerRequest, getSamplerInfo(samplerName));
+      int numPruned = selectedSegments.getNumPruned();
+      if (numPruned == 0) {
+        return Set.of();
+      }
+      // Built up rather than copied down: the count is already known and is usually a small fraction of the table's
+      // segments, so copying every selected segment only to remove most of them again would size the allocation to
+      // the table instead of to the answer.
+      Set<String> prunedSegments = Sets.newHashSetWithExpectedSize(numPruned);
+      for (String segment : selectedSegments._selected) {
+        if (!selectedSegments._surviving.contains(segment)) {
+          prunedSegments.add(segment);
         }
       }
-      return new ArrayList<>(selectedSegments);
+      return prunedSegments;
+    }
+  }
+
+  /// What one run of [RoutingEntry#selectThenPrune] decided: the segments selection offered, and the ones the pruners
+  /// left. Both are needed because the difference between them is the only sound proof that a segment cannot match.
+  private static class SelectedSegments {
+    final Set<String> _selected;
+    final Set<String> _surviving;
+
+    SelectedSegments(Set<String> selected, Set<String> surviving) {
+      _selected = selected;
+      _surviving = surviving;
+    }
+
+    int getNumPruned() {
+      return _selected.size() - _surviving.size();
     }
   }
 }

pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java

@@ -217,8 +217,93 @@ public List<String> getSegments(BrokerRequest brokerRequest, @Nullable String sa
     return combined.isEmpty() ? null : combined;
   }
 
+  /// Combines by *intersection* over the clusters that have routing for the table, which is the opposite of how
+  /// [#getSegments] combines and deliberately so: that returns segments that survive, and a segment survives if any
+  /// cluster keeps it, while this returns segments that are provably eliminated, and a proof only holds if every
+  /// cluster that could route the segment eliminated it. Unioning instead would let one cluster's pruners speak for a
+  /// segment another cluster would still have queried -- silently dropping matching data.
+  ///
+  /// Restricting the intersection to the clusters that have the table is what keeps it useful: the usual case is a
+  /// table in exactly one cluster, where the intersection is that cluster's own verdict. A cluster without the table
+  /// would otherwise contribute an empty set and reduce every answer to "nothing proven".
+  @Nullable
+  @Override
+  public Set<String> getPrunedSegments(BrokerRequest brokerRequest) {
+    String tableNameWithType = brokerRequest.getQuerySource().getTableName();
+    Set<String> combined = intersectPrunedSegments(null, _localClusterRoutingManager, brokerRequest,
+        tableNameWithType);
+    for (BaseBrokerRoutingManager remoteCluster : _remoteClusterRoutingManagers) {
+      combined = intersectPrunedSegments(combined, remoteCluster, brokerRequest, tableNameWithType);
+    }
+    // Still null when no cluster has the table at all, which is what the interface reports for a table that does not
+    // exist -- as opposed to an empty set, which is a cluster that ran the pruners and proved nothing.
+    return combined;
+  }
+
+  /// Folds one cluster's verdict into the running intersection, or returns an empty set to end it: once nothing is
+  /// proven, nothing downstream can make it provable again, and asking the remaining clusters would run a full
+  /// selection and pruner chain each for an answer that is already fixed. `null` means no cluster has answered yet.
+  @Nullable
+  private Set<String> intersectPrunedSegments(@Nullable Set<String> combined, BaseBrokerRoutingManager cluster,
+      BrokerRequest brokerRequest, String tableNameWithType) {
+    if (combined != null && combined.isEmpty()) {
+      return combined;
+    }
+    try {
+      // One lookup rather than routingExists-then-get: a table appearing between the two would let this skip a
+      // cluster that can route it, and the intersection would then claim segments eliminated that nobody asked about.
+      Set<String> prunedSegments = cluster.getPrunedSegments(brokerRequest);
+      if (prunedSegments == null) {
+        // This cluster has no routing for the table, so it eliminates nothing and constrains nothing.
+        return combined;
+      }
+      if (prunedSegments.isEmpty()) {
+        // This cluster proves nothing, so neither does the intersection.
+        return Set.of();
+      }
+      if (combined == null) {
+        return new HashSet<>(prunedSegments);
+      }
+      combined.retainAll(prunedSegments);
+      return combined;
+    } catch (Exception e) {
+      LOGGER.error("Error getting pruned segments from cluster routing manager for table {}", tableNameWithType, e);
+      // A cluster we could not ask may still have routed any of these segments, so prove nothing.
+      return Set.of();
+    }
+  }
+
+  /// Returns the partition info only when a single cluster has any, and `null` when more than one does.
+  ///
+  /// Unlike [#getRoutingTable], [#getSegments] and [#getServingInstances], this cannot union the clusters: the info is
+  /// a per-partition array of the servers holding every segment of that partition, and no server holds the segments
+  /// that live in another cluster. One cluster's array would make a partition served only by another cluster look like
+  /// a partition holding no data, and a colocated join treats such a partition as empty and silently drops its rows. So
+  /// a table spread over several clusters reports nothing and its callers fail. Expressing it properly needs the array
+  /// to carry each partition's cluster, which the current shape cannot do.
   @Override
   public TablePartitionReplicatedServersInfo getTablePartitionReplicatedServersInfo(String tableNameWithType) {
-    return findFirst(mgr -> mgr.getTablePartitionReplicatedServersInfo(tableNameWithType), tableNameWithType);
+    TablePartitionReplicatedServersInfo partitionInfo =
+        _localClusterRoutingManager.getTablePartitionReplicatedServersInfo(tableNameWithType);
+    for (BaseBrokerRoutingManager remoteCluster : _remoteClusterRoutingManagers) {
+      TablePartitionReplicatedServersInfo remotePartitionInfo;
+      try {
+        remotePartitionInfo = remoteCluster.getTablePartitionReplicatedServersInfo(tableNameWithType);
+      } catch (Exception e) {
+        LOGGER.error("Error getting table partition info from remote cluster routing manager for table {}",
+            tableNameWithType, e);
+        continue;
+      }
+      if (remotePartitionInfo == null) {
+        continue;
+      }
+      if (partitionInfo != null) {
+        LOGGER.warn("Found table partition info in multiple clusters for table: {}, returning null so that "
+            + "partition-aware routing is not attempted on a partial view", tableNameWithType);
+        return null;
+      }
+      partitionInfo = remotePartitionInfo;
+    }
+    return partitionInfo;
   }
 }

pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionMetadataManager.java

@@ -19,11 +19,13 @@
 package org.apache.pinot.broker.routing.segmentpartition;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeSet;
 import javax.annotation.Nullable;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.commons.lang3.tuple.Triple;
@@ -66,8 +68,13 @@ public class SegmentPartitionMetadataManager implements SegmentZkMetadataFetchLi
   private final Map<String, SegmentInfo> _segmentInfoMap = new HashMap<>();
 
   // computed value based on status change.
-  private transient TablePartitionInfo _tablePartitionInfo;
-  private transient TablePartitionReplicatedServersInfo _tablePartitionReplicatedServersInfo;
+  // NOTE: Volatile because they are written while the table's routing entry is built or updated, and read without any
+  // lock by unrelated threads (e.g. query planner threads). The writers are serialized by BaseBrokerRoutingManager,
+  // which holds the per-table routing table build lock around init() and around every subsequent update; this class'
+  // own 'synchronized' does not cover init(). Both graphs are effectively immutable once published, so the volatile
+  // write provides all the happens-before edge the readers need.
+  private volatile TablePartitionInfo _tablePartitionInfo;
+  private volatile TablePartitionReplicatedServersInfo _tablePartitionReplicatedServersInfo;
 
   public SegmentPartitionMetadataManager(String tableNameWithType, String partitionColumn, String partitionFunctionName,
       int numPartitions, long newSegmentExpirationMs) {
@@ -265,8 +272,13 @@ private void computeTablePartitionReplicatedServersInfo() {
               : segmentsReducingFullyReplicatedServers.subList(0, 10) + "...", _tableNameWithType);
     }
     // Process new segments
+    // Partitions whose segments are all excluded below hold data but end up without partition info. Track them so that
+    // consumers requiring a fully replicated server per partition can tell them apart from genuinely empty partitions.
+    Set<Integer> partitionsWithOnlyDeferredSegments = Set.of();
     if (!newSegmentInfoEntries.isEmpty()) {
       List<String> excludedNewSegments = new ArrayList<>();
+      // Sorted for deterministic reporting
+      Set<Integer> excludedNewSegmentPartitions = new TreeSet<>();
       for (Map.Entry<String, SegmentInfo> entry : newSegmentInfoEntries) {
         String segment = entry.getKey();
         SegmentInfo segmentInfo = entry.getValue();
@@ -284,6 +296,7 @@ private void computeTablePartitionReplicatedServersInfo() {
             partitionInfoMap[partitionId] = partitionInfo;
           } else {
             excludedNewSegments.add(segment);
+            excludedNewSegmentPartitions.add(partitionId);
           }
         } else {
           // If the new segment is not the first segment of a partition, add it only if it won't reduce the fully
@@ -295,6 +308,7 @@ private void computeTablePartitionReplicatedServersInfo() {
             partitionInfo._segments.add(segment);
           } else {
             excludedNewSegments.add(segment);
+            excludedNewSegmentPartitions.add(partitionId);
           }
         }
       }
@@ -303,10 +317,25 @@ private void computeTablePartitionReplicatedServersInfo() {
         LOGGER.info("Excluded {} new segments: {}... without all replicas available in table: {}", numSegments,
             numSegments <= 10 ? excludedNewSegments : excludedNewSegments.subList(0, 10) + "...", _tableNameWithType);
       }
+      // NOTE: Computed against the final partition info map, i.e. after the whole new segment pass, rather than latched
+      // when a segment is excluded: a partition can hold both an excluded new segment and one that ends up populating
+      // the partition info, and which of the two is visited first depends on the iteration order of _segmentInfoMap.
+      excludedNewSegmentPartitions.removeIf(partitionId -> partitionInfoMap[partitionId] != null);
+      if (!excludedNewSegmentPartitions.isEmpty()) {
+        // An unmodifiable view rather than Set.copyOf(): it enforces the accessor's effectively-immutable contract and
+        // keeps the sorted iteration order.
+        partitionsWithOnlyDeferredSegments = Collections.unmodifiableSet(excludedNewSegmentPartitions);
+        int numAffectedPartitions = excludedNewSegmentPartitions.size();
+        List<Integer> partitionsToLog = new ArrayList<>(excludedNewSegmentPartitions);
+        LOGGER.warn("Found {} partitions: {} without partition info because all their segments are new segments "
+                + "without all replicas available in table: {}", numAffectedPartitions,
+            numAffectedPartitions <= 10 ? partitionsToLog : partitionsToLog.subList(0, 10) + "...",
+            _tableNameWithType);
+      }
     }
     _tablePartitionReplicatedServersInfo =
         new TablePartitionReplicatedServersInfo(_tableNameWithType, _partitionColumn, _partitionFunctionName,
-            _numPartitions, partitionInfoMap, segmentsWithInvalidPartition);
+            _numPartitions, partitionInfoMap, segmentsWithInvalidPartition, partitionsWithOnlyDeferredSegments);
   }
 
   private void computeTablePartitionInfo() {