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