From 5a48ddb29070df73e3f86c1cc2506b0961270078 Mon Sep 17 00:00:00 2001 From: Yash Mayya Date: Wed, 5 Aug 2026 14:11:42 -0700 Subject: [PATCH 1/2] Support colocated joins when a partition holds no segments A colocated join failed outright when a partition of one of its tables held no segments, with "Failed to find any segment for table: X, partition: N". Worker ids came from a running counter over the partitions that held data, so skipping an empty one would shift every later partition down a slot. Two tables each dropping a different empty partition could then end up with equal worker counts and be wired 1-to-1 onto mismatched partitions, losing rows with no error, which is why the assignment refused to continue at all. The stages tied together by direct exchanges now share one ordered list of partition classes, dropping only the classes that hold no data on any member. A class the group keeps but a member holds no data for gets a worker with no segments, placed on a server borrowed from a member that does hold that class so the exchange stays in process. The two sides of every direct exchange assert that they agree on the list. The broker publishes the partitions whose only segments are new and have no online replica. Those hold data that no server can serve as a whole, so they keep failing rather than being read as empty. A worker with no segments is charged against the query thread estimate like any other worker: it is dispatched and does run a leaf operator. The estimate therefore over-counts by two threads per such worker, which is conservative and only affects colocated joins over a partition space that is largely unpopulated. Those queries failed outright before, so there is no earlier estimate to compare against. Aggregation merge identity is deliberately not covered here. A leaf that scans nothing emits one identity row for an aggregation with no GROUP BY, but that predates this change: a worker whose segments are all pruned on the server already does the same. Padding raises how many such rows reach the merge without introducing the dependency, and the one aggregation that is not a true merge identity fails only when every worker is empty, which this change does not newly reach. --- .../manager/MultiClusterRoutingManager.java | 31 +- .../SegmentPartitionMetadataManager.java | 35 +- .../MultiClusterRoutingManagerTest.java | 34 + .../SegmentPartitionMetadataManagerTest.java | 146 +++ .../TablePartitionReplicatedServersInfo.java | 31 + .../ColocatedJoinEmptyPartitionTest.java | 423 +++++++ .../physical/DispatchablePlanContext.java | 8 + .../physical/DispatchablePlanFragment.java | 6 +- .../physical/DispatchablePlanMetadata.java | 49 + .../physical/MailboxAssignmentVisitor.java | 50 +- .../routing/ColocationGroupAnalyzer.java | 253 ++++ .../query/routing/LeafPartitionHints.java | 118 ++ .../pinot/query/routing/WorkerManager.java | 620 +++++++-- .../pinot/query/QueryEnvironmentTestBase.java | 2 +- .../physical/DispatchableSubPlanTest.java | 8 +- .../MailboxAssignmentVisitorTest.java | 84 ++ .../physical/PinotDispatchPlannerTest.java | 28 + .../routing/ColocationGroupAnalyzerTest.java | 357 ++++++ .../query/routing/WorkerManagerTest.java | 1105 ++++++++++++++++- 19 files changed, 3277 insertions(+), 111 deletions(-) create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java create mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java create mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/query/routing/LeafPartitionHints.java create mode 100644 pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java index 17d62a7ffad1..312d7b0b7d10 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java @@ -217,8 +217,37 @@ public List 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; } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionMetadataManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionMetadataManager.java index 03a15564575b..1cf6ed6a4651 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionMetadataManager.java +++ b/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 _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 partitionsWithOnlyDeferredSegments = Set.of(); if (!newSegmentInfoEntries.isEmpty()) { List excludedNewSegments = new ArrayList<>(); + // Sorted for deterministic reporting + Set excludedNewSegmentPartitions = new TreeSet<>(); for (Map.Entry 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 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() { diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java index 7ba62cb47b73..8d8cdb34c262 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java +++ b/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 segments) { Map serverMap = new HashMap<>(); ServerInstance server = createMockServerInstance(serverName); diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionMetadataManagerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionMetadataManagerTest.java index 1c056ba9c74e..9901b71c96e3 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionMetadataManagerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionMetadataManagerTest.java @@ -19,6 +19,7 @@ package org.apache.pinot.broker.routing.segmentpartition; import com.google.common.collect.ImmutableSet; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -253,6 +254,8 @@ public void testPartitionMetadataManagerProcessingThroughSegmentChangesSinglePar assertEquals(partitionInfoMap[1]._fullyReplicatedServers, Set.of(SERVER_0)); assertEqualsNoOrder(partitionInfoMap[1]._segments.toArray(), new String[]{segment1, segment2}); assertTrue(tablePartitionReplicatedServersInfo.getSegmentsWithInvalidPartition().isEmpty()); + // Partition 0 is still served by segment0, so it is not a deferred empty partition + assertTrue(tablePartitionReplicatedServersInfo.getPartitionsWithOnlyDeferredSegments().isEmpty()); // Making all of them replicated will show full list, even for the new segment segmentAssignment.put(segment0, Map.of(SERVER_0, ONLINE, SERVER_1, ONLINE)); @@ -285,6 +288,149 @@ public void testPartitionMetadataManagerProcessingThroughSegmentChangesSinglePar assertEquals(tablePartitionReplicatedServersInfo.getSegmentsWithInvalidPartition().get(0), segmentInvalid); } + /// A partition whose only segments are new ones without all replicas available holds data that no single server can + /// serve as a whole, so it must be told apart from a genuinely empty partition (see + /// [TablePartitionReplicatedServersInfo#getPartitionsWithOnlyDeferredSegments()]). + @Test + public void testPartitionsWithOnlyDeferredSegments() { + ExternalView externalView = new ExternalView(OFFLINE_TABLE_NAME); + Map> segmentAssignment = externalView.getRecord().getMapFields(); + Set onlineSegments = new HashSet<>(); + // NOTE: Ideal state is not used in the current implementation. + IdealState idealState = new IdealState(OFFLINE_TABLE_NAME); + + SegmentPartitionMetadataManager partitionMetadataManager = + new SegmentPartitionMetadataManager(OFFLINE_TABLE_NAME, PARTITION_COLUMN, PARTITION_COLUMN_FUNC, NUM_PARTITIONS, + TimeUnit.MINUTES.toMillis(5)); + SegmentZkMetadataFetcher segmentZkMetadataFetcher = + new SegmentZkMetadataFetcher(OFFLINE_TABLE_NAME, _propertyStore); + segmentZkMetadataFetcher.register(partitionMetadataManager); + + // Initial state should be all empty + segmentZkMetadataFetcher.init(idealState, externalView, onlineSegments); + TablePartitionReplicatedServersInfo tablePartitionReplicatedServersInfo = + partitionMetadataManager.getTablePartitionReplicatedServersInfo(); + assertTrue(tablePartitionReplicatedServersInfo.getPartitionsWithOnlyDeferredSegments().isEmpty()); + + // A newly created segment without available replica as the only segment of partition 1 leaves the partition without + // partition info, and should be reported as a deferred empty partition. Partition 0 has no segment at all, and + // should not be reported. + long creationTimeMs = System.currentTimeMillis(); + String newSegmentWithoutReplica = "deferredSegment1"; + onlineSegments.add(newSegmentWithoutReplica); + setSegmentZKMetadata(newSegmentWithoutReplica, PARTITION_COLUMN_FUNC, NUM_PARTITIONS, 1, creationTimeMs); + segmentZkMetadataFetcher.onAssignmentChange(idealState, externalView, onlineSegments); + tablePartitionReplicatedServersInfo = partitionMetadataManager.getTablePartitionReplicatedServersInfo(); + TablePartitionReplicatedServersInfo.PartitionInfo[] partitionInfoMap = + tablePartitionReplicatedServersInfo.getPartitionInfoMap(); + assertNull(partitionInfoMap[0]); + assertNull(partitionInfoMap[1]); + assertEquals(tablePartitionReplicatedServersInfo.getPartitionsWithOnlyDeferredSegments(), Set.of(1)); + assertTrue(tablePartitionReplicatedServersInfo.getSegmentsWithInvalidPartition().isEmpty()); + + // Adding another newly created segment with all replicas available to partition 1 makes the partition servable. The + // first segment is still excluded, but the partition is no longer deferred empty. This holds regardless of the + // order the 2 new segments are processed in, which is why the deferred empty partitions are derived from the final + // partition info map instead of being latched when a segment is excluded. + String newSegmentWithReplicas = "deferredSegment2"; + onlineSegments.add(newSegmentWithReplicas); + segmentAssignment.put(newSegmentWithReplicas, Map.of(SERVER_0, ONLINE, SERVER_1, ONLINE)); + setSegmentZKMetadata(newSegmentWithReplicas, PARTITION_COLUMN_FUNC, NUM_PARTITIONS, 1, creationTimeMs); + segmentZkMetadataFetcher.onAssignmentChange(idealState, externalView, onlineSegments); + tablePartitionReplicatedServersInfo = partitionMetadataManager.getTablePartitionReplicatedServersInfo(); + partitionInfoMap = tablePartitionReplicatedServersInfo.getPartitionInfoMap(); + assertEquals(partitionInfoMap[1]._fullyReplicatedServers, Set.of(SERVER_0, SERVER_1)); + assertEquals(partitionInfoMap[1]._segments, List.of(newSegmentWithReplicas)); + assertTrue(tablePartitionReplicatedServersInfo.getPartitionsWithOnlyDeferredSegments().isEmpty()); + + // Bringing up the replicas of the first segment adds it to the partition info + segmentAssignment.put(newSegmentWithoutReplica, Map.of(SERVER_0, ONLINE, SERVER_1, ONLINE)); + segmentZkMetadataFetcher.onAssignmentChange(idealState, externalView, onlineSegments); + tablePartitionReplicatedServersInfo = partitionMetadataManager.getTablePartitionReplicatedServersInfo(); + partitionInfoMap = tablePartitionReplicatedServersInfo.getPartitionInfoMap(); + assertEqualsNoOrder(partitionInfoMap[1]._segments.toArray(), + new String[]{newSegmentWithoutReplica, newSegmentWithReplicas}); + assertTrue(tablePartitionReplicatedServersInfo.getPartitionsWithOnlyDeferredSegments().isEmpty()); + } + + /// A partition holding both a new segment without any online replica and a new segment with all of them is servable, + /// so it must never be reported -- whichever of the two the new-segment pass visits first, and one of them IS always + /// excluded (see the NOTE on the `removeIf` in [SegmentPartitionMetadataManager]). + /// + /// The pass walks a [HashMap] keyed by segment name, so the names decide the visit order. This test pins down a name + /// pair for each of the 2 orders and drives both announcement orders through the manager on top of that. + @Test + public void testPartitionsWithOnlyDeferredSegmentsAreOrderIndependent() { + for (boolean noReplicaVisitedFirst : List.of(true, false)) { + String[] segmentNames = findSegmentNamePair(noReplicaVisitedFirst); + for (boolean announceNoReplicaFirst : List.of(true, false)) { + assertPartitionHasNotOnlyDeferredSegments(segmentNames[0], segmentNames[1], noReplicaVisitedFirst, + announceNoReplicaFirst); + } + } + } + + /// Returns a `{noReplicaSegment, allReplicasSegment}` name pair that a [HashMap] holding exactly those 2 keys + /// iterates in the requested order. + private static String[] findSegmentNamePair(boolean noReplicaVisitedFirst) { + for (int i = 0; i < 1000; i++) { + String noReplicaSegment = "deferredNoReplica" + i; + String allReplicasSegment = "deferredAllReplicas" + i; + Map probe = new HashMap<>(); + probe.put(noReplicaSegment, noReplicaSegment); + probe.put(allReplicasSegment, allReplicasSegment); + if (probe.keySet().iterator().next().equals(noReplicaSegment) == noReplicaVisitedFirst) { + return new String[]{noReplicaSegment, allReplicasSegment}; + } + } + throw new AssertionError( + "Found no segment name pair iterated with the segment " + (noReplicaVisitedFirst ? "without" : "with") + + " replicas first"); + } + + /// Registers 2 new segments of partition 1 -- one without any online replica and one with all of them -- and asserts + /// that the partition ends up servable and is NOT reported as deferred. `announceNoReplicaFirst` picks which of the 2 + /// is announced first; `noReplicaVisitedFirst` only feeds the failure message (see [#findSegmentNamePair]). + private void assertPartitionHasNotOnlyDeferredSegments(String noReplicaSegment, String allReplicasSegment, + boolean noReplicaVisitedFirst, boolean announceNoReplicaFirst) { + ExternalView externalView = new ExternalView(OFFLINE_TABLE_NAME); + Map> segmentAssignment = externalView.getRecord().getMapFields(); + Set onlineSegments = new HashSet<>(); + // NOTE: Ideal state is not used in the current implementation. + IdealState idealState = new IdealState(OFFLINE_TABLE_NAME); + + SegmentPartitionMetadataManager partitionMetadataManager = + new SegmentPartitionMetadataManager(OFFLINE_TABLE_NAME, PARTITION_COLUMN, PARTITION_COLUMN_FUNC, NUM_PARTITIONS, + TimeUnit.MINUTES.toMillis(5)); + SegmentZkMetadataFetcher segmentZkMetadataFetcher = + new SegmentZkMetadataFetcher(OFFLINE_TABLE_NAME, _propertyStore); + segmentZkMetadataFetcher.register(partitionMetadataManager); + segmentZkMetadataFetcher.init(idealState, externalView, onlineSegments); + + long creationTimeMs = System.currentTimeMillis(); + setSegmentZKMetadata(noReplicaSegment, PARTITION_COLUMN_FUNC, NUM_PARTITIONS, 1, creationTimeMs); + setSegmentZKMetadata(allReplicasSegment, PARTITION_COLUMN_FUNC, NUM_PARTITIONS, 1, creationTimeMs); + // Only the second segment has replicas: the first one is absent from the external view altogether. + segmentAssignment.put(allReplicasSegment, Map.of(SERVER_0, ONLINE, SERVER_1, ONLINE)); + onlineSegments.add(announceNoReplicaFirst ? noReplicaSegment : allReplicasSegment); + segmentZkMetadataFetcher.onAssignmentChange(idealState, externalView, onlineSegments); + onlineSegments.add(announceNoReplicaFirst ? allReplicasSegment : noReplicaSegment); + segmentZkMetadataFetcher.onAssignmentChange(idealState, externalView, onlineSegments); + + TablePartitionReplicatedServersInfo tablePartitionReplicatedServersInfo = + partitionMetadataManager.getTablePartitionReplicatedServersInfo(); + String context = "with the segment without replicas visited " + (noReplicaVisitedFirst ? "first" : "second") + + " and announced " + (announceNoReplicaFirst ? "first" : "second"); + TablePartitionReplicatedServersInfo.PartitionInfo partitionInfo = + tablePartitionReplicatedServersInfo.getPartitionInfoMap()[1]; + assertNotNull(partitionInfo, "Partition 1 has no partition info " + context); + assertEquals(partitionInfo._fullyReplicatedServers, Set.of(SERVER_0, SERVER_1), context); + assertEquals(partitionInfo._segments, List.of(allReplicasSegment), context); + assertTrue(tablePartitionReplicatedServersInfo.getPartitionsWithOnlyDeferredSegments().isEmpty(), + "Servable partition reported as deferred empty " + context + ": " + + tablePartitionReplicatedServersInfo.getPartitionsWithOnlyDeferredSegments()); + } + private void setSegmentZKMetadata(String segment, String partitionFunction, int numPartitions, int partitionId, long creationTimeMs) { SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segment); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/TablePartitionReplicatedServersInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/TablePartitionReplicatedServersInfo.java index 706cb64ab8d1..42d47e6c249a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/TablePartitionReplicatedServersInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/TablePartitionReplicatedServersInfo.java @@ -20,6 +20,7 @@ import java.util.List; import java.util.Set; +import javax.annotation.Nullable; /// An advanced version of [TablePartitionInfo] that also contains information about the fully replicated servers @@ -31,16 +32,30 @@ public class TablePartitionReplicatedServersInfo { private final int _numPartitions; private final PartitionInfo[] _partitionInfoMap; private final List _segmentsWithInvalidPartition; + private final Set _partitionsWithOnlyDeferredSegments; + /// @deprecated Defaults [#getPartitionsWithOnlyDeferredSegments()] to empty, i.e. claims that no partition holds + /// deferred data, which is the unsafe direction (see that method). Use the overload and pass the real + /// set. + @Deprecated public TablePartitionReplicatedServersInfo(String tableNameWithType, String partitionColumn, String partitionFunctionName, int numPartitions, PartitionInfo[] partitionInfoMap, List segmentsWithInvalidPartition) { + this(tableNameWithType, partitionColumn, partitionFunctionName, numPartitions, partitionInfoMap, + segmentsWithInvalidPartition, Set.of()); + } + + public TablePartitionReplicatedServersInfo(String tableNameWithType, String partitionColumn, + String partitionFunctionName, int numPartitions, PartitionInfo[] partitionInfoMap, + List segmentsWithInvalidPartition, @Nullable Set partitionsWithOnlyDeferredSegments) { _tableNameWithType = tableNameWithType; _partitionColumn = partitionColumn; _partitionFunctionName = partitionFunctionName; _numPartitions = numPartitions; _partitionInfoMap = partitionInfoMap; _segmentsWithInvalidPartition = segmentsWithInvalidPartition; + _partitionsWithOnlyDeferredSegments = + partitionsWithOnlyDeferredSegments != null ? partitionsWithOnlyDeferredSegments : Set.of(); } public String getTableNameWithType() { @@ -67,6 +82,22 @@ public List getSegmentsWithInvalidPartition() { return _segmentsWithInvalidPartition; } + /// Returns the partitions that have no entry in [#getPartitionInfoMap()] *only* because all of their segments were + /// deferred: every one of them is a new segment (recently created or pushed) that does not have all of its replicas + /// online yet, so including it would leave the partition without a fully replicated server. + /// + /// A `null` slot in [#getPartitionInfoMap()] therefore has several causes: the partition genuinely holds no data, all + /// of its segments are deferred (this set), or its segments hold invalid partition metadata (see + /// [#getSegmentsWithInvalidPartition()]). Only the first is safe to read as empty. A consumer that needs one server + /// to scan a whole partition (e.g. a colocated join in the multi-stage engine) must fail the query on the others + /// rather than silently dropping their rows; one that scatters over all the servers holding the table (the regular + /// routing path) can ignore this set, because it picks the deferred segments up through the routing table. + /// + /// Empty when there is nothing to report, and never `null`: the consumers above read it without a null check. + public Set getPartitionsWithOnlyDeferredSegments() { + return _partitionsWithOnlyDeferredSegments; + } + public static class PartitionInfo { public final Set _fullyReplicatedServers; public final List _segments; diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java new file mode 100644 index 000000000000..3c75ccae35bb --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java @@ -0,0 +1,423 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.integration.tests.custom; + +import com.fasterxml.jackson.databind.JsonNode; +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.integration.tests.ClusterIntegrationTestUtils; +import org.apache.pinot.spi.config.table.ColumnPartitionConfig; +import org.apache.pinot.spi.config.table.SegmentPartitionConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.util.TestUtils; +import org.testng.annotations.AfterClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/// End-to-end coverage for a colocated join over a partitioned table whose declared partition count exceeds the +/// partitions that actually hold segments. +/// +/// Two offline tables are partitioned identically (`Modulo` over 8 partitions) on the join key, but each populates only +/// 3 of the 8 partitions, and they populate *different* ones: +/// +/// | table | populated partitions | +/// |---|---| +/// | left | 0, 1, 2 | +/// | right | 1, 2, 3 | +/// +/// The join keeps the union (classes 0..3) and drops the four classes neither side holds data in, so each side ends up +/// with one worker that has nothing to scan: the left table for class 3, the right table for class 0. What only an +/// end-to-end run can show is that a real server accepts and answers a leaf-stage request whose segment list is empty +/// for a genuinely partitioned table scan. +/// +/// The partition layout is supplied with explicit `tableOptions` hints rather than inferred, because hint inference is +/// off by default (`pinot.broker.multistage.infer.partition.hint`); the hints carry exactly what it would have +/// produced. The `is_colocated_by_join_keys` hint is spelled out for readability -- the exchange would be +/// pre-partitioned here without it, because the join key *is* the partition key. +/// +/// What these tests do NOT prove: the cross-server fallback in `WorkerManager#assignPaddedWorker`, which only fires +/// when the server borrowed from the peer does not host the empty worker's table at all. Both tables are replicated on +/// both servers of the shared cluster, so such a worker always lands on its peer's server here. +@Test(suiteName = "CustomClusterIntegrationTest") +public class ColocatedJoinEmptyPartitionTest extends CustomDataQueryClusterIntegrationTest { + private static final String LEFT_TABLE_NAME = "ColocatedJoinEmptyPartitionLeft"; + private static final String RIGHT_TABLE_NAME = "ColocatedJoinEmptyPartitionRight"; + + private static final String PARTITION_KEY_COLUMN = "partitionKey"; + private static final String METRIC_COLUMN = "metricValue"; + private static final String PARTITION_FUNCTION = "Modulo"; + + /// Deliberately larger than the number of partitions either table populates, which is what this test is about. + private static final int NUM_DECLARED_PARTITIONS = 8; + private static final List LEFT_POPULATED_PARTITIONS = List.of(0, 1, 2); + private static final List RIGHT_POPULATED_PARTITIONS = List.of(1, 2, 3); + /// The partition classes kept by a colocated join of the two tables, i.e. the union of the populated ones. + private static final int NUM_KEPT_CLASSES_FOR_JOIN = 4; + private static final int NUM_ROWS_PER_PARTITION = 2; + + private static final int LEFT_METRIC_MULTIPLIER = 10; + private static final int RIGHT_METRIC_MULTIPLIER = 100; + + private static final String COLOCATED_JOIN_HINT = "/*+ joinOptions(is_colocated_by_join_keys='true') */"; + private static final String TABLE_HINT = + String.format("/*+ tableOptions(partition_function='%s', partition_key='%s', partition_size='%d') */", + PARTITION_FUNCTION, PARTITION_KEY_COLUMN, NUM_DECLARED_PARTITIONS); + + /// Matches one worker's pre-partitioned mailbox send line of an `EXPLAIN IMPLEMENTATION PLAN` tree, e.g. + /// `[2]@localhost:1|[0] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[1]@localhost:1|[0]}`. Group 1 is the sender + /// worker id, group 2 the receiver list. + private static final Pattern PRE_PARTITIONED_SEND_PATTERN = + Pattern.compile("\\|\\[(\\d+)] MAIL_SEND\\([A-Z_]+\\)\\[PARTITIONED]->\\{([^}]*)}"); + + @Override + public String getTableName() { + return LEFT_TABLE_NAME; + } + + @Override + public Schema createSchema() { + return createSchemaForTable(LEFT_TABLE_NAME); + } + + @Override + public List createAvroFiles() { + // Not used: setUpTable builds one Avro file per populated partition, for each of the two tables. + return List.of(); + } + + @Override + protected long getCountStarResult() { + return (long) LEFT_POPULATED_PARTITIONS.size() * NUM_ROWS_PER_PARTITION; + } + + @Override + protected void setUpTable() + throws Exception { + setUpTable(LEFT_TABLE_NAME, LEFT_POPULATED_PARTITIONS, LEFT_METRIC_MULTIPLIER); + setUpTable(RIGHT_TABLE_NAME, RIGHT_POPULATED_PARTITIONS, RIGHT_METRIC_MULTIPLIER); + } + + @Override + protected void waitForAllDocsLoaded(long timeoutMs) { + long expectedNumDocs = getCountStarResult(); + for (String tableName : List.of(LEFT_TABLE_NAME, RIGHT_TABLE_NAME)) { + TestUtils.waitForCondition(aVoid -> getCurrentCountStarResult(tableName) == expectedNumDocs, 100L, timeoutMs, + "Failed to load " + expectedNumDocs + " documents into table: " + tableName); + } + } + + @Override + @AfterClass + public void tearDown() + throws IOException { + LOGGER.warn("Tearing down integration test class: {}", getClass().getSimpleName()); + dropOfflineTable(LEFT_TABLE_NAME); + dropOfflineTable(RIGHT_TABLE_NAME); + FileUtils.deleteDirectory(_tempDir); + LOGGER.warn("Finished tearing down integration test class: {}", getClass().getSimpleName()); + } + + /// The case where a real server has to answer a leaf-stage request with an empty segment list: the two tables + /// populate different subsets of the declared partitions, so each side ends up with a zero-segment worker. + @Test + public void testColocatedJoinWithEmptySegmentWorkers() + throws Exception { + setUseMultiStageQueryEngine(true); + String query = colocatedJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME); + + JsonNode response = queryBrokerHttpEndpoint(query); + assertNoExceptions(response); + + // Rows only join on the keys of the partitions both tables populate, i.e. 1 and 2. + List> expectedRows = new ArrayList<>(); + for (int partition : LEFT_POPULATED_PARTITIONS) { + if (!RIGHT_POPULATED_PARTITIONS.contains(partition)) { + continue; + } + for (int key : keysForPartition(partition)) { + expectedRows.add( + List.of((long) key, (long) key * LEFT_METRIC_MULTIPLIER, (long) key * RIGHT_METRIC_MULTIPLIER)); + } + } + assertRows(response, expectedRows); + + // Both leaves keep the union of the populated classes, so both have exactly one worker with nothing to scan. + assertLeafStages(response, 2, NUM_KEPT_CLASSES_FOR_JOIN, LEFT_POPULATED_PARTITIONS.size()); + assertDirectExchanges(query, 2, NUM_KEPT_CLASSES_FOR_JOIN); + } + + /// A self-join, where every kept class holds data on both sides: the plain worker-count reduction on its own, with + /// the leaves running 3 workers for 8 declared partitions and nothing padded. + @Test + public void testColocatedSelfJoinWithoutEmptySegmentWorkers() + throws Exception { + setUseMultiStageQueryEngine(true); + String query = colocatedJoinQuery(LEFT_TABLE_NAME, LEFT_TABLE_NAME); + + JsonNode response = queryBrokerHttpEndpoint(query); + assertNoExceptions(response); + + List> expectedRows = new ArrayList<>(); + for (int partition : LEFT_POPULATED_PARTITIONS) { + for (int key : keysForPartition(partition)) { + expectedRows.add( + List.of((long) key, (long) key * LEFT_METRIC_MULTIPLIER, (long) key * LEFT_METRIC_MULTIPLIER)); + } + } + assertRows(response, expectedRows); + + int numKeptClasses = LEFT_POPULATED_PARTITIONS.size(); + assertLeafStages(response, 2, numKeptClasses, LEFT_POPULATED_PARTITIONS.size()); + assertDirectExchanges(query, 2, numKeptClasses); + } + + /// Cross-checks the colocated result against the same join planned as a shuffle (no table hints), which rules out a + /// colocated plan that pairs the wrong partition classes and drops or duplicates rows with no error. It also pins + /// down that `fanOut` really tells the two plans apart, which the other tests rely on. + @Test + public void testColocatedJoinMatchesShuffledJoin() + throws Exception { + setUseMultiStageQueryEngine(true); + JsonNode colocatedResponse = queryBrokerHttpEndpoint(colocatedJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME)); + assertNoExceptions(colocatedResponse); + JsonNode shuffledResponse = queryBrokerHttpEndpoint(shuffledJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME)); + assertNoExceptions(shuffledResponse); + + assertEquals(colocatedResponse.get("resultTable").get("rows"), shuffledResponse.get("resultTable").get("rows"), + "Colocated and shuffled plans must return the same rows"); + + JsonNode shuffledStageStats = shuffledResponse.get("stageStats"); + assertNotNull(shuffledStageStats, "Missing stage stats in shuffled response: " + shuffledResponse); + List shuffledLeafStageSends = new ArrayList<>(); + collectLeafStageSends(shuffledStageStats, shuffledLeafStageSends); + assertEquals(shuffledLeafStageSends.size(), 2, + "Unexpected number of leaf stages in stage stats: " + shuffledStageStats.toPrettyString()); + for (JsonNode leafStageSend : shuffledLeafStageSends) { + assertTrue(leafStageSend.path("fanOut").asInt(-1) > 1, + "A shuffled leaf send must write more than one receive mailbox, otherwise the fanOut of 1 asserted for the " + + "colocated plan proves nothing. Stage stats: " + shuffledStageStats.toPrettyString()); + } + } + + private static String colocatedJoinQuery(String leftTableName, String rightTableName) { + return String.format( + "SELECT %s l.%s, l.%s, r.%s FROM %s %s AS l JOIN %s %s AS r ON l.%s = r.%s ORDER BY l.%s", + COLOCATED_JOIN_HINT, PARTITION_KEY_COLUMN, METRIC_COLUMN, METRIC_COLUMN, leftTableName, TABLE_HINT, + rightTableName, TABLE_HINT, PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN); + } + + private static String shuffledJoinQuery(String leftTableName, String rightTableName) { + return String.format("SELECT l.%s, l.%s, r.%s FROM %s AS l JOIN %s AS r ON l.%s = r.%s ORDER BY l.%s", + PARTITION_KEY_COLUMN, METRIC_COLUMN, METRIC_COLUMN, leftTableName, rightTableName, PARTITION_KEY_COLUMN, + PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN); + } + + private static void assertNoExceptions(JsonNode response) { + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions == null || exceptions.isEmpty(), "Query failed with exceptions: " + exceptions); + } + + /// Compares the result table against the expected rows, sorted by their first column to match the queries' `ORDER + /// BY`. + private static void assertRows(JsonNode response, List> unsortedExpectedRows) { + List> expectedRows = new ArrayList<>(unsortedExpectedRows); + expectedRows.sort(Comparator.comparingLong(row -> row.get(0))); + JsonNode resultTable = response.get("resultTable"); + assertNotNull(resultTable, "Missing result table in response: " + response); + JsonNode rows = resultTable.get("rows"); + assertNotNull(rows, "Missing rows in response: " + response); + assertEquals(rows.size(), expectedRows.size(), "Unexpected number of rows: " + rows); + for (int i = 0; i < expectedRows.size(); i++) { + List expectedRow = expectedRows.get(i); + JsonNode row = rows.get(i); + assertEquals(row.size(), expectedRow.size(), "Unexpected number of columns in row: " + row); + for (int j = 0; j < expectedRow.size(); j++) { + assertEquals(row.get(j).asLong(), (long) expectedRow.get(j), + "Unexpected value at row " + i + " column " + j + " in rows: " + rows); + } + } + } + + /// Asserts on every leaf stage of the executed plan, i.e. on every `MAILBOX_SEND` node of the `stageStats` tree whose + /// only child is a `LEAF` node. `expectedNumWorkers` is the number of partition classes the colocated group kept, + /// read from the send's summed `parallelism`; `expectedNumSegments` is lower than it exactly because some workers had + /// nothing to scan. + private static void assertLeafStages(JsonNode response, int expectedNumLeafStages, int expectedNumWorkers, + int expectedNumSegments) { + JsonNode stageStats = response.get("stageStats"); + assertNotNull(stageStats, "Missing stage stats in response: " + response); + List leafStageSends = new ArrayList<>(); + collectLeafStageSends(stageStats, leafStageSends); + assertEquals(leafStageSends.size(), expectedNumLeafStages, + "Unexpected number of leaf stages in stage stats: " + stageStats.toPrettyString()); + for (JsonNode leafStageSend : leafStageSends) { + assertEquals(leafStageSend.path("parallelism").asInt(-1), expectedNumWorkers, + "Unexpected leaf stage worker count, so the colocated group did not keep the expected partition classes. " + + "Stage stats: " + stageStats.toPrettyString()); + // A pre-partitioned send is wired 1-to-1, so each sender writes exactly one receive mailbox. A shuffle would make + // each sender write one mailbox per receiver worker. + assertEquals(leafStageSend.path("fanOut").asInt(-1), 1, + "Leaf stage send is not 1-to-1, so the plan fell back to a shuffle. Stage stats: " + + stageStats.toPrettyString()); + JsonNode leaf = leafStageSend.get("children").get(0); + assertEquals(leaf.path("numSegmentsQueried").asInt(-1), expectedNumSegments, + "Unexpected number of segments queried by the leaf stage. Stage stats: " + stageStats.toPrettyString()); + } + } + + private static void collectLeafStageSends(JsonNode node, List leafStageSends) { + JsonNode children = node.get("children"); + if ("MAILBOX_SEND".equals(node.path("type").asText()) && children != null && children.size() == 1 && "LEAF".equals( + children.get(0).path("type").asText())) { + leafStageSends.add(node); + return; + } + if (children != null) { + for (JsonNode child : children) { + collectLeafStageSends(child, leafStageSends); + } + } + } + + /// Asserts that the planner wired the leaf stages into direct (1-to-1) exchanges rather than shuffles, by reading the + /// physical plan: `MailboxSendNode#explain` marks a pre-partitioned send with `[PARTITIONED]`, and the physical + /// explain prints one such line per leaf worker together with the receiver mailboxes it targets. + private void assertDirectExchanges(String query, int expectedNumLeafStages, int expectedNumWorkers) + throws Exception { + JsonNode response = queryBrokerHttpEndpoint("EXPLAIN IMPLEMENTATION PLAN FOR " + query); + assertNoExceptions(response); + JsonNode rows = response.get("resultTable").get("rows"); + assertNotNull(rows, "Missing rows in explain response: " + response); + StringBuilder planBuilder = new StringBuilder(); + for (JsonNode row : rows) { + for (JsonNode cell : row) { + planBuilder.append(cell.asText()).append('\n'); + } + } + String plan = planBuilder.toString(); + assertFalse(plan.isEmpty(), "Empty implementation plan for query: " + query); + + int numPrePartitionedSends = 0; + Matcher matcher = PRE_PARTITIONED_SEND_PATTERN.matcher(plan); + while (matcher.find()) { + numPrePartitionedSends++; + String senderWorkerId = matcher.group(1); + String receivers = matcher.group(2); + // One receiver mailbox, and it is the receiver worker with the same id: that is the direct exchange. A shuffle + // would list every receiver worker here. + assertFalse(receivers.contains(","), + "Pre-partitioned send targets more than one receiver mailbox, so the exchange is not 1-to-1. Plan:\n" + plan); + assertTrue(receivers.endsWith("|[" + senderWorkerId + "]"), + "Pre-partitioned send from worker " + senderWorkerId + " targets receiver " + receivers + + " instead of the receiver worker with the same id. Plan:\n" + plan); + } + assertEquals(numPrePartitionedSends, expectedNumLeafStages * expectedNumWorkers, + "Unexpected number of pre-partitioned mailbox sends (one per leaf worker is expected) in plan:\n" + plan); + } + + private void setUpTable(String tableName, List populatedPartitions, int metricMultiplier) + throws Exception { + Schema schema = createSchemaForTable(tableName); + addSchema(schema); + TableConfig tableConfig = createTableConfigForTable(tableName); + addTableConfig(tableConfig); + + // The segment directories are shared across tables, and uploadSegments pushes everything it finds in the tar one. + TestUtils.ensureDirectoriesExistAndEmpty(_segmentDir, _tarDir); + int segmentIndex = 0; + for (int partition : populatedPartitions) { + // One segment per partition, so that every segment holds exactly one partition id (a segment spanning several has + // no usable partition metadata) and every partition has a fully replicated server. + File avroFile = createAvroFile(tableName, partition, metricMultiplier); + ClusterIntegrationTestUtils.buildSegmentFromAvro(avroFile, tableConfig, schema, segmentIndex++, _segmentDir, + _tarDir); + } + uploadSegments(tableName, _tarDir); + } + + private static Schema createSchemaForTable(String tableName) { + return new Schema.SchemaBuilder().setSchemaName(tableName) + .addSingleValueDimension(PARTITION_KEY_COLUMN, FieldSpec.DataType.INT) + .addMetric(METRIC_COLUMN, FieldSpec.DataType.INT) + .addDateTime(TIMESTAMP_FIELD_NAME, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + } + + private static TableConfig createTableConfigForTable(String tableName) { + return new TableConfigBuilder(TableType.OFFLINE).setTableName(tableName) + .setTimeColumnName(TIMESTAMP_FIELD_NAME) + // Replicate every segment on both servers of the shared cluster, so that each partition has both of them as + // fully replicated servers and a zero-segment worker deterministically lands on the server its peer picked. + .setNumReplicas(2) + .setSegmentPartitionConfig(new SegmentPartitionConfig( + Map.of(PARTITION_KEY_COLUMN, new ColumnPartitionConfig(PARTITION_FUNCTION, NUM_DECLARED_PARTITIONS)))) + .build(); + } + + private File createAvroFile(String tableName, int partition, int metricMultiplier) + throws IOException { + var avroSchema = SchemaBuilder.record("record") + .fields() + .name(PARTITION_KEY_COLUMN).type().intType().noDefault() + .name(METRIC_COLUMN).type().intType().noDefault() + .name(TIMESTAMP_FIELD_NAME).type().longType().noDefault() + .endRecord(); + File avroFile = new File(_tempDir, tableName + "_partition_" + partition + ".avro"); + try (DataFileWriter fileWriter = new DataFileWriter<>(new GenericDatumWriter<>(avroSchema))) { + fileWriter.create(avroSchema, avroFile); + for (int key : keysForPartition(partition)) { + GenericData.Record record = new GenericData.Record(avroSchema); + record.put(PARTITION_KEY_COLUMN, key); + record.put(METRIC_COLUMN, key * metricMultiplier); + record.put(TIMESTAMP_FIELD_NAME, 1_600_000_000_000L + key); + fileWriter.append(record); + } + } + return avroFile; + } + + /// Returns the join keys that land in the given partition: `Modulo` maps a key to `key % numPartitions`. + private static int[] keysForPartition(int partition) { + int[] keys = new int[NUM_ROWS_PER_PARTITION]; + for (int i = 0; i < NUM_ROWS_PER_PARTITION; i++) { + keys[i] = partition + i * NUM_DECLARED_PARTITIONS; + } + return keys; + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java index 635e9a1ab9bc..f50a3db9e1c3 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java @@ -60,6 +60,7 @@ public class DispatchablePlanContext { private final Map _dispatchablePlanMetadataMap = new HashMap<>(); private final Map _dispatchablePlanStageRootMap = new HashMap<>(); + private final Map _partitionTableInfoCache = new HashMap<>(); private long _numSegmentsPrunedByBroker; private int _leafStagesAssigned; private int _leafStagesEmpty; @@ -133,6 +134,13 @@ public Map getDispatchablePlanStageRootMap() { return _dispatchablePlanStageRootMap; } + /// The partition layout of each partitioned table scanned by this query, keyed by table name. Read from the routing + /// manager once per table so that the colocation pre-pass and every leaf stage scanning the same table (e.g. both + /// sides of a self-join) see one snapshot. The value is opaque here: [WorkerManager] builds and interprets it. + public Map getPartitionTableInfoCache() { + return _partitionTableInfoCache; + } + public long getNumSegmentsPrunedByBroker() { return _numSegmentsPrunedByBroker; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanFragment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanFragment.java index c4bd3462f0c6..2f8752b4669d 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanFragment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanFragment.java @@ -54,15 +54,17 @@ public DispatchablePlanFragment(PlanFragment planFragment) { } /// Returns a copy of `original` with its plan fragment root replaced by `newRoot`. - /// Worker metadata and server-instance mapping are shallow-copied so the new fragment is + /// Worker metadata, server-instance mapping and the worker-to-segments map are shallow-copied so the new fragment is /// independent of the original. public static DispatchablePlanFragment copyWithRoot(DispatchablePlanFragment original, PlanNode newRoot) { int fragmentId = original.getPlanFragment().getFragmentId(); - return new DispatchablePlanFragment( + DispatchablePlanFragment copy = new DispatchablePlanFragment( new PlanFragment(fragmentId, newRoot, List.of()), new ArrayList<>(original.getWorkerMetadataList()), new HashMap<>(original.getServerInstanceToWorkerIdMap()), new HashMap<>(original.getCustomProperties())); + copy.setWorkerIdToSegmentsMap(original.getWorkerIdToSegmentsMap()); + return copy; } public DispatchablePlanFragment(PlanFragment planFragment, List workerMetadataList, diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java index 995aa2261f5a..c47b434110b5 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java @@ -73,6 +73,10 @@ public class DispatchablePlanMetadata implements Serializable { private TimeBoundaryInfo _timeBoundaryInfo; private int _partitionParallelism = 1; private final Map> _tableToUnavailableSegmentsMap = new HashMap<>(); + // Broker-local, never serialized: see getPartitionClassIds() + private transient int[] _partitionClassIds; + // Broker-local, never serialized: see getPaddedClassCandidates() + private transient Map> _paddedClassCandidates; // Calculated in {@link MailboxAssignmentVisitor} // Map from workerId -> {planFragmentId -> mailboxes} @@ -173,6 +177,51 @@ public void setPartitionParallelism(int partitionParallelism) { _partitionParallelism = partitionParallelism; } + /// Returns the partition classes this stage's worker ids stand for, in worker-id order, or `null` when the worker ids + /// are not in partition-class space. + /// + /// A partition class is the set of partitions that share one worker: with a hinted partition size of `w`, class `j` + /// holds every partition `p` where `p % w == j`. Across a direct (1-to-1) exchange the worker id is the only carrier + /// of partition identity -- the wiring pairs sender worker `k` with receiver worker `k` and checks nothing about the + /// data behind them -- so equal worker counts are no evidence that two stages agree: had one dropped its empty class + /// 1 and the other its empty class 2, both would still have `w - 1` workers, and worker 1 would pair class 2 with + /// class 1, losing rows with no error. `WorkerManager` therefore computes one class list per colocated group, + /// dropping only the classes no member of the group holds data in, and shares that same array with every stage of it. + /// A leaf stage gets one worker per entry, i.e. worker `k` handles class `[k]`; an intermediate stage with a + /// partition parallelism of `p` gets `p` workers per entry, i.e. worker `k` handles class `[k / p]`, the same fan-out + /// the exchange performs. + /// + /// `null` means the worker ids are not partition classes (e.g. a stage assigned over candidate servers, or a + /// singleton reducer), or that the stage's group was not reduced, in which case worker `k` maps to class `k` as + /// before. + /// + /// Broker-local planning state: not serialized to the servers, and must not be mutated (the same array instance is + /// shared by every stage of the group). + @Nullable + public int[] getPartitionClassIds() { + return _partitionClassIds; + } + + public void setPartitionClassIds(@Nullable int[] partitionClassIds) { + _partitionClassIds = partitionClassIds; + } + + /// Returns the partition classes of [#getPartitionClassIds()] that this stage holds no data in, mapped to the servers + /// its colocated group expects the (empty) worker of that class to be picked from, or `null` when this stage has + /// nothing to pad. Only ever set together with [#getPartitionClassIds()], by the same producer; see + /// `WorkerManager#assignPaddedWorker`, which is where such a worker and its candidate servers are used. + /// + /// Broker-local planning state, like [#getPartitionClassIds()]: neither the map nor the server sets in it must be + /// mutated (the sets may be the ones the broker publishes its partition metadata with). + @Nullable + public Map> getPaddedClassCandidates() { + return _paddedClassCandidates; + } + + public void setPaddedClassCandidates(@Nullable Map> paddedClassCandidates) { + _paddedClassCandidates = paddedClassCandidates; + } + public Map> getTableToUnavailableSegmentsMap() { return _tableToUnavailableSegmentsMap; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitor.java index 41076e49d0ed..a1e0363e6086 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitor.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitor.java @@ -20,6 +20,7 @@ import com.google.common.base.Preconditions; import java.util.ArrayList; +import java.util.Arrays; import java.util.Comparator; import java.util.HashMap; import java.util.List; @@ -70,14 +71,14 @@ public Void process(PlanNode node, DispatchablePlanContext context) { } int parallelism = numReceivers / numSenders; computeDirectExchange(senderMailboxesMap, receiverMailboxesMap, senderStageId, receiverStageId, - senderServerMap, receiverServerMap, numSenders, parallelism); + senderServerMap, receiverServerMap, numSenders, parallelism, senderMetadata, receiverMetadata); } else if (senderMetadata.isPrePartitioned() && isDirectExchangeCompatible(senderMetadata, receiverMetadata)) { // Direct exchange: the data is already pre-partitioned, so send it 1-to-1 to the worker with the same worker // id (with parallelism, fan out each sender worker to a contiguous range of receiver workers). The // co-location handling is the same as SINGLETON, see computeDirectExchange. int parallelism = numReceivers / numSenders; computeDirectExchange(senderMailboxesMap, receiverMailboxesMap, senderStageId, receiverStageId, - senderServerMap, receiverServerMap, numSenders, parallelism); + senderServerMap, receiverServerMap, numSenders, parallelism, senderMetadata, receiverMetadata); } else { // For other exchange types, send the data to all the instances in the receiver fragment // TODO: Add support for more exchange types @@ -104,10 +105,16 @@ public Void process(PlanNode node, DispatchablePlanContext context) { /// partition to a different replica, leaving worker `i` on different servers. Rather than failing the query, we fall /// back to a cross-server send: the exchange stays correct because worker id still maps to the same partition on both /// sides, and we only lose locality (one extra network hop) until routing re-stabilizes. + /// + /// A sender worker with no segment to scan is wired like any other one: it is dispatched regardless, so leaving it + /// out of the receiver's mailbox map would strand its stage stats, and any error it reports, in a mailbox nobody + /// reads. private void computeDirectExchange(Map> senderMailboxesMap, Map> receiverMailboxesMap, Integer senderStageId, Integer receiverStageId, Map senderServerMap, Map receiverServerMap, - int numSenders, int parallelism) { + int numSenders, int parallelism, DispatchablePlanMetadata senderMetadata, + DispatchablePlanMetadata receiverMetadata) { + checkPartitionClassAgreement(senderMetadata, receiverMetadata, senderStageId, receiverStageId); if (parallelism == 1) { // 1-to-1 mapping for (int workerId = 0; workerId < numSenders; workerId++) { @@ -153,6 +160,29 @@ private void computeDirectExchange(Map> send } } + /// Fails when the two sides of a direct exchange do not agree on the partition classes their worker ids stand for. + /// [#computeDirectExchange] pairs sender worker `k` with receiver worker `k` and checks nothing about the data behind + /// them, so equal worker counts are no evidence of agreement -- see + /// [DispatchablePlanMetadata#getPartitionClassIds()]. `WorkerManager` shares one class list across every stage of a + /// colocated group, so this can only trip if that invariant regresses. + /// + /// A `null` list means that side's worker ids are not partition classes at all (e.g. a stage assigned over candidate + /// servers, or a singleton reducer) and makes no claim to compare against, so only two class-space sides are checked. + private static void checkPartitionClassAgreement(DispatchablePlanMetadata senderMetadata, + DispatchablePlanMetadata receiverMetadata, int senderStageId, int receiverStageId) { + int[] senderPartitionClassIds = senderMetadata.getPartitionClassIds(); + int[] receiverPartitionClassIds = receiverMetadata.getPartitionClassIds(); + if (senderPartitionClassIds == null || receiverPartitionClassIds == null) { + return; + } + Preconditions.checkState(Arrays.equals(senderPartitionClassIds, receiverPartitionClassIds), + "Partition class mismatch for the direct exchange from stage: %s to stage: %s, sender: %s vs receiver: %s", + senderStageId, receiverStageId, Arrays.toString(senderPartitionClassIds), + Arrays.toString(receiverPartitionClassIds)); + } + + /// Wires one sender worker of a direct exchange to the contiguous range of `parallelism` receiver workers it fans out + /// to. See [#computeDirectExchange]. private void computeDirectExchangeWithParallelism(Map> senderMailboxesMap, Map> receiverMailboxesMap, Integer senderStageId, Integer receiverStageId, int senderWorkerId, int receiverWorkerId, QueryServerInstance senderServer, QueryServerInstance receiverServer, @@ -178,12 +208,26 @@ private static boolean isDirectExchangeCompatible(DispatchablePlanMetadata sende if (numSenders * sender.getPartitionParallelism() != numReceivers) { return false; } + // A sender whose worker ids stand for partition classes may only be wired 1-to-1 to a receiver whose worker ids + // stand for the same ones: without a class list the receiver took its workers from the candidate servers, so equal + // worker counts would be a coincidence. The shuffle fallback is safe -- connectWorkers re-hashes across any worker + // count. + if (!Arrays.equals(sender.getPartitionClassIds(), receiver.getPartitionClassIds())) { + return false; + } if (sender.getPartitionFunction() == null) { return receiver.getPartitionFunction() == null; } return sender.getPartitionFunction().equalsIgnoreCase(receiver.getPartitionFunction()); } + /// Wires one side of a shuffled exchange: every worker of `stageId` (the source, sized by `serverMap`) becomes a + /// mailbox of every one of the `numWorkers` workers on the other side. + /// + /// NOTE: The source stage may have no worker at all (an empty or fully-pruned leaf), in which case every worker on + /// the other side still gets an entry holding an empty mailbox list -- that is what lets the other side's send + /// operator resolve this stage and its receive operator return end-of-stream at once, so do not short-circuit it + /// away. private void connectWorkers(int stageId, Map serverMap, Map> mailboxesMap, int numWorkers) { Map> serverToWorkerIdsMap = new HashMap<>(); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java new file mode 100644 index 000000000000..88b43fb61e7d --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java @@ -0,0 +1,253 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.routing; + +import java.util.ArrayDeque; +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.Queue; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.rel.RelDistribution; +import org.apache.pinot.query.planner.PlanFragment; +import org.apache.pinot.query.planner.physical.DispatchablePlanMetadata; +import org.apache.pinot.query.planner.plannode.MailboxSendNode; +import org.apache.pinot.query.planner.plannode.PlanNode; + + +/// Finds the groups of plan fragments that are tied together by direct (1-to-1) exchanges, so that [WorkerManager] can +/// give all the members of a group the same worker-id-to-partition-class mapping. Every member must drop exactly the +/// same classes, because the worker id is the only carrier of partition identity across such an exchange -- see +/// [DispatchablePlanMetadata#getPartitionClassIds()], which is what a group ends up sharing. +/// +/// The relation used to form the groups deliberately over-approximates: an edge is added for every send that *may* be +/// wired 1-to-1, that is a SINGLETON send (Pinot's representation of a local exchange) or a send from a pre-partitioned +/// stage, without checking the worker counts that ultimately decide it. That is always safe, because merging groups can +/// only shrink the set of classes a group is allowed to drop. A SINGLETON send counts even when the sender is not +/// marked pre-partitioned (a lookup join's local exchange, say), because the receiver still copies its worker map from +/// it. +/// +/// This class looks only at the plan shape and the table hints. Which classes actually hold data -- and therefore which +/// ones survive -- is resolved by [WorkerManager], which owns the routing information. +class ColocationGroupAnalyzer { + private ColocationGroupAnalyzer() { + } + + /// Returns the groups whose worker count may be reduced to the partition classes that survive. A group that does not + /// qualify (see [#toReducibleGroup]) is omitted entirely, keeping the existing assignment for every fragment in it. + static List findReducibleGroups(PlanFragment rootFragment, + Map metadataMap) { + Map fragmentMap = collectFragments(rootFragment); + Map parents = new HashMap<>(); + Set fragmentsWithUnsafePrePartitionedSend = new HashSet<>(); + Set fragmentsWithShuffledInput = new HashSet<>(); + for (PlanFragment fragment : fragmentMap.values()) { + PlanNode fragmentRoot = fragment.getFragmentRoot(); + if (!(fragmentRoot instanceof MailboxSendNode)) { + // Only the root (broker reduce) fragment, which has no send node and therefore no outgoing edge. + continue; + } + MailboxSendNode sendNode = (MailboxSendNode) fragmentRoot; + int senderFragmentId = fragment.getFragmentId(); + DispatchablePlanMetadata senderMetadata = metadataMap.get(senderFragmentId); + RelDistribution.Type distributionType = sendNode.getDistributionType(); + boolean prePartitioned = senderMetadata != null && senderMetadata.isPrePartitioned(); + if (distributionType != RelDistribution.Type.SINGLETON && !prePartitioned) { + // The data is shuffled, so the receiver re-hashes it across any worker count and the two sides need not agree + // on what a worker id stands for. Remember the receivers though: a shuffled sender hashes its rows over the + // receiver's worker count, so reducing that count moves a row to a different worker than the one the 1-to-1 + // side delivers that row's class to, and rows with the same key stop meeting. + for (int receiverFragmentId : sendNode.getReceiverStageIds()) { + fragmentsWithShuffledInput.add(receiverFragmentId); + } + continue; + } + if (prePartitioned && distributionType != RelDistribution.Type.SINGLETON + && distributionType != RelDistribution.Type.HASH_DISTRIBUTED) { + // A pre-partitioned BROADCAST (or RANDOM) send is wired 1-to-1 whenever the worker counts happen to line up, + // which is wrong for BROADCAST: the receiver would see one sender's slice instead of every row. Today an empty + // partition aborts such a plan, so leave the whole group alone rather than making that path reachable by + // reducing the worker count into a match. + fragmentsWithUnsafePrePartitionedSend.add(senderFragmentId); + } + for (int receiverFragmentId : sendNode.getReceiverStageIds()) { + union(parents, senderFragmentId, receiverFragmentId); + } + } + + // Bucket the fragments by the representative of their connected component. + Map> groupMembers = new HashMap<>(); + for (Integer fragmentId : fragmentMap.keySet()) { + groupMembers.computeIfAbsent(find(parents, fragmentId), k -> new ArrayList<>()).add(fragmentId); + } + + List reducibleGroups = new ArrayList<>(); + for (List members : groupMembers.values()) { + if (!Collections.disjoint(members, fragmentsWithUnsafePrePartitionedSend) + || !Collections.disjoint(members, fragmentsWithShuffledInput)) { + continue; + } + ColocationGroup group = toReducibleGroup(members, fragmentMap, metadataMap); + if (group != null) { + reducibleGroups.add(group); + } + } + return reducibleGroups; + } + + /// Collects every fragment reachable from the given root, keyed by fragment id. With spools the plan is a DAG rather + /// than a tree (the same fragment is a child of every receiver that reads the spool), so a fragment is collected + /// once. + private static Map collectFragments(PlanFragment rootFragment) { + Map fragmentMap = new HashMap<>(); + Queue pending = new ArrayDeque<>(); + pending.add(rootFragment); + while (!pending.isEmpty()) { + PlanFragment fragment = pending.poll(); + if (fragmentMap.put(fragment.getFragmentId(), fragment) != null) { + continue; + } + pending.addAll(fragment.getChildren()); + } + return fragmentMap; + } + + /// Classifies the members of one connected component and returns the group when its worker count may be reduced, or + /// `null` when it must keep today's assignment. A lone fragment is tied to nothing, so its worker ids owe no + /// agreement to another stage and it keeps that assignment; beyond that, and beyond holding a partitioned leaf to + /// reduce at all, a group qualifies only when: + /// + /// - all of its partitioned leaves share the same hinted partition size, function and parallelism, so a worker id + /// means the same class on all of them. The function matters as much as the size: class `j` of a `Murmur` + /// partitioned table and class `j` of a `HashCode` one hold different keys, so unioning their empty classes would + /// union two different class spaces; + /// - none of its leaves is assigned over servers rather than partitions. Such a leaf (the `is_colocated_by_join_keys` + /// escape hatch on a table without partition metadata) gets one worker per server, so changing the worker count of + /// its partitioned peers would change whether the exchange between them is wired 1-to-1. + @Nullable + private static ColocationGroup toReducibleGroup(List members, Map fragmentMap, + Map metadataMap) { + if (members.size() < 2) { + return null; + } + List partitionedLeafFragmentIds = new ArrayList<>(); + int partitionSize = -1; + int partitionParallelism = -1; + String partitionFunction = null; + for (Integer fragmentId : members) { + DispatchablePlanMetadata metadata = metadataMap.get(fragmentId); + if (metadata == null || !WorkerManager.isLeafPlan(metadata)) { + // An intermediate stage derives its worker map from a child (local exchange or pre-partitioned assignment) or + // is assigned over candidate servers. Either way it constrains no class, and WorkerManager copies the class + // list onto it when it derives its map from a member that has one. + continue; + } + PlanFragment fragment = fragmentMap.get(fragmentId); + if (WorkerManager.isLookupJoin(fragment.getChildren())) { + // The workers come from the single local exchange child, so the fragment's own table hints are ignored. + continue; + } + Map tableOptions = metadata.getTableOptions(); + if (tableOptions == null) { + return null; + } + if (LeafPartitionHints.isReplicated(tableOptions)) { + // Constrains no class either, see LeafPartitionHints#isReplicated. + continue; + } + LeafPartitionHints hints; + try { + hints = LeafPartitionHints.resolve(tableOptions); + } catch (IllegalStateException e) { + // Invalid hints. Leave the group alone so that the leaf assignment reports them. + return null; + } + if (hints.getPartitionKey() == null) { + return null; + } + String leafPartitionFunction = hints.getHintedPartitionFunction(); + if (partitionedLeafFragmentIds.isEmpty()) { + partitionSize = hints.getPartitionSize(); + partitionParallelism = hints.getPartitionParallelism(); + partitionFunction = leafPartitionFunction; + } else if (partitionSize != hints.getPartitionSize() + || partitionParallelism != hints.getPartitionParallelism() + || !isSamePartitionFunction(partitionFunction, leafPartitionFunction)) { + return null; + } + partitionedLeafFragmentIds.add(fragmentId); + } + if (partitionedLeafFragmentIds.isEmpty()) { + return null; + } + return new ColocationGroup(partitionSize, partitionedLeafFragmentIds); + } + + /// Compares two `partition_function` hints the way the rest of the engine compares partition function names: + /// case-insensitively, with a missing hint matching only another missing one (see + /// `MailboxAssignmentVisitor#isDirectExchangeCompatible`). Comparing the hints rather than the resolved names (see + /// [LeafPartitionHints#getPartitionFunction()]) is the stricter choice; it only leaves more groups alone, which costs + /// nothing but the reduction. + private static boolean isSamePartitionFunction(@Nullable String first, @Nullable String second) { + return first != null ? first.equalsIgnoreCase(second) : second == null; + } + + private static void union(Map parents, int first, int second) { + int firstRoot = find(parents, first); + int secondRoot = find(parents, second); + if (firstRoot != secondRoot) { + parents.put(firstRoot, secondRoot); + } + } + + private static int find(Map parents, int fragmentId) { + int root = fragmentId; + Integer parent = parents.get(root); + while (parent != null && parent != root) { + root = parent; + parent = parents.get(root); + } + // Path compression. + int current = fragmentId; + while (current != root) { + Integer next = parents.put(current, root); + assert next != null; + current = next; + } + return root; + } + + /// A set of plan fragments whose worker ids must all stand for the same partition class, together with the hinted + /// partition layout they share. + static class ColocationGroup { + /// The number of partition classes, and of workers before reduction, i.e. the hinted `partition_size`. + final int _partitionSize; + /// The members whose data decides which classes survive. + final List _partitionedLeafFragmentIds; + + ColocationGroup(int partitionSize, List partitionedLeafFragmentIds) { + _partitionSize = partitionSize; + _partitionedLeafFragmentIds = partitionedLeafFragmentIds; + } + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/LeafPartitionHints.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/LeafPartitionHints.java new file mode 100644 index 000000000000..03cfa3904850 --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/LeafPartitionHints.java @@ -0,0 +1,118 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.routing; + +import com.google.common.base.Preconditions; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.calcite.rel.hint.PinotHintOptions.TableHintOptions; + + +/// The partition layout hinted on one leaf stage's table, resolved in a single place so that every consumer resolves it +/// the same way. Agreement is load-bearing rather than cosmetic: [WorkerManager] gives worker `k` of a partitioned leaf +/// the `k`-th surviving partition class of the leaf's colocated group, and [ColocationGroupAnalyzer] decides that class +/// list from these same hints, so a `partition_size` resolved two ways would make worker `k` stand for a different +/// partition on each side of a 1-to-1 exchange. +class LeafPartitionHints { + private static final String DEFAULT_PARTITION_FUNCTION = "Murmur"; + + @Nullable + private final String _partitionKey; + private final int _partitionSize; + private final int _partitionParallelism; + @Nullable + private final String _hintedPartitionFunction; + + private LeafPartitionHints(@Nullable String partitionKey, int partitionSize, int partitionParallelism, + @Nullable String hintedPartitionFunction) { + _partitionKey = partitionKey; + _partitionSize = partitionSize; + _partitionParallelism = partitionParallelism; + _hintedPartitionFunction = hintedPartitionFunction; + } + + /// Resolves the partition hints of a leaf stage from its table hint options. A hint that cannot be used, including a + /// non-numeric `partition_size` or `partition_parallelism`, is reported as [IllegalStateException] so that a caller + /// which wants to degrade instead of failing (see [ColocationGroupAnalyzer]) has a single type to catch. + static LeafPartitionHints resolve(Map tableOptions) { + // Resolved for a non-partitioned leaf too, because it also sizes the workers of its local exchange. + int partitionParallelism = parsePositive(tableOptions, TableHintOptions.PARTITION_PARALLELISM, 1); + String partitionKey = tableOptions.get(TableHintOptions.PARTITION_KEY); + if (partitionKey == null) { + // Not a partitioned leaf, so the rest of the hints say nothing about it and are deliberately left unresolved. + return new LeafPartitionHints(null, -1, partitionParallelism, null); + } + int partitionSize = parsePositive(tableOptions, TableHintOptions.PARTITION_SIZE, -1); + Preconditions.checkState(partitionSize > 0, "'%s' must be provided for partition key: %s", + TableHintOptions.PARTITION_SIZE, partitionKey); + return new LeafPartitionHints(partitionKey, partitionSize, partitionParallelism, + tableOptions.get(TableHintOptions.PARTITION_FUNCTION)); + } + + /// Returns whether the given table hint options declare the table replicated across all workers. Such a leaf holds + /// every segment on every worker and takes its worker map from its peer, so no partition hint applies to it. + static boolean isReplicated(Map tableOptions) { + return Boolean.parseBoolean(tableOptions.get(TableHintOptions.IS_REPLICATED)); + } + + /// Returns the hinted partition key, or `null` when the leaf is not partitioned, in which case the partition size and + /// function are meaningless. + @Nullable + String getPartitionKey() { + return _partitionKey; + } + + /// Returns the number of partition classes, and of workers before any reduction, i.e. the hinted `partition_size`. + /// Positive when [#getPartitionKey()] is non-null, -1 otherwise. + int getPartitionSize() { + return _partitionSize; + } + + int getPartitionParallelism() { + return _partitionParallelism; + } + + /// Returns the partition function to use, i.e. the hinted one or `Murmur` when the hint is absent. + String getPartitionFunction() { + return _hintedPartitionFunction != null ? _hintedPartitionFunction : DEFAULT_PARTITION_FUNCTION; + } + + /// Returns the `partition_function` hint exactly as given, i.e. `null` when it is absent. Unlike + /// [#getPartitionFunction()], which fills the default in, so comparing two leaves through this never lets an omitted + /// hint match an explicit one. + @Nullable + String getHintedPartitionFunction() { + return _hintedPartitionFunction; + } + + private static int parsePositive(Map tableOptions, String option, int defaultValue) { + String value = tableOptions.get(option); + if (value == null) { + return defaultValue; + } + int parsed; + try { + parsed = Integer.parseInt(value); + } catch (NumberFormatException e) { + throw new IllegalStateException("'" + option + "' must be a positive integer, got: " + value); + } + Preconditions.checkState(parsed > 0, "'%s' must be positive, got: %s", option, parsed); + return parsed; + } +} 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 9ea02f0e9a43..a7031c5c4dc4 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 @@ -21,7 +21,9 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.collect.Maps; +import it.unimi.dsi.fastutil.ints.IntArrayList; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collection; import java.util.Comparator; import java.util.HashMap; @@ -31,11 +33,11 @@ import java.util.Map; import java.util.Random; import java.util.Set; +import java.util.TreeSet; import javax.annotation.Nullable; import org.apache.calcite.rel.RelDistribution; import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.collections4.MapUtils; -import org.apache.pinot.calcite.rel.hint.PinotHintOptions; import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType; import org.apache.pinot.calcite.rel.rules.ImmutableTableOptions; import org.apache.pinot.calcite.rel.rules.TableOptions; @@ -73,8 +75,6 @@ public class WorkerManager { private static final Random RANDOM = new Random(); // default shuffle method in v2 private static final String DEFAULT_SHUFFLE_PARTITION_FUNCTION = "AbsHashCodeSum"; - // default table partition function if not specified in hint - private static final String DEFAULT_TABLE_PARTITION_FUNCTION = "Murmur"; private final String _instanceId; private final String _hostName; @@ -111,25 +111,211 @@ public void assignWorkers(PlanFragment rootFragment, DispatchablePlanContext con metadata.setWorkerIdToServerInstanceMap( Map.of(0, new QueryServerInstance(_instanceId, _hostName, _port, _port))); + // Pre-pass: decide which partition classes get a worker, for every colocated group of fragments. It must run before + // any assignment: a group's leaves have to agree on the class list, and a leaf cannot see its peers while assigned. + assignPartitionClasses(rootFragment, context); + // Two-pass assignment: leaf stages must be assigned first so that the candidate server information // (_nonLookupTables or _leafServerInstances) is fully populated before intermediate stages use it. // Without this, literal-only stages (e.g. UNION ALL of constants) that are traversed before any table scan // would see an empty candidate set and fall back to all enabled servers across all tenants. + // Each pass gets its own visited set: with spools the plan is a DAG rather than a tree (the same PlanFragment is a + // child of every receiver that reads the spool), so a fragment must be assigned exactly once per pass, and one + // skipped by the first pass still has to be assigned by the second. + Set visitedInLeafPass = new HashSet<>(); for (PlanFragment child : rootFragment.getChildren()) { - assignWorkersToNonRootFragment(child, context, true); + assignWorkersToNonRootFragment(child, context, true, visitedInLeafPass); } + Set visitedInIntermediatePass = new HashSet<>(); for (PlanFragment child : rootFragment.getChildren()) { - assignWorkersToNonRootFragment(child, context, false); + assignWorkersToNonRootFragment(child, context, false, visitedInIntermediatePass); + } + } + + /// Decides which partition classes get a worker, for every colocated group of fragments that may be reduced, and + /// publishes the decision on each partitioned leaf of the group (see + /// [DispatchablePlanMetadata#getPartitionClassIds()] and [DispatchablePlanMetadata#getPaddedClassCandidates()]). + /// + /// A class survives when *any* member holds a segment in it: the union, not the intersection, because a class that + /// holds data for one member must keep its worker on every member or the members stop agreeing on what a worker id + /// stands for. A member holding no data in a surviving class gets a worker with no segments (see + /// [#assignPaddedWorker]). Emptiness is computed in class space (`0..partitionSize-1`) rather than over raw partition + /// ids because members may declare different partition counts; a member carrying no per-class visibility (replicated, + /// non-partitioned, or deriving its worker map from a peer) contributes nothing to the union. + /// + /// Marking no group keeps the assignment as it is without one: every class gets a worker, so a class holding no + /// segment fails the assignment instead of being dropped or padded. + private void assignPartitionClasses(PlanFragment rootFragment, DispatchablePlanContext context) { + Map metadataMap = context.getDispatchablePlanMetadataMap(); + Map partitionTableInfoCache = context.getPartitionTableInfoCache(); + for (ColocationGroupAnalyzer.ColocationGroup group : ColocationGroupAnalyzer.findReducibleGroups(rootFragment, + metadataMap)) { + int numWorkers = group._partitionSize; + List memberFragmentIds = group._partitionedLeafFragmentIds; + // The servers each member can scan each class on, in the same order as the member fragment ids. + List>> memberClassServers = new ArrayList<>(memberFragmentIds.size()); + // Allocated lazily, once the first member has checked the hint against its table: numWorkers is the raw hinted + // partition size, so sizing anything from it before that check would let a bogus hint allocate unboundedly. The + // check also bounds it by the table's partition count. + boolean[] survivingClasses = null; + boolean reducible = true; + for (Integer fragmentId : memberFragmentIds) { + DispatchablePlanMetadata metadata = metadataMap.get(fragmentId); + String tableName = metadata.getScannedTables().get(0); + // NOTE: A failure here is the same one the leaf assignment would hit for this table, only raised earlier. + PartitionTableInfo partitionTableInfo = + partitionTableInfoCache.computeIfAbsent(tableName, this::calculatePartitionTableInfo); + int numPartitions = partitionTableInfo._partitionInfoMap.length; + if (numPartitions == 0 || numPartitions % numWorkers != 0) { + // The table does not match the hinted partition size. Leave the group alone so that checkPartitionInfoMap + // reports it during the leaf assignment. + reducible = false; + break; + } + if (survivingClasses == null) { + survivingClasses = new boolean[numWorkers]; + } + List> classServers = collectClassServers(partitionTableInfo._partitionInfoMap, numWorkers); + boolean anyPopulated = false; + for (int classId = 0; classId < numWorkers; classId++) { + if (classServers.get(classId) != null) { + survivingClasses[classId] = true; + anyPopulated = true; + } + } + // A member holding no data at all leaves nothing to assign: no class to place its single empty worker in, and + // no server known to host the table to place it on. Check the deferred cause first though -- a table whose + // every partition is deferred also has no populated class, and reports far more actionably. That is the + // pre-pass' only deferred check: a group it marks gets no broker pruning, so the leaf assignment covers the + // rest. + if (!anyPopulated) { + checkNoPartitionsWithOnlyDeferredSegments(partitionTableInfo, tableName); + } + Preconditions.checkState(anyPopulated, + "Failed to find any segment in any partition for table: %s, which is required for a partitioned worker " + + "assignment", tableName); + memberClassServers.add(classServers); + } + if (!reducible) { + continue; + } + // The member list is never empty (see ColocationGroupAnalyzer#toReducibleGroup), so the loop allocated this, and + // the class list is never empty either: every member holds data in at least one class, and the union keeps it. + assert survivingClasses != null; + int[] partitionClassIds = toClassIds(survivingClasses); + Map>> padding = + computePadding(memberFragmentIds, memberClassServers, partitionClassIds); + if (padding.isEmpty() && partitionClassIds.length == numWorkers) { + // Worker k already stands for class k on every member: nothing to reduce, nothing to pad. Leaving the group + // unmarked also keeps broker pruning on for its leaves (see computePartitionsToKeep). A group that needs + // padding is marked even when it keeps every class, because a padded worker's id is its index in the class + // list. + continue; + } + // One shared array instance, so that the agreement check in MailboxAssignmentVisitor compares one list rather + // than copies of it. The padding goes on the same metadata: a padded worker's id only means something within the + // list. + for (Integer fragmentId : memberFragmentIds) { + DispatchablePlanMetadata metadata = metadataMap.get(fragmentId); + metadata.setPartitionClassIds(partitionClassIds); + metadata.setPaddedClassCandidates(padding.get(fragmentId)); + } + } + } + + /// Returns the servers that can scan each partition class of the given layout as a whole, in class-id order, or + /// `null` for a class that holds no segment at all. This is the intersection of the fully replicated servers of the + /// class's populated partitions, i.e. the candidate set its worker is picked from (see + /// [#assignMultiplePartitionsPerWorker]). An empty (rather than `null`) intersection means a class holding data that + /// no single server can scan as a whole, which the worker assignment reports. + /// + /// The returned sets must only be read: a class with a single populated partition (the common case) hands out that + /// partition's own set rather than a copy, and a copy is made only where an intersection has to be written. + private static List> collectClassServers(PartitionInfo[] partitionInfoMap, int numWorkers) { + int numPartitions = partitionInfoMap.length; + List> classServers = new ArrayList<>(numWorkers); + for (int classId = 0; classId < numWorkers; classId++) { + Set servers = null; + boolean copied = false; + for (int partitionId = classId; partitionId < numPartitions; partitionId += numWorkers) { + PartitionInfo partitionInfo = partitionInfoMap[partitionId]; + if (partitionInfo == null) { + continue; + } + if (servers == null) { + servers = partitionInfo._fullyReplicatedServers; + } else { + if (!copied) { + servers = new HashSet<>(servers); + copied = true; + } + servers.retainAll(partitionInfo._fullyReplicatedServers); + } + } + classServers.add(servers); + } + return classServers; + } + + /// Returns, for every member of a colocated group that holds no data in a class the group keeps, that class mapped to + /// the servers a peer holding data in it picks its own worker from (see [#assignPaddedWorker], which is where that + /// borrowed set is used). Keyed by fragment id and absent altogether for a member that needs no padding, so a + /// non-null entry is the signal that the leaf must pad. When several peers hold data in the class the first one in + /// member order is used; any of them keeps the exchange in process for that peer. + private static Map>> computePadding(List memberFragmentIds, + List>> memberClassServers, int[] partitionClassIds) { + Map>> padding = new HashMap<>(); + for (int memberIndex = 0; memberIndex < memberFragmentIds.size(); memberIndex++) { + List> classServers = memberClassServers.get(memberIndex); + Map> paddedClasses = null; + for (int classId : partitionClassIds) { + if (classServers.get(classId) != null) { + continue; + } + // A class is kept only because some member holds data in it, so there is always such a peer. + Set peerServers = null; + for (List> peerClassServers : memberClassServers) { + peerServers = peerClassServers.get(classId); + if (peerServers != null) { + break; + } + } + if (paddedClasses == null) { + paddedClasses = new HashMap<>(); + } + paddedClasses.put(classId, peerServers); + } + if (paddedClasses != null) { + padding.put(memberFragmentIds.get(memberIndex), paddedClasses); + } + } + return padding; + } + + /// Returns the ids of the set classes, ascending. The order is part of the mapping: worker `k` handles the class at + /// index `k`, so all the members of a group must walk the list the same way. + private static int[] toClassIds(boolean[] survivingClasses) { + IntArrayList classIds = new IntArrayList(survivingClasses.length); + for (int classId = 0; classId < survivingClasses.length; classId++) { + if (survivingClasses[classId]) { + classIds.add(classId); + } } + return classIds.toIntArray(); } /// Post-order traversal that assigns workers to either leaf or intermediate fragments. /// @param leafOnly when true, only leaf fragments are assigned; when false, only intermediate fragments are assigned + /// @param visitedFragmentIds the fragment ids already traversed in this pass; a spooled fragment is reachable from + /// multiple receivers and must be assigned only once private void assignWorkersToNonRootFragment(PlanFragment fragment, DispatchablePlanContext context, - boolean leafOnly) { + boolean leafOnly, Set visitedFragmentIds) { + if (!visitedFragmentIds.add(fragment.getFragmentId())) { + return; + } List children = fragment.getChildren(); for (PlanFragment child : children) { - assignWorkersToNonRootFragment(child, context, leafOnly); + assignWorkersToNonRootFragment(child, context, leafOnly, visitedFragmentIds); } Map metadataMap = context.getDispatchablePlanMetadataMap(); DispatchablePlanMetadata metadata = metadataMap.get(fragment.getFragmentId()); @@ -146,6 +332,8 @@ private void assignWorkersToNonRootFragment(PlanFragment fragment, DispatchableP Map workerIdToServerInstanceMap = assignWorkersForLocalExchange(childMetadata); metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap); metadata.setPartitionFunction(childMetadata.getPartitionFunction()); + // The worker map comes from the child, so the worker ids stand for the same partition classes as the child's. + metadata.setPartitionClassIds(childMetadata.getPartitionClassIds()); // Fake a segments map so that the worker can be correctly identified as leaf stage Map> segmentsMap = Map.of(TableType.OFFLINE.name(), List.of()); Map>> workerIdToSegmentsMap = @@ -162,7 +350,7 @@ private void assignWorkersToNonRootFragment(PlanFragment fragment, DispatchableP } } - private boolean isLookupJoin(List children) { + static boolean isLookupJoin(List children) { if (children.size() != 1) { return false; } @@ -210,13 +398,17 @@ private Map assignWorkersForLocalExchange(Dispatch } } - private static boolean isLeafPlan(DispatchablePlanMetadata metadata) { + static boolean isLeafPlan(DispatchablePlanMetadata metadata) { return metadata.getScannedTables().size() == 1; } // -------------------------------------------------------------------------- // Intermediate stage assign logic // -------------------------------------------------------------------------- + + /// Assigns the workers of an intermediate (non table scanning) fragment. An override must copy the partition class + /// list of the child it derives its worker map from (see [DispatchablePlanMetadata#getPartitionClassIds()]); not + /// copying it costs the colocation of the exchange (the data is shuffled) but never correctness. protected void assignWorkersToIntermediateFragment(PlanFragment fragment, DispatchablePlanContext context) { List children = fragment.getChildren(); Map metadataMap = context.getDispatchablePlanMetadataMap(); @@ -247,6 +439,8 @@ protected void assignWorkersToIntermediateFragment(PlanFragment fragment, Dispat DispatchablePlanMetadata firstChildMetadata = metadataMap.get(children.get(0).getFragmentId()); metadata.setWorkerIdToServerInstanceMap(assignWorkersForLocalExchange(firstChildMetadata)); metadata.setPartitionFunction(firstChildMetadata.getPartitionFunction()); + // isPrePartitionAssignment verified that the children all agree on the classes their worker ids stand for. + metadata.setPartitionClassIds(firstChildMetadata.getPartitionClassIds()); return; } @@ -325,12 +519,17 @@ protected void assignWorkersToIntermediateFragment(PlanFragment fragment, Dispat } childMetadata.setWorkerIdToServerInstanceMap(childWorkerIdToServerInstanceMap); childMetadata.setWorkerIdToSegmentsMap(childWorkerIdToSegmentsMap); + // With a local exchange peer the worker map is copied from it, so the classes come along; without one it comes + // from the candidate servers, whose worker ids are not classes at all. + childMetadata.setPartitionClassIds( + localExchangeChildMetadata != null ? localExchangeChildMetadata.getPartitionClassIds() : null); } } metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap); if (localExchangeChildMetadata != null) { metadata.setPartitionFunction(localExchangeChildMetadata.getPartitionFunction()); + metadata.setPartitionClassIds(localExchangeChildMetadata.getPartitionClassIds()); } else { metadata.setPartitionFunction(DEFAULT_SHUFFLE_PARTITION_FUNCTION); } @@ -347,11 +546,17 @@ private boolean isPrePartitionAssignment(List children, // 2. Pick the most colocate assignment instead of picking the first children String partitionFunction = null; int partitionCount = 0; + // The children are wired 1-to-1 to this stage, so they must also agree on the class each worker id stands for. A + // mismatch means the plan does not form one colocated group; shuffle rather than mispair the classes. + int[] partitionClassIds = metadataMap.get(children.get(0).getFragmentId()).getPartitionClassIds(); for (PlanFragment child : children) { DispatchablePlanMetadata childMetadata = metadataMap.get(child.getFragmentId()); if (!childMetadata.isPrePartitioned()) { return false; } + if (!Arrays.equals(partitionClassIds, childMetadata.getPartitionClassIds())) { + return false; + } if (partitionFunction == null) { partitionFunction = childMetadata.getPartitionFunction(); } else if (!partitionFunction.equalsIgnoreCase(childMetadata.getPartitionFunction())) { @@ -473,26 +678,22 @@ private void assignWorkersToLeafFragment(PlanFragment fragment, DispatchablePlan Map tableOptions = metadata.getTableOptions(); if (tableOptions != null) { - if (Boolean.parseBoolean(tableOptions.get(PinotHintOptions.TableHintOptions.IS_REPLICATED))) { + if (LeafPartitionHints.isReplicated(tableOptions)) { setSegmentsForReplicatedLeafFragment(metadata, context); return; } - String partitionParallelismStr = tableOptions.get(PinotHintOptions.TableHintOptions.PARTITION_PARALLELISM); - int partitionParallelism = partitionParallelismStr != null ? Integer.parseInt(partitionParallelismStr) : 1; - Preconditions.checkState(partitionParallelism > 0, "'%s' must be positive: %s, got: %s", - PinotHintOptions.TableHintOptions.PARTITION_PARALLELISM, partitionParallelism); - metadata.setPartitionParallelism(partitionParallelism); + LeafPartitionHints partitionHints = LeafPartitionHints.resolve(tableOptions); + metadata.setPartitionParallelism(partitionHints.getPartitionParallelism()); - String partitionKey = tableOptions.get(PinotHintOptions.TableHintOptions.PARTITION_KEY); - if (partitionKey != null) { + if (partitionHints.getPartitionKey() != null) { // Broker pruning: build a filter-bearing routing query (null when disabled/unsupported) so the partitioned // assignment can drop partitions with no matching segments. Reuses the same gate as the non-partitioned path. - // Skip pre-partitioned leaves up front: pruning is disabled for them (see computePartitionsToKeep), so don't - // spend planning time building the routing query, e.g. for colocated-join leaves. - PinotQuery routingPinotQuery = metadata.isPrePartitioned() ? null + // Skip pre-partitioned leaves and leaves of a reduced colocated group up front: pruning is disabled for them + // (see computePartitionsToKeep), so don't spend planning time building the routing query. + PinotQuery routingPinotQuery = metadata.isPrePartitioned() || metadata.getPartitionClassIds() != null ? null : extractRoutingQuery(fragment.getFragmentRoot(), metadata.getScannedTables().get(0), context); - assignWorkersToPartitionedLeafFragment(metadata, context, partitionKey, tableOptions, routingPinotQuery); + assignWorkersToPartitionedLeafFragment(metadata, context, partitionHints, routingPinotQuery); updateContextForLeafStage(metadata, context); return; } @@ -743,6 +944,11 @@ private void setSegmentsForReplicatedLeafFragment(DispatchablePlanMetadata metad } /// Extension point to filter the non-replicated leaf-stage per-worker segment assignment; no-op by default. + /// + /// An override must treat the assignment it is handed as read-only and publish its result by replacing the per-worker + /// segment lists (or the whole map) on the metadata, rather than by editing them in place: part of what the + /// assignment is built from is the broker's published partition metadata, shared across queries and read concurrently + /// by other planning threads. What is handed over is nevertheless kept safe to edit in place. protected void filterLeafStageSegments(DispatchablePlanContext context, DispatchablePlanMetadata metadata) { } @@ -928,24 +1134,21 @@ private static void transferToServerInstanceLogicalSegmentsMap(String physicalTa // -------------------------------------------------------------------------- // Partitioned leaf stage assignment // -------------------------------------------------------------------------- + + /// Assigns one worker per partition class of a leaf that scans a partitioned table. private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata metadata, - DispatchablePlanContext context, String partitionKey, Map tableOptions, + DispatchablePlanContext context, LeafPartitionHints partitionHints, @Nullable PinotQuery routingPinotQuery) { // when partition key exist, we assign workers for leaf-stage in partitioned fashion. - - String numPartitionsStr = tableOptions.get(PinotHintOptions.TableHintOptions.PARTITION_SIZE); - Preconditions.checkState(numPartitionsStr != null, "'%s' must be provided for partition key: %s", - PinotHintOptions.TableHintOptions.PARTITION_SIZE, partitionKey); - int numWorkers = Integer.parseInt(numPartitionsStr); - Preconditions.checkState(numWorkers > 0, "'%s' must be positive, got: %s", - PinotHintOptions.TableHintOptions.PARTITION_SIZE, numWorkers); - - String partitionFunction = tableOptions.getOrDefault(PinotHintOptions.TableHintOptions.PARTITION_FUNCTION, - DEFAULT_TABLE_PARTITION_FUNCTION); + String partitionKey = partitionHints.getPartitionKey(); + assert partitionKey != null; + int numWorkers = partitionHints.getPartitionSize(); + String partitionFunction = partitionHints.getPartitionFunction(); String tableName = metadata.getScannedTables().get(0); - // calculates the partition table info using the routing manager - PartitionTableInfo partitionTableInfo = calculatePartitionTableInfo(tableName); + // calculates the partition table info using the routing manager, reusing this query's cached snapshot + PartitionTableInfo partitionTableInfo = + context.getPartitionTableInfoCache().computeIfAbsent(tableName, this::calculatePartitionTableInfo); // verifies that the partition table obtained from routing manager is compatible with the hint options checkPartitionInfoMap(partitionTableInfo, tableName, partitionKey, partitionFunction, numWorkers); @@ -953,6 +1156,20 @@ private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata met int numPartitions = partitionInfoMap.length; assert numPartitions % numWorkers == 0; int numPartitionsPerWorker = numPartitions / numWorkers; + // The partition classes that get a worker, one per worker in worker-id order, or null to give every class a worker. + int[] partitionClassIds = metadata.getPartitionClassIds(); + if (partitionClassIds != null) { + // The list is resolved from the same hints by the same LeafPartitionHints, so it is a non-empty ascending + // subsequence of 0..numWorkers-1; a mismatch would index outside the partition info map below. + Preconditions.checkState( + partitionClassIds.length > 0 && partitionClassIds[partitionClassIds.length - 1] < numWorkers, + "Invalid partition classes: %s for table: %s with hinted partition size: %s", + Arrays.toString(partitionClassIds), tableName, numWorkers); + } + // The classes to pad, if any, resolved once for the whole leaf (see PaddingInfo). + Map> paddedClassCandidates = metadata.getPaddedClassCandidates(); + PaddingInfo paddingInfo = paddedClassCandidates != null ? new PaddingInfo(paddedClassCandidates, + collectHostingServers(partitionInfoMap)) : null; // Broker pruning: the partitions to keep (null means keep all). Partitions absent from the set are skipped below. Set partitionsToKeep = @@ -962,18 +1179,30 @@ private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata met if (numSegmentsPrunedByBroker > 0) { context.addNumSegmentsPrunedByBroker(numSegmentsPrunedByBroker); } + } else { + // Every partition needs a worker here (pruning is off for a pre-partitioned leaf and for a leaf of a reduced + // colocated group), so a partition that holds data without a fully replicated server has nowhere to go: report it + // rather than dropping (or padding away) its rows. The other cause of a data-holding partition without an entry, + // segments with invalid partition metadata, is rejected while the partition table info is built. + // TODO: With pruning active a deferred partition is simply absent from partitionsToKeep and skipped, dropping its + // rows for a query whose filter does match it. Checking it there instead would fail every query on the + // table while any segment is new; deciding it per query needs the deferred segment names, not just their + // ids. + checkNoPartitionsWithOnlyDeferredSegments(partitionTableInfo, tableName); } Map workerIdToServerInstanceMap = new HashMap<>(); Map>> workerIdToSegmentsMap = new HashMap<>(); if (numPartitionsPerWorker == 1) { - assignOnePartitionPerWorker(tableName, context.getRequestId(), partitionInfoMap, partitionsToKeep, - _routingManager.getEnabledServerInstanceMap(), workerIdToServerInstanceMap, workerIdToSegmentsMap); - } else { - assignMultiplePartitionsPerWorker(tableName, context.getRequestId(), numPartitionsPerWorker, partitionInfoMap, - partitionsToKeep, _routingManager.getEnabledServerInstanceMap(), workerIdToServerInstanceMap, + assignOnePartitionPerWorker(tableName, context.getRequestId(), partitionInfoMap, partitionClassIds, + partitionsToKeep, paddingInfo, _routingManager.getEnabledServerInstanceMap(), workerIdToServerInstanceMap, workerIdToSegmentsMap); + } else { + assignMultiplePartitionsPerWorker(tableName, context.getRequestId(), numWorkers, partitionInfoMap, + partitionClassIds, partitionsToKeep, paddingInfo, _routingManager.getEnabledServerInstanceMap(), + workerIdToServerInstanceMap, workerIdToSegmentsMap); } + checkLeafWorkerAssignment(tableName, workerIdToServerInstanceMap, workerIdToSegmentsMap); metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap); metadata.setWorkerIdToSegmentsMap(workerIdToSegmentsMap); metadata.setTimeBoundaryInfo(partitionTableInfo._timeBoundaryInfo); @@ -987,9 +1216,10 @@ private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata met /// /// - broker pruning is disabled or the leaf shape is unsupported (the routing query is `null`), or there is /// no filter to prune with; - /// - the leaf feeds a pre-partitioned (1-to-1 direct) exchange -- dropping/compacting workers would misalign - /// sender/receiver worker ids in `MailboxAssignmentVisitor`. A non-pre-partitioned leaf is shuffled via - /// `connectWorkers`, which re-hashes across any worker count, so pruning is safe there; + /// - the leaf feeds a pre-partitioned (1-to-1 direct) exchange, or it belongs to a colocated group that agreed on a + /// partition class list -- dropping/compacting workers would misalign sender/receiver worker ids in + /// `MailboxAssignmentVisitor`. A non-pre-partitioned leaf is shuffled via `connectWorkers`, which re-hashes across + /// any worker count, so pruning is safe there; /// - routing fails (pruning is best-effort); /// - every partition would be pruned -- an empty worker map would break exchanges in a multi-leaf plan (the /// all-leaves-empty short-circuit does not fire for a partially-empty plan), and the server-side filter still @@ -1010,7 +1240,8 @@ private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata met @Nullable private Set computePartitionsToKeep(@Nullable PinotQuery routingPinotQuery, DispatchablePlanMetadata metadata, long requestId, PartitionInfo[] partitionInfoMap) { - if (routingPinotQuery == null || routingPinotQuery.getFilterExpression() == null || metadata.isPrePartitioned()) { + if (routingPinotQuery == null || routingPinotQuery.getFilterExpression() == null || metadata.isPrePartitioned() + || metadata.getPartitionClassIds() != null) { return null; } Map routingTableMap; @@ -1069,63 +1300,100 @@ private static long countPrunedSegments(PartitionInfo[] partitionInfoMap, Set partitionsToKeep, Map enabledServerInstanceMap, + @Nullable int[] partitionClassIds, @Nullable Set partitionsToKeep, @Nullable PaddingInfo paddingInfo, + Map enabledServerInstanceMap, Map workerIdToServerInstanceMap, Map>> workerIdToSegmentsMap) { - int numPartitions = partitionInfoMap.length; - int workerId = 0; - for (int i = 0; i < numPartitions; i++) { - // Skip partitions pruned by the broker filter. Empty partitions are never in partitionsToKeep, so under pruning - // they are skipped here too; the precondition below only fires when pruning is inactive (partitionsToKeep null). - if (partitionsToKeep != null && !partitionsToKeep.contains(i)) { + int[] partitionIds = selectPartitionsToAssign(partitionInfoMap.length, partitionClassIds, partitionsToKeep); + for (int workerId = 0; workerId < partitionIds.length; workerId++) { + int partitionId = partitionIds[workerId]; + PartitionInfo partitionInfo = partitionInfoMap[partitionId]; + if (partitionInfo == null) { + // Pad a class the colocated group keeps but this table has no data in, see assignPaddedWorker. + // TODO: Currently we don't support the case when a partition doesn't contain any segment outside of a colocated + // group, where there is nothing to keep the worker ids aligned with. The reason is that the leaf stage + // won't be able to directly return empty response. + Preconditions.checkState(paddingInfo != null && paddingInfo._classCandidates.containsKey(partitionId), + "Failed to find any segment for table: %s, partition: %s", tableName, partitionId); + assignPaddedWorker(tableName, requestId, partitionId, paddingInfo, enabledServerInstanceMap, workerId, + workerIdToServerInstanceMap, workerIdToSegmentsMap); continue; } - PartitionInfo partitionInfo = partitionInfoMap[i]; - // TODO: Currently we don't support the case when a partition doesn't contain any segment. The reason is that - // the leaf stage won't be able to directly return empty response. - Preconditions.checkState(partitionInfo != null, "Failed to find any segment for table: %s, partition: %s", - tableName, i); // NOTE: Pick worker based on the request id plus the partition id (not a running counter) so that the same worker // is picked across different table scans when the segments for the same partition are colocated, and so - // that skipping pruned partitions does not shift the server assignment of the surviving ones. + // that skipping pruned or empty partitions does not shift the server assignment of the surviving ones. ServerInstance serverInstance = - pickEnabledServer(partitionInfo._fullyReplicatedServers, enabledServerInstanceMap, requestId + i); + pickEnabledServer(partitionInfo._fullyReplicatedServers, enabledServerInstanceMap, requestId + partitionId); Preconditions.checkState(serverInstance != null, - "Failed to find enabled fully replicated server for table: %s, partition: %s", tableName, i); + "Failed to find enabled fully replicated server for table: %s, partition: %s", tableName, partitionId); workerIdToServerInstanceMap.put(workerId, new QueryServerInstance(serverInstance)); + // NOTE: Copy the segment lists. Unlike the multiple-partitions-per-worker path (which merges into fresh lists), + // these are the broker's published metadata, shared across queries and never to be mutated (see + // filterLeafStageSegments). workerIdToSegmentsMap.put(workerId, - getSegmentsMap(partitionInfo._offlineSegments, partitionInfo._realtimeSegments)); - workerId++; + getSegmentsMap(copySegments(partitionInfo._offlineSegments), copySegments(partitionInfo._realtimeSegments))); } } + /// Returns the partitions to assign in worker-id order, i.e. worker `k` gets the partition at index `k`. + /// + /// For a leaf in a colocated group (`partitionClassIds` non-null) this is the group's surviving class list itself: + /// the worker id must be the position in that list, not a running counter, so that worker `k` stands for the same + /// class on every member of the group. Otherwise it is every partition, minus the ones broker pruning dropped. The + /// returned array may be the class list shared by the whole group, so the caller must only read it. + private static int[] selectPartitionsToAssign(int numPartitions, @Nullable int[] partitionClassIds, + @Nullable Set partitionsToKeep) { + if (partitionClassIds != null) { + return partitionClassIds; + } + if (partitionsToKeep == null) { + int[] partitionIds = new int[numPartitions]; + for (int partitionId = 0; partitionId < numPartitions; partitionId++) { + partitionIds[partitionId] = partitionId; + } + return partitionIds; + } + IntArrayList partitionIds = new IntArrayList(partitionsToKeep.size()); + for (int partitionId = 0; partitionId < numPartitions; partitionId++) { + if (partitionsToKeep.contains(partitionId)) { + partitionIds.add(partitionId); + } + } + return partitionIds.toIntArray(); + } + /// Round-robin partitions to workers, where each worker gets numPartitionsPerWorker partitions. This setup works only /// if all segments for these partitions are assigned to the same group of servers. This is useful when user wants to /// colocate tables with different partition count, but same partition function. /// E.g. when there are 16 partitions for table A and 4 partitions for table B, we may assign 16 partitions for table /// A to 4 workers, where partition 0, 4, 8, 12 goes to worker 0, partition 1, 5, 9, 13 goes to worker 1, etc. - private void assignMultiplePartitionsPerWorker(String tableName, long requestId, int numPartitionsPerWorker, - PartitionInfo[] partitionInfoMap, @Nullable Set partitionsToKeep, - Map enabledServerInstanceMap, + /// + /// The worker index is already the partition class id here, so when `partitionClassIds` is non-null only the classes + /// in that list get a worker, in that order, padding the ones this table holds no data in (see + /// [#selectPartitionsToAssign], which makes the same decision on the one-partition-per-worker path). + private void assignMultiplePartitionsPerWorker(String tableName, long requestId, int numWorkers, + PartitionInfo[] partitionInfoMap, @Nullable int[] partitionClassIds, @Nullable Set partitionsToKeep, + @Nullable PaddingInfo paddingInfo, Map enabledServerInstanceMap, Map workerIdToServerInstanceMap, Map>> workerIdToSegmentsMap) { int numPartitions = partitionInfoMap.length; - assert numPartitions % numPartitionsPerWorker == 0; - int numWorkers = numPartitions / numPartitionsPerWorker; + int numPartitionsPerWorker = numPartitions / numWorkers; + int numClasses = partitionClassIds != null ? partitionClassIds.length : numWorkers; int workerId = 0; - for (int i = 0; i < numWorkers; i++) { + for (int classIndex = 0; classIndex < numClasses; classIndex++) { + int classId = partitionClassIds != null ? partitionClassIds[classIndex] : classIndex; Set fullyReplicatedServers = null; List offlineSegments = null; List realtimeSegments = null; - for (int j = i; j < numPartitions; j += numWorkers) { - if (partitionsToKeep != null && !partitionsToKeep.contains(j)) { + for (int partitionId = classId; partitionId < numPartitions; partitionId += numWorkers) { + if (partitionsToKeep != null && !partitionsToKeep.contains(partitionId)) { // Partition pruned by the broker filter. continue; } - PartitionInfo partitionInfo = partitionInfoMap[j]; + PartitionInfo partitionInfo = partitionInfoMap[partitionId]; if (partitionInfo == null) { continue; } @@ -1149,28 +1417,128 @@ private void assignMultiplePartitionsPerWorker(String tableName, long requestId, } } } - // Without broker pruning we don't support a worker whose partitions all lack segments, because the leaf stage - // can't directly return an empty response. With pruning active a fully-pruned worker is legitimate and skipped. if (fullyReplicatedServers == null) { + // Pad a class the colocated group keeps but this table has no data in, see assignPaddedWorker. + if (paddingInfo != null && paddingInfo._classCandidates.containsKey(classId)) { + assignPaddedWorker(tableName, requestId, classId, paddingInfo, enabledServerInstanceMap, workerId, + workerIdToServerInstanceMap, workerIdToSegmentsMap); + workerId++; + continue; + } + // Without broker pruning we don't support a worker whose partitions all lack segments, because the leaf stage + // can't directly return an empty response. With pruning active a fully-pruned worker is legitimate and skipped. Preconditions.checkState(partitionsToKeep != null, - "Failed to find any segment for table: %s, worker: %s, partitions per worker: %s", tableName, i, - numPartitionsPerWorker); + "Failed to find any segment for table: %s, partition class: %s, partitions per worker: %s", tableName, + classId, numPartitionsPerWorker); continue; } - // NOTE: Pick worker based on the request id plus the worker index (not a running counter) so that the same worker - // is picked across different table scans when the segments for the same partition are colocated, and so - // that skipping fully-pruned workers does not shift the server assignment of the surviving ones. + // NOTE: Pick worker based on the request id plus the partition class id (not a running counter) so that the same + // worker is picked across different table scans when the segments for the same partition are colocated, and + // so that skipping fully-pruned or dropped classes does not shift the assignment of the surviving ones. ServerInstance serverInstance = - pickEnabledServer(fullyReplicatedServers, enabledServerInstanceMap, requestId + i); + pickEnabledServer(fullyReplicatedServers, enabledServerInstanceMap, requestId + classId); Preconditions.checkState(serverInstance != null, - "Failed to find enabled fully replicated server for table: %s, worker: %s, partitions per worker: %s", - tableName, i, numPartitionsPerWorker); + "Failed to find enabled fully replicated server for table: %s, partition class: %s, partitions per worker: " + + "%s", tableName, classId, numPartitionsPerWorker); workerIdToServerInstanceMap.put(workerId, new QueryServerInstance(serverInstance)); workerIdToSegmentsMap.put(workerId, getSegmentsMap(offlineSegments, realtimeSegments)); workerId++; } } + /// Assigns a worker with no segments to scan, for a partition class that this table holds no data in while its + /// colocated group keeps it because a peer does hold data there (see [#computePadding]). Without it this member would + /// have fewer workers than its peers, and the 1-to-1 exchange between them would either pair the wrong classes or + /// degrade to a shuffle. + /// + /// The server is picked from the candidate set the peer picks its own worker for this class from, with the same seed, + /// so that the empty worker lands on the peer's server and the exchange stays in process: [#pickEnabledServer] sorts + /// the candidates and starts at `seed % size`, so the set is what decides the pick. A server outside the ones that + /// provably host this table cannot be used at all (it may have no table data manager for it, which it reports as a + /// missing table), so fall back to the servers that do host it and accept one cross-server send. + /// + /// Exactly one [TableType] key is emitted: the one the chosen server provably has a table data manager for (see + /// [#collectHostingServers]), because the server resolves one data manager per key in the map and fails the query + /// when it is missing. The segment list is mutable because the leaf-stage segment filters may edit the lists they are + /// handed (see [#filterLeafStageSegments]). + private static void assignPaddedWorker(String tableName, long requestId, int classId, PaddingInfo paddingInfo, + Map enabledServerInstanceMap, int workerId, + Map workerIdToServerInstanceMap, + Map>> workerIdToSegmentsMap) { + Set peerServers = paddingInfo._classCandidates.get(classId); + // The callers only pad a class the colocated group decided to pad, so there is always a candidate set for it. + assert peerServers != null; + Map hostingServers = paddingInfo._hostingServers; + ServerInstance serverInstance = pickEnabledServer(peerServers, enabledServerInstanceMap, requestId + classId); + String tableType = serverInstance != null ? hostingServers.get(serverInstance.getInstanceId()) : null; + if (tableType == null) { + serverInstance = pickEnabledServer(hostingServers.keySet(), enabledServerInstanceMap, requestId + classId); + Preconditions.checkState(serverInstance != null, + "Failed to find an enabled server hosting table: %s for the empty worker of partition class: %s", tableName, + classId); + // Non-null because the server was picked from the hosting map's own key set. + tableType = hostingServers.get(serverInstance.getInstanceId()); + } + workerIdToServerInstanceMap.put(workerId, new QueryServerInstance(serverInstance)); + workerIdToSegmentsMap.put(workerId, Map.of(tableType, new ArrayList<>())); + } + + /// Returns every server that provably hosts the given table -- the union of the fully replicated servers over its + /// populated partitions -- mapped to the [TableType] name to hand a worker placed on that server. A server outside + /// this map is not known to host the table at all, so it cannot be given a worker for it. The table type of a server + /// is taken from the first populated partition it hosts, so it is one the server provably has a data manager for. + private static Map collectHostingServers(PartitionInfo[] partitionInfoMap) { + Map hostingServers = new HashMap<>(); + for (PartitionInfo partitionInfo : partitionInfoMap) { + if (partitionInfo == null) { + continue; + } + String tableType = partitionInfo._offlineSegments != null ? TableType.OFFLINE.name() : TableType.REALTIME.name(); + for (String server : partitionInfo._fullyReplicatedServers) { + hostingServers.putIfAbsent(server, tableType); + } + } + return hostingServers; + } + + /// Validates the worker assignment computed for a partitioned leaf fragment before it is published on the + /// [DispatchablePlanMetadata]. Both invariants hold by construction today; the checks exist so that a regression + /// fails here, naming the table and the offending worker id, instead of much later: + /// + /// - the worker ids must be exactly `0..numWorkers-1`, because + /// [DispatchablePlanContext#constructDispatchablePlanFragmentMap] indexes a `WorkerMetadata[]` sized from the + /// server map by worker id, where a gap leaves a null entry; + /// - every worker must have a segments map keyed by 1 or 2 [TableType] names, with non-null lists, because the server + /// splits the request on the number of entries and resolves one table data manager per key: an unexpected key + /// becomes an opaque server-side failure. + @VisibleForTesting + static void checkLeafWorkerAssignment(String tableName, + Map workerIdToServerInstanceMap, + Map>> workerIdToSegmentsMap) { + int numWorkers = workerIdToServerInstanceMap.size(); + Preconditions.checkState(workerIdToSegmentsMap.size() == numWorkers, + "Got %s workers but %s worker segment entries for table: %s", numWorkers, workerIdToSegmentsMap.size(), + tableName); + for (int workerId = 0; workerId < numWorkers; workerId++) { + Preconditions.checkState(workerIdToServerInstanceMap.containsKey(workerId), + "Missing server instance for worker: %s (num workers: %s) for table: %s", workerId, numWorkers, tableName); + Map> segmentsMap = workerIdToSegmentsMap.get(workerId); + Preconditions.checkState(segmentsMap != null, "Missing segments for worker: %s (num workers: %s) for table: %s", + workerId, numWorkers, tableName); + int numTableTypes = segmentsMap.size(); + Preconditions.checkState(numTableTypes == 1 || numTableTypes == 2, + "Expected 1 or 2 table types for worker: %s, got: %s for table: %s", workerId, numTableTypes, tableName); + for (Map.Entry> entry : segmentsMap.entrySet()) { + String tableType = entry.getKey(); + Preconditions.checkState( + TableType.OFFLINE.name().equals(tableType) || TableType.REALTIME.name().equals(tableType), + "Unexpected table type: %s for worker: %s for table: %s", tableType, workerId, tableName); + Preconditions.checkState(entry.getValue() != null, + "Null segment list for table type: %s, worker: %s for table: %s", tableType, workerId, tableName); + } + } + } + @Nullable public TableOptions inferTableOptions(String tableName) { try { @@ -1213,6 +1581,11 @@ private PartitionTableInfo calculatePartitionTableInfo(String tableName) { verifyCompatibility(offlineTpi, realtimeTpi); + // This branch builds the merged partition info map itself instead of going through + // PartitionTableInfo.fromTablePartitionInfo, so it runs the check (on both sides) itself. + checkNoSegmentsWithInvalidPartition(offlineTpi); + checkNoSegmentsWithInvalidPartition(realtimeTpi); + TablePartitionReplicatedServersInfo.PartitionInfo[] offlinePartitionInfoMap = offlineTpi.getPartitionInfoMap(); TablePartitionReplicatedServersInfo.PartitionInfo[] realtimePartitionInfoMap = realtimeTpi.getPartitionInfoMap(); @@ -1242,8 +1615,17 @@ private PartitionTableInfo calculatePartitionTableInfo(String tableName) { partitionInfoMap[i] = new PartitionInfo(fullyReplicatedServers, offlinePartitionInfo._segments, realtimePartitionInfo._segments); } + // Union the two sides, then keep only the partitions the merged map has no entry for: a partition one side + // deferred but the other still serves as a whole does get a worker, so reporting it would fail a query the + // other side can answer on its own. A TreeSet keeps the broker's sorted order, so the error message is + // deterministic. + Set partitionsWithOnlyDeferredSegments = + new TreeSet<>(offlineTpi.getPartitionsWithOnlyDeferredSegments()); + partitionsWithOnlyDeferredSegments.addAll(realtimeTpi.getPartitionsWithOnlyDeferredSegments()); + partitionsWithOnlyDeferredSegments.removeIf( + partitionId -> partitionId < partitionInfoMap.length && partitionInfoMap[partitionId] != null); return new PartitionTableInfo(offlineTpi.getPartitionColumn(), offlineTpi.getPartitionFunctionName(), - partitionInfoMap, timeBoundaryInfo); + partitionInfoMap, timeBoundaryInfo, partitionsWithOnlyDeferredSegments); } else if (offlineRoutingExists) { return getOfflinePartitionTableInfo(offlineTableName); } else { @@ -1273,10 +1655,44 @@ private static void verifyCompatibility(TablePartitionReplicatedServersInfo offl offlineTpi.getPartitionFunctionName(), realtimeTpi.getPartitionFunctionName()); } + /// Rejects a table that has segments whose partition metadata is invalid (e.g. a segment holding multiple partition + /// ids for the partition column). Such segments are not represented in the partition info map at all, so a + /// partitioned assignment would silently omit their rows. + /// + /// Throws [IllegalStateException] rather than using [Preconditions] so that the implicit table hint path + /// ([#inferTableOptions]) keeps degrading quietly to a non-partitioned (shuffled) plan. + private static void checkNoSegmentsWithInvalidPartition(TablePartitionReplicatedServersInfo tpi) { + int numSegmentsWithInvalidPartition = tpi.getSegmentsWithInvalidPartition().size(); + if (numSegmentsWithInvalidPartition > 0) { + throw new IllegalStateException("Find " + numSegmentsWithInvalidPartition + + " segments with invalid partition for table: " + tpi.getTableNameWithType()); + } + } + + /// Rejects a table that has partitions holding data which no single server can serve as a whole right now, because + /// every segment of the partition is new and does not have all of its replicas online yet (see + /// [TablePartitionReplicatedServersInfo#getPartitionsWithOnlyDeferredSegments()], which is also where the other + /// causes of a partition without an entry in the partition info map are listed). + /// + /// The partitioned assignment needs one worker to scan a whole partition and the multi-stage engine has no + /// optional-segment mechanism to fall back on, so the only alternatives are failing here or silently omitting the + /// partition's rows. Only called where every partition needs a worker, i.e. where broker pruning is inactive. + private static void checkNoPartitionsWithOnlyDeferredSegments(PartitionTableInfo partitionTableInfo, + String tableNameWithType) { + Set partitionsWithOnlyDeferredSegments = partitionTableInfo._partitionsWithOnlyDeferredSegments; + Preconditions.checkState(partitionsWithOnlyDeferredSegments.isEmpty(), + "Failed to find a fully replicated server for partitions: %s of table: %s, because all of their segments are " + + "new and don't have all replicas online yet", partitionsWithOnlyDeferredSegments, tableNameWithType); + } + /// Verifies that the partition info maps from the table partition info are compatible with the information supplied /// as arguments. private void checkPartitionInfoMap(PartitionTableInfo partitionTableInfo, String tableNameWithType, String partitionKey, String partitionFunction, int numPartitions) { + // Must be checked first: the modulo check below passes trivially for an empty partition info map, leaving the + // caller with 0 partitions per worker. + Preconditions.checkState(partitionTableInfo._partitionInfoMap.length > 0, + "Failed to find any partition for table: %s", tableNameWithType); Preconditions.checkState(partitionTableInfo._partitionKey.equals(partitionKey), "Partition key: %s does not match partition column: %s for table: %s", partitionKey, partitionTableInfo._partitionKey, tableNameWithType); @@ -1303,29 +1719,46 @@ private PartitionTableInfo getRealtimePartitionTableInfo(String realtimeTableNam return PartitionTableInfo.fromTablePartitionInfo(realtimeTpi, TableType.REALTIME); } - private static class PartitionTableInfo { + /// What one partitioned leaf needs to pad the partition classes its colocated group keeps but its own table holds no + /// data in. Resolved once per leaf: both members depend only on the partition layout, so a padded worker would + /// otherwise re-scan it, which is quadratic for a wide table joined to one with few populated classes. + private static class PaddingInfo { + /// See [DispatchablePlanMetadata#getPaddedClassCandidates()]. + final Map> _classCandidates; + /// See [#collectHostingServers]. + final Map _hostingServers; + + PaddingInfo(Map> classCandidates, Map hostingServers) { + _classCandidates = classCandidates; + _hostingServers = hostingServers; + } + } + + /// The partition layout of one table, as the worker assignment needs it. Public only so that the per-query cache of + /// these can live on [DispatchablePlanContext]; its contents stay internal to the worker assignment. + public static class PartitionTableInfo { final String _partitionKey; final String _partitionFunction; final PartitionInfo[] _partitionInfoMap; @Nullable final TimeBoundaryInfo _timeBoundaryInfo; + /// Partitions with no entry in `_partitionInfoMap` even though they hold data. See + /// [TablePartitionReplicatedServersInfo#getPartitionsWithOnlyDeferredSegments()]. + final Set _partitionsWithOnlyDeferredSegments; PartitionTableInfo(String partitionKey, String partitionFunction, PartitionInfo[] partitionInfoMap, - @Nullable TimeBoundaryInfo timeBoundaryInfo) { + @Nullable TimeBoundaryInfo timeBoundaryInfo, Set partitionsWithOnlyDeferredSegments) { _partitionKey = partitionKey; _partitionFunction = partitionFunction; _partitionInfoMap = partitionInfoMap; _timeBoundaryInfo = timeBoundaryInfo; + _partitionsWithOnlyDeferredSegments = partitionsWithOnlyDeferredSegments; } static PartitionTableInfo fromTablePartitionInfo( TablePartitionReplicatedServersInfo tablePartitionReplicatedServersInfo, TableType tableType) { - if (!tablePartitionReplicatedServersInfo.getSegmentsWithInvalidPartition().isEmpty()) { - throw new IllegalStateException( - "Find " + tablePartitionReplicatedServersInfo.getSegmentsWithInvalidPartition().size() - + " segments with invalid partition"); - } + checkNoSegmentsWithInvalidPartition(tablePartitionReplicatedServersInfo); int numPartitions = tablePartitionReplicatedServersInfo.getNumPartitions(); TablePartitionReplicatedServersInfo.PartitionInfo[] tablePartitionInfoMap = tablePartitionReplicatedServersInfo @@ -1349,7 +1782,8 @@ static PartitionTableInfo fromTablePartitionInfo( } } return new PartitionTableInfo(tablePartitionReplicatedServersInfo.getPartitionColumn(), - tablePartitionReplicatedServersInfo.getPartitionFunctionName(), workerPartitionInfoMap, null); + tablePartitionReplicatedServersInfo.getPartitionFunctionName(), workerPartitionInfoMap, null, + tablePartitionReplicatedServersInfo.getPartitionsWithOnlyDeferredSegments()); } } @@ -1390,6 +1824,12 @@ private static ServerInstance pickEnabledServer(Set candidates, return null; } + /// Copies a segment list published by the broker so that the planner never hands out (or mutates) the shared one. + @Nullable + private static List copySegments(@Nullable List segments) { + return segments != null ? new ArrayList<>(segments) : null; + } + private static Map> getSegmentsMap(@Nullable List offlineSegments, @Nullable List realtimeSegments) { if (offlineSegments != null) { diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java index e26dd6222e2c..4843e8766e27 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java @@ -333,7 +333,7 @@ public static QueryEnvironment getQueryEnvironment(int reducerPort, int port1, i } TablePartitionReplicatedServersInfo tablePartitionReplicatedServersInfo = new TablePartitionReplicatedServersInfo(tableNameWithType, partitionColumn, "Hashcode", numPartitions, - partitionIdToInfoMap, List.of()); + partitionIdToInfoMap, List.of(), Set.of()); partitionInfoMap.put(tableNameWithType, tablePartitionReplicatedServersInfo); } } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/DispatchableSubPlanTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/DispatchableSubPlanTest.java index 3cbcc517f3e1..f488d11411b7 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/DispatchableSubPlanTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/DispatchableSubPlanTest.java @@ -29,6 +29,7 @@ import org.apache.pinot.query.planner.plannode.ValueNode; import org.testng.annotations.Test; +import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertSame; import static org.testng.Assert.assertTrue; @@ -71,18 +72,21 @@ public void testIsAllLeafStagesEmptyNoTables() { } @Test - public void testCopyWithRootPreservesFragmentId() { + public void testCopyWithRootPreservesFragmentIdAndSegmentsMap() { ValueNode oldRoot = new ValueNode(0, new DataSchema(new String[0], new ColumnDataType[0]), PlanNode.NodeHint.EMPTY, List.of(), List.of()); PlanFragment fragment = new PlanFragment(0, oldRoot, List.of()); DispatchablePlanFragment original = new DispatchablePlanFragment(fragment); + Map>> workerIdToSegmentsMap = Map.of(0, Map.of("OFFLINE", List.of("segment0"))); + original.setWorkerIdToSegmentsMap(workerIdToSegmentsMap); ValueNode newRoot = new ValueNode(0, new DataSchema(new String[0], new ColumnDataType[0]), PlanNode.NodeHint.EMPTY, List.of(), List.of()); DispatchablePlanFragment copy = DispatchablePlanFragment.copyWithRoot(original, newRoot); - org.testng.Assert.assertEquals(copy.getPlanFragment().getFragmentId(), 0); + assertEquals(copy.getPlanFragment().getFragmentId(), 0); assertSame(copy.getPlanFragment().getFragmentRoot(), newRoot); assertSame(original.getPlanFragment().getFragmentRoot(), oldRoot); + assertEquals(copy.getWorkerIdToSegmentsMap(), workerIdToSegmentsMap); } } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitorTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitorTest.java index 81699e74340d..33476c72f878 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitorTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitorTest.java @@ -37,6 +37,7 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; @@ -133,6 +134,65 @@ public void testSingletonWithParallelismAllowsCrossServer() { assertEquals(singleMailbox(receiver.getWorkerIdToMailboxesMap(), 1, SENDER_STAGE).getHostname(), "host_A"); } + /// A mismatch means the one-class-list-per-colocated-group invariant regressed (see + /// [DispatchablePlanMetadata#getPartitionClassIds()]) and must be reported rather than pairing one class with + /// another. + @Test(expectedExceptions = IllegalStateException.class, + expectedExceptionsMessageRegExp = ".*Partition class mismatch.*\\[0, 2\\].*\\[0, 3\\].*") + public void testDirectExchangeRejectsMismatchedPartitionClasses() { + DispatchablePlanMetadata sender = metadata(Map.of(0, server("A"), 1, server("B"))); + sender.setPartitionClassIds(new int[]{0, 2}); + DispatchablePlanMetadata receiver = metadata(Map.of(0, server("A"), 1, server("B"))); + receiver.setPartitionClassIds(new int[]{0, 3}); + process(singletonSendNode(List.of()), sender, receiver); + } + + /// A receiver with no class list took its workers from the candidate servers, so matching worker counts are a + /// coincidence and the exchange must fall back to a shuffle rather than pair the two 1-to-1. + @Test + public void testPrePartitionedSendWithoutMatchingClassesFallsBackToShuffle() { + DispatchablePlanMetadata sender = prePartitionedSender(); + sender.setPartitionClassIds(new int[]{0, 2}); + DispatchablePlanMetadata receiver = metadata(Map.of(0, server("A"), 1, server("B"))); + process(hashSendNode(), sender, receiver); + + // Shuffled: every receiver worker reads from every sender worker, rather than only from the one with its own id. + assertEquals(expandedWorkerIds(receiver.getWorkerIdToMailboxesMap(), 0, SENDER_STAGE), List.of(0, 1)); + assertEquals(expandedWorkerIds(receiver.getWorkerIdToMailboxesMap(), 1, SENDER_STAGE), List.of(0, 1)); + } + + /// The control for the test above: with both sides in the same class space the very same shapes are wired 1-to-1. + @Test + public void testPrePartitionedSendWithMatchingClassesIsDirect() { + DispatchablePlanMetadata sender = prePartitionedSender(); + sender.setPartitionClassIds(new int[]{0, 2}); + DispatchablePlanMetadata receiver = metadata(Map.of(0, server("A"), 1, server("B"))); + receiver.setPartitionClassIds(new int[]{0, 2}); + receiver.setPartitionFunction("absHashCodeSum"); + process(hashSendNode(), sender, receiver); + + assertEquals(expandedWorkerIds(receiver.getWorkerIdToMailboxesMap(), 0, SENDER_STAGE), List.of(0)); + assertEquals(expandedWorkerIds(receiver.getWorkerIdToMailboxesMap(), 1, SENDER_STAGE), List.of(1)); + } + + /// A receiver stage with no worker at all (an empty or fully pruned leaf, while another leaf of the plan is not) must + /// still leave every sender worker an entry holding an empty mailbox list, or the sender's `WorkerMetadata` carries a + /// null mailbox map and fails while the dispatch request is serialized. + @Test + public void testShuffleToReceiverWithoutWorkersKeepsEmptySenderEntry() { + DispatchablePlanMetadata sender = metadata(Map.of(0, server("A"), 1, server("B"))); + DispatchablePlanMetadata receiver = metadata(Map.of()); + process(hashSendNode(), sender, receiver); + + for (int workerId = 0; workerId < 2; workerId++) { + MailboxInfos mailboxInfos = sender.getWorkerIdToMailboxesMap().get(workerId).get(RECEIVER_STAGE); + assertNotNull(mailboxInfos, "Missing entry for worker: " + workerId); + assertTrue(mailboxInfos.getMailboxInfos().isEmpty(), String.valueOf(mailboxInfos.getMailboxInfos())); + } + // Nothing to receive on: the receiver has no worker to hold an entry. + assertTrue(receiver.getWorkerIdToMailboxesMap().isEmpty()); + } + private static QueryServerInstance server(String id) { return new QueryServerInstance(id, "host_" + id, 1, 1); } @@ -143,12 +203,36 @@ private static DispatchablePlanMetadata metadata(Map keys) { DataSchema dataSchema = new DataSchema(new String[]{"col"}, new ColumnDataType[]{ColumnDataType.INT}); return new MailboxSendNode(SENDER_STAGE, dataSchema, List.of(), RECEIVER_STAGE, PinotRelExchangeType.PIPELINE_BREAKER, RelDistribution.Type.SINGLETON, keys, false, null, false, "absHashCode"); } + private static MailboxSendNode hashSendNode() { + DataSchema dataSchema = new DataSchema(new String[]{"col"}, new ColumnDataType[]{ColumnDataType.INT}); + return new MailboxSendNode(SENDER_STAGE, dataSchema, List.of(), RECEIVER_STAGE, PinotRelExchangeType.STREAMING, + RelDistribution.Type.HASH_DISTRIBUTED, List.of(0), false, null, false, "absHashCode"); + } + + /// The sender worker ids the given receiver worker reads from, in mailbox order. + private static List expandedWorkerIds(Map> mailboxesMap, int workerId, + int stageId) { + List workerIds = new ArrayList<>(); + for (MailboxInfo mailboxInfo : mailboxesMap.get(workerId).get(stageId).getMailboxInfos()) { + workerIds.addAll(mailboxInfo.getWorkerIds()); + } + return workerIds; + } + private static void process(MailboxSendNode sendNode, DispatchablePlanMetadata sender, DispatchablePlanMetadata receiver) { DispatchablePlanContext context = Mockito.mock(DispatchablePlanContext.class); diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/PinotDispatchPlannerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/PinotDispatchPlannerTest.java index a394d508b0ea..08abe33540df 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/PinotDispatchPlannerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/PinotDispatchPlannerTest.java @@ -22,6 +22,7 @@ import java.util.List; import java.util.Map; import org.apache.pinot.query.QueryEnvironmentTestBase; +import org.apache.pinot.query.planner.plannode.MailboxSendNode; import org.apache.pinot.query.planner.plannode.PlanNode; import org.testng.annotations.Test; @@ -104,6 +105,33 @@ public void testRewriteReduceStageWithJoinInlinesAllBranches() { assertAllStageIdsAreZero(root); } + /// `WorkerManager` walks the plan twice (leaves first, then intermediate stages), and with a spool the same + /// `PlanFragment` is a child of every receiver that reads it. Sharing one visited set across both passes would leave + /// a spooled intermediate fragment with no workers: the leaf pass visits it and skips it as a non-leaf, then the + /// intermediate pass skips it as visited. + @Test + public void testSpooledIntermediateStageGetsWorkers() { + DispatchableSubPlan subPlan = _queryEnvironment.planQuery("SET useSpools=true; " + + "WITH mySpool AS (SELECT col1, SUM(col3) AS s FROM a GROUP BY col1) " + + "SELECT 1 FROM mySpool AS a1 JOIN b ON a1.col1 = b.col1 JOIN mySpool AS a2 ON a2.col1 = b.col1"); + + // The spool is only useful if some fragment really is read by more than one receiver. + assertTrue(hasMultiReceiverSend(subPlan), "Query did not produce a spool: " + subPlan.getQueryStageMap().keySet()); + for (Map.Entry entry : subPlan.getQueryStageMap().entrySet()) { + assertFalse(entry.getValue().getWorkerMetadataList().isEmpty(), "No worker for stage: " + entry.getKey()); + } + } + + private static boolean hasMultiReceiverSend(DispatchableSubPlan subPlan) { + for (DispatchablePlanFragment fragment : subPlan.getQueryStageMap().values()) { + PlanNode root = fragment.getPlanFragment().getFragmentRoot(); + if (root instanceof MailboxSendNode && ((MailboxSendNode) root).isMultiSend()) { + return true; + } + } + return false; + } + private static void assertAllStageIdsAreZero(PlanNode node) { assertEquals(node.getStageId(), 0); for (PlanNode input : node.getInputs()) { diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java new file mode 100644 index 000000000000..3a779db87312 --- /dev/null +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java @@ -0,0 +1,357 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.routing; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.rel.RelDistribution; +import org.apache.pinot.calcite.rel.hint.PinotHintOptions; +import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.query.planner.PlanFragment; +import org.apache.pinot.query.planner.physical.DispatchablePlanMetadata; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; +import org.apache.pinot.query.planner.plannode.MailboxSendNode; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +/// Tests the plan-shape classification [ColocationGroupAnalyzer] does. Which partition classes actually survive is +/// decided by [WorkerManager] and covered by `WorkerManagerTest`. +public class ColocationGroupAnalyzerTest { + private static final DataSchema SCHEMA = + new DataSchema(new String[]{"col1"}, new ColumnDataType[]{ColumnDataType.INT}); + private static final String HASH_FUNCTION = "absHashCodeSum"; + + /// The plan shape a colocated join takes: both leaves are pre-partitioned and send 1-to-1 to the join stage, so they + /// and the stages they feed form one reducible group. + @Test + public void testGroupWithOnlyPrePartitionedSendsIsReducible() { + List groups = + ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap(true)); + + assertEquals(groups.size(), 1); + assertEquals(groups.get(0)._partitionSize, 4); + assertEquals(Set.copyOf(groups.get(0)._partitionedLeafFragmentIds), Set.of(2, 3)); + } + + /// A member that also receives a shuffled send must keep today's worker count, or that sender's rows land on + /// different workers than the 1-to-1 side's; see ColocationGroupAnalyzer#findReducibleGroups. + @Test + public void testGroupWithAShuffledSendIntoAMemberIsNotReducible() { + List groups = + ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap(false)); + + assertTrue(groups.isEmpty(), String.valueOf(groups.size())); + } + + /// A SINGLETON send ties the two stages together even when the sender is not marked pre-partitioned, because the + /// receiver still copies its worker map from the sender. + @Test + public void testSingletonSendFormsAnEdgeWithoutPrePartitioning() { + Map metadataMap = metadataMap(false); + // Both leaves send SINGLETON, and neither is pre-partitioned. + metadataMap.get(2).setPrePartitioned(false); + List groups = + ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(RelDistribution.Type.SINGLETON), metadataMap); + + assertEquals(groups.size(), 1); + assertEquals(Set.copyOf(groups.get(0)._partitionedLeafFragmentIds), Set.of(2, 3)); + } + + /// Reducing the worker count must not turn mismatched counts into a match for a pre-partitioned BROADCAST send, which + /// would then be wired 1-to-1; see ColocationGroupAnalyzer#findReducibleGroups. + @Test + public void testGroupWithPrePartitionedBroadcastSendIsNotReducible() { + List groups = ColocationGroupAnalyzer.findReducibleGroups( + twoLeafPlan(RelDistribution.Type.HASH_DISTRIBUTED, RelDistribution.Type.BROADCAST_DISTRIBUTED), + metadataMap(true)); + + assertTrue(groups.isEmpty(), String.valueOf(groups.size())); + } + + /// A lone fragment is tied to nothing, so its worker ids owe nothing to another stage and its assignment is kept. + @Test + public void testLoneFragmentComponentIsNotReducible() { + // The single leaf is not pre-partitioned and shuffles into the reduce stage, so no edge is formed at all and the + // leaf ends up in a component of its own. + PlanFragment leaf = new PlanFragment(1, sendNode(1, 0, RelDistribution.Type.HASH_DISTRIBUTED), List.of()); + PlanFragment root = new PlanFragment(0, receiveNode(1), List.of(leaf)); + Map metadataMap = new HashMap<>(); + metadataMap.put(0, new DispatchablePlanMetadata()); + metadataMap.put(1, partitionedLeafMetadata("tableA", false, "4", null)); + + assertTrue(ColocationGroupAnalyzer.findReducibleGroups(root, metadataMap).isEmpty()); + } + + /// Leaves that disagree on the hinted partition size cannot share a worker-id-to-class mapping: worker `k` would + /// stand for class `k mod 4` on one and `k mod 8` on the other. Same for the parallelism, which sizes the derived + /// stages. + @Test + public void testGroupWithMismatchedPartitionSizeIsNotReducible() { + Map metadataMap = metadataMap(true); + metadataMap.put(3, partitionedLeafMetadata("tableB", true, "8", null)); + + assertTrue(ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap).isEmpty()); + } + + @Test + public void testGroupWithMismatchedPartitionParallelismIsNotReducible() { + Map metadataMap = metadataMap(true); + metadataMap.put(2, partitionedLeafMetadata("tableA", true, "4", "2")); + metadataMap.put(3, partitionedLeafMetadata("tableB", true, "4", "3")); + + assertTrue(ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap).isEmpty()); + } + + /// The control for the two tests above: the same shape agreeing on a partition parallelism above 1 is reducible. + @Test + public void testGroupWithMatchingPartitionParallelismIsReducible() { + Map metadataMap = metadataMap(true); + metadataMap.put(2, partitionedLeafMetadata("tableA", true, "4", "2")); + metadataMap.put(3, partitionedLeafMetadata("tableB", true, "4", "2")); + + assertEquals(ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap).size(), 1); + } + + /// Agreeing on the partition size is not enough: two functions put different keys in class `j`, see + /// ColocationGroupAnalyzer#toReducibleGroup. + @Test + public void testGroupWithMismatchedPartitionFunctionIsNotReducible() { + Map metadataMap = metadataMap(true); + metadataMap.put(2, partitionedLeafMetadata("tableA", true, "4", null, "Murmur")); + metadataMap.put(3, partitionedLeafMetadata("tableB", true, "4", null, "HashCode")); + + assertTrue(ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap).isEmpty()); + } + + /// An omitted partition function hint is not resolved to the default here, so it does not match an explicit one. + @Test + public void testGroupWithOnlyOneLeafHintingAPartitionFunctionIsNotReducible() { + Map metadataMap = metadataMap(true); + metadataMap.put(3, partitionedLeafMetadata("tableB", true, "4", null, "Murmur")); + + assertTrue(ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap).isEmpty()); + } + + /// The control for the two tests above: function names are compared case-insensitively, as elsewhere in the engine. + @Test + public void testGroupWithMatchingPartitionFunctionIsReducible() { + Map metadataMap = metadataMap(true); + metadataMap.put(2, partitionedLeafMetadata("tableA", true, "4", null, "Murmur")); + metadataMap.put(3, partitionedLeafMetadata("tableB", true, "4", null, "murmur")); + + assertEquals(ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap).size(), 1); + } + + /// A leaf with no table hints, or none declaring a partition key, is assigned over servers rather than partitions -- + /// the `is_colocated_by_join_keys` escape hatch, which must keep working; see + /// ColocationGroupAnalyzer#toReducibleGroup. + @Test + public void testGroupWithALeafWithoutTableOptionsIsNotReducible() { + Map metadataMap = metadataMap(true); + DispatchablePlanMetadata noHints = new DispatchablePlanMetadata(); + noHints.addScannedTable("tableB"); + noHints.setPrePartitioned(true); + metadataMap.put(3, noHints); + + assertTrue(ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap).isEmpty()); + } + + @Test + public void testGroupWithALeafWithoutPartitionKeyIsNotReducible() { + Map metadataMap = metadataMap(true); + DispatchablePlanMetadata noPartitionKey = new DispatchablePlanMetadata(); + noPartitionKey.addScannedTable("tableB"); + noPartitionKey.setTableOptions(Map.of(PinotHintOptions.TableHintOptions.PARTITION_SIZE, "4")); + noPartitionKey.setPrePartitioned(true); + metadataMap.put(3, noPartitionKey); + + assertTrue(ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap).isEmpty()); + } + + /// An invalid partition size is left for the leaf assignment to report, rather than being interpreted here. + @Test + public void testGroupWithInvalidPartitionSizeIsNotReducible() { + for (String partitionSize : new String[]{"0", "-4", "four"}) { + Map metadataMap = metadataMap(true); + metadataMap.put(3, partitionedLeafMetadata("tableB", true, partitionSize, null)); + + assertTrue(ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap).isEmpty(), partitionSize); + } + } + + @Test + public void testGroupWithInvalidPartitionParallelismIsNotReducible() { + Map metadataMap = metadataMap(true); + metadataMap.put(3, partitionedLeafMetadata("tableB", true, "4", "0")); + + assertTrue(ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap).isEmpty()); + } + + /// A replicated leaf constrains no class (see LeafPartitionHints#isReplicated), so a group mixing one with a + /// partitioned fact table stays reducible -- without that, its missing partition key would reject the whole group. + @Test + public void testReplicatedLeafDoesNotBlockTheGroup() { + Map metadataMap = metadataMap(true); + DispatchablePlanMetadata replicated = new DispatchablePlanMetadata(); + replicated.addScannedTable("dimTable"); + replicated.setTableOptions(Map.of(PinotHintOptions.TableHintOptions.IS_REPLICATED, "true")); + replicated.setPrePartitioned(true); + metadataMap.put(3, replicated); + + List groups = + ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(), metadataMap); + + assertEquals(groups.size(), 1); + // Only the partitioned leaf decides which classes survive. + assertEquals(groups.get(0)._partitionedLeafFragmentIds, List.of(2)); + } + + /// A lookup join's workers come from its single local exchange child, so its own hints (a different partition size + /// here) are not a constraint on the group. + @Test + public void testLookupJoinMemberIsIgnored() { + PlanFragment localExchangeChild = new PlanFragment(2, sendNode(2, 1, RelDistribution.Type.SINGLETON), List.of()); + PlanFragment lookupJoin = + new PlanFragment(1, sendNode(1, 0, RelDistribution.Type.SINGLETON), List.of(localExchangeChild)); + PlanFragment root = new PlanFragment(0, receiveNode(1), List.of(lookupJoin)); + Map metadataMap = new HashMap<>(); + metadataMap.put(0, new DispatchablePlanMetadata()); + // The lookup join stage scans the dimension table itself, with hints of its own. + metadataMap.put(1, partitionedLeafMetadata("dimTable", false, "8", null)); + metadataMap.put(2, partitionedLeafMetadata("tableA", false, "4", null)); + + List groups = ColocationGroupAnalyzer.findReducibleGroups(root, + metadataMap); + + assertEquals(groups.size(), 1); + assertEquals(groups.get(0)._partitionSize, 4); + assertEquals(groups.get(0)._partitionedLeafFragmentIds, List.of(2)); + } + + /// A group of intermediate stages only has nothing to reduce: only a partitioned leaf's data decides the classes. + @Test + public void testGroupWithoutAPartitionedLeafIsNotReducible() { + PlanFragment intermediate = new PlanFragment(1, sendNode(1, 0, RelDistribution.Type.SINGLETON), List.of()); + PlanFragment root = new PlanFragment(0, receiveNode(1), List.of(intermediate)); + Map metadataMap = new HashMap<>(); + metadataMap.put(0, new DispatchablePlanMetadata()); + metadataMap.put(1, new DispatchablePlanMetadata()); + + assertTrue(ColocationGroupAnalyzer.findReducibleGroups(root, metadataMap).isEmpty()); + } + + /// With a spool the same fragment is a child of every receiver that reads it, and its send node lists all of them: + /// every receiver must end up in the spooled sender's group, and the sender must be visited only once. + @Test + public void testSpooledFragmentTiesEveryReceiverIntoOneGroup() { + PlanFragment spooledLeaf = new PlanFragment(3, + new MailboxSendNode(3, SCHEMA, List.of(), List.of(1, 2), PinotRelExchangeType.STREAMING, + RelDistribution.Type.HASH_DISTRIBUTED, List.of(0), false, null, false, HASH_FUNCTION), List.of()); + PlanFragment firstReceiver = + new PlanFragment(1, sendNode(1, 0, RelDistribution.Type.SINGLETON), List.of(spooledLeaf)); + PlanFragment secondReceiver = + new PlanFragment(2, sendNode(2, 0, RelDistribution.Type.SINGLETON), List.of(spooledLeaf)); + PlanFragment root = new PlanFragment(0, receiveNode(1), List.of(firstReceiver, secondReceiver)); + Map metadataMap = new HashMap<>(); + metadataMap.put(0, new DispatchablePlanMetadata()); + metadataMap.put(1, new DispatchablePlanMetadata()); + metadataMap.put(2, new DispatchablePlanMetadata()); + metadataMap.put(3, partitionedLeafMetadata("tableA", true, "4", null)); + + List groups = + ColocationGroupAnalyzer.findReducibleGroups(root, metadataMap); + + // One group, and the spooled leaf is listed once rather than once per receiver. + assertEquals(groups.size(), 1); + assertEquals(groups.get(0)._partitionedLeafFragmentIds, List.of(3)); + } + + /// Builds a 4 stage plan: 2 partitioned leaves (stages 2 and 3) sending to a join stage (stage 1), which sends + /// SINGLETON to the broker reduce stage (stage 0). + private static PlanFragment twoLeafPlan() { + return twoLeafPlan(RelDistribution.Type.HASH_DISTRIBUTED); + } + + private static PlanFragment twoLeafPlan(RelDistribution.Type leafDistributionType) { + return twoLeafPlan(leafDistributionType, leafDistributionType); + } + + /// Same as [#twoLeafPlan()], with the distribution type of each leaf's send. + private static PlanFragment twoLeafPlan(RelDistribution.Type firstLeafDistributionType, + RelDistribution.Type secondLeafDistributionType) { + PlanFragment firstLeaf = new PlanFragment(2, sendNode(2, 1, firstLeafDistributionType), List.of()); + PlanFragment secondLeaf = new PlanFragment(3, sendNode(3, 1, secondLeafDistributionType), List.of()); + PlanFragment joinFragment = + new PlanFragment(1, sendNode(1, 0, RelDistribution.Type.SINGLETON), List.of(firstLeaf, secondLeaf)); + return new PlanFragment(0, receiveNode(1), List.of(joinFragment)); + } + + private static MailboxReceiveNode receiveNode(int senderStageId) { + return new MailboxReceiveNode(0, SCHEMA, senderStageId, PinotRelExchangeType.STREAMING, + RelDistribution.Type.SINGLETON, null, null, false, false, null); + } + + private static MailboxSendNode sendNode(int stageId, int receiverStageId, RelDistribution.Type distributionType) { + return new MailboxSendNode(stageId, SCHEMA, List.of(), receiverStageId, PinotRelExchangeType.STREAMING, + distributionType, List.of(0), false, null, false, HASH_FUNCTION); + } + + /// The metadata for [#twoLeafPlan()]. The second leaf is pre-partitioned -- i.e. its hash send may be wired 1-to-1 -- + /// only when `prePartitionSecondLeaf` is set, otherwise its send is a plain shuffle into the join stage. + private static Map metadataMap(boolean prePartitionSecondLeaf) { + Map metadataMap = new HashMap<>(); + metadataMap.put(0, new DispatchablePlanMetadata()); + metadataMap.put(1, new DispatchablePlanMetadata()); + metadataMap.put(2, partitionedLeafMetadata("tableA", true, "4", null)); + metadataMap.put(3, partitionedLeafMetadata("tableB", prePartitionSecondLeaf, "4", null)); + return metadataMap; + } + + private static DispatchablePlanMetadata partitionedLeafMetadata(String tableName, boolean prePartitioned, + String partitionSize, @Nullable String partitionParallelism) { + return partitionedLeafMetadata(tableName, prePartitioned, partitionSize, partitionParallelism, null); + } + + /// A hint of `null` is left out of the table options altogether, i.e. the leaf does not declare that option. + private static DispatchablePlanMetadata partitionedLeafMetadata(String tableName, boolean prePartitioned, + String partitionSize, @Nullable String partitionParallelism, @Nullable String partitionFunction) { + DispatchablePlanMetadata metadata = new DispatchablePlanMetadata(); + metadata.addScannedTable(tableName); + Map tableOptions = new HashMap<>(); + tableOptions.put(PinotHintOptions.TableHintOptions.PARTITION_KEY, "col1"); + tableOptions.put(PinotHintOptions.TableHintOptions.PARTITION_SIZE, partitionSize); + if (partitionParallelism != null) { + tableOptions.put(PinotHintOptions.TableHintOptions.PARTITION_PARALLELISM, partitionParallelism); + } + if (partitionFunction != null) { + tableOptions.put(PinotHintOptions.TableHintOptions.PARTITION_FUNCTION, partitionFunction); + } + metadata.setTableOptions(tableOptions); + metadata.setPrePartitioned(prePartitioned); + return metadata; + } +} diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java index c5071f4aef93..60fe0a7b93e5 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java @@ -44,6 +44,7 @@ import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; import org.apache.pinot.query.planner.physical.DispatchableSubPlan; import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.utils.CommonConstants; @@ -56,9 +57,12 @@ import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotEquals; import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNotSame; import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; /// Tests for [WorkerManager]. @@ -835,12 +839,1061 @@ public void testBrokerPruningPartitionedLeafHybridTable() { } } + @Test + public void testHybridPartitionedLeafRejectsSegmentsWithInvalidPartition() { + // The hybrid branch merges the two sides' maps itself instead of going through + // PartitionTableInfo.fromTablePartitionInfo, so it must run the invalid-partition check on both. Here only the + // realtime side has such a segment. + QueryEnvironment queryEnvironment = newHybridPartitionedQueryEnvironment(List.of("segO2"), List.of("segR1"), + List.of(), List.of("segRbad")); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SELECT col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ " + + "WHERE col1 = 'foo'")) { + // QueryEnvironment wraps the planning failure, so assert on the cause. + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("segments with invalid partition"), cause.getMessage()); + assertTrue(cause.getMessage().contains("testTable_REALTIME"), cause.getMessage()); + } + } + + @Test + public void testPartitionedLeafRejectsPartitionWithOnlyDeferredSegments() { + // Partition 2 has no entry in the partition info map, but not because it is empty: all of its segments are + // deferred, so no single server can scan it whole and padding it would silently drop its rows. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment(Set.of(2), Set.of(2), Set.of(2), Set.of()); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("Failed to find a fully replicated server for partitions: [2]"), + cause.getMessage()); + // The message names the scanned (raw) table name. + assertTrue(cause.getMessage().contains("of table: " + COLOCATED_TABLE_A), cause.getMessage()); + } + } + + @Test + public void testPlainPartitionedLeafRejectsPartitionWithOnlyDeferredSegmentsWithoutPruning() { + // A plain partitioned leaf, outside any colocated group, so only the check at the assignment site can fire. Pruning + // is off, so partition 3 needs a worker and has nowhere to go: padding it would drop the held-back segments' rows. + QueryEnvironment queryEnvironment = + newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4, 1, List.of("seg2"), List.of(), 0, false, Set.of(3), + Set.of(3)); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=false; SELECT col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ " + + "WHERE col1 = 'foo'")) { + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("Failed to find a fully replicated server for partitions: [3]"), + cause.getMessage()); + assertTrue(cause.getMessage().contains("of table: " + PARTITIONED_TABLE), cause.getMessage()); + } + } + + @Test + public void testPlainPartitionedLeafWithPartitionWithOnlyDeferredSegmentsStillPrunes() { + // Same layout, with broker pruning active and a filter that only matches partition 2. The deferred partition gets + // no worker either way, so the query must keep planning: failing every query on the table while a segment is new + // would be a bigger regression than the rows this one cannot see. + QueryEnvironment queryEnvironment = + newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4, 1, List.of("seg2"), List.of(), 0, false, Set.of(3), + Set.of(3)); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; SELECT col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ " + + "WHERE col1 = 'foo'")) { + DispatchablePlanFragment leaf = leafFragment(compiledQuery.planQuery(0).getQueryPlan()); + assertNotNull(leaf); + assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 1); + assertEquals(assignedSegments(leaf), List.of("seg2")); + } + } + + @Test + public void testPartitionedLeafRejectsTableWithoutAnyPartition() { + // An empty partition info map passes the "partitions must be a multiple of the hinted partition size" check + // trivially, leaving 0 partitions per worker, so it has to be rejected on its own. + QueryEnvironment queryEnvironment = newPartitionedQueryEnvironment(new int[0], 4, List.of(), 0); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SELECT col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ " + + "WHERE col1 = 'foo'")) { + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("Failed to find any partition for table: " + PARTITIONED_TABLE), + cause.getMessage()); + } + } + + @Test + public void testPartitionedLeafPublishesACopyOfTheBrokerSegmentList() { + // The lists of a one-partition-per-worker assignment come from the broker's published metadata and are handed to + // filterLeafStageSegments, which may edit them in place, so they must be copied first. + QueryEnvironment queryEnvironment = newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4, List.of(), 0); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=false; SELECT col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ " + + "WHERE col1 = 'foo'")) { + DispatchablePlanFragment leaf = leafFragment(compiledQuery.planQuery(0).getQueryPlan()); + assertNotNull(leaf); + List segments = leaf.getWorkerIdToSegmentsMap().get(0).get(TableType.OFFLINE.name()); + assertEquals(segments, List.of("seg0")); + segments.add("mutated"); + // Planning the same query again must see the broker's original list, not the mutation above. + DispatchablePlanFragment leafAgain = leafFragment(compiledQuery.planQuery(1).getQueryPlan()); + assertNotNull(leafAgain); + List segmentsAgain = leafAgain.getWorkerIdToSegmentsMap().get(0).get(TableType.OFFLINE.name()); + assertEquals(segmentsAgain, List.of("seg0")); + assertNotSame(segmentsAgain, segments); + } + } + + @Test + public void testHybridPartitionedLeafRejectsOfflineSegmentsWithInvalidPartition() { + // The mirror of testHybridPartitionedLeafRejectsSegmentsWithInvalidPartition: the merged map is built from both + // sides, so the check has to run on both. Here only the offline side has such a segment. + QueryEnvironment queryEnvironment = newHybridPartitionedQueryEnvironment(List.of("segO2"), List.of("segR1"), + List.of("segObad"), List.of()); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SELECT col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ " + + "WHERE col1 = 'foo'")) { + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("segments with invalid partition"), cause.getMessage()); + assertTrue(cause.getMessage().contains(PARTITIONED_TABLE_OFFLINE), cause.getMessage()); + } + } + + @Test + public void testHybridPartitionedLeafKeepsPartitionDeferredOnOneSideOnly() { + // Partition 3's offline segments were all held back, but the realtime side still serves the whole partition, so the + // merged map has an entry for it. Reporting it would fail a query the realtime side can answer on its own. + QueryEnvironment queryEnvironment = newHybridPartitionedQueryEnvironment(List.of(), List.of(), List.of(), List.of(), + Set.of(3), Set.of(3)); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=false; SELECT col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ " + + "WHERE col1 = 'foo'")) { + DispatchablePlanFragment leaf = leafFragment(compiledQuery.planQuery(0).getQueryPlan()); + assertNotNull(leaf); + assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 4); + // Worker 3 is realtime-only, the other 3 workers carry both table types. + assertEquals(leaf.getWorkerIdToSegmentsMap().get(3).keySet(), Set.of(TableType.REALTIME.name())); + assertEquals(assignedSegments(leaf, 3), List.of("segR3")); + assertEquals(leaf.getWorkerIdToSegmentsMap().get(0).keySet(), + Set.of(TableType.OFFLINE.name(), TableType.REALTIME.name())); + } + } + + @Test + public void testColocatedJoinDropsEmptyPartitionOnBothSides() { + // Partition 3 holds no segment on either side of the colocated join. Both leaves must drop it, keeping the same + // worker id -> partition mapping so that the 1-to-1 exchange between them stays correct. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment(Set.of(3), Set.of(3), Set.of(), Set.of()); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + List leafFragments = leafFragments(dispatchableSubPlan); + assertEquals(leafFragments.size(), 2); + for (DispatchablePlanFragment leaf : leafFragments) { + // 3 workers instead of 4, one per surviving partition, and partition 3's (absent) segment is not assigned. + assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 3); + String segmentPrefix = leaf.getTableName().startsWith(COLOCATED_TABLE_A) ? "a_seg" : "b_seg"; + assertEquals(new HashSet<>(assignedSegments(leaf)), + Set.of(segmentPrefix + "0", segmentPrefix + "1", segmentPrefix + "2")); + } + // Worker k of both leaves must hold partition k's segment and live on partition k's server, otherwise the 1-to-1 + // exchange would pair rows of different partitions. + Map workerIdToServerA = workerIdToServer(leafFragments.get(0)); + Map workerIdToServerB = workerIdToServer(leafFragments.get(1)); + assertEquals(workerIdToServerA, workerIdToServerB); + for (int workerId = 0; workerId < 3; workerId++) { + assertEquals(assignedSegments(leafFragments.get(0), workerId).size(), 1); + // Partition p lives on server p, i.e. localhost:p+1 (see newColocatedJoinQueryEnvironment). + assertTrue(workerIdToServerA.get(workerId).endsWith("_" + (workerId + 1)), workerIdToServerA.get(workerId)); + } + // The join stage takes its workers from the leaves, so it must be reduced along with them, and it must land on + // their servers: the exchange is still a 1-to-1 local exchange rather than a shuffle across all the servers. + DispatchablePlanFragment joinFragment = joinFragment(dispatchableSubPlan); + assertEquals(joinFragment.getWorkerMetadataList().size(), 3); + assertEquals(workerIdToServer(joinFragment), workerIdToServerA); + } + } + + @Test + public void testColocatedJoinPadsPartitionEmptyOnOneSide() { + // Table A is empty in partition 3 but table B is not, so the group keeps the class and table A gets a worker with + // no segments for it. Here the server holding B's partition 3 does not host table A at all (each partition lives on + // its own server), so that worker cannot be placed with its peer and falls back to a server that does host A. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment(Set.of(3), Set.of(), Set.of(), Set.of()); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + List leafFragments = leafFragments(dispatchableSubPlan); + assertEquals(leafFragments.size(), 2); + DispatchablePlanFragment leafA = leafFragments.get(0).getTableName().startsWith(COLOCATED_TABLE_A) + ? leafFragments.get(0) : leafFragments.get(1); + DispatchablePlanFragment leafB = leafA == leafFragments.get(0) ? leafFragments.get(1) : leafFragments.get(0); + // Both sides keep all 4 workers, one per partition class, so worker k still stands for partition k on both. + assertEquals(leafA.getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafB.getWorkerIdToSegmentsMap().size(), 4); + for (int workerId = 0; workerId < 3; workerId++) { + assertEquals(assignedSegments(leafA, workerId), List.of("a_seg" + workerId)); + } + assertEquals(assignedSegments(leafB, 3), List.of("b_seg3")); + // A single table type key mapped to an empty (mutable) list, on a server that hosts table A rather than on + // partition 3's server (localhost_4), which only hosts table B's partition 3. + Map> emptyWorkerSegmentsMap = leafA.getWorkerIdToSegmentsMap().get(3); + assertEquals(emptyWorkerSegmentsMap.keySet(), Set.of(TableType.OFFLINE.name())); + assertEquals(emptyWorkerSegmentsMap.get(TableType.OFFLINE.name()), List.of()); + emptyWorkerSegmentsMap.get(TableType.OFFLINE.name()).add("mutable"); + String emptyWorkerServer = workerIdToServer(leafA).get(3); + assertTrue(Set.of("_1", "_2", "_3").stream().anyMatch(emptyWorkerServer::endsWith), emptyWorkerServer); + // The join stage takes its workers from a leaf, so it keeps all 4 workers as well. + assertEquals(joinFragment(dispatchableSubPlan).getWorkerMetadataList().size(), 4); + } + } + + @Test + public void testColocatedJoinPadsWorkerOnPeerServer() { + // Same as above, but every server hosts every partition of both tables, so the empty worker can borrow both the + // candidate servers and the seed of the peer holding partition 3, which is what keeps the exchange in process. + QueryEnvironment queryEnvironment = + newColocatedJoinQueryEnvironment(Set.of(3), Set.of(), Set.of(), Set.of(), true); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + List leafFragments = leafFragments(dispatchableSubPlan); + assertEquals(leafFragments.size(), 2); + DispatchablePlanFragment leafA = leafFragments.get(0).getTableName().startsWith(COLOCATED_TABLE_A) + ? leafFragments.get(0) : leafFragments.get(1); + DispatchablePlanFragment leafB = leafA == leafFragments.get(0) ? leafFragments.get(1) : leafFragments.get(0); + assertEquals(leafA.getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafB.getWorkerIdToSegmentsMap().size(), 4); + assertEquals(assignedSegments(leafA, 3), List.of()); + assertEquals(assignedSegments(leafB, 3), List.of("b_seg3")); + // Every worker of the padding side lands on the same server as its peer, the empty one included. + assertEquals(workerIdToServer(leafA), workerIdToServer(leafB)); + DispatchablePlanFragment joinFragment = joinFragment(dispatchableSubPlan); + assertEquals(workerIdToServer(joinFragment), workerIdToServer(leafA)); + // The empty worker is wired like any other one: it has an outbound mailbox and the join worker reading its class + // has an inbound one. Dropping either would write its end-of-stream block, and any error, where nobody reads. + int leafAFragmentId = leafA.getPlanFragment().getFragmentId(); + int joinFragmentId = joinFragment.getPlanFragment().getFragmentId(); + assertNotNull(leafA.getWorkerMetadataList().get(3).getMailboxInfosMap().get(joinFragmentId)); + assertNotNull(joinFragment.getWorkerMetadataList().get(3).getMailboxInfosMap().get(leafAFragmentId)); + // Landing with the peer is what keeps the exchange in process: the sender and the join worker reading it share + // one local mailbox rather than a cross-server pair. + MailboxInfos mailboxInfos = + joinFragment.getWorkerMetadataList().get(2).getMailboxInfosMap().get(leafAFragmentId); + assertNotNull(mailboxInfos); + assertTrue(mailboxInfos instanceof SharedMailboxInfos, String.valueOf(mailboxInfos)); + assertEquals(mailboxInfos.getMailboxInfos().size(), 1); + assertEquals(mailboxInfos.getMailboxInfos().get(0).getHostname(), "localhost"); + assertEquals(mailboxInfos.getMailboxInfos().get(0).getWorkerIds(), List.of(2)); + } + } + + @Test + public void testColocatedJoinPadsWorkerOnPeerServerRatherThanItsOwn() { + // Borrowing the peer's candidate servers is the only thing that keeps the empty worker with its peer, and here the + // peer's set is a strict subset of the servers hosting table A, so the two resolve differently: + // - table A holds partitions 0..2 on servers 1, 2 and {3, 4}, so its own candidate set is all 4 servers; + // - table B's partition 3 lives on server 1 alone, so borrowing lands the empty worker there; + // - picking from table A's own set would land it on server 3 instead. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).emptyPartitions(Set.of(3)) + .partitionServerIndexes(Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2, 3))), + new ColocatedTableSpec(4, false) + .partitionServerIndexes(Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2), 3, Set.of(0)))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + // Worker 3 is the empty one: it stands for class 3, which table A holds no data in. + assertEquals(assignedSegments(leafA, 3), List.of()); + String peerServer = getServerInstance("localhost", 1).getInstanceId(); + assertEquals(workerIdToServer(leafB).get(3), peerServer); + // The discriminating assertions: on the peer's server, not on the one table A's own candidate set resolves to. + assertEquals(workerIdToServer(leafA).get(3), peerServer); + assertNotEquals(workerIdToServer(leafA).get(3), getServerInstance("localhost", 3).getInstanceId()); + } + } + + @Test + public void testColocatedJoinPadsRealtimeWorkerWithRealtimeSegmentsMap() { + // The one table type key an empty worker emits must be one the chosen server actually has a table data manager for. + // Every other colocated test registers offline tables only, so this covers the realtime branch. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, true).emptyPartitions(Set.of(3)), new ColocatedTableSpec(4, true), + TableType.REALTIME); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + Map> emptyWorkerSegmentsMap = leafA.getWorkerIdToSegmentsMap().get(3); + assertEquals(emptyWorkerSegmentsMap.keySet(), Set.of(TableType.REALTIME.name())); + assertEquals(emptyWorkerSegmentsMap.get(TableType.REALTIME.name()), List.of()); + } + } + + @Test + public void testColocatedNonEquiJoinIsNotReduced() { + // A non-equi colocated join sends one side BROADCAST with prePartitioned set, which reducing the worker count must + // not wire 1-to-1 (see ColocationGroupAnalyzer#findReducibleGroups), so the group keeps today's assignment -- and + // today's assignment rejects the empty partition. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment(Set.of(3), Set.of(3), Set.of(), Set.of()); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_NON_EQUI_JOIN_QUERY)) { + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("Failed to find any segment for table"), cause.getMessage()); + } + } + + @Test + public void testColocatedJoinPadsClassEmptyOnOneSideWithMultiplePartitionsPerWorker() { + // 8 partitions per table over a hinted partition size of 4, so worker k handles the class {k, k + 4}. Table A holds + // no segment in either partition of class 3, so it pads that class while table B keeps its 2 segments there. + QueryEnvironment queryEnvironment = + newColocatedJoinQueryEnvironment(Set.of(3, 7), Set.of(), Set.of(), Set.of(), true, 8); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + List leafFragments = leafFragments(dispatchableSubPlan); + assertEquals(leafFragments.size(), 2); + DispatchablePlanFragment leafA = leafFragments.get(0).getTableName().startsWith(COLOCATED_TABLE_A) + ? leafFragments.get(0) : leafFragments.get(1); + DispatchablePlanFragment leafB = leafA == leafFragments.get(0) ? leafFragments.get(1) : leafFragments.get(0); + assertEquals(leafA.getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafB.getWorkerIdToSegmentsMap().size(), 4); + // The workers of the surviving classes are not shifted by the empty one, which keeps its own class index. + for (int workerId = 0; workerId < 3; workerId++) { + assertEquals(new HashSet<>(assignedSegments(leafA, workerId)), + Set.of("a_seg" + workerId, "a_seg" + (workerId + 4))); + } + assertEquals(assignedSegments(leafA, 3), List.of()); + assertEquals(new HashSet<>(assignedSegments(leafB, 3)), Set.of("b_seg3", "b_seg7")); + // Same shape as on the one-partition-per-worker path: one table type key, mapped to a mutable empty list. + Map> emptyWorkerSegmentsMap = leafA.getWorkerIdToSegmentsMap().get(3); + assertEquals(emptyWorkerSegmentsMap.keySet(), Set.of(TableType.OFFLINE.name())); + assertEquals(emptyWorkerSegmentsMap.get(TableType.OFFLINE.name()), List.of()); + emptyWorkerSegmentsMap.get(TableType.OFFLINE.name()).add("mutable"); + // Every server hosts every partition here, so the empty worker lands with its peer. + assertEquals(workerIdToServer(leafA), workerIdToServer(leafB)); + } + } + + @Test + public void testColocatedJoinRejectsFullyEmptyTable() { + QueryEnvironment queryEnvironment = + newColocatedJoinQueryEnvironment(Set.of(0, 1, 2, 3), Set.of(), Set.of(), Set.of()); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("Failed to find any segment in any partition for table: " + + COLOCATED_TABLE_A), cause.getMessage()); + } + } + + @Test + public void testColocatedJoinAlignsWorkersWhenEmptyClassesDiffer() { + // The case a naive "skip the empty partition" fix gets wrong: table A holds no segment in partition 1 and table B + // holds none in partition 2, so skipping what each side is missing would leave both with 3 workers and mispair them + // (see DispatchablePlanMetadata#getPartitionClassIds). Taking the union keeps both partitions on both sides. + QueryEnvironment queryEnvironment = + newColocatedJoinQueryEnvironment(Set.of(1), Set.of(2), Set.of(), Set.of(), true); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + // Neither side dropped a class the other kept, so the worker counts agree for the right reason. + assertEquals(leafA.getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafB.getWorkerIdToSegmentsMap().size(), 4); + // Each side scans its own partition on the 3 workers it has data for, at that partition's structural index. + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0), 2, Set.of(2), 3, Set.of(3))); + assertEquals(workerIdToPartitions(leafB, "b_seg"), Map.of(0, Set.of(0), 1, Set.of(1), 3, Set.of(3))); + // Each side pads the class only the other carries: A pads worker 1, B pads worker 2. + assertEquals(assignedSegments(leafA, 1), List.of()); + assertEquals(assignedSegments(leafB, 2), List.of()); + // The mapping itself, not just the counts: the two sides must agree on every worker id, and cover them all. + Map workerIdToClass = new HashMap<>(); + mergeWorkerIdToClass(workerIdToClass, leafA, "a_seg", 4); + mergeWorkerIdToClass(workerIdToClass, leafB, "b_seg", 4); + assertEquals(workerIdToClass, Map.of(0, 0, 1, 1, 2, 2, 3, 3)); + // Both empty workers land with their peer, so the join still runs on the leaves' servers. + assertEquals(workerIdToServer(leafA), workerIdToServer(leafB)); + assertEquals(workerIdToServer(joinFragment(dispatchableSubPlan)), workerIdToServer(leafA)); + } + } + + @Test + public void testColocatedJoinReducesFanOutToPopulatedClasses() { + // 8 declared partition classes but only 2 populated (partitions 0 and 5), on both sides. Each leaf gets 2 workers + // instead of 8, and the query is only dispatched to the 2 servers holding those classes: the fan-out follows from + // the reduced class list, because the dispatched server set is built from the worker -> server map. + Set emptyPartitions = Set.of(1, 2, 3, 4, 6, 7); + QueryEnvironment queryEnvironment = + newColocatedJoinQueryEnvironment(emptyPartitions, emptyPartitions, Set.of(), Set.of(), false, 8); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(colocatedJoinQuery(8))) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + assertEquals(leafA.getWorkerIdToSegmentsMap().size(), 2); + assertEquals(leafB.getWorkerIdToSegmentsMap().size(), 2); + // Worker 0 -> class 0, worker 1 -> class 5: the worker id is the index in the surviving class list. + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0), 1, Set.of(5))); + assertEquals(workerIdToPartitions(leafB, "b_seg"), Map.of(0, Set.of(0), 1, Set.of(5))); + // Partition p is hosted by server p % 4, so only servers 1 and 2 hold the surviving classes. Nothing in the plan + // may be dispatched to the other 2 servers. + String server1 = getServerInstance("localhost", 1).getInstanceId(); + String server2 = getServerInstance("localhost", 2).getInstanceId(); + assertEquals(new HashSet<>(workerIdToServer(leafA).values()), Set.of(server1, server2)); + assertEquals(new HashSet<>(workerIdToServer(leafB).values()), Set.of(server1, server2)); + assertEquals(joinFragment(dispatchableSubPlan).getWorkerMetadataList().size(), 2); + assertEquals(dispatchedServers(dispatchableSubPlan), Set.of(server1, server2)); + } + } + + @Test + public void testColocatedJoinReducedGroupIgnoresBrokerPruning() { + // A reduced group's worker id is a position in the group's surviving class list, not a running counter over what a + // filter leaves behind, so broker pruning has to be off for its leaves. This shape is the only one that reaches + // that gate: a join written with is_colocated_by_join_keys marks its leaves pre-partitioned and is gated one step + // earlier (see testBrokerPruningPartitionedLeafSkippedForColocatedJoin), while a fact table joined with a + // replicated dimension table over an explicit local exchange is not marked pre-partitioned. + // + // The fact table's class 3 is empty, so the group is reduced to [0, 1, 2], and the filter leaves only class 0. Were + // pruning left on, assignMultiplePartitionsPerWorker would find no segment for classes 1 and 2 and skip them + // WITHOUT consuming a worker id, leaving the leaf one worker while its class list still claimed three. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(8, false).emptyPartitions(Set.of(3, 7)).survivingSegments(List.of("a_seg0", "a_seg4")), + new ColocatedTableSpec(8, false)); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile("SET useBrokerPruning=true; " + + replicatedDimensionJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + // Nothing was pruned, even though the filtered routing query would have dropped 2 of the 3 surviving classes. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + // One worker per surviving class, holding both partitions of that class, and no worker dropped or padded. + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0, 4), 1, Set.of(1, 5), 2, Set.of(2, 6))); + assertEquals(leafA.getWorkerIdToSegmentsMap().keySet(), Set.of(0, 1, 2)); + Map workerIdToClass = new HashMap<>(); + mergeWorkerIdToClass(workerIdToClass, leafA, "a_seg", 4); + assertEquals(workerIdToClass, Map.of(0, 0, 1, 1, 2, 2)); + // The replicated leaf and the join derive their workers from the fact leaf, so they follow it class for class. + assertEquals(leafB.getWorkerIdToSegmentsMap().keySet(), leafA.getWorkerIdToSegmentsMap().keySet()); + assertEquals(workerIdToServer(leafB), workerIdToServer(leafA)); + assertEquals(joinFragment(dispatchableSubPlan).getWorkerMetadataList().size(), 3); + } + } + + @Test + public void testColocatedJoinReducedGroupWithReplicatedLeaf() { + // The most common colocated shape: a partitioned fact table joined with a replicated dimension table over a local + // exchange. The fact table is the group's only source of classes -- a replicated one says nothing about which of + // them hold data (see LeafPartitionHints#isReplicated) -- so its empty class 3 is dropped rather than padded. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment(Set.of(3), Set.of(), Set.of(), Set.of(), true); + try (QueryEnvironment.CompiledQuery compiledQuery = + queryEnvironment.compile(replicatedDimensionJoinQuery(4))) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + // The fact leaf keeps one worker per surviving class and pads nothing. + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2))); + assertEquals(leafA.getWorkerIdToSegmentsMap().keySet(), Set.of(0, 1, 2)); + // The replicated leaf follows the reduced fact leaf: same worker ids, each scanning the whole dimension table on + // the same server as its fact-table peer. + assertEquals(leafB.getWorkerIdToSegmentsMap().keySet(), leafA.getWorkerIdToSegmentsMap().keySet()); + for (Integer workerId : leafB.getWorkerIdToSegmentsMap().keySet()) { + assertEquals(new HashSet<>(assignedSegments(leafB, workerId)), + Set.of("b_seg0", "b_seg1", "b_seg2", "b_seg3")); + } + assertEquals(workerIdToServer(leafB), workerIdToServer(leafA)); + // The join keeps the same 3 workers on the same servers, so both exchanges into it stay 1-to-1 and in process: + // each join worker reads a single local mailbox holding its own worker id from each side. + DispatchablePlanFragment joinFragment = joinFragment(dispatchableSubPlan); + assertEquals(joinFragment.getWorkerMetadataList().size(), 3); + assertEquals(workerIdToServer(joinFragment), workerIdToServer(leafA)); + int leafAFragmentId = leafA.getPlanFragment().getFragmentId(); + int leafBFragmentId = leafB.getPlanFragment().getFragmentId(); + for (int workerId = 0; workerId < 3; workerId++) { + Map mailboxInfosMap = + joinFragment.getWorkerMetadataList().get(workerId).getMailboxInfosMap(); + for (int senderFragmentId : List.of(leafAFragmentId, leafBFragmentId)) { + MailboxInfos mailboxInfos = mailboxInfosMap.get(senderFragmentId); + assertNotNull(mailboxInfos, "No mailbox for sender: " + senderFragmentId + " on worker: " + workerId); + assertTrue(mailboxInfos instanceof SharedMailboxInfos, String.valueOf(mailboxInfos)); + assertEquals(mailboxInfos.getMailboxInfos().size(), 1); + assertEquals(mailboxInfos.getMailboxInfos().get(0).getWorkerIds(), List.of(workerId)); + } + } + } + } + + @Test + public void testColocatedJoinReducedGroupWithPartitionParallelism() { + // With partition_parallelism = p the leaf still gets one worker per surviving class while the stage reading it gets + // p workers per sender, i.e. join worker k handles the class at index k / p. Class 3 is empty on both sides, so + // that arithmetic runs over the reduced list [0, 1, 2] rather than over 0..partitionSize-1. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment(Set.of(3), Set.of(3), Set.of(), Set.of()); + String tableHint = colocatedTableHint(4, 2); + try (QueryEnvironment.CompiledQuery compiledQuery = + queryEnvironment.compile(colocatedJoinQuery(tableHint, tableHint))) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2))); + assertEquals(workerIdToPartitions(leafB, "b_seg"), Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2))); + // 3 surviving classes x parallelism 2. + DispatchablePlanFragment joinFragment = joinFragment(dispatchableSubPlan); + assertEquals(joinFragment.getWorkerMetadataList().size(), 6); + int leafAFragmentId = leafA.getPlanFragment().getFragmentId(); + int joinFragmentId = joinFragment.getPlanFragment().getFragmentId(); + // Receiver k reads sender k / 2, and runs on that sender's server. The two receivers of a sender share its single + // local mailbox, hence SharedMailboxInfos. + for (int workerId = 0; workerId < 6; workerId++) { + MailboxInfos mailboxInfos = + joinFragment.getWorkerMetadataList().get(workerId).getMailboxInfosMap().get(leafAFragmentId); + assertNotNull(mailboxInfos, "No mailbox for table A's leaf on worker: " + workerId); + assertTrue(mailboxInfos instanceof SharedMailboxInfos, String.valueOf(mailboxInfos)); + assertEquals(mailboxInfos.getMailboxInfos().size(), 1); + assertEquals(mailboxInfos.getMailboxInfos().get(0).getWorkerIds(), List.of(workerId / 2)); + assertEquals(workerIdToServer(joinFragment).get(workerId), workerIdToServer(leafA).get(workerId / 2)); + } + // And the other way round: sender k fans out to the contiguous receiver range [2k, 2k + 1]. + for (int workerId = 0; workerId < 3; workerId++) { + MailboxInfos mailboxInfos = + leafA.getWorkerMetadataList().get(workerId).getMailboxInfosMap().get(joinFragmentId); + assertNotNull(mailboxInfos, "No mailbox for the join stage on worker: " + workerId); + assertEquals(mailboxInfos.getMailboxInfos().size(), 1); + assertEquals(mailboxInfos.getMailboxInfos().get(0).getWorkerIds(), List.of(2 * workerId, 2 * workerId + 1)); + } + } + } + + @Test + public void testPartitionedLeafRejectsPartitionWithOnlyDeferredSegmentsWithMultiplePartitionsPerWorker() { + // Worker 3 covers the partition class {3, 7}, where partition 3 is genuinely empty but partition 7 has no entry + // only because all of its segments are deferred, so the class cannot be padded: that would drop partition 7's rows. + QueryEnvironment queryEnvironment = + newColocatedJoinQueryEnvironment(Set.of(3, 7), Set.of(), Set.of(7), Set.of(), true, 8); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("Failed to find a fully replicated server for partitions: [7]"), + cause.getMessage()); + assertTrue(cause.getMessage().contains("of table: " + COLOCATED_TABLE_A), cause.getMessage()); + } + } + + @Test + public void testColocatedJoinRejectsSegmentsWithInvalidPartition() { + // Segments with invalid partition metadata are absent from the partition info map altogether, and unlike an empty + // partition there is nothing to pad: their rows may belong to any partition. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, true).segmentsWithInvalidPartition(List.of("a_segBad")), + new ColocatedTableSpec(4, true)); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("1 segments with invalid partition for table: " + + COLOCATED_TABLE_A_OFFLINE), cause.getMessage()); + } + } + + @Test + public void testColocatedJoinRejectsPartitionWithoutFullyReplicatedServer() { + // Partition 2 of table A holds a segment, but no single server holds the whole partition. There is an entry, so + // nothing to pad, and it must keep failing at the server-pick precondition instead of being reduced away. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).partitionsWithoutFullyReplicatedServer(Set.of(2)), + new ColocatedTableSpec(4, false)); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("Failed to find enabled fully replicated server for table: " + + COLOCATED_TABLE_A), cause.getMessage()); + assertTrue(cause.getMessage().contains("partition: 2"), cause.getMessage()); + } + } + + @Test + public void testColocatedJoinRejectsClassWithoutFullyReplicatedServerWithMultiplePartitionsPerWorker() { + // Same, on the several-partitions-per-worker path: worker 3 covers {3, 7}, where partition 3 is empty and partition + // 7 has no fully replicated server. The class holds data, so it is not padded, and no server can scan it whole. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(8, false).emptyPartitions(Set.of(3)) + .partitionsWithoutFullyReplicatedServer(Set.of(7)), + new ColocatedTableSpec(8, false)); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(COLOCATED_JOIN_QUERY)) { + RuntimeException e = expectThrows(RuntimeException.class, () -> compiledQuery.planQuery(0)); + Throwable cause = e.getCause(); + assertTrue(cause instanceof IllegalStateException, String.valueOf(cause)); + assertTrue(cause.getMessage().contains("Failed to find enabled fully replicated server for table: " + + COLOCATED_TABLE_A), cause.getMessage()); + assertTrue(cause.getMessage().contains("partition class: 3"), cause.getMessage()); + } + } + + // --------------------------------------------------------------------------- + // Partitioned leaf assignment shape invariants + // --------------------------------------------------------------------------- + + @Test + public void testCheckLeafWorkerAssignmentRejectsSparseWorkerIds() { + // DispatchablePlanContext sizes a WorkerMetadata[] from the server map and indexes it by worker id, so a gap would + // leave a null entry (and an out-of-range id would throw an ArrayIndexOutOfBoundsException) there instead of here. + Map serverMap = Map.of(0, queryServerInstance(1), 2, queryServerInstance(2)); + Map>> segmentsMap = + Map.of(0, offlineSegments("seg0"), 2, offlineSegments("seg2")); + IllegalStateException e = expectThrows(IllegalStateException.class, + () -> WorkerManager.checkLeafWorkerAssignment("testTable", serverMap, segmentsMap)); + assertTrue(e.getMessage().contains("Missing server instance for worker: 1"), e.getMessage()); + } + + @Test + public void testCheckLeafWorkerAssignmentRejectsKeySetMismatch() { + Map serverMap = Map.of(0, queryServerInstance(1), 1, queryServerInstance(2)); + Map>> segmentsMap = + Map.of(0, offlineSegments("seg0"), 5, offlineSegments("seg5")); + IllegalStateException e = expectThrows(IllegalStateException.class, + () -> WorkerManager.checkLeafWorkerAssignment("testTable", serverMap, segmentsMap)); + assertTrue(e.getMessage().contains("Missing segments for worker: 1"), e.getMessage()); + } + + @Test + public void testCheckLeafWorkerAssignmentRejectsNullSegmentList() { + Map> nullList = new HashMap<>(); + nullList.put(TableType.OFFLINE.name(), null); + IllegalStateException e = expectThrows(IllegalStateException.class, + () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + Map.of(0, nullList))); + assertTrue(e.getMessage().contains("Null segment list for table type: OFFLINE"), e.getMessage()); + } + + @Test + public void testCheckLeafWorkerAssignmentRejectsEmptyTableTypeMap() { + // The server splits the request on the number of entries in this map, so a worker with no table type at all would + // produce no server request. + IllegalStateException e = expectThrows(IllegalStateException.class, + () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + Map.of(0, Map.of()))); + assertTrue(e.getMessage().contains("Expected 1 or 2 table types for worker: 0, got: 0"), e.getMessage()); + } + + @Test + public void testCheckLeafWorkerAssignmentRejectsThreeTableTypeMap() { + Map> threeTypes = new HashMap<>(); + threeTypes.put(TableType.OFFLINE.name(), List.of()); + threeTypes.put(TableType.REALTIME.name(), List.of()); + threeTypes.put("HYBRID", List.of()); + IllegalStateException e = expectThrows(IllegalStateException.class, + () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + Map.of(0, threeTypes))); + assertTrue(e.getMessage().contains("Expected 1 or 2 table types for worker: 0, got: 3"), e.getMessage()); + } + + @Test + public void testCheckLeafWorkerAssignmentRejectsUnknownTableType() { + // The server resolves one table data manager per key in this map, and reports a missing table for an unknown one. + IllegalStateException e = expectThrows(IllegalStateException.class, + () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + Map.of(0, Map.of("HYBRID", List.of())))); + assertTrue(e.getMessage().contains("Unexpected table type: HYBRID for worker: 0"), e.getMessage()); + } + + @Test + public void testCheckLeafWorkerAssignmentAcceptsHybridAndEmptySegmentWorkers() { + // The two shapes the partitioned assignment produces: a hybrid worker with both table types, and one with a single + // table type mapped to an empty list. + Map> hybridSegments = new HashMap<>(); + hybridSegments.put(TableType.OFFLINE.name(), List.of("segO0")); + hybridSegments.put(TableType.REALTIME.name(), List.of("segR0")); + WorkerManager.checkLeafWorkerAssignment("testTable", + Map.of(0, queryServerInstance(1), 1, queryServerInstance(2)), + Map.of(0, hybridSegments, 1, Map.of(TableType.OFFLINE.name(), new ArrayList<>()))); + } + + private static QueryServerInstance queryServerInstance(int port) { + return new QueryServerInstance(getServerInstance("localhost", port)); + } + + private static Map> offlineSegments(String... segments) { + return Map.of(TableType.OFFLINE.name(), List.of(segments)); + } + + private static final String COLOCATED_TABLE_A = "tableA"; + private static final String COLOCATED_TABLE_A_OFFLINE = "tableA_OFFLINE"; + private static final String COLOCATED_TABLE_B = "tableB"; + private static final String COLOCATED_TABLE_B_OFFLINE = "tableB_OFFLINE"; + private static final String COLOCATED_TABLE_HINT = colocatedTableHint(4); + private static final String COLOCATED_JOIN_QUERY = colocatedJoinQuery(4); + private static final String COLOCATED_NON_EQUI_JOIN_QUERY = + "SELECT /*+ joinOptions(is_colocated_by_join_keys='true') */ " + COLOCATED_TABLE_A + ".col2, " + + COLOCATED_TABLE_B + ".col2 FROM " + COLOCATED_TABLE_A + " " + COLOCATED_TABLE_HINT + "JOIN " + + COLOCATED_TABLE_B + " " + COLOCATED_TABLE_HINT + "ON " + COLOCATED_TABLE_A + ".col3 < " + + COLOCATED_TABLE_B + ".col3"; + + private static String colocatedTableHint(int partitionSize) { + return "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='" + partitionSize + + "') */ "; + } + + /// Same as [#colocatedTableHint(int)], with an explicit `partition_parallelism`. + private static String colocatedTableHint(int partitionSize, int partitionParallelism) { + return "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='" + partitionSize + + "', partition_parallelism='" + partitionParallelism + "') */ "; + } + + /// A colocated (equi) join of "tableA" and "tableB", both hinted with the given `partition_size`. + private static String colocatedJoinQuery(int partitionSize) { + String tableHint = colocatedTableHint(partitionSize); + return colocatedJoinQuery(tableHint, tableHint); + } + + /// Same as [#colocatedJoinQuery(int)], with the table hint of each side given explicitly so that the two sides can + /// differ (e.g. an explicit partition parallelism on both). + private static String colocatedJoinQuery(String tableHintA, String tableHintB) { + return "SELECT /*+ joinOptions(is_colocated_by_join_keys='true') */ " + COLOCATED_TABLE_A + ".col2, " + + COLOCATED_TABLE_B + ".col2 FROM " + COLOCATED_TABLE_A + " " + tableHintA + "JOIN " + COLOCATED_TABLE_B + " " + + tableHintB + "ON " + COLOCATED_TABLE_A + ".col1 = " + COLOCATED_TABLE_B + ".col1"; + } + + /// A join of the partitioned fact table "tableA" with "tableB" hinted replicated, both sides sent over a local + /// exchange. This is the shape [#colocatedJoinQuery(int)] cannot express: `is_colocated_by_join_keys` claims both + /// sides are partitioned by the join key, while a replicated table is simply present in full on every worker. + private static String replicatedDimensionJoinQuery(int partitionSize) { + return "SELECT /*+ joinOptions(left_distribution_type='local', right_distribution_type='local') */ " + + COLOCATED_TABLE_A + ".col2, " + COLOCATED_TABLE_B + ".col2 FROM " + COLOCATED_TABLE_A + " " + + colocatedTableHint(partitionSize) + "JOIN " + COLOCATED_TABLE_B + + " /*+ tableOptions(is_replicated='true') */ ON " + COLOCATED_TABLE_A + ".col1 = " + COLOCATED_TABLE_B + + ".col1"; + } + + /// Builds a QueryEnvironment for two offline partitioned tables "tableA" and "tableB" (function Hashcode on col1, 4 + /// partitions each), for colocated join tests. Partition `p` of table `t` holds one segment `"{t}_seg{p}"` fully + /// replicated on server `p`, unless `p` is in that table's `emptyPartitions`, in which case it has no entry in the + /// partition info map at all. The `deferredPartitions` are the ones the broker reports as absent only because all of + /// their segments are new and not fully online yet. + private static QueryEnvironment newColocatedJoinQueryEnvironment(Set emptyPartitionsA, + Set emptyPartitionsB, Set deferredPartitionsA, Set deferredPartitionsB) { + return newColocatedJoinQueryEnvironment(emptyPartitionsA, emptyPartitionsB, deferredPartitionsA, + deferredPartitionsB, false); + } + + /// Same as [#newColocatedJoinQueryEnvironment(Set, Set, Set, Set)], except that when + /// `everyServerHostsEveryPartition` is set each partition is fully replicated on all the servers instead of only on + /// its own one. + private static QueryEnvironment newColocatedJoinQueryEnvironment(Set emptyPartitionsA, + Set emptyPartitionsB, Set deferredPartitionsA, Set deferredPartitionsB, + boolean everyServerHostsEveryPartition) { + return newColocatedJoinQueryEnvironment(emptyPartitionsA, emptyPartitionsB, deferredPartitionsA, + deferredPartitionsB, everyServerHostsEveryPartition, 4); + } + + /// Same as [#newColocatedJoinQueryEnvironment(Set, Set, Set, Set, boolean)], with the number of partitions of each + /// table. There are always 4 servers, so with more partitions than that, partition `p` lives on server `p % 4` and + /// several partitions share a worker. + private static QueryEnvironment newColocatedJoinQueryEnvironment(Set emptyPartitionsA, + Set emptyPartitionsB, Set deferredPartitionsA, Set deferredPartitionsB, + boolean everyServerHostsEveryPartition, int numPartitionsPerTable) { + return newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(numPartitionsPerTable, everyServerHostsEveryPartition).emptyPartitions(emptyPartitionsA) + .partitionsWithOnlyDeferredSegments(deferredPartitionsA), + new ColocatedTableSpec(numPartitionsPerTable, everyServerHostsEveryPartition).emptyPartitions(emptyPartitionsB) + .partitionsWithOnlyDeferredSegments(deferredPartitionsB)); + } + + /// Same as [#newColocatedJoinQueryEnvironment(Set, Set, Set, Set, boolean, int)], taking the full layout of each + /// table so that a test can also make a partition unservable or give it invalid partition metadata. + private static QueryEnvironment newColocatedJoinQueryEnvironment(ColocatedTableSpec specA, + ColocatedTableSpec specB) { + return newColocatedJoinQueryEnvironment(specA, specB, TableType.OFFLINE); + } + + /// Same as [#newColocatedJoinQueryEnvironment(ColocatedTableSpec, ColocatedTableSpec)], with the table type both + /// tables are registered under, so that the realtime-only shape can be covered too. + private static QueryEnvironment newColocatedJoinQueryEnvironment(ColocatedTableSpec specA, ColocatedTableSpec specB, + TableType tableType) { + int numServers = 4; + ServerInstance[] servers = new ServerInstance[numServers]; + Map enabledServers = new HashMap<>(); + for (int i = 0; i < numServers; i++) { + servers[i] = getServerInstance("localhost", i + 1); + enabledServers.put(servers[i].getInstanceId(), servers[i]); + } + String tableAWithType = COLOCATED_TABLE_A + "_" + tableType.name(); + String tableBWithType = COLOCATED_TABLE_B + "_" + tableType.name(); + Map partitionInfoByTable = new HashMap<>(); + partitionInfoByTable.put(tableAWithType, + colocatedTablePartitionInfo(tableAWithType, "a_seg", servers, specA)); + partitionInfoByTable.put(tableBWithType, + colocatedTablePartitionInfo(tableBWithType, "b_seg", servers, specB)); + Map routingTableByTable = new HashMap<>(); + if (specA._survivingSegments != null) { + routingTableByTable.put(tableAWithType, colocatedRoutingTable(servers, "a_seg", specA._survivingSegments)); + } + if (specB._survivingSegments != null) { + routingTableByTable.put(tableBWithType, colocatedRoutingTable(servers, "b_seg", specB._survivingSegments)); + } + PartitionedRoutingManager routingManager = + new PartitionedRoutingManager(enabledServers, partitionInfoByTable, routingTableByTable, false); + + Map tableNameMap = new HashMap<>(); + tableNameMap.put(tableAWithType, tableAWithType); + tableNameMap.put(COLOCATED_TABLE_A, COLOCATED_TABLE_A); + tableNameMap.put(tableBWithType, tableBWithType); + tableNameMap.put(COLOCATED_TABLE_B, COLOCATED_TABLE_B); + TableCache tableCache = mock(TableCache.class); + when(tableCache.getTableNameMap()).thenReturn(tableNameMap); + when(tableCache.getActualTableName(anyString())).thenAnswer(inv -> tableNameMap.get(inv.getArgument(0))); + when(tableCache.getSchema(anyString())).thenAnswer( + inv -> getSchemaBuilder(inv.getArgument(0, String.class)).build()); + when(tableCache.getTableConfig(anyString())).thenReturn(mock(TableConfig.class)); + + WorkerManager workerManager = new WorkerManager("Broker_localhost", "localhost", 5, routingManager); + return new QueryEnvironment(QueryEnvironment.configBuilder() + .requestId(-1L) + .database(CommonConstants.DEFAULT_DATABASE) + .tableCache(tableCache) + .workerManager(workerManager) + .build()); + } + + private static TablePartitionReplicatedServersInfo colocatedTablePartitionInfo(String tableNameWithType, + String segmentPrefix, ServerInstance[] servers, ColocatedTableSpec spec) { + Set allServers = new HashSet<>(); + for (ServerInstance server : servers) { + allServers.add(server.getInstanceId()); + } + int numPartitions = spec._numPartitions; + TablePartitionReplicatedServersInfo.PartitionInfo[] partitionInfoMap = + new TablePartitionReplicatedServersInfo.PartitionInfo[numPartitions]; + for (int p = 0; p < numPartitions; p++) { + if (!spec._emptyPartitions.contains(p)) { + Set partitionServers; + if (spec._partitionsWithoutFullyReplicatedServer.contains(p)) { + // The partition holds a segment, but no single server holds all of it. + partitionServers = Set.of(); + } else if (spec._partitionServerIndexes.containsKey(p)) { + partitionServers = new HashSet<>(); + for (Integer serverIndex : spec._partitionServerIndexes.get(p)) { + partitionServers.add(servers[serverIndex].getInstanceId()); + } + } else { + partitionServers = spec._everyServerHostsEveryPartition ? allServers + : Set.of(servers[p % servers.length].getInstanceId()); + } + // Mutable, like the lists the broker publishes: the assignment must hand out a copy rather than this instance. + partitionInfoMap[p] = new TablePartitionReplicatedServersInfo.PartitionInfo(partitionServers, + new ArrayList<>(List.of(segmentPrefix + p))); + } + } + return new TablePartitionReplicatedServersInfo(tableNameWithType, "col1", "Hashcode", numPartitions, + partitionInfoMap, spec._segmentsWithInvalidPartition, spec._partitionsWithOnlyDeferredSegments); + } + + /// Buckets the given surviving segments onto the server hosting their partition (partition `p` lives on server + /// `p % 4`), i.e. builds what the routing manager returns for one colocated table's filtered routing query. + private static RoutingTable colocatedRoutingTable(ServerInstance[] servers, String segmentPrefix, + List survivingSegments) { + Map> serverToSegmentList = new HashMap<>(); + for (String segment : survivingSegments) { + int partition = Integer.parseInt(segment.substring(segmentPrefix.length())); + serverToSegmentList.computeIfAbsent(servers[partition % servers.length], k -> new ArrayList<>()).add(segment); + } + Map serverToSegments = new HashMap<>(); + serverToSegmentList.forEach((server, segments) -> serverToSegments.put(server, + new SegmentsToQuery(segments, List.of()))); + return new RoutingTable(serverToSegments, List.of(), 0); + } + + /// How one side of a colocated join is laid out, for [#newColocatedJoinQueryEnvironment(ColocatedTableSpec, + /// ColocatedTableSpec)]. Partition `p` holds one segment named after the table's prefix, fully replicated on server + /// `p % 4` (or on all 4 servers when `everyServerHostsEveryPartition` is set), unless a set below says otherwise. + private static class ColocatedTableSpec { + final int _numPartitions; + final boolean _everyServerHostsEveryPartition; + /// Partitions with no entry at all in the partition info map, i.e. holding no segment. + Set _emptyPartitions = Set.of(); + /// Partitions the broker reports as having no entry only because all their segments are new and not fully online. + Set _partitionsWithOnlyDeferredSegments = Set.of(); + /// Partitions with an entry but no fully replicated server, i.e. holding a segment that no single server has whole. + Set _partitionsWithoutFullyReplicatedServer = Set.of(); + /// Segments the broker reports as holding invalid partition metadata (absent from the partition info map). + List _segmentsWithInvalidPartition = List.of(); + /// Overrides the servers of individual partitions, by server index, so that a partition's servers can be a strict + /// subset of the ones hosting the table as a whole. + Map> _partitionServerIndexes = Map.of(); + /// The segments the routing manager reports as surviving the filtered routing query, i.e. what broker pruning would + /// keep. Null when the table gets no routing table at all, so that a routing call a test did not expect returns + /// null rather than a silently empty answer. + @Nullable + List _survivingSegments; + + ColocatedTableSpec(int numPartitions, boolean everyServerHostsEveryPartition) { + _numPartitions = numPartitions; + _everyServerHostsEveryPartition = everyServerHostsEveryPartition; + } + + ColocatedTableSpec emptyPartitions(Set emptyPartitions) { + _emptyPartitions = emptyPartitions; + return this; + } + + ColocatedTableSpec partitionsWithOnlyDeferredSegments(Set partitionsWithOnlyDeferredSegments) { + _partitionsWithOnlyDeferredSegments = partitionsWithOnlyDeferredSegments; + return this; + } + + ColocatedTableSpec partitionsWithoutFullyReplicatedServer(Set partitionsWithoutFullyReplicatedServer) { + _partitionsWithoutFullyReplicatedServer = partitionsWithoutFullyReplicatedServer; + return this; + } + + ColocatedTableSpec segmentsWithInvalidPartition(List segmentsWithInvalidPartition) { + _segmentsWithInvalidPartition = segmentsWithInvalidPartition; + return this; + } + + ColocatedTableSpec partitionServerIndexes(Map> partitionServerIndexes) { + _partitionServerIndexes = partitionServerIndexes; + return this; + } + + ColocatedTableSpec survivingSegments(List survivingSegments) { + _survivingSegments = survivingSegments; + return this; + } + } + + /// Returns the only fragment below the reduce stage with neither segments nor children, i.e. the join stage. + private static DispatchablePlanFragment joinFragment(DispatchableSubPlan dispatchableSubPlan) { + for (Map.Entry entry : dispatchableSubPlan.getQueryStageMap().entrySet()) { + if (entry.getKey() != 0 && entry.getValue().getWorkerIdToSegmentsMap().isEmpty()) { + return entry.getValue(); + } + } + throw new AssertionError("Found no join fragment in: " + dispatchableSubPlan.getQueryStageMap().keySet()); + } + + /// Returns the leaf fragment scanning the given table. + private static DispatchablePlanFragment leafFragmentForTable(DispatchableSubPlan dispatchableSubPlan, + String tableName) { + for (DispatchablePlanFragment leafFragment : leafFragments(dispatchableSubPlan)) { + if (leafFragment.getTableName().startsWith(tableName)) { + return leafFragment; + } + } + throw new AssertionError("Found no leaf fragment for table: " + tableName); + } + + /// Returns every server instance id the plan is dispatched to, over all the stages but the broker reduce root. + private static Set dispatchedServers(DispatchableSubPlan dispatchableSubPlan) { + Set servers = new HashSet<>(); + for (Map.Entry entry : dispatchableSubPlan.getQueryStageMap().entrySet()) { + if (entry.getKey() != 0) { + for (QueryServerInstance server : entry.getValue().getServerInstances()) { + servers.add(server.getInstanceId()); + } + } + } + return servers; + } + + /// Maps each worker id of the given leaf to the partitions it scans, decoded from the `{segmentPrefix}{partition}` + /// segment names. A worker with no segment is absent from the result: nothing but its index in the class list the + /// colocated group shares says which class it stands for. + private static Map> workerIdToPartitions(DispatchablePlanFragment leafFragment, + String segmentPrefix) { + Map> workerIdToPartitions = new HashMap<>(); + for (Map.Entry>> entry : leafFragment.getWorkerIdToSegmentsMap().entrySet()) { + Set partitions = new HashSet<>(); + for (List segments : entry.getValue().values()) { + for (String segment : segments) { + assertTrue(segment.startsWith(segmentPrefix), "Unexpected segment: " + segment); + partitions.add(Integer.parseInt(segment.substring(segmentPrefix.length()))); + } + } + if (!partitions.isEmpty()) { + workerIdToPartitions.put(entry.getKey(), partitions); + } + } + return workerIdToPartitions; + } + + /// Folds the worker id -> partition class mapping of one leaf into `workerIdToClass`, failing when this leaf + /// contradicts what another leaf of the same colocated group already recorded for a worker id. A worker with no + /// segment contributes nothing, so the map is only filled in from the sides that hold data. + private static void mergeWorkerIdToClass(Map workerIdToClass, + DispatchablePlanFragment leafFragment, String segmentPrefix, int partitionSize) { + for (Map.Entry> entry : workerIdToPartitions(leafFragment, segmentPrefix).entrySet()) { + Integer workerId = entry.getKey(); + Set partitionClasses = new HashSet<>(); + for (Integer partition : entry.getValue()) { + partitionClasses.add(partition % partitionSize); + } + assertEquals(partitionClasses.size(), 1, + "Worker: " + workerId + " scans several partition classes: " + partitionClasses); + int partitionClass = partitionClasses.iterator().next(); + Integer recorded = workerIdToClass.put(workerId, partitionClass); + assertTrue(recorded == null || recorded == partitionClass, "Worker: " + workerId + " stands for partition class: " + + recorded + " on one side of the exchange and: " + partitionClass + " on the other"); + } + } + + private static Map workerIdToServer(DispatchablePlanFragment leafFragment) { + Map workerIdToServer = new HashMap<>(); + for (Map.Entry> entry + : leafFragment.getServerInstanceToWorkerIdMap().entrySet()) { + for (Integer workerId : entry.getValue()) { + workerIdToServer.put(workerId, entry.getKey().getInstanceId()); + } + } + return workerIdToServer; + } + + private static List assignedSegments(DispatchablePlanFragment leafFragment, int workerId) { + List segments = new ArrayList<>(); + leafFragment.getWorkerIdToSegmentsMap().get(workerId).values().forEach(segments::addAll); + return segments; + } + /// Builds a QueryEnvironment for a hybrid partitioned table "testTable" (function Hashcode on col1, 4 partitions). /// Partition `p` holds offline segment `"segO{p}"` and realtime segment `"segR{p}"`, both fully /// replicated on server `p`. The given surviving segment lists are what the [RoutingManager] returns for /// the filtered routing query of each table type. private static QueryEnvironment newHybridPartitionedQueryEnvironment(List survivingOfflineSegments, List survivingRealtimeSegments) { + return newHybridPartitionedQueryEnvironment(survivingOfflineSegments, survivingRealtimeSegments, List.of(), + List.of()); + } + + /// Same as [#newHybridPartitionedQueryEnvironment(List, List)], with the segments reported as having invalid + /// partition metadata for each table type. + private static QueryEnvironment newHybridPartitionedQueryEnvironment(List survivingOfflineSegments, + List survivingRealtimeSegments, List offlineSegmentsWithInvalidPartition, + List realtimeSegmentsWithInvalidPartition) { + return newHybridPartitionedQueryEnvironment(survivingOfflineSegments, survivingRealtimeSegments, + offlineSegmentsWithInvalidPartition, realtimeSegmentsWithInvalidPartition, Set.of(), Set.of()); + } + + /// Same as [#newHybridPartitionedQueryEnvironment(List, List, List, List)], with the OFFLINE partitions that have no + /// entry in the offline partition info map and, among those, the ones the broker reports as absent only because all + /// of their segments are new and not fully online yet. The REALTIME side always has an entry for every partition. + private static QueryEnvironment newHybridPartitionedQueryEnvironment(List survivingOfflineSegments, + List survivingRealtimeSegments, List offlineSegmentsWithInvalidPartition, + List realtimeSegmentsWithInvalidPartition, Set emptyOfflinePartitions, + Set offlinePartitionsWithOnlyDeferredSegments) { int numPartitions = 4; ServerInstance[] servers = new ServerInstance[numPartitions]; Map enabledServers = new HashMap<>(); @@ -854,16 +1907,20 @@ private static QueryEnvironment newHybridPartitionedQueryEnvironment(List partitionServers = Set.of(servers[p].getInstanceId()); - offlinePartitions[p] = new TablePartitionReplicatedServersInfo.PartitionInfo(partitionServers, - List.of("segO" + p)); + if (!emptyOfflinePartitions.contains(p)) { + offlinePartitions[p] = new TablePartitionReplicatedServersInfo.PartitionInfo(partitionServers, + new ArrayList<>(List.of("segO" + p))); + } realtimePartitions[p] = new TablePartitionReplicatedServersInfo.PartitionInfo(partitionServers, - List.of("segR" + p)); + new ArrayList<>(List.of("segR" + p))); } String realtimeTableName = PARTITIONED_TABLE + "_REALTIME"; TablePartitionReplicatedServersInfo offlineInfo = new TablePartitionReplicatedServersInfo( - PARTITIONED_TABLE_OFFLINE, "col1", "Hashcode", numPartitions, offlinePartitions, List.of()); + PARTITIONED_TABLE_OFFLINE, "col1", "Hashcode", numPartitions, offlinePartitions, + offlineSegmentsWithInvalidPartition, offlinePartitionsWithOnlyDeferredSegments); TablePartitionReplicatedServersInfo realtimeInfo = new TablePartitionReplicatedServersInfo( - realtimeTableName, "col1", "Hashcode", numPartitions, realtimePartitions, List.of()); + realtimeTableName, "col1", "Hashcode", numPartitions, realtimePartitions, + realtimeSegmentsWithInvalidPartition, Set.of()); PartitionedRoutingManager routingManager = new PartitionedRoutingManager(enabledServers, Map.of(PARTITIONED_TABLE_OFFLINE, offlineInfo, realtimeTableName, realtimeInfo), @@ -915,6 +1972,18 @@ private static QueryEnvironment newPartitionedQueryEnvironment(int[] serverIdxPe private static QueryEnvironment newPartitionedQueryEnvironment(int[] serverIdxPerPartition, int numServers, int replicasPerPartition, List survivingSegments, List unavailableSegments, int reportedPrunedByRouting, boolean throwOnRouting) { + return newPartitionedQueryEnvironment(serverIdxPerPartition, numServers, replicasPerPartition, survivingSegments, + unavailableSegments, reportedPrunedByRouting, throwOnRouting, Set.of(), Set.of()); + } + + /// Same as [#newPartitionedQueryEnvironment(int[], int, int, List, List, int, boolean)], with the partitions that + /// have no entry in the partition info map at all (`emptyPartitions`) and, among those, the ones the broker reports + /// as absent only because all of their segments are new and not fully online yet + /// (`partitionsWithOnlyDeferredSegments`). + private static QueryEnvironment newPartitionedQueryEnvironment(int[] serverIdxPerPartition, int numServers, + int replicasPerPartition, List survivingSegments, List unavailableSegments, + int reportedPrunedByRouting, boolean throwOnRouting, Set emptyPartitions, + Set partitionsWithOnlyDeferredSegments) { int numPartitions = serverIdxPerPartition.length; ServerInstance[] servers = new ServerInstance[numServers]; Map enabledServers = new HashMap<>(); @@ -925,15 +1994,20 @@ private static QueryEnvironment newPartitionedQueryEnvironment(int[] serverIdxPe TablePartitionReplicatedServersInfo.PartitionInfo[] partitionInfoMap = new TablePartitionReplicatedServersInfo.PartitionInfo[numPartitions]; for (int p = 0; p < numPartitions; p++) { + if (emptyPartitions.contains(p)) { + continue; + } Set fullyReplicatedServers = new HashSet<>(); for (int r = 0; r < replicasPerPartition; r++) { fullyReplicatedServers.add(servers[serverIdxPerPartition[p] + r].getInstanceId()); } - partitionInfoMap[p] = - new TablePartitionReplicatedServersInfo.PartitionInfo(fullyReplicatedServers, List.of("seg" + p)); + // Mutable, like the lists the broker publishes: the assignment must hand out a copy rather than this instance. + partitionInfoMap[p] = new TablePartitionReplicatedServersInfo.PartitionInfo(fullyReplicatedServers, + new ArrayList<>(List.of("seg" + p))); } TablePartitionReplicatedServersInfo tablePartitionInfo = new TablePartitionReplicatedServersInfo( - PARTITIONED_TABLE_OFFLINE, "col1", "Hashcode", numPartitions, partitionInfoMap, List.of()); + PARTITIONED_TABLE_OFFLINE, "col1", "Hashcode", numPartitions, partitionInfoMap, List.of(), + partitionsWithOnlyDeferredSegments); // Model the pruned routing table: surviving segments bucketed onto their owning server. Map> serverToSegmentList = new HashMap<>(); @@ -1392,10 +2466,23 @@ public RoutingTable getRoutingTable(BrokerRequest brokerRequest, String tableNam return getRoutingTable(brokerRequest, requestId); } + /// Only the replicated leaf path reads this (a table hinted `is_replicated` holds every segment on every worker), + /// so answer it from the same partition layout the partitioned path reads. @Nullable @Override public List getSegments(BrokerRequest brokerRequest) { - return List.of(); + TablePartitionReplicatedServersInfo partitionInfo = + _partitionInfoByTable.get(brokerRequest.getQuerySource().getTableName()); + if (partitionInfo == null) { + return List.of(); + } + List segments = new ArrayList<>(); + for (TablePartitionReplicatedServersInfo.PartitionInfo entry : partitionInfo.getPartitionInfoMap()) { + if (entry != null) { + segments.addAll(entry._segments); + } + } + return segments; } @Override From ec4ac46693a1fc82dea69df8cc08c66583e9e04e Mon Sep 17 00:00:00 2001 From: Yash Mayya Date: Thu, 6 Aug 2026 16:47:16 -0700 Subject: [PATCH 2/2] Support broker segment pruning for colocated joins --- .../manager/BaseBrokerRoutingManager.java | 86 +++- .../manager/MultiClusterRoutingManager.java | 56 +++ .../manager/BrokerRoutingManagerTest.java | 167 ++++++- .../MultiClusterRoutingManagerTest.java | 86 ++++ .../pinot/core/routing/RoutingManager.java | 32 ++ .../ColocatedJoinEmptyPartitionTest.java | 176 ++++++- .../physical/DispatchablePlanContext.java | 9 + .../routing/ColocationGroupAnalyzer.java | 19 +- .../routing/PlanNodeRoutingQueryBuilder.java | 21 + .../pinot/query/routing/WorkerManager.java | 349 ++++++++++---- .../routing/ColocationGroupAnalyzerTest.java | 15 +- .../query/routing/WorkerManagerTest.java | 431 ++++++++++++++++-- .../queries/ExplainPhysicalPlans.json | 24 +- .../pinot/spi/utils/CommonConstants.java | 5 + 14 files changed, 1285 insertions(+), 191 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java index e0ccad18e369..fa94094787e0 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java +++ b/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 getSegments(BrokerRequest brokerRequest, @Nullable String sa return routingEntry.getSegments(brokerRequest, samplerName); } + @Nullable + @Override + public Set 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 selectedSegments = segmentSelector.select(brokerRequest); - int numTotalSelectedSegments = selectedSegments.size(); + Set 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 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 getSegments(BrokerRequest brokerRequest, @Nullable String samplerName) { - SamplerInfo samplerInfo = getSamplerInfo(samplerName); - SegmentSelector segmentSelector = samplerInfo != null ? samplerInfo._segmentSelector : _segmentSelector; - Set 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 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 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 _selected; + final Set _surviving; + + SelectedSegments(Set selected, Set surviving) { + _selected = selected; + _surviving = surviving; + } + + int getNumPruned() { + return _selected.size() - _surviving.size(); } } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java index 312d7b0b7d10..9f04a4ca3d40 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java @@ -217,6 +217,62 @@ public List 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 getPrunedSegments(BrokerRequest brokerRequest) { + String tableNameWithType = brokerRequest.getQuerySource().getTableName(); + Set 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 intersectPrunedSegments(@Nullable Set 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 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 diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java index bc4134ebb615..492463d5c21c 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java @@ -19,9 +19,10 @@ package org.apache.pinot.broker.routing.manager; import java.lang.reflect.Constructor; -import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.Consumer; import org.apache.helix.AccessOption; import org.apache.helix.BaseDataAccessor; @@ -40,6 +41,8 @@ import org.apache.pinot.broker.routing.segmentselector.SegmentSelector; import org.apache.pinot.broker.routing.timeboundary.TimeBoundaryManager; import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.common.request.QuerySource; import org.apache.pinot.core.routing.TablePartitionInfo; import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; @@ -59,9 +62,11 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNull; import static org.testng.Assert.assertSame; import static org.testng.Assert.assertTrue; @@ -241,9 +246,162 @@ public void testSamplerContextSharesTimeBoundaryAndPartitionMetadata() assertSame(_routingManager.getTablePartitionReplicatedServersInfo(TEST_TABLE), expectedReplicatedServersInfo); } + @Test + public void testGetPrunedSegmentsIsExactlySelectedMinusSurvivors() + throws Exception { + SegmentSelector segmentSelector = selectorOf("seg1", "seg2", "seg3", "seg4"); + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, segmentSelector, List.of(prunerDropping("seg1", "seg3")), + mock(InstanceSelector.class))); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of("seg1", "seg3")); + } + + @Test + public void testGetPrunedSegmentsChainsEveryPruner() + throws Exception { + SegmentPruner firstPruner = prunerDropping("seg1"); + SegmentPruner secondPruner = prunerDropping("seg3"); + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, selectorOf("seg1", "seg2", "seg3"), List.of(firstPruner, secondPruner), + mock(InstanceSelector.class))); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of("seg1", "seg3")); + + // The second pruner judges what the first left, so consulting only the last one would lose "seg1". + ArgumentCaptor> captor = ArgumentCaptor.captor(); + verify(secondPruner).prune(any(), captor.capture()); + assertEquals(captor.getValue(), Set.of("seg2", "seg3")); + } + + /// Nothing pruned must read as "proved nothing", not as "proved every selected segment empty" -- the latter would + /// let a caller treat a table that fully matches the filter as a table with no matching rows. + @Test + public void testGetPrunedSegmentsIsEmptyWhenNothingWasPruned() + throws Exception { + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, selectorOf("seg1", "seg2"), List.of(prunerDropping()), + mock(InstanceSelector.class))); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of()); + } + + @Test + public void testGetPrunedSegmentsIsEmptyWhenSelectionIsEmpty() + throws Exception { + SegmentPruner pruner = prunerDropping("seg1"); + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, selectorOf(), List.of(pruner), mock(InstanceSelector.class))); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of()); + // An empty selection is answered without asking anyone, so the empty result cannot have come from a pruner. + verify(pruner, never()).prune(any(), any()); + } + + /// A table this broker has no routing for is `null`, not an empty set: it eliminated nothing because it would have + /// routed nothing, which is a different claim from "the pruners ran and proved nothing". + @Test + public void testGetPrunedSegmentsIsNullForUnknownTable() { + assertNull(_routingManager.getPrunedSegments(brokerRequest("noSuchTable_OFFLINE"))); + } + + /// The whole point of the API: only presence in the result is a proof. A segment the selector never offered is + /// absent from the survivors for a reason that has nothing to do with the filter -- here the selector withheld + /// "seg3" -- and reporting it would let a caller skip a segment that may well hold matching rows. + @Test + public void testGetPrunedSegmentsDoesNotReportASegmentTheSelectorNeverOffered() + throws Exception { + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, selectorOf("seg1", "seg2"), List.of(prunerDropping("seg1", "seg3")), + mock(InstanceSelector.class))); + + Set prunedSegments = _routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)); + + assertEquals(prunedSegments, Set.of("seg1")); + assertFalse(prunedSegments.contains("seg3")); + } + + /// Instance selection is what makes routing depend on the request id and on which replicas are up; keeping it out + /// is what makes this deterministic enough to plan on. + @Test + public void testGetPrunedSegmentsNeverConsultsInstanceSelection() + throws Exception { + InstanceSelector instanceSelector = mock(InstanceSelector.class); + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, selectorOf("seg1", "seg2"), List.of(prunerDropping("seg1")), instanceSelector)); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of("seg1")); + verifyNoInteractions(instanceSelector); + } + + /// A pruner that edits the set it was handed leaves nothing to take a difference against. That has to degrade to + /// "proved nothing" rather than to a wrong proof. + @Test + public void testGetPrunedSegmentsIsEmptyWhenAPrunerEditsInPlace() + throws Exception { + SegmentPruner pruner = mock(SegmentPruner.class); + when(pruner.prune(any(), any())).thenAnswer(invocation -> { + Set segments = invocation.getArgument(1); + segments.remove("seg1"); + return segments; + }); + SegmentSelector segmentSelector = mock(SegmentSelector.class); + when(segmentSelector.select(any())).thenReturn(new HashSet<>(Set.of("seg1", "seg2"))); + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, segmentSelector, List.of(pruner), mock(InstanceSelector.class))); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of()); + } + + private static BrokerRequest brokerRequest(String tableNameWithType) { + QuerySource querySource = new QuerySource(); + querySource.setTableName(tableNameWithType); + BrokerRequest brokerRequest = new BrokerRequest(); + brokerRequest.setQuerySource(querySource); + return brokerRequest; + } + + private static SegmentSelector selectorOf(String... segments) { + SegmentSelector segmentSelector = mock(SegmentSelector.class); + when(segmentSelector.select(any())).thenReturn(Set.of(segments)); + return segmentSelector; + } + + /// Mirrors [org.apache.pinot.broker.routing.segmentpruner.EmptySegmentPruner]: a fresh set when it prunes + /// something, the very set it was handed when it does not. + private static SegmentPruner prunerDropping(String... segments) { + Set droppedSegments = Set.of(segments); + SegmentPruner segmentPruner = mock(SegmentPruner.class); + when(segmentPruner.prune(any(), any())).thenAnswer(invocation -> { + Set candidateSegments = invocation.getArgument(1); + if (droppedSegments.stream().noneMatch(candidateSegments::contains)) { + return candidateSegments; + } + Set survivingSegments = new HashSet<>(candidateSegments); + survivingSegments.removeAll(droppedSegments); + return survivingSegments; + }); + return segmentPruner; + } + private static Object createRoutingEntry(String tableNameWithType, TimeBoundaryManager timeBoundaryManager, SegmentPartitionMetadataManager partitionMetadataManager, Map samplerInfos) throws Exception { + return createRoutingEntry(tableNameWithType, mock(SegmentSelector.class), List.of(), mock(InstanceSelector.class), + timeBoundaryManager, partitionMetadataManager, samplerInfos); + } + + private static Object createRoutingEntry(String tableNameWithType, SegmentSelector segmentSelector, + List segmentPruners, InstanceSelector instanceSelector) + throws Exception { + return createRoutingEntry(tableNameWithType, segmentSelector, segmentPruners, instanceSelector, + mock(TimeBoundaryManager.class), mock(SegmentPartitionMetadataManager.class), Map.of()); + } + + private static Object createRoutingEntry(String tableNameWithType, SegmentSelector segmentSelector, + List segmentPruners, InstanceSelector instanceSelector, TimeBoundaryManager timeBoundaryManager, + SegmentPartitionMetadataManager partitionMetadataManager, Map samplerInfos) + throws Exception { Class routingEntryClass = Class.forName(BaseBrokerRoutingManager.class.getName() + "$RoutingEntry"); Constructor constructor = routingEntryClass.getDeclaredConstructor(String.class, String.class, String.class, SegmentPreSelector.class, SegmentSelector.class, List.class, InstanceSelector.class, int.class, int.class, @@ -251,10 +409,9 @@ private static Object createRoutingEntry(String tableNameWithType, TimeBoundaryM Map.class, boolean.class); constructor.setAccessible(true); return constructor.newInstance(tableNameWithType, "/IDEALSTATES/" + tableNameWithType, - "/EXTERNALVIEW/" + tableNameWithType, mock(SegmentPreSelector.class), mock(SegmentSelector.class), - Collections.emptyList(), mock(InstanceSelector.class), 1, 1, - mock(SegmentZkMetadataFetcher.class), timeBoundaryManager, partitionMetadataManager, null, samplerInfos, - false); + "/EXTERNALVIEW/" + tableNameWithType, mock(SegmentPreSelector.class), segmentSelector, segmentPruners, + instanceSelector, 1, 1, mock(SegmentZkMetadataFetcher.class), timeBoundaryManager, partitionMetadataManager, + null, samplerInfos, false); } @SuppressWarnings({"rawtypes", "unchecked"}) diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java index 8d8cdb34c262..0ea105a56e16 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java @@ -44,6 +44,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; @@ -327,6 +328,91 @@ public void testGetTablePartitionInfoIgnoresAFailingRemoteCluster() { assertEquals(_multiClusterRoutingManager.getTablePartitionReplicatedServersInfo(TEST_TABLE), partitionInfo); } + /// The usual case: the table lives in exactly one cluster, so the intersection is that cluster's own verdict. + @Test + public void testGetPrunedSegmentsReturnsTheSoleClusterVerdictVerbatim() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + when(_localClusterRoutingManager.getPrunedSegments(brokerRequest)) + .thenReturn(Set.of("seg1", "seg2")); + withoutTheTable(brokerRequest, _remoteClusterRoutingManager1, _remoteClusterRoutingManager2); + + assertEquals(_multiClusterRoutingManager.getPrunedSegments(brokerRequest), Set.of("seg1", "seg2")); + } + + /// Intersection, not union: unioning would let one cluster's pruners speak for a segment another cluster would + /// still have queried, so the planner would skip data that matches -- a silent wrong answer rather than a slow one. + /// Here only "seg2" is eliminated everywhere; "seg1" and "seg3" each survive in one cluster. + @Test + public void testGetPrunedSegmentsIntersectsRatherThanUnions() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + when(_localClusterRoutingManager.getPrunedSegments(brokerRequest)) + .thenReturn(Set.of("seg1", "seg2")); + when(_remoteClusterRoutingManager1.getPrunedSegments(brokerRequest)) + .thenReturn(Set.of("seg2", "seg3")); + withoutTheTable(brokerRequest, _remoteClusterRoutingManager2); + + Set prunedSegments = _multiClusterRoutingManager.getPrunedSegments(brokerRequest); + + assertEquals(prunedSegments, Set.of("seg2")); + assertFalse(prunedSegments.contains("seg1")); + assertFalse(prunedSegments.contains("seg3")); + } + + @Test + public void testGetPrunedSegmentsIsEmptyWhenAClusterProvesNothing() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + when(_localClusterRoutingManager.getPrunedSegments(brokerRequest)).thenReturn(Set.of("seg1")); + when(_remoteClusterRoutingManager1.getPrunedSegments(brokerRequest)).thenReturn(Set.of()); + withoutTheTable(brokerRequest, _remoteClusterRoutingManager2); + + assertEquals(_multiClusterRoutingManager.getPrunedSegments(brokerRequest), Set.of()); + } + + /// A cluster that does not have the table constrains nothing. Were its absence folded into the same empty set the + /// pruners use for "proved nothing", every answer would collapse to "nothing proven" in the common deployment. + @Test + public void testGetPrunedSegmentsSkipsAClusterWithoutTheTable() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + when(_localClusterRoutingManager.getPrunedSegments(brokerRequest)) + .thenReturn(Set.of("seg1", "seg2")); + // A cluster without the table reports null rather than an empty verdict, which is what keeps it from collapsing + // the intersection to "nothing proven". + withoutTheTable(brokerRequest, _remoteClusterRoutingManager1, _remoteClusterRoutingManager2); + + assertEquals(_multiClusterRoutingManager.getPrunedSegments(brokerRequest), Set.of("seg1", "seg2")); + } + + @Test + public void testGetPrunedSegmentsIsNullWhenNoClusterHasTheTable() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + withoutTheTable(brokerRequest, _localClusterRoutingManager, _remoteClusterRoutingManager1, + _remoteClusterRoutingManager2); + + assertNull(_multiClusterRoutingManager.getPrunedSegments(brokerRequest)); + } + + /// The deliberate opposite of [#testGetTablePartitionInfoIgnoresAFailingRemoteCluster]: a cluster we could not ask + /// may still have routed any of these segments, so its silence cannot be read as agreement. + @Test + public void testGetPrunedSegmentsIsEmptyWhenAClusterThrows() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + when(_localClusterRoutingManager.getPrunedSegments(brokerRequest)).thenReturn(Set.of("seg1")); + when(_remoteClusterRoutingManager1.getPrunedSegments(brokerRequest)) + .thenThrow(new RuntimeException("remote cluster is down")); + withoutTheTable(brokerRequest, _remoteClusterRoutingManager2); + + assertEquals(_multiClusterRoutingManager.getPrunedSegments(brokerRequest), Set.of()); + } + + /// Makes the given clusters report that they have no routing for the table. Worth spelling out in every pruning + /// test: an unstubbed mock hands back an empty set, which means "ran the pruners and proved nothing" and would + /// collapse the intersection for a reason the test did not intend. + private static void withoutTheTable(BrokerRequest brokerRequest, BaseBrokerRoutingManager... clusters) { + for (BaseBrokerRoutingManager cluster : clusters) { + when(cluster.getPrunedSegments(brokerRequest)).thenReturn(null); + } + } + private RoutingTable createRoutingTable(String serverName, List segments) { Map serverMap = new HashMap<>(); ServerInstance server = createMockServerInstance(serverName); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java index 9f33b868b519..53198ede031a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java @@ -89,6 +89,38 @@ default List getSegments(BrokerRequest brokerRequest, @Nullable String s return getSegments(brokerRequest); } + /// Returns the segments that the segment pruners *provably eliminated* for the given broker request, i.e. the + /// selected segments that cannot hold a row matching the request's filter. Absence from the returned set means + /// nothing: a segment may be missing because it matches, because it was never selected, or because the table does + /// not exist. Only presence is a proof. + /// + /// This is the complement of [#getSegments(BrokerRequest)], and the distinction is what makes it usable as a + /// planning-time emptiness proof. Deciding "this segment does not match" from *absence* of a survivor conflates + /// pruning with the several innocent reasons a segment can be missing from a routing result -- it was classified as + /// optional by instance selection, its server left the enabled server map, or it entered the partition metadata + /// before it became selectable -- and each of those would silently drop matching data. Deciding it from presence in + /// this set cannot. + /// + /// Instance selection deliberately takes no part, so the result depends only on the request and the pruners, never + /// on a request id or on which replica a query happens to pick. + /// + /// Returns `null` if the table does not exist, as [#getSegments(BrokerRequest)] does, which is not the same as an + /// empty set: a broker that does not have the table eliminated nothing because it would have routed nothing, while + /// an empty set is a broker that ran the pruners and proved nothing. A caller combining several brokers' verdicts + /// has to tell those apart, and gets both from this one lookup rather than from a separate existence check that + /// could disagree with it. + /// + /// The returned set is for reading only: an implementation may hand back an immutable or a shared set, so a caller + /// that needs to modify it must copy it first. + /// + /// The default implementation returns an empty set, i.e. proves nothing about any segment. Note that this is a + /// statement about segments only: a caller may still act on emptiness it can see for itself, such as a partition + /// that lists no segment at all. + @Nullable + default Set getPrunedSegments(BrokerRequest brokerRequest) { + return Set.of(); + } + /// Validate routing exist for a table /// /// @param tableNameWithType the name of the table. diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java index 3c75ccae35bb..37809b8cbf35 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java @@ -27,6 +27,8 @@ import java.util.Map; import java.util.regex.Matcher; import java.util.regex.Pattern; +import java.util.stream.Collectors; +import java.util.stream.Stream; import org.apache.avro.SchemaBuilder; import org.apache.avro.file.DataFileWriter; import org.apache.avro.generic.GenericData; @@ -34,6 +36,7 @@ import org.apache.commons.io.FileUtils; import org.apache.pinot.integration.tests.ClusterIntegrationTestUtils; import org.apache.pinot.spi.config.table.ColumnPartitionConfig; +import org.apache.pinot.spi.config.table.RoutingConfig; import org.apache.pinot.spi.config.table.SegmentPartitionConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; @@ -66,6 +69,11 @@ /// end-to-end run can show is that a real server accepts and answers a leaf-stage request whose segment list is empty /// for a genuinely partitioned table scan. /// +/// The same fixture covers the class reduction a filter buys. The tables configure the `partition` segment pruner, so +/// a restriction on the join key lets the broker eliminate segments before planning, and a class every member +/// eliminates leaves the group: `partitionKey IN (1, 2)` keeps only classes 1 and 2, halving the leaf worker count. +/// A restriction matching nothing anywhere falls back to the populated classes and returns an empty result. +/// /// The partition layout is supplied with explicit `tableOptions` hints rather than inferred, because hint inference is /// off by default (`pinot.broker.multistage.infer.partition.hint`); the hints carry exactly what it would have /// produced. The `is_colocated_by_join_keys` hint is spelled out for readability -- the exchange would be @@ -91,6 +99,12 @@ public class ColocatedJoinEmptyPartitionTest extends CustomDataQueryClusterInteg private static final int NUM_KEPT_CLASSES_FOR_JOIN = 4; private static final int NUM_ROWS_PER_PARTITION = 2; + /// The partition keys the pruning tests filter on. They land in partitions 1 and 2, the only ones both tables + /// populate, so classes 0 (left only) and 3 (right only) are eliminated on every member of the colocated group. + private static final List FILTERED_KEYS = List.of(1, 2); + /// A key in a partition neither table populates, so the filter matches nothing anywhere. + private static final List UNMATCHED_KEYS = List.of(5); + private static final int LEFT_METRIC_MULTIPLIER = 10; private static final int RIGHT_METRIC_MULTIPLIER = 100; @@ -221,31 +235,143 @@ public void testColocatedJoinMatchesShuffledJoin() assertEquals(colocatedResponse.get("resultTable").get("rows"), shuffledResponse.get("resultTable").get("rows"), "Colocated and shuffled plans must return the same rows"); + assertShuffledLeafStages(shuffledResponse, 2); + } - JsonNode shuffledStageStats = shuffledResponse.get("stageStats"); - assertNotNull(shuffledStageStats, "Missing stage stats in shuffled response: " + shuffledResponse); - List shuffledLeafStageSends = new ArrayList<>(); - collectLeafStageSends(shuffledStageStats, shuffledLeafStageSends); - assertEquals(shuffledLeafStageSends.size(), 2, - "Unexpected number of leaf stages in stage stats: " + shuffledStageStats.toPrettyString()); - for (JsonNode leafStageSend : shuffledLeafStageSends) { - assertTrue(leafStageSend.path("fanOut").asInt(-1) > 1, - "A shuffled leaf send must write more than one receive mailbox, otherwise the fanOut of 1 asserted for the " - + "colocated plan proves nothing. Stage stats: " + shuffledStageStats.toPrettyString()); - } + /// The class reduction broker pruning buys: a filter that every member of the colocated group can prune with drops + /// the classes all of them eliminate, so the leaves run fewer workers and scan fewer segments than the union of the + /// populated classes. What the reduced width must not cost is the 1-to-1 wiring, which is what the two exchange + /// assertions are for. + @Test + public void testColocatedJoinPrunesClassesEveryMemberFilters() + throws Exception { + setUseMultiStageQueryEngine(true); + String query = colocatedJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME, bothSidesFilter(FILTERED_KEYS)); + + JsonNode response = queryBrokerHttpEndpoint(query); + assertNoExceptions(response); + assertRows(response, expectedFilteredRows(FILTERED_KEYS)); + + // Classes 0 and 3 go: the left table's only class-0 segment is pruned by its own filter and the right table never + // populated that class, and symmetrically for class 3. Both tables populate every surviving class with exactly one + // segment, so both leaves scan one segment per kept class. + List keptClasses = keptClassesFor(FILTERED_KEYS); + assertLeafStages(response, 2, keptClasses.size(), keptClasses.size()); + assertDirectExchanges(query, 2, keptClasses.size()); + + // One segment per class that a table populated but the filter did not keep: the left table's class-0 segment and + // the right table's class-3 one. + long expectedNumPrunedSegments = + Stream.concat(LEFT_POPULATED_PARTITIONS.stream(), RIGHT_POPULATED_PARTITIONS.stream()) + .filter(partition -> !keptClasses.contains(partition)) + .count(); + assertEquals(response.path("numSegmentsPrunedByBroker").asLong(-1), expectedNumPrunedSegments, + "Unexpected number of broker-pruned segments in response: " + response); + } + + /// The same differential check as [#testColocatedJoinMatchesShuffledJoin], at the reduced width. A plan that quietly + /// fell back to a shuffle would return the right rows too, so this pairs with the `fanOut` and `[PARTITIONED]` + /// assertions rather than replacing them; what it rules out is a reduction that pairs the wrong classes and drops or + /// duplicates rows with no error. The shuffled plan reaches its answer by an independent route: with no table hints + /// its leaves are assigned per server and pruned per segment, not per partition class. + @Test + public void testFilteredColocatedJoinMatchesFilteredShuffledJoin() + throws Exception { + setUseMultiStageQueryEngine(true); + String whereClause = bothSidesFilter(FILTERED_KEYS); + JsonNode colocatedResponse = queryBrokerHttpEndpoint( + colocatedJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME, whereClause)); + assertNoExceptions(colocatedResponse); + // Pin the rows rather than only comparing the two plans: were the filter to stop matching anything, both sides + // would return nothing and the comparison below would still pass while proving nothing at all. + assertRows(colocatedResponse, expectedFilteredRows(FILTERED_KEYS)); + JsonNode shuffledResponse = queryBrokerHttpEndpoint( + shuffledJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME, whereClause)); + assertNoExceptions(shuffledResponse); + + assertEquals(colocatedResponse.get("resultTable").get("rows"), shuffledResponse.get("resultTable").get("rows"), + "Class-reduced colocated and shuffled plans must return the same rows"); + assertShuffledLeafStages(shuffledResponse, 2); + } + + /// A filter every member prunes every segment with: the group is left with no surviving class at all. It must fall + /// back to its populated classes and let the servers return the empty result, because a zero-worker leaf has no + /// handling on a 1-to-1 exchange -- an empty answer, not an error. + @Test + public void testColocatedJoinWithFilterMatchingNothing() + throws Exception { + setUseMultiStageQueryEngine(true); + assertTrue(keptClassesFor(UNMATCHED_KEYS).isEmpty(), + "The filter must match no partition either table populates, otherwise this is not the all-pruned fallback"); + String query = colocatedJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME, bothSidesFilter(UNMATCHED_KEYS)); + + JsonNode response = queryBrokerHttpEndpoint(query); + assertNoExceptions(response); + assertRows(response, List.of()); + + // Planned exactly as if there were no filter, and nothing is reported as pruned: the fallback dropped no class the + // group would otherwise have kept. + assertLeafStages(response, 2, NUM_KEPT_CLASSES_FOR_JOIN, LEFT_POPULATED_PARTITIONS.size()); + assertDirectExchanges(query, 2, NUM_KEPT_CLASSES_FOR_JOIN); + assertEquals(response.path("numSegmentsPrunedByBroker").asLong(-1), 0L, + "Unexpected number of broker-pruned segments in response: " + response); } private static String colocatedJoinQuery(String leftTableName, String rightTableName) { + return colocatedJoinQuery(leftTableName, rightTableName, ""); + } + + private static String colocatedJoinQuery(String leftTableName, String rightTableName, String whereClause) { return String.format( - "SELECT %s l.%s, l.%s, r.%s FROM %s %s AS l JOIN %s %s AS r ON l.%s = r.%s ORDER BY l.%s", + "SELECT %s l.%s, l.%s, r.%s FROM %s %s AS l JOIN %s %s AS r ON l.%s = r.%s %s ORDER BY l.%s", COLOCATED_JOIN_HINT, PARTITION_KEY_COLUMN, METRIC_COLUMN, METRIC_COLUMN, leftTableName, TABLE_HINT, - rightTableName, TABLE_HINT, PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN); + rightTableName, TABLE_HINT, PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN, whereClause, PARTITION_KEY_COLUMN); } private static String shuffledJoinQuery(String leftTableName, String rightTableName) { - return String.format("SELECT l.%s, l.%s, r.%s FROM %s AS l JOIN %s AS r ON l.%s = r.%s ORDER BY l.%s", + return shuffledJoinQuery(leftTableName, rightTableName, ""); + } + + private static String shuffledJoinQuery(String leftTableName, String rightTableName, String whereClause) { + return String.format("SELECT l.%s, l.%s, r.%s FROM %s AS l JOIN %s AS r ON l.%s = r.%s %s ORDER BY l.%s", PARTITION_KEY_COLUMN, METRIC_COLUMN, METRIC_COLUMN, leftTableName, rightTableName, PARTITION_KEY_COLUMN, - PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN); + PARTITION_KEY_COLUMN, whereClause, PARTITION_KEY_COLUMN); + } + + /// A partition-key restriction spelled out on both sides of the join rather than on one and left to Calcite's + /// transitive inference, so that the leaf of each member carries it whatever the planner decides to push down. + private static String bothSidesFilter(List keys) { + String keyList = keys.stream().map(String::valueOf).collect(Collectors.joining(", ")); + return String.format("WHERE l.%s IN (%s) AND r.%s IN (%s)", PARTITION_KEY_COLUMN, keyList, PARTITION_KEY_COLUMN, + keyList); + } + + /// The partition classes a colocated join of the two tables keeps under the given partition-key restriction. A class + /// survives when at least one member still holds a segment its own filter leaves, i.e. when some restricted key + /// hashes into a partition that member populates. One partition per class here, since the declared partition count + /// is the hinted partition size. + private static List keptClassesFor(List keys) { + return keys.stream() + .map(key -> key % NUM_DECLARED_PARTITIONS) + .filter(partition -> LEFT_POPULATED_PARTITIONS.contains(partition) + || RIGHT_POPULATED_PARTITIONS.contains(partition)) + .distinct() + .sorted() + .collect(Collectors.toList()); + } + + /// The rows the given partition-key restriction leaves in an inner join of the two tables: a key survives when both + /// tables populate the partition it hashes into, and every populated partition holds every one of its keys. + private static List> expectedFilteredRows(List keys) { + List> expectedRows = new ArrayList<>(); + for (int key : keys) { + int partition = key % NUM_DECLARED_PARTITIONS; + if (LEFT_POPULATED_PARTITIONS.contains(partition) && RIGHT_POPULATED_PARTITIONS.contains(partition)) { + expectedRows.add( + List.of((long) key, (long) key * LEFT_METRIC_MULTIPLIER, (long) key * RIGHT_METRIC_MULTIPLIER)); + } + } + return expectedRows; } private static void assertNoExceptions(JsonNode response) { @@ -301,6 +427,22 @@ private static void assertLeafStages(JsonNode response, int expectedNumLeafStage } } + /// Asserts that every leaf stage of a shuffled plan writes more than one receive mailbox. This is the control that + /// gives the `fanOut` of 1 asserted for a colocated plan its meaning. + private static void assertShuffledLeafStages(JsonNode response, int expectedNumLeafStages) { + JsonNode stageStats = response.get("stageStats"); + assertNotNull(stageStats, "Missing stage stats in shuffled response: " + response); + List leafStageSends = new ArrayList<>(); + collectLeafStageSends(stageStats, leafStageSends); + assertEquals(leafStageSends.size(), expectedNumLeafStages, + "Unexpected number of leaf stages in stage stats: " + stageStats.toPrettyString()); + for (JsonNode leafStageSend : leafStageSends) { + assertTrue(leafStageSend.path("fanOut").asInt(-1) > 1, + "A shuffled leaf send must write more than one receive mailbox, otherwise the fanOut of 1 asserted for the " + + "colocated plan proves nothing. Stage stats: " + stageStats.toPrettyString()); + } + } + private static void collectLeafStageSends(JsonNode node, List leafStageSends) { JsonNode children = node.get("children"); if ("MAILBOX_SEND".equals(node.path("type").asText()) && children != null && children.size() == 1 && "LEAF".equals( @@ -387,6 +529,10 @@ private static TableConfig createTableConfigForTable(String tableName) { .setNumReplicas(2) .setSegmentPartitionConfig(new SegmentPartitionConfig( Map.of(PARTITION_KEY_COLUMN, new ColumnPartitionConfig(PARTITION_FUNCTION, NUM_DECLARED_PARTITIONS)))) + // Without this the broker builds no partition pruner at all (SegmentPrunerFactory only reads the routing + // config), a filter on the partition key would prune nothing, and the filtered tests below would assert the + // unfiltered worker count. The unfiltered tests are unaffected: with no filter there is nothing to prune. + .setRoutingConfig(new RoutingConfig(null, List.of(RoutingConfig.PARTITION_SEGMENT_PRUNER_TYPE), null, null)) .build(); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java index f50a3db9e1c3..585576dfa992 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java @@ -61,6 +61,7 @@ public class DispatchablePlanContext { private final Map _dispatchablePlanMetadataMap = new HashMap<>(); private final Map _dispatchablePlanStageRootMap = new HashMap<>(); private final Map _partitionTableInfoCache = new HashMap<>(); + private final Map> _prunedSegmentsCache = new HashMap<>(); private long _numSegmentsPrunedByBroker; private int _leafStagesAssigned; private int _leafStagesEmpty; @@ -141,6 +142,14 @@ public Map getPartitionTableInfoCache( return _partitionTableInfoCache; } + /// The segments the broker's pruners provably eliminated for each leaf fragment, keyed by fragment id. Keyed by + /// fragment rather than by table because the verdict depends on the leaf's own filter, and the two sides of a + /// self-join scan one table under two different ones. Cached because the colocation pre-pass and the leaf + /// assignment both need it, and each entry costs a routing call. + public Map> getPrunedSegmentsCache() { + return _prunedSegmentsCache; + } + public long getNumSegmentsPrunedByBroker() { return _numSegmentsPrunedByBroker; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java index 88b43fb61e7d..0d5508f6fc90 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java @@ -150,7 +150,7 @@ private static ColocationGroup toReducibleGroup(List members, Map partitionedLeafFragmentIds = new ArrayList<>(); + List partitionedLeafFragments = new ArrayList<>(); int partitionSize = -1; int partitionParallelism = -1; String partitionFunction = null; @@ -186,7 +186,7 @@ private static ColocationGroup toReducibleGroup(List members, Map members, Map parents, int fragmentId) { static class ColocationGroup { /// The number of partition classes, and of workers before reduction, i.e. the hinted `partition_size`. final int _partitionSize; - /// The members whose data decides which classes survive. - final List _partitionedLeafFragmentIds; + /// The members whose data decides which classes survive. The whole fragment rather than its id because deciding + /// survival reads each member's own filter off its leaf stage tree, see `WorkerManager#assignPartitionClasses`. + final List _partitionedLeafFragments; - ColocationGroup(int partitionSize, List partitionedLeafFragmentIds) { + ColocationGroup(int partitionSize, List partitionedLeafFragments) { _partitionSize = partitionSize; - _partitionedLeafFragmentIds = partitionedLeafFragmentIds; + _partitionedLeafFragments = partitionedLeafFragments; } } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java index d7b32f3e4b37..0b6a4ff65c3f 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java @@ -85,6 +85,27 @@ public static PinotQuery createPinotQueryForRouting(String tableName, PlanNode l return pinotQuery; } + /// Whether [#createPinotQueryForRouting] can fold the given leaf stage tree, i.e. whether it holds no multi-input + /// node. Lets a caller skip the attempt rather than pay a thrown-and-caught exception per query, which a colocated + /// semi-join's probe leaf -- the one holding the join -- would otherwise do on every query. It lives here so that it + /// walks the tree the same way [#accumulateBottomToTop] does; the two disagreeing would either bring the exception + /// back or, worse, silently refuse shapes that fold perfectly well. + /// + /// A `true` result is not a promise that the fold succeeds: the tree may still be missing a table scan, which only + /// the fold itself detects. + public static boolean canBuildRoutingQuery(PlanNode leafStageRoot) { + List inputs = leafStageRoot.getInputs(); + if (inputs.size() > 1) { + return false; + } + for (PlanNode input : inputs) { + if (!canBuildRoutingQuery(input)) { + return false; + } + } + return true; + } + private static void accumulateBottomToTop(PlanNode root, List parentNodes) { Preconditions.checkState(root.getInputs().size() <= 1, "Leaf stage nodes should have at most one input, found: %s", root.getInputs().size()); 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 a7031c5c4dc4..2bae257eee97 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 @@ -136,12 +136,24 @@ public void assignWorkers(PlanFragment rootFragment, DispatchablePlanContext con /// publishes the decision on each partitioned leaf of the group (see /// [DispatchablePlanMetadata#getPartitionClassIds()] and [DispatchablePlanMetadata#getPaddedClassCandidates()]). /// - /// A class survives when *any* member holds a segment in it: the union, not the intersection, because a class that - /// holds data for one member must keep its worker on every member or the members stop agreeing on what a worker id - /// stands for. A member holding no data in a surviving class gets a worker with no segments (see - /// [#assignPaddedWorker]). Emptiness is computed in class space (`0..partitionSize-1`) rather than over raw partition - /// ids because members may declare different partition counts; a member carrying no per-class visibility (replicated, - /// non-partitioned, or deriving its worker map from a peer) contributes nothing to the union. + /// A class survives when *any* member holds a segment in it that the query's filter does not provably exclude: the + /// union, not the intersection, because a class that holds matching data for one member must keep its worker on + /// every member or the members stop agreeing on what a worker id stands for. Dropping a class the group as a whole + /// has no matching row in is what turns broker pruning into fewer workers, and therefore fewer dispatched servers, + /// for a colocated join. The union direction is also what keeps this layer free of relational semantics: no input + /// row anywhere in a class means no output row attributable to it, whatever operator sits above, whereas dropping a + /// class only one side filtered away would be wrong for a RIGHT or FULL join, a union, or an anti-join. + /// + /// A member holding no data in a surviving class gets a worker with no segments (see [#assignPaddedWorker]). A + /// member that holds data the filter excludes does *not*: it keeps every segment of the class. Survival is decided + /// per class rather than per partition on purpose -- a class's worker is placed on the servers shared by all of its + /// populated partitions, so dropping some of them on one member and not on another would move that member's worker + /// off its peer's server and turn an in-process exchange into a network one. Filtering inside a surviving class is a + /// possible follow-up. + /// + /// Emptiness is computed in class space (`0..partitionSize-1`) rather than over raw partition ids because members + /// may declare different partition counts; a member carrying no per-class visibility (replicated, non-partitioned, + /// or deriving its worker map from a peer) contributes nothing to the union. /// /// Marking no group keeps the assignment as it is without one: every class gets a worker, so a class holding no /// segment fails the assignment instead of being dropped or padded. @@ -151,67 +163,108 @@ private void assignPartitionClasses(PlanFragment rootFragment, DispatchablePlanC for (ColocationGroupAnalyzer.ColocationGroup group : ColocationGroupAnalyzer.findReducibleGroups(rootFragment, metadataMap)) { int numWorkers = group._partitionSize; - List memberFragmentIds = group._partitionedLeafFragmentIds; - // The servers each member can scan each class on, in the same order as the member fragment ids. - List>> memberClassServers = new ArrayList<>(memberFragmentIds.size()); + List memberFragments = group._partitionedLeafFragments; + List memberFragmentIds = new ArrayList<>(memberFragments.size()); + // The servers each member can scan each class on, in the same order as the member fragments. + List>> memberClassServers = new ArrayList<>(memberFragments.size()); + // The partition layouts, same order again, kept only to count what a filter-dropped class cost each member. + List memberPartitionInfoMaps = new ArrayList<>(memberFragments.size()); // Allocated lazily, once the first member has checked the hint against its table: numWorkers is the raw hinted // partition size, so sizing anything from it before that check would let a bogus hint allocate unboundedly. The // check also bounds it by the table's partition count. - boolean[] survivingClasses = null; + // populatedClasses ignores the filter and is what decides padding and the all-pruned fallback; matchingClasses + // is the same union taken over the segments the filter leaves, and is a subset of it. + boolean[] populatedClasses = null; + boolean[] matchingClasses = null; boolean reducible = true; - for (Integer fragmentId : memberFragmentIds) { - DispatchablePlanMetadata metadata = metadataMap.get(fragmentId); + for (PlanFragment fragment : memberFragments) { + DispatchablePlanMetadata metadata = metadataMap.get(fragment.getFragmentId()); String tableName = metadata.getScannedTables().get(0); // NOTE: A failure here is the same one the leaf assignment would hit for this table, only raised earlier. PartitionTableInfo partitionTableInfo = partitionTableInfoCache.computeIfAbsent(tableName, this::calculatePartitionTableInfo); - int numPartitions = partitionTableInfo._partitionInfoMap.length; + PartitionInfo[] partitionInfoMap = partitionTableInfo._partitionInfoMap; + int numPartitions = partitionInfoMap.length; if (numPartitions == 0 || numPartitions % numWorkers != 0) { // The table does not match the hinted partition size. Leave the group alone so that checkPartitionInfoMap // reports it during the leaf assignment. reducible = false; break; } - if (survivingClasses == null) { - survivingClasses = new boolean[numWorkers]; + if (populatedClasses == null) { + populatedClasses = new boolean[numWorkers]; + matchingClasses = new boolean[numWorkers]; } - List> classServers = collectClassServers(partitionTableInfo._partitionInfoMap, numWorkers); + List> classServers = collectClassServers(partitionInfoMap, numWorkers); + Set prunedSegments = getPrunedSegments(fragment, tableName, context); boolean anyPopulated = false; for (int classId = 0; classId < numWorkers; classId++) { if (classServers.get(classId) != null) { - survivingClasses[classId] = true; + populatedClasses[classId] = true; anyPopulated = true; + if (prunedSegments == null) { + // No filter to prune this member with, so it contributes every class it holds data in. + matchingClasses[classId] = true; + } } } // A member holding no data at all leaves nothing to assign: no class to place its single empty worker in, and // no server known to host the table to place it on. Check the deferred cause first though -- a table whose // every partition is deferred also has no populated class, and reports far more actionably. That is the - // pre-pass' only deferred check: a group it marks gets no broker pruning, so the leaf assignment covers the - // rest. + // pre-pass' only deferred check: a group it marks gets no broker pruning at the leaf, so the leaf assignment + // covers the rest. + // + // NOTE: This reads the unfiltered population on purpose. A member whose filter matches nothing holds data + // all the same, and failing it here would turn a correct empty result into a query error. if (!anyPopulated) { checkNoPartitionsWithOnlyDeferredSegments(partitionTableInfo, tableName); } Preconditions.checkState(anyPopulated, "Failed to find any segment in any partition for table: %s, which is required for a partitioned worker " + "assignment", tableName); + if (prunedSegments != null) { + markClassesWithMatchingData(partitionInfoMap, numWorkers, prunedSegments, matchingClasses); + } memberClassServers.add(classServers); + memberPartitionInfoMaps.add(partitionInfoMap); + memberFragmentIds.add(fragment.getFragmentId()); } if (!reducible) { continue; } - // The member list is never empty (see ColocationGroupAnalyzer#toReducibleGroup), so the loop allocated this, and - // the class list is never empty either: every member holds data in at least one class, and the union keeps it. - assert survivingClasses != null; + // The member list is never empty (see ColocationGroupAnalyzer#toReducibleGroup), so the loop allocated these, + // and the class list is never empty either: every member holds data in at least one class, and the union keeps + // it. + assert populatedClasses != null && matchingClasses != null; + // A group the filter empties keeps all of its populated classes, mirroring the leaf-level fallback in + // computePartitionsToKeep: a zero-worker leaf has no handling on a 1-to-1 exchange, and the server-side filter + // still returns the correct empty result from an unreduced plan. + boolean[] survivingClasses = anyTrue(matchingClasses) ? matchingClasses : populatedClasses; int[] partitionClassIds = toClassIds(survivingClasses); Map>> padding = computePadding(memberFragmentIds, memberClassServers, partitionClassIds); if (padding.isEmpty() && partitionClassIds.length == numWorkers) { // Worker k already stands for class k on every member: nothing to reduce, nothing to pad. Leaving the group - // unmarked also keeps broker pruning on for its leaves (see computePartitionsToKeep). A group that needs + // unmarked also keeps leaf-level broker pruning on for the members that are eligible for it (see + // computePartitionsToKeep), which prunes at partition rather than class granularity. A group that needs // padding is marked even when it keeps every class, because a padded worker's id is its index in the class // list. continue; } + // Report what the filter cost, not what the class reduction did: a class no member holds data in is empty + // rather than pruned, and it is already dropped above without being counted. Derived from the decision actually + // taken, so the all-pruned fallback above reports nothing. + long numPrunedSegments = 0; + for (int classId = 0; classId < numWorkers; classId++) { + if (populatedClasses[classId] && !survivingClasses[classId]) { + for (PartitionInfo[] partitionInfoMap : memberPartitionInfoMaps) { + numPrunedSegments += countSegmentsInClass(partitionInfoMap, numWorkers, classId); + } + } + } + if (numPrunedSegments > 0) { + context.addNumSegmentsPrunedByBroker(numPrunedSegments); + } // One shared array instance, so that the agreement check in MailboxAssignmentVisitor compares one list rather // than copies of it. The padding goes on the same metadata: a padded worker's id only means something within the // list. @@ -223,6 +276,48 @@ private void assignPartitionClasses(PlanFragment rootFragment, DispatchablePlanC } } + private static long countSegmentsInClass(PartitionInfo[] partitionInfoMap, int numWorkers, int classId) { + long numSegments = 0; + for (int partitionId = classId; partitionId < partitionInfoMap.length; partitionId += numWorkers) { + PartitionInfo partitionInfo = partitionInfoMap[partitionId]; + if (partitionInfo != null) { + numSegments += CollectionUtils.size(partitionInfo._offlineSegments) + + CollectionUtils.size(partitionInfo._realtimeSegments); + } + } + return numSegments; + } + + /// Sets, for every class holding at least one segment the given pruned set does not cover, the corresponding entry + /// of `matchingClasses`. Only presence in the pruned set is a proof (see [RoutingManager#getPrunedSegments]), so a + /// partition with no segment listed at all -- which the hybrid layout allows for one of the two table types -- is + /// read as matching rather than as empty. + private static void markClassesWithMatchingData(PartitionInfo[] partitionInfoMap, int numWorkers, + Set prunedSegments, boolean[] matchingClasses) { + int numPartitions = partitionInfoMap.length; + for (int classId = 0; classId < numWorkers; classId++) { + if (matchingClasses[classId]) { + continue; + } + for (int partitionId = classId; partitionId < numPartitions; partitionId += numWorkers) { + PartitionInfo partitionInfo = partitionInfoMap[partitionId]; + if (partitionInfo != null && !allSegmentsPruned(partitionInfo, prunedSegments)) { + matchingClasses[classId] = true; + break; + } + } + } + } + + private static boolean anyTrue(boolean[] flags) { + for (boolean flag : flags) { + if (flag) { + return true; + } + } + return false; + } + /// Returns the servers that can scan each partition class of the given layout as a whole, in class-id order, or /// `null` for a class that holds no segment at all. This is the intersection of the fully replicated servers of the /// class's populated partitions, i.e. the candidate set its worker is picked from (see @@ -687,13 +782,14 @@ private void assignWorkersToLeafFragment(PlanFragment fragment, DispatchablePlan metadata.setPartitionParallelism(partitionHints.getPartitionParallelism()); if (partitionHints.getPartitionKey() != null) { - // Broker pruning: build a filter-bearing routing query (null when disabled/unsupported) so the partitioned - // assignment can drop partitions with no matching segments. Reuses the same gate as the non-partitioned path. - // Skip pre-partitioned leaves and leaves of a reduced colocated group up front: pruning is disabled for them - // (see computePartitionsToKeep), so don't spend planning time building the routing query. - PinotQuery routingPinotQuery = metadata.isPrePartitioned() || metadata.getPartitionClassIds() != null ? null - : extractRoutingQuery(fragment.getFragmentRoot(), metadata.getScannedTables().get(0), context); - assignWorkersToPartitionedLeafFragment(metadata, context, partitionHints, routingPinotQuery); + // Broker pruning: the segments the pruners provably eliminated (empty when disabled/unsupported) so the + // partitioned assignment can drop partitions holding none of them. Skip the lookup for a pre-partitioned leaf + // and for a leaf of a marked colocated group: leaf-level pruning is disabled for both (see + // computePartitionsToKeep) because the group's shared class list already carries their verdict, so asking + // would only cost planning time. + Set prunedSegments = metadata.isPrePartitioned() || metadata.getPartitionClassIds() != null ? null + : getPrunedSegments(fragment, metadata.getScannedTables().get(0), context); + assignWorkersToPartitionedLeafFragment(metadata, context, partitionHints, prunedSegments); updateContextForLeafStage(metadata, context); return; } @@ -867,6 +963,86 @@ private Map getRoutingTable(PinotQuery pinotQuery, long re } } + /// Returns the segments the broker's pruners provably eliminated for the given leaf fragment, or `null` when there + /// was no filter to prune with at all: broker pruning off, an unsupported leaf shape, a filterless leaf, or a + /// routing failure (pruning is best-effort and never fails a query that would otherwise route). + /// + /// The empty set and `null` mean different things and callers depend on it. An empty set is "a filter ran and + /// proved nothing", which still lets the partitioned assignment skip a partition holding no segment at all -- + /// behaviour that predates this and that a query with an empty partition relies on to plan. + /// + /// Memoised per fragment for the query, because the colocation pre-pass and the leaf assignment both ask for it and + /// each answer costs a routing call. Keyed by fragment rather than by table: the two sides of a self-join scan one + /// table under two different filters. + @Nullable + private Set getPrunedSegments(PlanFragment fragment, String tableName, DispatchablePlanContext context) { + // Not computeIfAbsent: null is a meaningful answer here and would be recomputed on every call. + Map> prunedSegmentsCache = context.getPrunedSegmentsCache(); + Integer fragmentId = fragment.getFragmentId(); + if (prunedSegmentsCache.containsKey(fragmentId)) { + return prunedSegmentsCache.get(fragmentId); + } + Set prunedSegments = computePrunedSegments(fragment, tableName, context); + prunedSegmentsCache.put(fragmentId, prunedSegments); + return prunedSegments; + } + + @Nullable + private Set computePrunedSegments(PlanFragment fragment, String tableName, + DispatchablePlanContext context) { + PinotQuery routingPinotQuery = extractRoutingQuery(fragment.getFragmentRoot(), tableName, context); + if (routingPinotQuery == null || routingPinotQuery.getFilterExpression() == null) { + return null; + } + try { + TableType tableType = TableNameBuilder.getTableTypeFromTableName(routingPinotQuery.getDataSource() + .getTableName()); + if (tableType != null) { + return getPrunedSegmentsHelper(routingPinotQuery); + } + // A raw table name may resolve to either or both physical tables. Segment names are unique across them, and a + // segment is pruned by the table that holds it, so the two verdicts simply add up. Only merge when both prove + // something: a table of one type alone is the common case, and copying its verdict to union it with an empty + // set would allocate a second set over every segment name for nothing. + Set offlinePrunedSegments = getPrunedSegmentsHelper(routingPinotQuery, TableType.OFFLINE); + Set realtimePrunedSegments = getPrunedSegmentsHelper(routingPinotQuery, TableType.REALTIME); + if (offlinePrunedSegments.isEmpty()) { + return realtimePrunedSegments; + } + if (realtimePrunedSegments.isEmpty()) { + return offlinePrunedSegments; + } + Set prunedSegments = new HashSet<>(offlinePrunedSegments); + prunedSegments.addAll(realtimePrunedSegments); + return prunedSegments; + } catch (RuntimeException e) { + // Pruning is best-effort: never fail a query that would otherwise route successfully unpruned. + LOGGER.warn("Broker pruning skipped for table {} due to routing failure", tableName, e); + return null; + } + } + + /// A table the routing manager does not have is reported as `null` there; here it is simply a table that proves + /// nothing, which is the same thing this path does with a table whose pruners eliminated no segment. + private Set getPrunedSegmentsHelper(PinotQuery pinotQuery) { + Set prunedSegments = + _routingManager.getPrunedSegments(CalciteSqlCompiler.convertToBrokerRequest(pinotQuery)); + return prunedSegments != null ? prunedSegments : Set.of(); + } + + private Set getPrunedSegmentsHelper(PinotQuery pinotQuery, TableType tableType) { + return getPrunedSegmentsHelper(withTableType(pinotQuery, tableType)); + } + + /// Returns a copy of the given routing query aimed at one physical table, so that a query written against a raw + /// table name can be routed against each type in turn. + private static PinotQuery withTableType(PinotQuery pinotQuery, TableType tableType) { + PinotQuery copy = pinotQuery.deepCopy(); + copy.getDataSource().setTableName(TableNameBuilder.forType(tableType).tableNameWithType( + TableNameBuilder.extractRawTableName(pinotQuery.getDataSource().getTableName()))); + return copy; + } + /// Builds a [PinotQuery] from the leaf stage tree for broker-side segment pruning on the logical planner path. /// Returns `null` if broker pruning is disabled or the leaf stage shape is unsupported. @Nullable @@ -878,6 +1054,9 @@ private PinotQuery extractRoutingQuery(PlanNode leafStageRoot, String tableName, if (!useBrokerPruning) { return null; } + if (!PlanNodeRoutingQueryBuilder.canBuildRoutingQuery(leafStageRoot)) { + return null; + } try { PinotQuery pinotQuery = PlanNodeRoutingQueryBuilder.createPinotQueryForRouting(tableName, leafStageRoot, false); Map queryOptions = context.getPlannerContext().getOptions(); @@ -911,10 +1090,7 @@ private RoutingTable getRoutingTableHelper(PinotQuery pinotQuery, long requestId @Nullable private RoutingTable getRoutingTableHelper(PinotQuery pinotQuery, long requestId, TableType tableType) { - PinotQuery copy = pinotQuery.deepCopy(); - copy.getDataSource().setTableName(TableNameBuilder.forType(tableType).tableNameWithType( - TableNameBuilder.extractRawTableName(pinotQuery.getDataSource().getTableName()))); - return getRoutingTableHelper(copy, requestId); + return getRoutingTableHelper(withTableType(pinotQuery, tableType), requestId); } // -------------------------------------------------------------------------- @@ -1137,8 +1313,7 @@ private static void transferToServerInstanceLogicalSegmentsMap(String physicalTa /// Assigns one worker per partition class of a leaf that scans a partitioned table. private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata metadata, - DispatchablePlanContext context, LeafPartitionHints partitionHints, - @Nullable PinotQuery routingPinotQuery) { + DispatchablePlanContext context, LeafPartitionHints partitionHints, @Nullable Set prunedSegments) { // when partition key exist, we assign workers for leaf-stage in partitioned fashion. String partitionKey = partitionHints.getPartitionKey(); assert partitionKey != null; @@ -1172,8 +1347,7 @@ private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata met collectHostingServers(partitionInfoMap)) : null; // Broker pruning: the partitions to keep (null means keep all). Partitions absent from the set are skipped below. - Set partitionsToKeep = - computePartitionsToKeep(routingPinotQuery, metadata, context.getRequestId(), partitionInfoMap); + Set partitionsToKeep = computePartitionsToKeep(prunedSegments, metadata, partitionInfoMap); if (partitionsToKeep != null) { long numSegmentsPrunedByBroker = countPrunedSegments(partitionInfoMap, partitionsToKeep); if (numSegmentsPrunedByBroker > 0) { @@ -1202,35 +1376,42 @@ private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata met partitionClassIds, partitionsToKeep, paddingInfo, _routingManager.getEnabledServerInstanceMap(), workerIdToServerInstanceMap, workerIdToSegmentsMap); } - checkLeafWorkerAssignment(tableName, workerIdToServerInstanceMap, workerIdToSegmentsMap); + checkLeafWorkerAssignment(tableName, partitionClassIds, workerIdToServerInstanceMap, workerIdToSegmentsMap); metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap); metadata.setWorkerIdToSegmentsMap(workerIdToSegmentsMap); metadata.setTimeBoundaryInfo(partitionTableInfo._timeBoundaryInfo); metadata.setPartitionFunction(partitionFunction); } - /// Broker pruning for the partitioned leaf path. Returns the set of partition ids that still have at least one - /// segment matching the query filter, or `null` to keep all partitions. + /// Broker pruning for the partitioned leaf path. Returns the set of partition ids that are not provably empty for + /// this query, or `null` to keep all partitions. + /// + /// Note that an *empty* pruned set is not the same as an absent one and does not return `null` here: a filter that + /// proved nothing still leaves this the job of skipping a partition holding no segment at all, which is what lets a + /// table with an empty partition plan rather than fail on a worker it cannot place. /// /// Returns `null` (no pruning) when any of the following hold: /// - /// - broker pruning is disabled or the leaf shape is unsupported (the routing query is `null`), or there is - /// no filter to prune with; + /// - there was no filter to prune with at all (see [#getPrunedSegments]) -- broker pruning is disabled, the leaf + /// shape is unsupported, the leaf carries no filter, or routing failed (pruning is best-effort); /// - the leaf feeds a pre-partitioned (1-to-1 direct) exchange, or it belongs to a colocated group that agreed on a - /// partition class list -- dropping/compacting workers would misalign sender/receiver worker ids in - /// `MailboxAssignmentVisitor`. A non-pre-partitioned leaf is shuffled via `connectWorkers`, which re-hashes across - /// any worker count, so pruning is safe there; - /// - routing fails (pruning is best-effort); + /// partition class list. Both get their verdict from the group instead, in `assignPartitionClasses`, because it is + /// the only place that sees every member before any of them is assigned: a leaf deciding on its own would drop a + /// class its peer keeps, and the two would stop agreeing on what a worker id stands for. A non-pre-partitioned, + /// unmarked leaf is shuffled via `connectWorkers`, which re-hashes across any worker count, so it can decide alone + /// -- and at partition rather than class granularity; /// - every partition would be pruned -- an empty worker map would break exchanges in a multi-leaf plan (the /// all-leaves-empty short-circuit does not fire for a partially-empty plan), and the server-side filter still /// yields the correct empty result unpruned. /// - /// Partition survival is decided by routing the filter-bearing query through the [RoutingManager] (the same - /// mechanism the non-partitioned path uses), so the segment-level pruners judge survival using each segment's own - /// partition metadata. This is correct for every partition function and configuration, unlike recomputing the - /// partition id from the table-level function name (which lacks the per-segment function config). A partition is - /// dropped only when every one of its segments was pruned; a segment that merely became unavailable keeps its - /// partition alive so matching data is never silently dropped. + /// A partition is dropped only when every one of its segments is in the *provably pruned* set (see + /// [RoutingManager#getPrunedSegments]), never because a segment failed to appear somewhere. That direction is the + /// whole point: absence from a routing result has innocent causes -- a segment classified as optional by instance + /// selection, one whose server left the enabled server map, one that entered the partition metadata before it became + /// selectable -- and each would otherwise be read as "this partition is empty" and silently drop matching rows. It + /// also makes the verdict independent of the request id, so two leaves scanning one table under one filter cannot + /// disagree. Judging by pruner verdict rather than by recomputing the partition id from the table-level function + /// name is also what keeps it correct for every partition function and per-segment function config. /// /// Note that pruning here is partition-level, not segment-level: a surviving partition dispatches all of its /// segments, including ones the pruners eliminated (the server-side pruners drop those again cheaply). This keeps @@ -1238,37 +1419,15 @@ private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata met /// workers -- at the cost of a lower pruning ceiling than the non-partitioned path for partitions with mixed-match /// segments. Segment-level pruning within surviving partitions is a possible follow-up. @Nullable - private Set computePartitionsToKeep(@Nullable PinotQuery routingPinotQuery, - DispatchablePlanMetadata metadata, long requestId, PartitionInfo[] partitionInfoMap) { - if (routingPinotQuery == null || routingPinotQuery.getFilterExpression() == null || metadata.isPrePartitioned() - || metadata.getPartitionClassIds() != null) { - return null; - } - Map routingTableMap; - try { - routingTableMap = getRoutingTable(routingPinotQuery, requestId); - } catch (RuntimeException e) { - // Pruning is best-effort: never fail a query that would otherwise route successfully unpruned. - LOGGER.warn("Broker pruning skipped for partitioned table {} due to routing failure", - routingPinotQuery.getDataSource().getTableName(), e); + private static Set computePartitionsToKeep(@Nullable Set prunedSegments, + DispatchablePlanMetadata metadata, PartitionInfo[] partitionInfoMap) { + if (prunedSegments == null || metadata.isPrePartitioned() || metadata.getPartitionClassIds() != null) { return null; } - if (routingTableMap.isEmpty()) { - return null; - } - Set matchedSegments = new HashSet<>(); - for (RoutingTable routingTable : routingTableMap.values()) { - for (SegmentsToQuery segmentsToQuery : routingTable.getServerInstanceToSegmentsMap().values()) { - matchedSegments.addAll(segmentsToQuery.getSegments()); - } - // Keep a partition alive if any of its segments is merely unavailable (rather than pruned) so we never drop data. - matchedSegments.addAll(routingTable.getUnavailableSegments()); - } Set partitionsToKeep = new HashSet<>(); for (int i = 0; i < partitionInfoMap.length; i++) { PartitionInfo partitionInfo = partitionInfoMap[i]; - if (partitionInfo != null && (containsAny(partitionInfo._offlineSegments, matchedSegments) || containsAny( - partitionInfo._realtimeSegments, matchedSegments))) { + if (partitionInfo != null && !allSegmentsPruned(partitionInfo, prunedSegments)) { partitionsToKeep.add(i); } } @@ -1276,18 +1435,28 @@ private Set computePartitionsToKeep(@Nullable PinotQuery routingPinotQu return partitionsToKeep.isEmpty() ? null : partitionsToKeep; } - private static boolean containsAny(@Nullable List segments, Set matchedSegments) { + /// Returns whether every segment of the given partition is provably pruned. Vacuously true for a partition listing + /// no segment at all, which has no rows to contribute either way -- note that a partition holding data the broker + /// cannot route yet has no entry in the map rather than an empty one, so it is not this case (see + /// [#checkNoPartitionsWithOnlyDeferredSegments]). + private static boolean allSegmentsPruned(PartitionInfo partitionInfo, Set prunedSegments) { + return allPruned(partitionInfo._offlineSegments, prunedSegments) + && allPruned(partitionInfo._realtimeSegments, prunedSegments); + } + + private static boolean allPruned(@Nullable List segments, Set prunedSegments) { if (segments != null) { for (String segment : segments) { - if (matchedSegments.contains(segment)) { - return true; + if (!prunedSegments.contains(segment)) { + return false; } } } - return false; + return true; } - /// Counts the segments in partitions dropped by broker pruning (those absent from `partitionsToKeep`). + /// Counts the segments in partitions dropped by broker pruning (those absent from `partitionsToKeep`). A partition + /// dropped for holding no segment rather than for being pruned contributes nothing, so it is not miscounted. private static long countPrunedSegments(PartitionInfo[] partitionInfoMap, Set partitionsToKeep) { long numPrunedSegments = 0; for (int i = 0; i < partitionInfoMap.length; i++) { @@ -1510,12 +1679,20 @@ private static Map collectHostingServers(PartitionInfo[] partiti /// server map by worker id, where a gap leaves a null entry; /// - every worker must have a segments map keyed by 1 or 2 [TableType] names, with non-null lists, because the server /// splits the request on the number of entries and resolves one table data manager per key: an unexpected key - /// becomes an opaque server-side failure. + /// becomes an opaque server-side failure; + /// - a leaf of a colocated group must produce exactly one worker per class of the group's shared list, because a + /// worker id *is* an index into that list. Nothing downstream can catch a leaf that skipped one: + /// `MailboxAssignmentVisitor#checkPartitionClassAgreement` compares the shared array against itself, so a member + /// that quietly assigned fewer workers than the array claims still agrees with its peers on the array while + /// disagreeing with them on what every worker id after the gap means. @VisibleForTesting - static void checkLeafWorkerAssignment(String tableName, + static void checkLeafWorkerAssignment(String tableName, @Nullable int[] partitionClassIds, Map workerIdToServerInstanceMap, Map>> workerIdToSegmentsMap) { int numWorkers = workerIdToServerInstanceMap.size(); + Preconditions.checkState(partitionClassIds == null || partitionClassIds.length == numWorkers, + "Got %s workers for partition classes: %s of table: %s", numWorkers, + partitionClassIds != null ? Arrays.toString(partitionClassIds) : null, tableName); Preconditions.checkState(workerIdToSegmentsMap.size() == numWorkers, "Got %s workers but %s worker segment entries for table: %s", numWorkers, workerIdToSegmentsMap.size(), tableName); diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java index 3a779db87312..5908900e33c4 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java @@ -22,6 +22,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.stream.Collectors; import javax.annotation.Nullable; import org.apache.calcite.rel.RelDistribution; import org.apache.pinot.calcite.rel.hint.PinotHintOptions; @@ -54,7 +55,7 @@ public void testGroupWithOnlyPrePartitionedSendsIsReducible() { assertEquals(groups.size(), 1); assertEquals(groups.get(0)._partitionSize, 4); - assertEquals(Set.copyOf(groups.get(0)._partitionedLeafFragmentIds), Set.of(2, 3)); + assertEquals(Set.copyOf(fragmentIds(groups.get(0))), Set.of(2, 3)); } /// A member that also receives a shuffled send must keep today's worker count, or that sender's rows land on @@ -78,7 +79,7 @@ public void testSingletonSendFormsAnEdgeWithoutPrePartitioning() { ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(RelDistribution.Type.SINGLETON), metadataMap); assertEquals(groups.size(), 1); - assertEquals(Set.copyOf(groups.get(0)._partitionedLeafFragmentIds), Set.of(2, 3)); + assertEquals(Set.copyOf(fragmentIds(groups.get(0))), Set.of(2, 3)); } /// Reducing the worker count must not turn mismatched counts into a match for a pre-partitioned BROADCAST send, which @@ -227,7 +228,7 @@ public void testReplicatedLeafDoesNotBlockTheGroup() { assertEquals(groups.size(), 1); // Only the partitioned leaf decides which classes survive. - assertEquals(groups.get(0)._partitionedLeafFragmentIds, List.of(2)); + assertEquals(fragmentIds(groups.get(0)), List.of(2)); } /// A lookup join's workers come from its single local exchange child, so its own hints (a different partition size @@ -249,7 +250,7 @@ public void testLookupJoinMemberIsIgnored() { assertEquals(groups.size(), 1); assertEquals(groups.get(0)._partitionSize, 4); - assertEquals(groups.get(0)._partitionedLeafFragmentIds, List.of(2)); + assertEquals(fragmentIds(groups.get(0)), List.of(2)); } /// A group of intermediate stages only has nothing to reduce: only a partitioned leaf's data decides the classes. @@ -287,7 +288,11 @@ public void testSpooledFragmentTiesEveryReceiverIntoOneGroup() { // One group, and the spooled leaf is listed once rather than once per receiver. assertEquals(groups.size(), 1); - assertEquals(groups.get(0)._partitionedLeafFragmentIds, List.of(3)); + assertEquals(fragmentIds(groups.get(0)), List.of(3)); + } + + private static List fragmentIds(ColocationGroupAnalyzer.ColocationGroup group) { + return group._partitionedLeafFragments.stream().map(PlanFragment::getFragmentId).collect(Collectors.toList()); } /// Builds a 4 stage plan: 2 partitioned leaves (stages 2 and 3) sending to a join stage (stage 1), which sends diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java index 60fe0a7b93e5..b84ee7a982ee 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java @@ -720,10 +720,14 @@ public void testBuildLogicalTableRoutingRequestWithoutFilterUsesSelectStar() { } @Test - public void testBrokerPruningPartitionedLeafSkippedForColocatedJoin() { - // A pre-partitioned leaf (here both sides of a colocated self-join) feeds a 1-to-1 direct exchange wired by worker - // id. Compacting a side's workers for pruned partitions could pair mismatched partitions across the exchange, so - // pruning is skipped for any pre-partitioned leaf and all partitions stay assigned on every scan. + public void testBrokerPruningColocatedJoinDropsClassesBothSidesPrune() { + // A pre-partitioned leaf feeds a 1-to-1 direct exchange wired by worker id, so no leaf may decide on its own what + // to drop. The colocation pre-pass decides for the whole group instead: a class survives when ANY member still + // holds a segment its own filter leaves, so dropping one is safe without knowing which operator sits above. + // + // This is the shape the reduction is worth having for. The filter is on the partition key, so it transfers across + // the join equality to both sides, both prune every class but 2, and the group's class list shrinks to [2]: one + // worker per leaf, and the query is dispatched to the single server holding partition 2 instead of all four. QueryEnvironment queryEnvironment = newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4, List.of("seg2"), 3); try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( @@ -733,13 +737,76 @@ public void testBrokerPruningPartitionedLeafSkippedForColocatedJoin() { + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ t2 " + "ON t1.col1 = t2.col1 WHERE t1.col1 = 'foo'")) { DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); - assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); List leafFragments = leafFragments(dispatchableSubPlan); - assertFalse(leafFragments.isEmpty()); + assertEquals(leafFragments.size(), 2); + for (DispatchablePlanFragment leaf : leafFragments) { + assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 1); + // Worker 0 stands for class 2 on both members, and carries that class's whole segment list. + assertEquals(assignedSegments(leaf, 0), List.of("seg2")); + } + // Three classes dropped, one segment each, counted once per member of the group rather than from the routing + // table's own self-reported count (the fixture reports 3). + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 6); + // The point of the exercise: partition 2 lives only on server 3, so that is the whole dispatched set. + assertEquals(dispatchedServers(dispatchableSubPlan), Set.of(getServerInstance("localhost", 3).getInstanceId())); + } + } + + @Test + public void testBrokerPruningSelfJoinKeepsTheTwoSidesVerdictsApart() { + // Two leaves scanning ONE table under two different filters. The pruning verdict is memoised per leaf fragment, + // and this is what says so: keyed by table instead, one side would inherit the other's verdict and the union would + // be computed from one filter applied twice. + // + // Left keeps only class 1, right only class 2, so the union is {1, 2} -- a result neither side's verdict produces + // on its own, and neither does either verdict applied to both sides ({1} or {2}). + QueryEnvironment queryEnvironment = + newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4, 1, List.of(), List.of(), 0, false, Set.of(), Set.of(), + Map.of("left", List.of("seg1"), "right", List.of("seg2"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; SELECT t1.col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ t1 " + + "JOIN testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ t2 " + + "ON t1.col1 = t2.col1 WHERE t1.col2 = 'left' AND t2.col2 = 'right'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + List leafFragments = leafFragments(dispatchableSubPlan); + assertEquals(leafFragments.size(), 2); + Map workerIdToClass = new HashMap<>(); for (DispatchablePlanFragment leaf : leafFragments) { - assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 4, "Expected all partitions assigned for a pruned-gated " - + "colocated join leaf"); + // Both classes survive on both sides, and a class the group keeps dispatches all of its segments even on the + // member whose own filter excluded them. + assertEquals(workerIdToPartitions(leaf, "seg"), Map.of(0, Set.of(1), 1, Set.of(2))); + mergeWorkerIdToClass(workerIdToClass, leaf, "seg", 4); } + assertEquals(workerIdToClass, Map.of(0, 1, 1, 2)); + // Classes 0 and 3 dropped, one segment each, on each of the two leaves. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 4); + } + } + + @Test + public void testBrokerPruningColocatedJoinKeepsClassOnlyOneSidePrunes() { + // The union rule, and the reason it is not an intersection: only t1 is filtered, t2 still holds every class, so + // every class keeps its worker and the fan-out is unchanged. Dropping the classes t1's filter empties would be + // wrong the moment the operator above is a RIGHT or FULL join, a union, or an anti-join -- and the worker + // assignment deliberately knows nothing about which one it is. A one-sided filter buying nothing is the price. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).survivingSegments(List.of("a_seg1")), + new ColocatedTableSpec(4, false).survivingSegments(List.of("b_seg0", "b_seg1", "b_seg2", "b_seg3"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + assertEquals(leafA.getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafB.getWorkerIdToSegmentsMap().size(), 4); + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + // A kept class dispatches all of its segments on every member, including the ones the member's own filter + // excluded, so worker 0 of the filtered side still scans a_seg0. + assertEquals(workerIdToPartitions(leafA, "a_seg"), + Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2), 3, Set.of(3))); + assertEquals(workerIdToServer(leafA), workerIdToServer(leafB)); } } @@ -914,6 +981,29 @@ public void testPlainPartitionedLeafWithPartitionWithOnlyDeferredSegmentsStillPr } } + @Test + public void testPlainPartitionedLeafWithEmptyPartitionPlansWhenNothingIsPruned() { + // Partition 3 holds no segment at all and the filter prunes nothing, so the pruning verdict is an EMPTY set rather + // than an absent one. It still has to be acted on: without a verdict every partition gets a worker, and the one + // with no segment has no server to place it on, which fails a query that plans perfectly well with 3 workers. + // Hence the verdict distinguishes "a filter ran and proved nothing" from "there was no filter". + QueryEnvironment queryEnvironment = + newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4, 1, List.of("seg0", "seg1", "seg2"), List.of(), 0, + false, Set.of(3), Set.of()); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; SELECT col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ " + + "WHERE col1 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan); + assertNotNull(leaf); + assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 3); + assertEquals(assignedSegments(leaf), List.of("seg0", "seg1", "seg2")); + // Nothing was pruned, only skipped for holding nothing, so nothing is reported as pruned either. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + } + } + @Test public void testPartitionedLeafRejectsTableWithoutAnyPartition() { // An empty partition info map passes the "partitions must be a multiple of the hinted partition size" check @@ -1260,36 +1350,176 @@ public void testColocatedJoinReducesFanOutToPopulatedClasses() { } @Test - public void testColocatedJoinReducedGroupIgnoresBrokerPruning() { - // A reduced group's worker id is a position in the group's surviving class list, not a running counter over what a - // filter leaves behind, so broker pruning has to be off for its leaves. This shape is the only one that reaches - // that gate: a join written with is_colocated_by_join_keys marks its leaves pre-partitioned and is gated one step - // earlier (see testBrokerPruningPartitionedLeafSkippedForColocatedJoin), while a fact table joined with a - // replicated dimension table over an explicit local exchange is not marked pre-partitioned. + public void testBrokerPruningColocatedJoinDropsClassWhosePartitionsListNoSegment() { + // A partition whose entry lists no segment holds no rows, so its class is dropped even though the pruners proved + // nothing about any segment -- emptiness the planner can see for itself, not a pruning verdict. It is reported as + // such: numSegmentsPrunedByBroker stays 0 because no segment was pruned. + // + // The broker does not publish this shape today (a partition's entry is created together with its first segment), + // so this pins the behaviour rather than describing something reachable. It is what the golden physical plans + // record for a leaf whose fixture builds partitions this way. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).partitionsWithoutSegments(Set.of(2, 3)) + .survivingSegments(List.of("a_seg0", "a_seg1")), + new ColocatedTableSpec(4, false).partitionsWithoutSegments(Set.of(2, 3)) + .survivingSegments(List.of("b_seg0", "b_seg1"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo' AND " + + COLOCATED_TABLE_B + ".col2 = 'bar'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0), 1, Set.of(1))); + assertEquals(leafA.getWorkerIdToSegmentsMap().keySet(), Set.of(0, 1)); + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + } + } + + @Test + public void testBrokerPruningColocatedJoinAllPrunedKeepsEveryPopulatedClass() { + // Both sides prune everything, so the filtered union is empty and the group falls back to its populated classes. + // A group reduced to zero workers would leave a 1-to-1 exchange with no worker to wire on either side, and the + // server-side filter returns the same empty result from the unreduced plan anyway. So the most selective query + // gets the least reduction, which is the same trade the leaf-level fallback makes. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).survivingSegments(List.of()), + new ColocatedTableSpec(4, false).survivingSegments(List.of())); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + } + } + + @Test + public void testBrokerPruningColocatedJoinDisabledByQueryOption() { + // The existing useBrokerPruning switch is the kill switch for this too: with it off no routing query is built, so + // nothing is provably pruned and every populated class keeps its worker. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).survivingSegments(List.of("a_seg1")), + new ColocatedTableSpec(4, false).survivingSegments(List.of("b_seg1"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=false; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + } + } + + @Test + public void testBrokerPruningColocatedJoinFallsBackOnRoutingFailure() { + // Pruning is best-effort on this path too: a routing call that throws must leave the group with every populated + // class rather than fail a query that would otherwise plan. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).survivingSegments(List.of("a_seg1")), + new ColocatedTableSpec(4, false).survivingSegments(List.of("b_seg1")), TableType.OFFLINE, true); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + } + } + + @Test + public void testBrokerPruningColocatedJoinKeepsClassWithUnavailableSegment() { + // An unavailable segment was selected and not pruned, so the broker cannot prove its class empty. Class 2 keeps + // its worker on the strength of a_seg2 being unavailable rather than eliminated; only class 3, which both sides + // prune, is dropped. Reading "absent from the routing table" as "pruned" instead would silently drop the rows a + // transient outage hid. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).survivingSegments(List.of("a_seg0", "a_seg1")) + .unavailableSegments(List.of("a_seg2")), + new ColocatedTableSpec(4, false).survivingSegments(List.of("b_seg0", "b_seg1"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo' AND " + + COLOCATED_TABLE_B + ".col2 = 'bar'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2))); + assertEquals(workerIdToPartitions(leafB, "b_seg"), Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2))); + // Only class 3 dropped: one segment on each side. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 2); + } + } + + @Test + public void testBrokerPruningColocatedJoinStillPadsAMemberWithNoDataInASurvivingClass() { + // Emptiness padding and filter reduction have to compose. Table A holds nothing in class 3 while B does, and B's + // filter keeps class 3, so the class survives and A gets an empty worker for it -- placed on the server B picks + // for that class, so the exchange stays in process. Classes 1 and 2 are dropped because BOTH sides prune them. + // + // Padding is decided from unfiltered presence on purpose: a member that holds data the filter excludes dispatches + // it rather than being padded, so a surviving class always has the same segment list it would have unpruned. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).emptyPartitions(Set.of(3)).survivingSegments(List.of("a_seg0")), + new ColocatedTableSpec(4, false).survivingSegments(List.of("b_seg0", "b_seg3"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo' AND " + + COLOCATED_TABLE_B + ".col2 = 'bar'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + // Class list [0, 3]: worker 1 of A is the padded one, so it has an entry with no segment in it. + assertEquals(leafA.getWorkerIdToSegmentsMap().keySet(), Set.of(0, 1)); + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0))); + assertEquals(workerIdToPartitions(leafB, "b_seg"), Map.of(0, Set.of(0), 1, Set.of(3))); + Map workerIdToClass = new HashMap<>(); + mergeWorkerIdToClass(workerIdToClass, leafA, "a_seg", 4); + mergeWorkerIdToClass(workerIdToClass, leafB, "b_seg", 4); + assertEquals(workerIdToClass, Map.of(0, 0, 1, 3)); + // Worker 0 stands on the same server on both sides, which is what the 1-to-1 exchange is for. + assertEquals(workerIdToServer(leafA).get(0), workerIdToServer(leafB).get(0)); + // Worker 1 is A's padded one. It would rather land on B's server for class 3 to keep the exchange in process, + // but here A holds no data on that server at all -- partition 3 lives only on server 4 and A is empty there -- + // and a server with no data manager for the table fails the query outright, so it falls back to a server that + // provably hosts A and accepts one cross-server send. + assertEquals(workerIdToServer(leafB).get(1), getServerInstance("localhost", 4).getInstanceId()); + assertTrue(Set.of(getServerInstance("localhost", 1).getInstanceId(), + getServerInstance("localhost", 2).getInstanceId(), getServerInstance("localhost", 3).getInstanceId()) + .contains(workerIdToServer(leafA).get(1)), workerIdToServer(leafA).toString()); + // Classes 1 and 2 dropped, one segment each on each side. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 4); + } + } + + @Test + public void testColocatedJoinReducedGroupPrunesWholeClasses() { + // Emptiness reduction and filter reduction compose, on the one shape whose partitioned leaf is not marked + // pre-partitioned: a fact table joined with a replicated dimension table over an explicit local exchange. + // + // 8 partitions over 4 classes, so class c holds partitions c and c+4. The fact table's class 3 is empty, and the + // filter leaves only a_seg0 and a_seg4, which are both class 0. Emptiness drops class 3, the filter drops classes + // 1 and 2, and the class list ends up [0] -- a single worker holding the whole of class 0. // - // The fact table's class 3 is empty, so the group is reduced to [0, 1, 2], and the filter leaves only class 0. Were - // pruning left on, assignMultiplePartitionsPerWorker would find no segment for classes 1 and 2 and skip them - // WITHOUT consuming a worker id, leaving the leaf one worker while its class list still claimed three. + // This is the shape that used to be gated: reducing to [0] by dropping classes from the SHARED list keeps the + // worker id equal to its index in that list. Compacting at the leaf instead -- which is what + // assignMultiplePartitionsPerWorker would do if a per-leaf verdict ever reached it -- would leave the leaf one + // worker while its class list still claimed three. QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( new ColocatedTableSpec(8, false).emptyPartitions(Set.of(3, 7)).survivingSegments(List.of("a_seg0", "a_seg4")), new ColocatedTableSpec(8, false)); try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile("SET useBrokerPruning=true; " + replicatedDimensionJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); - // Nothing was pruned, even though the filtered routing query would have dropped 2 of the 3 surviving classes. - assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + // Classes 1 and 2, two segments each. Class 3 is empty, not pruned, so it is not counted. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 4); DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); - // One worker per surviving class, holding both partitions of that class, and no worker dropped or padded. - assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0, 4), 1, Set.of(1, 5), 2, Set.of(2, 6))); - assertEquals(leafA.getWorkerIdToSegmentsMap().keySet(), Set.of(0, 1, 2)); + // One worker for the one surviving class, holding both of its partitions, and no worker dropped or padded. + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0, 4))); + assertEquals(leafA.getWorkerIdToSegmentsMap().keySet(), Set.of(0)); Map workerIdToClass = new HashMap<>(); mergeWorkerIdToClass(workerIdToClass, leafA, "a_seg", 4); - assertEquals(workerIdToClass, Map.of(0, 0, 1, 1, 2, 2)); + assertEquals(workerIdToClass, Map.of(0, 0)); // The replicated leaf and the join derive their workers from the fact leaf, so they follow it class for class. assertEquals(leafB.getWorkerIdToSegmentsMap().keySet(), leafA.getWorkerIdToSegmentsMap().keySet()); assertEquals(workerIdToServer(leafB), workerIdToServer(leafA)); - assertEquals(joinFragment(dispatchableSubPlan).getWorkerMetadataList().size(), 3); + assertEquals(joinFragment(dispatchableSubPlan).getWorkerMetadataList().size(), 1); } } @@ -1456,7 +1686,7 @@ public void testCheckLeafWorkerAssignmentRejectsSparseWorkerIds() { Map>> segmentsMap = Map.of(0, offlineSegments("seg0"), 2, offlineSegments("seg2")); IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", serverMap, segmentsMap)); + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, serverMap, segmentsMap)); assertTrue(e.getMessage().contains("Missing server instance for worker: 1"), e.getMessage()); } @@ -1466,7 +1696,7 @@ public void testCheckLeafWorkerAssignmentRejectsKeySetMismatch() { Map>> segmentsMap = Map.of(0, offlineSegments("seg0"), 5, offlineSegments("seg5")); IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", serverMap, segmentsMap)); + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, serverMap, segmentsMap)); assertTrue(e.getMessage().contains("Missing segments for worker: 1"), e.getMessage()); } @@ -1475,7 +1705,7 @@ public void testCheckLeafWorkerAssignmentRejectsNullSegmentList() { Map> nullList = new HashMap<>(); nullList.put(TableType.OFFLINE.name(), null); IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, Map.of(0, queryServerInstance(1)), Map.of(0, nullList))); assertTrue(e.getMessage().contains("Null segment list for table type: OFFLINE"), e.getMessage()); } @@ -1485,7 +1715,7 @@ public void testCheckLeafWorkerAssignmentRejectsEmptyTableTypeMap() { // The server splits the request on the number of entries in this map, so a worker with no table type at all would // produce no server request. IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, Map.of(0, queryServerInstance(1)), Map.of(0, Map.of()))); assertTrue(e.getMessage().contains("Expected 1 or 2 table types for worker: 0, got: 0"), e.getMessage()); } @@ -1497,7 +1727,7 @@ public void testCheckLeafWorkerAssignmentRejectsThreeTableTypeMap() { threeTypes.put(TableType.REALTIME.name(), List.of()); threeTypes.put("HYBRID", List.of()); IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, Map.of(0, queryServerInstance(1)), Map.of(0, threeTypes))); assertTrue(e.getMessage().contains("Expected 1 or 2 table types for worker: 0, got: 3"), e.getMessage()); } @@ -1506,11 +1736,30 @@ public void testCheckLeafWorkerAssignmentRejectsThreeTableTypeMap() { public void testCheckLeafWorkerAssignmentRejectsUnknownTableType() { // The server resolves one table data manager per key in this map, and reports a missing table for an unknown one. IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, Map.of(0, queryServerInstance(1)), Map.of(0, Map.of("HYBRID", List.of())))); assertTrue(e.getMessage().contains("Unexpected table type: HYBRID for worker: 0"), e.getMessage()); } + @Test + public void testCheckLeafWorkerAssignmentRejectsFewerWorkersThanPartitionClasses() { + // A worker id of a colocated leaf IS an index into the group's shared class list, so a member that assigned fewer + // workers than the list has renumbered every class after the gap. checkPartitionClassAgreement cannot see it: it + // compares the shared array against itself, and both sides still hold the same instance. + IllegalStateException e = expectThrows(IllegalStateException.class, + () -> WorkerManager.checkLeafWorkerAssignment("testTable", new int[]{0, 2, 5}, + Map.of(0, queryServerInstance(1), 1, queryServerInstance(2)), + Map.of(0, offlineSegments("seg0"), 1, offlineSegments("seg2")))); + assertTrue(e.getMessage().contains("Got 2 workers for partition classes: [0, 2, 5]"), e.getMessage()); + } + + @Test + public void testCheckLeafWorkerAssignmentAcceptsOneWorkerPerPartitionClass() { + WorkerManager.checkLeafWorkerAssignment("testTable", new int[]{0, 2}, + Map.of(0, queryServerInstance(1), 1, queryServerInstance(2)), + Map.of(0, offlineSegments("seg0"), 1, offlineSegments("seg2"))); + } + @Test public void testCheckLeafWorkerAssignmentAcceptsHybridAndEmptySegmentWorkers() { // The two shapes the partitioned assignment produces: a hybrid worker with both table types, and one with a single @@ -1518,7 +1767,7 @@ public void testCheckLeafWorkerAssignmentAcceptsHybridAndEmptySegmentWorkers() { Map> hybridSegments = new HashMap<>(); hybridSegments.put(TableType.OFFLINE.name(), List.of("segO0")); hybridSegments.put(TableType.REALTIME.name(), List.of("segR0")); - WorkerManager.checkLeafWorkerAssignment("testTable", + WorkerManager.checkLeafWorkerAssignment("testTable", null, Map.of(0, queryServerInstance(1), 1, queryServerInstance(2)), Map.of(0, hybridSegments, 1, Map.of(TableType.OFFLINE.name(), new ArrayList<>()))); } @@ -1624,6 +1873,13 @@ private static QueryEnvironment newColocatedJoinQueryEnvironment(ColocatedTableS /// tables are registered under, so that the realtime-only shape can be covered too. private static QueryEnvironment newColocatedJoinQueryEnvironment(ColocatedTableSpec specA, ColocatedTableSpec specB, TableType tableType) { + return newColocatedJoinQueryEnvironment(specA, specB, tableType, false); + } + + /// Same as [#newColocatedJoinQueryEnvironment(ColocatedTableSpec, ColocatedTableSpec, TableType)], with a routing + /// manager that throws on every routing call, to exercise the best-effort fallback. + private static QueryEnvironment newColocatedJoinQueryEnvironment(ColocatedTableSpec specA, ColocatedTableSpec specB, + TableType tableType, boolean throwOnRouting) { int numServers = 4; ServerInstance[] servers = new ServerInstance[numServers]; Map enabledServers = new HashMap<>(); @@ -1640,13 +1896,15 @@ private static QueryEnvironment newColocatedJoinQueryEnvironment(ColocatedTableS colocatedTablePartitionInfo(tableBWithType, "b_seg", servers, specB)); Map routingTableByTable = new HashMap<>(); if (specA._survivingSegments != null) { - routingTableByTable.put(tableAWithType, colocatedRoutingTable(servers, "a_seg", specA._survivingSegments)); + routingTableByTable.put(tableAWithType, + colocatedRoutingTable(servers, "a_seg", specA._survivingSegments, specA._unavailableSegments)); } if (specB._survivingSegments != null) { - routingTableByTable.put(tableBWithType, colocatedRoutingTable(servers, "b_seg", specB._survivingSegments)); + routingTableByTable.put(tableBWithType, + colocatedRoutingTable(servers, "b_seg", specB._survivingSegments, specB._unavailableSegments)); } PartitionedRoutingManager routingManager = - new PartitionedRoutingManager(enabledServers, partitionInfoByTable, routingTableByTable, false); + new PartitionedRoutingManager(enabledServers, partitionInfoByTable, routingTableByTable, throwOnRouting); Map tableNameMap = new HashMap<>(); tableNameMap.put(tableAWithType, tableAWithType); @@ -1694,8 +1952,9 @@ private static TablePartitionReplicatedServersInfo colocatedTablePartitionInfo(S : Set.of(servers[p % servers.length].getInstanceId()); } // Mutable, like the lists the broker publishes: the assignment must hand out a copy rather than this instance. - partitionInfoMap[p] = new TablePartitionReplicatedServersInfo.PartitionInfo(partitionServers, - new ArrayList<>(List.of(segmentPrefix + p))); + List segments = spec._partitionsWithoutSegments.contains(p) ? new ArrayList<>() + : new ArrayList<>(List.of(segmentPrefix + p)); + partitionInfoMap[p] = new TablePartitionReplicatedServersInfo.PartitionInfo(partitionServers, segments); } } return new TablePartitionReplicatedServersInfo(tableNameWithType, "col1", "Hashcode", numPartitions, @@ -1705,7 +1964,7 @@ private static TablePartitionReplicatedServersInfo colocatedTablePartitionInfo(S /// Buckets the given surviving segments onto the server hosting their partition (partition `p` lives on server /// `p % 4`), i.e. builds what the routing manager returns for one colocated table's filtered routing query. private static RoutingTable colocatedRoutingTable(ServerInstance[] servers, String segmentPrefix, - List survivingSegments) { + List survivingSegments, List unavailableSegments) { Map> serverToSegmentList = new HashMap<>(); for (String segment : survivingSegments) { int partition = Integer.parseInt(segment.substring(segmentPrefix.length())); @@ -1714,7 +1973,7 @@ private static RoutingTable colocatedRoutingTable(ServerInstance[] servers, Stri Map serverToSegments = new HashMap<>(); serverToSegmentList.forEach((server, segments) -> serverToSegments.put(server, new SegmentsToQuery(segments, List.of()))); - return new RoutingTable(serverToSegments, List.of(), 0); + return new RoutingTable(serverToSegments, new ArrayList<>(unavailableSegments), 0); } /// How one side of a colocated join is laid out, for [#newColocatedJoinQueryEnvironment(ColocatedTableSpec, @@ -1739,6 +1998,8 @@ private static class ColocatedTableSpec { /// null rather than a silently empty answer. @Nullable List _survivingSegments; + List _unavailableSegments = List.of(); + Set _partitionsWithoutSegments = Set.of(); ColocatedTableSpec(int numPartitions, boolean everyServerHostsEveryPartition) { _numPartitions = numPartitions; @@ -1774,6 +2035,21 @@ ColocatedTableSpec survivingSegments(List survivingSegments) { _survivingSegments = survivingSegments; return this; } + + /// Partitions that get an entry listing no segment, as opposed to no entry at all. The broker never publishes + /// this shape today -- a partition's entry is created together with its first segment -- so it exists only to pin + /// what the assignment does if one ever appears. + ColocatedTableSpec partitionsWithoutSegments(Set partitionsWithoutSegments) { + _partitionsWithoutSegments = partitionsWithoutSegments; + return this; + } + + /// Segments the routing table reports as unavailable. They were selected and not pruned, so the broker cannot + /// prove them empty and their partition class has to survive. + ColocatedTableSpec unavailableSegments(List unavailableSegments) { + _unavailableSegments = unavailableSegments; + return this; + } } /// Returns the only fragment below the reduce stage with neither segments nor children, i.e. the join stage. @@ -1984,6 +2260,18 @@ private static QueryEnvironment newPartitionedQueryEnvironment(int[] serverIdxPe int replicasPerPartition, List survivingSegments, List unavailableSegments, int reportedPrunedByRouting, boolean throwOnRouting, Set emptyPartitions, Set partitionsWithOnlyDeferredSegments) { + return newPartitionedQueryEnvironment(serverIdxPerPartition, numServers, replicasPerPartition, survivingSegments, + unavailableSegments, reportedPrunedByRouting, throwOnRouting, emptyPartitions, + partitionsWithOnlyDeferredSegments, Map.of()); + } + + /// Same again, with the surviving segments a leaf sees when its own filter carries a given string literal. Lets two + /// leaves scanning the SAME table be given different verdicts, which is the only way to tell a per-leaf pruning + /// verdict from a per-table one. + private static QueryEnvironment newPartitionedQueryEnvironment(int[] serverIdxPerPartition, int numServers, + int replicasPerPartition, List survivingSegments, List unavailableSegments, + int reportedPrunedByRouting, boolean throwOnRouting, Set emptyPartitions, + Set partitionsWithOnlyDeferredSegments, Map> survivingSegmentsByFilterLiteral) { int numPartitions = serverIdxPerPartition.length; ServerInstance[] servers = new ServerInstance[numServers]; Map enabledServers = new HashMap<>(); @@ -2025,6 +2313,7 @@ private static QueryEnvironment newPartitionedQueryEnvironment(int[] serverIdxPe PartitionedRoutingManager routingManager = new PartitionedRoutingManager(enabledServers, Map.of(PARTITIONED_TABLE_OFFLINE, tablePartitionInfo), Map.of(PARTITIONED_TABLE_OFFLINE, prunedRoutingTable), throwOnRouting); + survivingSegmentsByFilterLiteral.forEach(routingManager::survivingSegmentsForFilterLiteral); Map tableNameMap = new HashMap<>(); tableNameMap.put(PARTITIONED_TABLE_OFFLINE, PARTITIONED_TABLE_OFFLINE); @@ -2424,6 +2713,9 @@ private static class PartitionedRoutingManager implements RoutingManager { private final Map _enabledServers; private final Map _partitionInfoByTable; private final Map _routingTableByTable; + /// Surviving segments for a leaf whose filter carries the given string literal, which is how a test gives two + /// leaves scanning ONE table two different verdicts. Falls back to the per-table routing table when absent. + private final Map> _survivingSegmentsByFilterLiteral = new HashMap<>(); private final boolean _throwOnRouting; @Nullable private final TimeBoundaryInfo _timeBoundaryInfo; @@ -2471,8 +2763,67 @@ public RoutingTable getRoutingTable(BrokerRequest brokerRequest, String tableNam @Nullable @Override public List getSegments(BrokerRequest brokerRequest) { - TablePartitionReplicatedServersInfo partitionInfo = - _partitionInfoByTable.get(brokerRequest.getQuerySource().getTableName()); + return allSegments(brokerRequest.getQuerySource().getTableName()); + } + + /// Derives the pruned set the way `BaseBrokerRoutingManager` does -- everything selection offered minus what the + /// pruners kept -- from the same configured routing table. Unavailable segments were selected and not pruned, so + /// they stay out of it and keep their partition alive. A table with no routing table registered proves nothing, + /// which is what an unexpected routing call should look like. + @Override + public Set getPrunedSegments(BrokerRequest brokerRequest) { + if (_throwOnRouting) { + throw new RuntimeException("Simulated routing failure"); + } + validatePrunableFilter(brokerRequest.getPinotQuery().getFilterExpression()); + String tableNameWithType = brokerRequest.getQuerySource().getTableName(); + Set prunedSegments = new HashSet<>(allSegments(tableNameWithType)); + String filterLiteral = firstStringLiteral(brokerRequest.getPinotQuery().getFilterExpression()); + List survivingSegments = + filterLiteral != null ? _survivingSegmentsByFilterLiteral.get(filterLiteral) : null; + if (survivingSegments != null) { + prunedSegments.removeAll(survivingSegments); + return prunedSegments; + } + RoutingTable routingTable = _routingTableByTable.get(tableNameWithType); + if (routingTable == null) { + return Set.of(); + } + for (SegmentsToQuery segmentsToQuery : routingTable.getServerInstanceToSegmentsMap().values()) { + prunedSegments.removeAll(segmentsToQuery.getSegments()); + } + prunedSegments.removeAll(routingTable.getUnavailableSegments()); + return prunedSegments; + } + + PartitionedRoutingManager survivingSegmentsForFilterLiteral(String filterLiteral, List segments) { + _survivingSegmentsByFilterLiteral.put(filterLiteral, segments); + return this; + } + + /// The first string literal in the filter, which the tests use as a stand-in for "which filter is this". + @Nullable + private static String firstStringLiteral(@Nullable Expression expression) { + if (expression == null) { + return null; + } + if (expression.getLiteral() != null && expression.getLiteral().isSetStringValue()) { + return expression.getLiteral().getStringValue(); + } + Function function = expression.getFunctionCall(); + if (function != null) { + for (Expression operand : function.getOperands()) { + String literal = firstStringLiteral(operand); + if (literal != null) { + return literal; + } + } + } + return null; + } + + private List allSegments(String tableNameWithType) { + TablePartitionReplicatedServersInfo partitionInfo = _partitionInfoByTable.get(tableNameWithType); if (partitionInfo == null) { return List.of(); } diff --git a/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json b/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json index 20ef9aed09ee..5f6c4b84e9d5 100644 --- a/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json +++ b/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json @@ -421,21 +421,15 @@ "\n ├── [2]@localhost:1|[1] PROJECT", "\n │ └── [2]@localhost:1|[1] TABLE SCAN (a) null", "\n └── [2]@localhost:1|[1] MAIL_RECEIVE(BROADCAST_DISTRIBUTED)", - "\n ├── [3]@localhost:2|[2] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2, 3]} (Subtree Omitted)", - "\n ├── [3]@localhost:2|[3] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2, 3]} (Subtree Omitted)", - "\n ├── [3]@localhost:1|[0] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2, 3]} (Subtree Omitted)", - "\n └── [3]@localhost:1|[1] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2, 3]}", - "\n └── [3]@localhost:1|[1] PROJECT", - "\n └── [3]@localhost:1|[1] FILTER", - "\n └── [3]@localhost:1|[1] AGGREGATE_FINAL", - "\n └── [3]@localhost:1|[1] MAIL_RECEIVE(HASH_DISTRIBUTED)", - "\n ├── [4]@localhost:2|[2] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[3]@localhost:2|[2]} (Subtree Omitted)", - "\n ├── [4]@localhost:2|[3] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[3]@localhost:2|[3]} (Subtree Omitted)", - "\n ├── [4]@localhost:1|[0] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[3]@localhost:1|[0]} (Subtree Omitted)", - "\n └── [4]@localhost:1|[1] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[3]@localhost:1|[1]}", - "\n └── [4]@localhost:1|[1] AGGREGATE_LEAF", - "\n └── [4]@localhost:1|[1] FILTER", - "\n └── [4]@localhost:1|[1] TABLE SCAN (b) null", + "\n └── [3]@localhost:1|[0] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2, 3]}", + "\n └── [3]@localhost:1|[0] PROJECT", + "\n └── [3]@localhost:1|[0] FILTER", + "\n └── [3]@localhost:1|[0] AGGREGATE_FINAL", + "\n └── [3]@localhost:1|[0] MAIL_RECEIVE(HASH_DISTRIBUTED)", + "\n └── [4]@localhost:1|[0] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[3]@localhost:1|[0]}", + "\n └── [4]@localhost:1|[0] AGGREGATE_LEAF", + "\n └── [4]@localhost:1|[0] FILTER", + "\n └── [4]@localhost:1|[0] TABLE SCAN (b) null", "\n" ] }, diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index e00598b582e7..30a70bd8da9f 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -625,6 +625,11 @@ public static class Broker { /// Separated from [#CONFIG_OF_USE_BROKER_PRUNING] so the two paths can be rolled out independently; both /// default to enabled now that all logical-planner leaf paths (non-partitioned, partitioned, logical tables) /// support broker pruning. Actual pruning still requires segment pruners to be configured on the table. + /// + /// On a colocated join this governs more than which segments are dispatched: a partition class that every member + /// of the colocated group prunes away is dropped from the group's shared class list, so the leaves and the stages + /// derived from them run fewer workers and the query is dispatched to fewer servers. Turning it off restores one + /// worker per populated class. public static final String CONFIG_OF_LOGICAL_PLANNER_USE_BROKER_PRUNING = "pinot.broker.multistage.logical.planner.use.broker.pruning"; public static final boolean DEFAULT_LOGICAL_PLANNER_USE_BROKER_PRUNING = true;