From 45447e784b218e7502e564eb8cc0576c1891f1a7 Mon Sep 17 00:00:00 2001 From: Yash Mayya Date: Wed, 5 Aug 2026 14:11:42 -0700 Subject: [PATCH] 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 | 1175 ++++++++++++++++- 19 files changed, 3347 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..84ce1a4e108e 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,1131 @@ 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 testColocatedJoinOnNonPartitionKeyAlignsWorkersWhenEmptyClassesDiffer() { + // Same as above, joining on a column that is not the hinted partition key (see + // #colocatedJoinQueryOnNonPartitionKey). Worker assignment reads the hinted partition key and the broker's + // partition info, never the join condition, so this must resolve identically: which class a row sits in is decided + // by the column the data was partitioned on, whatever column the join then matches on. + QueryEnvironment queryEnvironment = + newColocatedJoinQueryEnvironment(Set.of(1), Set.of(2), Set.of(), Set.of(), true); + try (QueryEnvironment.CompiledQuery compiledQuery = + queryEnvironment.compile(colocatedJoinQueryOnNonPartitionKey(4))) { + 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), 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))); + assertEquals(assignedSegments(leafA, 1), List.of()); + assertEquals(assignedSegments(leafB, 2), List.of()); + 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)); + assertEquals(workerIdToServer(leafA), workerIdToServer(leafB)); + // Both exchanges into the join stay 1-to-1 and in process rather than degrading to a shuffle. + DispatchablePlanFragment joinFragment = joinFragment(dispatchableSubPlan); + assertEquals(workerIdToServer(joinFragment), workerIdToServer(leafA)); + int leafAFragmentId = leafA.getPlanFragment().getFragmentId(); + int leafBFragmentId = leafB.getPlanFragment().getFragmentId(); + for (int workerId = 0; workerId < 4; 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 testColocatedJoinOnNonPartitionKeyReducesFanOut() { + // The fan-out reduction below, on a join key that is not the hinted partition key. + 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(colocatedJoinQueryOnNonPartitionKey(8))) { + 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(5))); + assertEquals(workerIdToPartitions(leafB, "b_seg"), Map.of(0, Set.of(0), 1, Set.of(5))); + String server1 = getServerInstance("localhost", 1).getInstanceId(); + String server2 = getServerInstance("localhost", 2).getInstanceId(); + assertEquals(dispatchedServers(dispatchableSubPlan), Set.of(server1, server2)); + } + } + + @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"; + } + + /// Same as [#colocatedJoinQuery(int)], joining on "col2" while both tables are hinted partitioned on "col1". This is + /// `is_colocated_by_join_keys` in its intended form: the user asserts that the join key is laid out the same way as + /// the hinted partition key, which Pinot has no metadata to verify. + private static String colocatedJoinQueryOnNonPartitionKey(int partitionSize) { + String tableHint = colocatedTableHint(partitionSize); + return "SELECT /*+ joinOptions(is_colocated_by_join_keys='true') */ " + COLOCATED_TABLE_A + ".col1, " + + COLOCATED_TABLE_B + ".col1 FROM " + COLOCATED_TABLE_A + " " + tableHint + "JOIN " + COLOCATED_TABLE_B + " " + + tableHint + "ON " + COLOCATED_TABLE_A + ".col2 = " + COLOCATED_TABLE_B + ".col2"; + } + + /// 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 +1977,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 +2042,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 +2064,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 +2536,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