Support colocated joins when a partition holds no segments

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

@@ -217,8 +217,37 @@ public List<String> getSegments(BrokerRequest brokerRequest, @Nullable String sa
     return combined.isEmpty() ? null : combined;
   }
 
+  /// 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() {

pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java

@@ -28,6 +28,7 @@
 import org.apache.pinot.common.request.QuerySource;
 import org.apache.pinot.core.routing.RoutingTable;
 import org.apache.pinot.core.routing.SegmentsToQuery;
+import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo;
 import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo;
 import org.apache.pinot.core.transport.ServerInstance;
 import org.mockito.Mock;
@@ -293,6 +294,39 @@ private BrokerRequest createMockBrokerRequest(String tableName) {
     return brokerRequest;
   }
 
+  @Test
+  public void testGetTablePartitionInfoReturnsTheSingleClusterThatHasIt() {
+    TablePartitionReplicatedServersInfo partitionInfo = mock(TablePartitionReplicatedServersInfo.class);
+    when(_localClusterRoutingManager.getTablePartitionReplicatedServersInfo(TEST_TABLE)).thenReturn(null);
+    when(_remoteClusterRoutingManager1.getTablePartitionReplicatedServersInfo(TEST_TABLE)).thenReturn(partitionInfo);
+    when(_remoteClusterRoutingManager2.getTablePartitionReplicatedServersInfo(TEST_TABLE)).thenReturn(null);
+
+    assertEquals(_multiClusterRoutingManager.getTablePartitionReplicatedServersInfo(TEST_TABLE), partitionInfo);
+  }
+
+  /// A partial view would make a partition served only by another cluster look empty, so nothing is reported at all.
+  @Test
+  public void testGetTablePartitionInfoReturnsNullWhenSeveralClustersHaveIt() {
+    when(_localClusterRoutingManager.getTablePartitionReplicatedServersInfo(TEST_TABLE))
+        .thenReturn(mock(TablePartitionReplicatedServersInfo.class));
+    when(_remoteClusterRoutingManager1.getTablePartitionReplicatedServersInfo(TEST_TABLE))
+        .thenReturn(mock(TablePartitionReplicatedServersInfo.class));
+    when(_remoteClusterRoutingManager2.getTablePartitionReplicatedServersInfo(TEST_TABLE)).thenReturn(null);
+
+    assertNull(_multiClusterRoutingManager.getTablePartitionReplicatedServersInfo(TEST_TABLE));
+  }
+
+  @Test
+  public void testGetTablePartitionInfoIgnoresAFailingRemoteCluster() {
+    TablePartitionReplicatedServersInfo partitionInfo = mock(TablePartitionReplicatedServersInfo.class);
+    when(_localClusterRoutingManager.getTablePartitionReplicatedServersInfo(TEST_TABLE)).thenReturn(partitionInfo);
+    when(_remoteClusterRoutingManager1.getTablePartitionReplicatedServersInfo(TEST_TABLE))
+        .thenThrow(new RuntimeException("remote cluster is down"));
+    when(_remoteClusterRoutingManager2.getTablePartitionReplicatedServersInfo(TEST_TABLE)).thenReturn(null);
+
+    assertEquals(_multiClusterRoutingManager.getTablePartitionReplicatedServersInfo(TEST_TABLE), partitionInfo);
+  }
+
   private RoutingTable createRoutingTable(String serverName, List<String> segments) {
     Map<ServerInstance, SegmentsToQuery> serverMap = new HashMap<>();
     ServerInstance server = createMockServerInstance(serverName);