diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java index 631b67f63cc2..269243cb5b64 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java @@ -20,6 +20,7 @@ import com.google.common.base.Preconditions; import com.google.common.collect.Maps; +import com.google.common.collect.Sets; import com.google.common.util.concurrent.ThreadFactoryBuilder; import java.util.ArrayList; import java.util.HashMap; @@ -1235,6 +1236,16 @@ public List getSegments(BrokerRequest brokerRequest, @Nullable String sa return routingEntry.getSegments(brokerRequest, samplerName); } + @Nullable + @Override + public Set getPrunedSegments(BrokerRequest brokerRequest) { + RoutingEntry routingEntry = _routingEntryMap.get(brokerRequest.getQuerySource().getTableName()); + if (routingEntry == null) { + return null; + } + return routingEntry.getPrunedSegments(brokerRequest, extractSamplerName(brokerRequest)); + } + private static String normalizeSamplerName(String samplerName) { return samplerName.trim().toLowerCase(Locale.ROOT); } @@ -1510,22 +1521,33 @@ void refreshSegment(String segment) { } } - InstanceSelector.SelectionResult calculateRouting(BrokerRequest brokerRequest, long requestId, - @Nullable String samplerName) { - SamplerInfo samplerInfo = getSamplerInfo(samplerName); + /// Runs selection and then the pruner chain, which is the one place that decides what a query sees. Every caller + /// goes through here on purpose: the routing table, the plain segment list and the planner's emptiness proof must + /// all be judged by the same selector and the same pruners. A second copy of this sequence that drifted would let + /// the planner prove a partition empty that a real query would still have scanned, and that loses rows with no + /// error anywhere. + private SelectedSegments selectThenPrune(BrokerRequest brokerRequest, @Nullable SamplerInfo samplerInfo) { SegmentSelector segmentSelector = samplerInfo != null ? samplerInfo._segmentSelector : _segmentSelector; - InstanceSelector instanceSelector = samplerInfo != null ? samplerInfo._instanceSelector : _instanceSelector; Set selectedSegments = segmentSelector.select(brokerRequest); - int numTotalSelectedSegments = selectedSegments.size(); + Set survivingSegments = selectedSegments; if (!selectedSegments.isEmpty()) { for (SegmentPruner segmentPruner : _segmentPruners) { - selectedSegments = segmentPruner.prune(brokerRequest, selectedSegments); + survivingSegments = segmentPruner.prune(brokerRequest, survivingSegments); } } - int numPrunedSegments = numTotalSelectedSegments - selectedSegments.size(); - if (!selectedSegments.isEmpty()) { + return new SelectedSegments(selectedSegments, survivingSegments); + } + + InstanceSelector.SelectionResult calculateRouting(BrokerRequest brokerRequest, long requestId, + @Nullable String samplerName) { + SamplerInfo samplerInfo = getSamplerInfo(samplerName); + InstanceSelector instanceSelector = samplerInfo != null ? samplerInfo._instanceSelector : _instanceSelector; + SelectedSegments selectedSegments = selectThenPrune(brokerRequest, samplerInfo); + Set survivingSegments = selectedSegments._surviving; + int numPrunedSegments = selectedSegments.getNumPruned(); + if (!survivingSegments.isEmpty()) { InstanceSelector.SelectionResult selectionResult = - instanceSelector.select(brokerRequest, new ArrayList<>(selectedSegments), requestId); + instanceSelector.select(brokerRequest, new ArrayList<>(survivingSegments), requestId); selectionResult.setNumPrunedSegments(numPrunedSegments); return selectionResult; } else { @@ -1535,15 +1557,47 @@ InstanceSelector.SelectionResult calculateRouting(BrokerRequest brokerRequest, l } List getSegments(BrokerRequest brokerRequest, @Nullable String samplerName) { - SamplerInfo samplerInfo = getSamplerInfo(samplerName); - SegmentSelector segmentSelector = samplerInfo != null ? samplerInfo._segmentSelector : _segmentSelector; - Set selectedSegments = segmentSelector.select(brokerRequest); - if (!selectedSegments.isEmpty()) { - for (SegmentPruner segmentPruner : _segmentPruners) { - selectedSegments = segmentPruner.prune(brokerRequest, selectedSegments); + return new ArrayList<>(selectThenPrune(brokerRequest, getSamplerInfo(samplerName))._surviving); + } + + /// See [RoutingManager#getPrunedSegments]. The sampler is honoured for the same reason the query path honours it: + /// a narrower selection only ever shrinks what this can prove, never widens it. + /// + /// The pruners return a new set rather than editing the one they are handed, so taking the difference costs + /// nothing unless something was actually pruned. If one ever did edit in place the two sets would be the same + /// object, the difference would come out empty, and this would fall back to proving nothing -- the safe direction. + Set getPrunedSegments(BrokerRequest brokerRequest, @Nullable String samplerName) { + SelectedSegments selectedSegments = selectThenPrune(brokerRequest, getSamplerInfo(samplerName)); + int numPruned = selectedSegments.getNumPruned(); + if (numPruned == 0) { + return Set.of(); + } + // Built up rather than copied down: the count is already known and is usually a small fraction of the table's + // segments, so copying every selected segment only to remove most of them again would size the allocation to + // the table instead of to the answer. + Set prunedSegments = Sets.newHashSetWithExpectedSize(numPruned); + for (String segment : selectedSegments._selected) { + if (!selectedSegments._surviving.contains(segment)) { + prunedSegments.add(segment); } } - return new ArrayList<>(selectedSegments); + return prunedSegments; + } + } + + /// What one run of [RoutingEntry#selectThenPrune] decided: the segments selection offered, and the ones the pruners + /// left. Both are needed because the difference between them is the only sound proof that a segment cannot match. + private static class SelectedSegments { + final Set _selected; + final Set _surviving; + + SelectedSegments(Set selected, Set surviving) { + _selected = selected; + _surviving = surviving; + } + + int getNumPruned() { + return _selected.size() - _surviving.size(); } } } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java index 312d7b0b7d10..9f04a4ca3d40 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java @@ -217,6 +217,62 @@ public List getSegments(BrokerRequest brokerRequest, @Nullable String sa return combined.isEmpty() ? null : combined; } + /// Combines by *intersection* over the clusters that have routing for the table, which is the opposite of how + /// [#getSegments] combines and deliberately so: that returns segments that survive, and a segment survives if any + /// cluster keeps it, while this returns segments that are provably eliminated, and a proof only holds if every + /// cluster that could route the segment eliminated it. Unioning instead would let one cluster's pruners speak for a + /// segment another cluster would still have queried -- silently dropping matching data. + /// + /// Restricting the intersection to the clusters that have the table is what keeps it useful: the usual case is a + /// table in exactly one cluster, where the intersection is that cluster's own verdict. A cluster without the table + /// would otherwise contribute an empty set and reduce every answer to "nothing proven". + @Nullable + @Override + public Set getPrunedSegments(BrokerRequest brokerRequest) { + String tableNameWithType = brokerRequest.getQuerySource().getTableName(); + Set combined = intersectPrunedSegments(null, _localClusterRoutingManager, brokerRequest, + tableNameWithType); + for (BaseBrokerRoutingManager remoteCluster : _remoteClusterRoutingManagers) { + combined = intersectPrunedSegments(combined, remoteCluster, brokerRequest, tableNameWithType); + } + // Still null when no cluster has the table at all, which is what the interface reports for a table that does not + // exist -- as opposed to an empty set, which is a cluster that ran the pruners and proved nothing. + return combined; + } + + /// Folds one cluster's verdict into the running intersection, or returns an empty set to end it: once nothing is + /// proven, nothing downstream can make it provable again, and asking the remaining clusters would run a full + /// selection and pruner chain each for an answer that is already fixed. `null` means no cluster has answered yet. + @Nullable + private Set intersectPrunedSegments(@Nullable Set combined, BaseBrokerRoutingManager cluster, + BrokerRequest brokerRequest, String tableNameWithType) { + if (combined != null && combined.isEmpty()) { + return combined; + } + try { + // One lookup rather than routingExists-then-get: a table appearing between the two would let this skip a + // cluster that can route it, and the intersection would then claim segments eliminated that nobody asked about. + Set prunedSegments = cluster.getPrunedSegments(brokerRequest); + if (prunedSegments == null) { + // This cluster has no routing for the table, so it eliminates nothing and constrains nothing. + return combined; + } + if (prunedSegments.isEmpty()) { + // This cluster proves nothing, so neither does the intersection. + return Set.of(); + } + if (combined == null) { + return new HashSet<>(prunedSegments); + } + combined.retainAll(prunedSegments); + return combined; + } catch (Exception e) { + LOGGER.error("Error getting pruned segments from cluster routing manager for table {}", tableNameWithType, e); + // A cluster we could not ask may still have routed any of these segments, so prove nothing. + return Set.of(); + } + } + /// Returns the partition info only when a single cluster has any, and `null` when more than one does. /// /// Unlike [#getRoutingTable], [#getSegments] and [#getServingInstances], this cannot union the clusters: the info is diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java index 32be5f024cfd..75c2f3570ba3 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java @@ -19,9 +19,10 @@ package org.apache.pinot.broker.routing.manager; import java.lang.reflect.Constructor; -import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.Consumer; import org.apache.helix.AccessOption; import org.apache.helix.BaseDataAccessor; @@ -45,6 +46,8 @@ import org.apache.pinot.broker.routing.timeboundary.TimeBoundaryManager; import org.apache.pinot.common.metrics.BrokerGauge; import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.common.request.QuerySource; import org.apache.pinot.common.utils.config.TableConfigSerDeUtils; import org.apache.pinot.core.routing.TablePartitionInfo; import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo; @@ -73,9 +76,11 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNull; import static org.testng.Assert.assertSame; import static org.testng.Assert.assertTrue; @@ -258,17 +263,170 @@ public void testSamplerContextSharesTimeBoundaryAndPartitionMetadata() assertSame(_routingManager.getTablePartitionReplicatedServersInfo(TEST_TABLE), expectedReplicatedServersInfo); } + @Test + public void testGetPrunedSegmentsIsExactlySelectedMinusSurvivors() + throws Exception { + SegmentSelector segmentSelector = selectorOf("seg1", "seg2", "seg3", "seg4"); + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, segmentSelector, List.of(prunerDropping("seg1", "seg3")), + mock(InstanceSelector.class))); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of("seg1", "seg3")); + } + + @Test + public void testGetPrunedSegmentsChainsEveryPruner() + throws Exception { + SegmentPruner firstPruner = prunerDropping("seg1"); + SegmentPruner secondPruner = prunerDropping("seg3"); + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, selectorOf("seg1", "seg2", "seg3"), List.of(firstPruner, secondPruner), + mock(InstanceSelector.class))); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of("seg1", "seg3")); + + // The second pruner judges what the first left, so consulting only the last one would lose "seg1". + ArgumentCaptor> captor = ArgumentCaptor.captor(); + verify(secondPruner).prune(any(), captor.capture()); + assertEquals(captor.getValue(), Set.of("seg2", "seg3")); + } + + /// Nothing pruned must read as "proved nothing", not as "proved every selected segment empty" -- the latter would + /// let a caller treat a table that fully matches the filter as a table with no matching rows. + @Test + public void testGetPrunedSegmentsIsEmptyWhenNothingWasPruned() + throws Exception { + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, selectorOf("seg1", "seg2"), List.of(prunerDropping()), + mock(InstanceSelector.class))); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of()); + } + + @Test + public void testGetPrunedSegmentsIsEmptyWhenSelectionIsEmpty() + throws Exception { + SegmentPruner pruner = prunerDropping("seg1"); + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, selectorOf(), List.of(pruner), mock(InstanceSelector.class))); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of()); + // An empty selection is answered without asking anyone, so the empty result cannot have come from a pruner. + verify(pruner, never()).prune(any(), any()); + } + + /// A table this broker has no routing for is `null`, not an empty set: it eliminated nothing because it would have + /// routed nothing, which is a different claim from "the pruners ran and proved nothing". + @Test + public void testGetPrunedSegmentsIsNullForUnknownTable() { + assertNull(_routingManager.getPrunedSegments(brokerRequest("noSuchTable_OFFLINE"))); + } + + /// The whole point of the API: only presence in the result is a proof. A segment the selector never offered is + /// absent from the survivors for a reason that has nothing to do with the filter -- here the selector withheld + /// "seg3" -- and reporting it would let a caller skip a segment that may well hold matching rows. + @Test + public void testGetPrunedSegmentsDoesNotReportASegmentTheSelectorNeverOffered() + throws Exception { + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, selectorOf("seg1", "seg2"), List.of(prunerDropping("seg1", "seg3")), + mock(InstanceSelector.class))); + + Set prunedSegments = _routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)); + + assertEquals(prunedSegments, Set.of("seg1")); + assertFalse(prunedSegments.contains("seg3")); + } + + /// Instance selection is what makes routing depend on the request id and on which replicas are up; keeping it out + /// is what makes this deterministic enough to plan on. + @Test + public void testGetPrunedSegmentsNeverConsultsInstanceSelection() + throws Exception { + InstanceSelector instanceSelector = mock(InstanceSelector.class); + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, selectorOf("seg1", "seg2"), List.of(prunerDropping("seg1")), instanceSelector)); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of("seg1")); + verifyNoInteractions(instanceSelector); + } + + /// A pruner that edits the set it was handed leaves nothing to take a difference against. That has to degrade to + /// "proved nothing" rather than to a wrong proof. + @Test + public void testGetPrunedSegmentsIsEmptyWhenAPrunerEditsInPlace() + throws Exception { + SegmentPruner pruner = mock(SegmentPruner.class); + when(pruner.prune(any(), any())).thenAnswer(invocation -> { + Set segments = invocation.getArgument(1); + segments.remove("seg1"); + return segments; + }); + SegmentSelector segmentSelector = mock(SegmentSelector.class); + when(segmentSelector.select(any())).thenReturn(new HashSet<>(Set.of("seg1", "seg2"))); + putRoutingEntry(TEST_TABLE, + createRoutingEntry(TEST_TABLE, segmentSelector, List.of(pruner), mock(InstanceSelector.class))); + + assertEquals(_routingManager.getPrunedSegments(brokerRequest(TEST_TABLE)), Set.of()); + } + + private static BrokerRequest brokerRequest(String tableNameWithType) { + QuerySource querySource = new QuerySource(); + querySource.setTableName(tableNameWithType); + BrokerRequest brokerRequest = new BrokerRequest(); + brokerRequest.setQuerySource(querySource); + return brokerRequest; + } + + private static SegmentSelector selectorOf(String... segments) { + SegmentSelector segmentSelector = mock(SegmentSelector.class); + when(segmentSelector.select(any())).thenReturn(Set.of(segments)); + return segmentSelector; + } + + /// Mirrors [org.apache.pinot.broker.routing.segmentpruner.EmptySegmentPruner]: a fresh set when it prunes + /// something, the very set it was handed when it does not. + private static SegmentPruner prunerDropping(String... segments) { + Set droppedSegments = Set.of(segments); + SegmentPruner segmentPruner = mock(SegmentPruner.class); + when(segmentPruner.prune(any(), any())).thenAnswer(invocation -> { + Set candidateSegments = invocation.getArgument(1); + if (droppedSegments.stream().noneMatch(candidateSegments::contains)) { + return candidateSegments; + } + Set survivingSegments = new HashSet<>(candidateSegments); + survivingSegments.removeAll(droppedSegments); + return survivingSegments; + }); + return segmentPruner; + } + private static Object createRoutingEntry(String tableNameWithType, TimeBoundaryManager timeBoundaryManager, SegmentPartitionMetadataManager partitionMetadataManager, Map samplerInfos) throws Exception { - return createRoutingEntry(tableNameWithType, timeBoundaryManager, partitionMetadataManager, samplerInfos, - mock(InstanceSelector.class), false); + return createRoutingEntry(tableNameWithType, mock(SegmentSelector.class), List.of(), + mock(InstanceSelector.class), timeBoundaryManager, partitionMetadataManager, samplerInfos, false); } private static Object createRoutingEntry(String tableNameWithType, TimeBoundaryManager timeBoundaryManager, SegmentPartitionMetadataManager partitionMetadataManager, Map samplerInfos, InstanceSelector instanceSelector, boolean disabled) throws Exception { + return createRoutingEntry(tableNameWithType, mock(SegmentSelector.class), List.of(), instanceSelector, + timeBoundaryManager, partitionMetadataManager, samplerInfos, disabled); + } + + private static Object createRoutingEntry(String tableNameWithType, SegmentSelector segmentSelector, + List segmentPruners, InstanceSelector instanceSelector) + throws Exception { + return createRoutingEntry(tableNameWithType, segmentSelector, segmentPruners, instanceSelector, + mock(TimeBoundaryManager.class), mock(SegmentPartitionMetadataManager.class), Map.of(), false); + } + + private static Object createRoutingEntry(String tableNameWithType, SegmentSelector segmentSelector, + List segmentPruners, InstanceSelector instanceSelector, TimeBoundaryManager timeBoundaryManager, + SegmentPartitionMetadataManager partitionMetadataManager, Map samplerInfos, boolean disabled) + throws Exception { Class routingEntryClass = Class.forName(BaseBrokerRoutingManager.class.getName() + "$RoutingEntry"); Constructor constructor = routingEntryClass.getDeclaredConstructor(String.class, String.class, String.class, SegmentPreSelector.class, SegmentSelector.class, List.class, InstanceSelector.class, int.class, int.class, @@ -276,8 +434,8 @@ private static Object createRoutingEntry(String tableNameWithType, TimeBoundaryM Map.class, boolean.class); constructor.setAccessible(true); return constructor.newInstance(tableNameWithType, "/IDEALSTATES/" + tableNameWithType, - "/EXTERNALVIEW/" + tableNameWithType, mock(SegmentPreSelector.class), mock(SegmentSelector.class), - Collections.emptyList(), instanceSelector, 1, 1, + "/EXTERNALVIEW/" + tableNameWithType, mock(SegmentPreSelector.class), segmentSelector, segmentPruners, + instanceSelector, 1, 1, mock(SegmentZkMetadataFetcher.class), timeBoundaryManager, partitionMetadataManager, null, samplerInfos, disabled); } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java index 8d8cdb34c262..0ea105a56e16 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManagerTest.java @@ -44,6 +44,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; @@ -327,6 +328,91 @@ public void testGetTablePartitionInfoIgnoresAFailingRemoteCluster() { assertEquals(_multiClusterRoutingManager.getTablePartitionReplicatedServersInfo(TEST_TABLE), partitionInfo); } + /// The usual case: the table lives in exactly one cluster, so the intersection is that cluster's own verdict. + @Test + public void testGetPrunedSegmentsReturnsTheSoleClusterVerdictVerbatim() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + when(_localClusterRoutingManager.getPrunedSegments(brokerRequest)) + .thenReturn(Set.of("seg1", "seg2")); + withoutTheTable(brokerRequest, _remoteClusterRoutingManager1, _remoteClusterRoutingManager2); + + assertEquals(_multiClusterRoutingManager.getPrunedSegments(brokerRequest), Set.of("seg1", "seg2")); + } + + /// Intersection, not union: unioning would let one cluster's pruners speak for a segment another cluster would + /// still have queried, so the planner would skip data that matches -- a silent wrong answer rather than a slow one. + /// Here only "seg2" is eliminated everywhere; "seg1" and "seg3" each survive in one cluster. + @Test + public void testGetPrunedSegmentsIntersectsRatherThanUnions() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + when(_localClusterRoutingManager.getPrunedSegments(brokerRequest)) + .thenReturn(Set.of("seg1", "seg2")); + when(_remoteClusterRoutingManager1.getPrunedSegments(brokerRequest)) + .thenReturn(Set.of("seg2", "seg3")); + withoutTheTable(brokerRequest, _remoteClusterRoutingManager2); + + Set prunedSegments = _multiClusterRoutingManager.getPrunedSegments(brokerRequest); + + assertEquals(prunedSegments, Set.of("seg2")); + assertFalse(prunedSegments.contains("seg1")); + assertFalse(prunedSegments.contains("seg3")); + } + + @Test + public void testGetPrunedSegmentsIsEmptyWhenAClusterProvesNothing() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + when(_localClusterRoutingManager.getPrunedSegments(brokerRequest)).thenReturn(Set.of("seg1")); + when(_remoteClusterRoutingManager1.getPrunedSegments(brokerRequest)).thenReturn(Set.of()); + withoutTheTable(brokerRequest, _remoteClusterRoutingManager2); + + assertEquals(_multiClusterRoutingManager.getPrunedSegments(brokerRequest), Set.of()); + } + + /// A cluster that does not have the table constrains nothing. Were its absence folded into the same empty set the + /// pruners use for "proved nothing", every answer would collapse to "nothing proven" in the common deployment. + @Test + public void testGetPrunedSegmentsSkipsAClusterWithoutTheTable() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + when(_localClusterRoutingManager.getPrunedSegments(brokerRequest)) + .thenReturn(Set.of("seg1", "seg2")); + // A cluster without the table reports null rather than an empty verdict, which is what keeps it from collapsing + // the intersection to "nothing proven". + withoutTheTable(brokerRequest, _remoteClusterRoutingManager1, _remoteClusterRoutingManager2); + + assertEquals(_multiClusterRoutingManager.getPrunedSegments(brokerRequest), Set.of("seg1", "seg2")); + } + + @Test + public void testGetPrunedSegmentsIsNullWhenNoClusterHasTheTable() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + withoutTheTable(brokerRequest, _localClusterRoutingManager, _remoteClusterRoutingManager1, + _remoteClusterRoutingManager2); + + assertNull(_multiClusterRoutingManager.getPrunedSegments(brokerRequest)); + } + + /// The deliberate opposite of [#testGetTablePartitionInfoIgnoresAFailingRemoteCluster]: a cluster we could not ask + /// may still have routed any of these segments, so its silence cannot be read as agreement. + @Test + public void testGetPrunedSegmentsIsEmptyWhenAClusterThrows() { + BrokerRequest brokerRequest = createMockBrokerRequest(TEST_TABLE); + when(_localClusterRoutingManager.getPrunedSegments(brokerRequest)).thenReturn(Set.of("seg1")); + when(_remoteClusterRoutingManager1.getPrunedSegments(brokerRequest)) + .thenThrow(new RuntimeException("remote cluster is down")); + withoutTheTable(brokerRequest, _remoteClusterRoutingManager2); + + assertEquals(_multiClusterRoutingManager.getPrunedSegments(brokerRequest), Set.of()); + } + + /// Makes the given clusters report that they have no routing for the table. Worth spelling out in every pruning + /// test: an unstubbed mock hands back an empty set, which means "ran the pruners and proved nothing" and would + /// collapse the intersection for a reason the test did not intend. + private static void withoutTheTable(BrokerRequest brokerRequest, BaseBrokerRoutingManager... clusters) { + for (BaseBrokerRoutingManager cluster : clusters) { + when(cluster.getPrunedSegments(brokerRequest)).thenReturn(null); + } + } + private RoutingTable createRoutingTable(String serverName, List segments) { Map serverMap = new HashMap<>(); ServerInstance server = createMockServerInstance(serverName); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java index 9f33b868b519..53198ede031a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java @@ -89,6 +89,38 @@ default List getSegments(BrokerRequest brokerRequest, @Nullable String s return getSegments(brokerRequest); } + /// Returns the segments that the segment pruners *provably eliminated* for the given broker request, i.e. the + /// selected segments that cannot hold a row matching the request's filter. Absence from the returned set means + /// nothing: a segment may be missing because it matches, because it was never selected, or because the table does + /// not exist. Only presence is a proof. + /// + /// This is the complement of [#getSegments(BrokerRequest)], and the distinction is what makes it usable as a + /// planning-time emptiness proof. Deciding "this segment does not match" from *absence* of a survivor conflates + /// pruning with the several innocent reasons a segment can be missing from a routing result -- it was classified as + /// optional by instance selection, its server left the enabled server map, or it entered the partition metadata + /// before it became selectable -- and each of those would silently drop matching data. Deciding it from presence in + /// this set cannot. + /// + /// Instance selection deliberately takes no part, so the result depends only on the request and the pruners, never + /// on a request id or on which replica a query happens to pick. + /// + /// Returns `null` if the table does not exist, as [#getSegments(BrokerRequest)] does, which is not the same as an + /// empty set: a broker that does not have the table eliminated nothing because it would have routed nothing, while + /// an empty set is a broker that ran the pruners and proved nothing. A caller combining several brokers' verdicts + /// has to tell those apart, and gets both from this one lookup rather than from a separate existence check that + /// could disagree with it. + /// + /// The returned set is for reading only: an implementation may hand back an immutable or a shared set, so a caller + /// that needs to modify it must copy it first. + /// + /// The default implementation returns an empty set, i.e. proves nothing about any segment. Note that this is a + /// statement about segments only: a caller may still act on emptiness it can see for itself, such as a partition + /// that lists no segment at all. + @Nullable + default Set getPrunedSegments(BrokerRequest brokerRequest) { + return Set.of(); + } + /// Validate routing exist for a table /// /// @param tableNameWithType the name of the table. diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java index 3c75ccae35bb..37809b8cbf35 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ColocatedJoinEmptyPartitionTest.java @@ -27,6 +27,8 @@ import java.util.Map; import java.util.regex.Matcher; import java.util.regex.Pattern; +import java.util.stream.Collectors; +import java.util.stream.Stream; import org.apache.avro.SchemaBuilder; import org.apache.avro.file.DataFileWriter; import org.apache.avro.generic.GenericData; @@ -34,6 +36,7 @@ import org.apache.commons.io.FileUtils; import org.apache.pinot.integration.tests.ClusterIntegrationTestUtils; import org.apache.pinot.spi.config.table.ColumnPartitionConfig; +import org.apache.pinot.spi.config.table.RoutingConfig; import org.apache.pinot.spi.config.table.SegmentPartitionConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; @@ -66,6 +69,11 @@ /// end-to-end run can show is that a real server accepts and answers a leaf-stage request whose segment list is empty /// for a genuinely partitioned table scan. /// +/// The same fixture covers the class reduction a filter buys. The tables configure the `partition` segment pruner, so +/// a restriction on the join key lets the broker eliminate segments before planning, and a class every member +/// eliminates leaves the group: `partitionKey IN (1, 2)` keeps only classes 1 and 2, halving the leaf worker count. +/// A restriction matching nothing anywhere falls back to the populated classes and returns an empty result. +/// /// The partition layout is supplied with explicit `tableOptions` hints rather than inferred, because hint inference is /// off by default (`pinot.broker.multistage.infer.partition.hint`); the hints carry exactly what it would have /// produced. The `is_colocated_by_join_keys` hint is spelled out for readability -- the exchange would be @@ -91,6 +99,12 @@ public class ColocatedJoinEmptyPartitionTest extends CustomDataQueryClusterInteg private static final int NUM_KEPT_CLASSES_FOR_JOIN = 4; private static final int NUM_ROWS_PER_PARTITION = 2; + /// The partition keys the pruning tests filter on. They land in partitions 1 and 2, the only ones both tables + /// populate, so classes 0 (left only) and 3 (right only) are eliminated on every member of the colocated group. + private static final List FILTERED_KEYS = List.of(1, 2); + /// A key in a partition neither table populates, so the filter matches nothing anywhere. + private static final List UNMATCHED_KEYS = List.of(5); + private static final int LEFT_METRIC_MULTIPLIER = 10; private static final int RIGHT_METRIC_MULTIPLIER = 100; @@ -221,31 +235,143 @@ public void testColocatedJoinMatchesShuffledJoin() assertEquals(colocatedResponse.get("resultTable").get("rows"), shuffledResponse.get("resultTable").get("rows"), "Colocated and shuffled plans must return the same rows"); + assertShuffledLeafStages(shuffledResponse, 2); + } - JsonNode shuffledStageStats = shuffledResponse.get("stageStats"); - assertNotNull(shuffledStageStats, "Missing stage stats in shuffled response: " + shuffledResponse); - List shuffledLeafStageSends = new ArrayList<>(); - collectLeafStageSends(shuffledStageStats, shuffledLeafStageSends); - assertEquals(shuffledLeafStageSends.size(), 2, - "Unexpected number of leaf stages in stage stats: " + shuffledStageStats.toPrettyString()); - for (JsonNode leafStageSend : shuffledLeafStageSends) { - assertTrue(leafStageSend.path("fanOut").asInt(-1) > 1, - "A shuffled leaf send must write more than one receive mailbox, otherwise the fanOut of 1 asserted for the " - + "colocated plan proves nothing. Stage stats: " + shuffledStageStats.toPrettyString()); - } + /// The class reduction broker pruning buys: a filter that every member of the colocated group can prune with drops + /// the classes all of them eliminate, so the leaves run fewer workers and scan fewer segments than the union of the + /// populated classes. What the reduced width must not cost is the 1-to-1 wiring, which is what the two exchange + /// assertions are for. + @Test + public void testColocatedJoinPrunesClassesEveryMemberFilters() + throws Exception { + setUseMultiStageQueryEngine(true); + String query = colocatedJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME, bothSidesFilter(FILTERED_KEYS)); + + JsonNode response = queryBrokerHttpEndpoint(query); + assertNoExceptions(response); + assertRows(response, expectedFilteredRows(FILTERED_KEYS)); + + // Classes 0 and 3 go: the left table's only class-0 segment is pruned by its own filter and the right table never + // populated that class, and symmetrically for class 3. Both tables populate every surviving class with exactly one + // segment, so both leaves scan one segment per kept class. + List keptClasses = keptClassesFor(FILTERED_KEYS); + assertLeafStages(response, 2, keptClasses.size(), keptClasses.size()); + assertDirectExchanges(query, 2, keptClasses.size()); + + // One segment per class that a table populated but the filter did not keep: the left table's class-0 segment and + // the right table's class-3 one. + long expectedNumPrunedSegments = + Stream.concat(LEFT_POPULATED_PARTITIONS.stream(), RIGHT_POPULATED_PARTITIONS.stream()) + .filter(partition -> !keptClasses.contains(partition)) + .count(); + assertEquals(response.path("numSegmentsPrunedByBroker").asLong(-1), expectedNumPrunedSegments, + "Unexpected number of broker-pruned segments in response: " + response); + } + + /// The same differential check as [#testColocatedJoinMatchesShuffledJoin], at the reduced width. A plan that quietly + /// fell back to a shuffle would return the right rows too, so this pairs with the `fanOut` and `[PARTITIONED]` + /// assertions rather than replacing them; what it rules out is a reduction that pairs the wrong classes and drops or + /// duplicates rows with no error. The shuffled plan reaches its answer by an independent route: with no table hints + /// its leaves are assigned per server and pruned per segment, not per partition class. + @Test + public void testFilteredColocatedJoinMatchesFilteredShuffledJoin() + throws Exception { + setUseMultiStageQueryEngine(true); + String whereClause = bothSidesFilter(FILTERED_KEYS); + JsonNode colocatedResponse = queryBrokerHttpEndpoint( + colocatedJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME, whereClause)); + assertNoExceptions(colocatedResponse); + // Pin the rows rather than only comparing the two plans: were the filter to stop matching anything, both sides + // would return nothing and the comparison below would still pass while proving nothing at all. + assertRows(colocatedResponse, expectedFilteredRows(FILTERED_KEYS)); + JsonNode shuffledResponse = queryBrokerHttpEndpoint( + shuffledJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME, whereClause)); + assertNoExceptions(shuffledResponse); + + assertEquals(colocatedResponse.get("resultTable").get("rows"), shuffledResponse.get("resultTable").get("rows"), + "Class-reduced colocated and shuffled plans must return the same rows"); + assertShuffledLeafStages(shuffledResponse, 2); + } + + /// A filter every member prunes every segment with: the group is left with no surviving class at all. It must fall + /// back to its populated classes and let the servers return the empty result, because a zero-worker leaf has no + /// handling on a 1-to-1 exchange -- an empty answer, not an error. + @Test + public void testColocatedJoinWithFilterMatchingNothing() + throws Exception { + setUseMultiStageQueryEngine(true); + assertTrue(keptClassesFor(UNMATCHED_KEYS).isEmpty(), + "The filter must match no partition either table populates, otherwise this is not the all-pruned fallback"); + String query = colocatedJoinQuery(LEFT_TABLE_NAME, RIGHT_TABLE_NAME, bothSidesFilter(UNMATCHED_KEYS)); + + JsonNode response = queryBrokerHttpEndpoint(query); + assertNoExceptions(response); + assertRows(response, List.of()); + + // Planned exactly as if there were no filter, and nothing is reported as pruned: the fallback dropped no class the + // group would otherwise have kept. + assertLeafStages(response, 2, NUM_KEPT_CLASSES_FOR_JOIN, LEFT_POPULATED_PARTITIONS.size()); + assertDirectExchanges(query, 2, NUM_KEPT_CLASSES_FOR_JOIN); + assertEquals(response.path("numSegmentsPrunedByBroker").asLong(-1), 0L, + "Unexpected number of broker-pruned segments in response: " + response); } private static String colocatedJoinQuery(String leftTableName, String rightTableName) { + return colocatedJoinQuery(leftTableName, rightTableName, ""); + } + + private static String colocatedJoinQuery(String leftTableName, String rightTableName, String whereClause) { return String.format( - "SELECT %s l.%s, l.%s, r.%s FROM %s %s AS l JOIN %s %s AS r ON l.%s = r.%s ORDER BY l.%s", + "SELECT %s l.%s, l.%s, r.%s FROM %s %s AS l JOIN %s %s AS r ON l.%s = r.%s %s ORDER BY l.%s", COLOCATED_JOIN_HINT, PARTITION_KEY_COLUMN, METRIC_COLUMN, METRIC_COLUMN, leftTableName, TABLE_HINT, - rightTableName, TABLE_HINT, PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN); + rightTableName, TABLE_HINT, PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN, whereClause, PARTITION_KEY_COLUMN); } private static String shuffledJoinQuery(String leftTableName, String rightTableName) { - return String.format("SELECT l.%s, l.%s, r.%s FROM %s AS l JOIN %s AS r ON l.%s = r.%s ORDER BY l.%s", + return shuffledJoinQuery(leftTableName, rightTableName, ""); + } + + private static String shuffledJoinQuery(String leftTableName, String rightTableName, String whereClause) { + return String.format("SELECT l.%s, l.%s, r.%s FROM %s AS l JOIN %s AS r ON l.%s = r.%s %s ORDER BY l.%s", PARTITION_KEY_COLUMN, METRIC_COLUMN, METRIC_COLUMN, leftTableName, rightTableName, PARTITION_KEY_COLUMN, - PARTITION_KEY_COLUMN, PARTITION_KEY_COLUMN); + PARTITION_KEY_COLUMN, whereClause, PARTITION_KEY_COLUMN); + } + + /// A partition-key restriction spelled out on both sides of the join rather than on one and left to Calcite's + /// transitive inference, so that the leaf of each member carries it whatever the planner decides to push down. + private static String bothSidesFilter(List keys) { + String keyList = keys.stream().map(String::valueOf).collect(Collectors.joining(", ")); + return String.format("WHERE l.%s IN (%s) AND r.%s IN (%s)", PARTITION_KEY_COLUMN, keyList, PARTITION_KEY_COLUMN, + keyList); + } + + /// The partition classes a colocated join of the two tables keeps under the given partition-key restriction. A class + /// survives when at least one member still holds a segment its own filter leaves, i.e. when some restricted key + /// hashes into a partition that member populates. One partition per class here, since the declared partition count + /// is the hinted partition size. + private static List keptClassesFor(List keys) { + return keys.stream() + .map(key -> key % NUM_DECLARED_PARTITIONS) + .filter(partition -> LEFT_POPULATED_PARTITIONS.contains(partition) + || RIGHT_POPULATED_PARTITIONS.contains(partition)) + .distinct() + .sorted() + .collect(Collectors.toList()); + } + + /// The rows the given partition-key restriction leaves in an inner join of the two tables: a key survives when both + /// tables populate the partition it hashes into, and every populated partition holds every one of its keys. + private static List> expectedFilteredRows(List keys) { + List> expectedRows = new ArrayList<>(); + for (int key : keys) { + int partition = key % NUM_DECLARED_PARTITIONS; + if (LEFT_POPULATED_PARTITIONS.contains(partition) && RIGHT_POPULATED_PARTITIONS.contains(partition)) { + expectedRows.add( + List.of((long) key, (long) key * LEFT_METRIC_MULTIPLIER, (long) key * RIGHT_METRIC_MULTIPLIER)); + } + } + return expectedRows; } private static void assertNoExceptions(JsonNode response) { @@ -301,6 +427,22 @@ private static void assertLeafStages(JsonNode response, int expectedNumLeafStage } } + /// Asserts that every leaf stage of a shuffled plan writes more than one receive mailbox. This is the control that + /// gives the `fanOut` of 1 asserted for a colocated plan its meaning. + private static void assertShuffledLeafStages(JsonNode response, int expectedNumLeafStages) { + JsonNode stageStats = response.get("stageStats"); + assertNotNull(stageStats, "Missing stage stats in shuffled response: " + response); + List leafStageSends = new ArrayList<>(); + collectLeafStageSends(stageStats, leafStageSends); + assertEquals(leafStageSends.size(), expectedNumLeafStages, + "Unexpected number of leaf stages in stage stats: " + stageStats.toPrettyString()); + for (JsonNode leafStageSend : leafStageSends) { + assertTrue(leafStageSend.path("fanOut").asInt(-1) > 1, + "A shuffled leaf send must write more than one receive mailbox, otherwise the fanOut of 1 asserted for the " + + "colocated plan proves nothing. Stage stats: " + stageStats.toPrettyString()); + } + } + private static void collectLeafStageSends(JsonNode node, List leafStageSends) { JsonNode children = node.get("children"); if ("MAILBOX_SEND".equals(node.path("type").asText()) && children != null && children.size() == 1 && "LEAF".equals( @@ -387,6 +529,10 @@ private static TableConfig createTableConfigForTable(String tableName) { .setNumReplicas(2) .setSegmentPartitionConfig(new SegmentPartitionConfig( Map.of(PARTITION_KEY_COLUMN, new ColumnPartitionConfig(PARTITION_FUNCTION, NUM_DECLARED_PARTITIONS)))) + // Without this the broker builds no partition pruner at all (SegmentPrunerFactory only reads the routing + // config), a filter on the partition key would prune nothing, and the filtered tests below would assert the + // unfiltered worker count. The unfiltered tests are unaffected: with no filter there is nothing to prune. + .setRoutingConfig(new RoutingConfig(null, List.of(RoutingConfig.PARTITION_SEGMENT_PRUNER_TYPE), null, null)) .build(); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java index f50a3db9e1c3..585576dfa992 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java @@ -61,6 +61,7 @@ public class DispatchablePlanContext { private final Map _dispatchablePlanMetadataMap = new HashMap<>(); private final Map _dispatchablePlanStageRootMap = new HashMap<>(); private final Map _partitionTableInfoCache = new HashMap<>(); + private final Map> _prunedSegmentsCache = new HashMap<>(); private long _numSegmentsPrunedByBroker; private int _leafStagesAssigned; private int _leafStagesEmpty; @@ -141,6 +142,14 @@ public Map getPartitionTableInfoCache( return _partitionTableInfoCache; } + /// The segments the broker's pruners provably eliminated for each leaf fragment, keyed by fragment id. Keyed by + /// fragment rather than by table because the verdict depends on the leaf's own filter, and the two sides of a + /// self-join scan one table under two different ones. Cached because the colocation pre-pass and the leaf + /// assignment both need it, and each entry costs a routing call. + public Map> getPrunedSegmentsCache() { + return _prunedSegmentsCache; + } + public long getNumSegmentsPrunedByBroker() { return _numSegmentsPrunedByBroker; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java index 88b43fb61e7d..0d5508f6fc90 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/ColocationGroupAnalyzer.java @@ -150,7 +150,7 @@ private static ColocationGroup toReducibleGroup(List members, Map partitionedLeafFragmentIds = new ArrayList<>(); + List partitionedLeafFragments = new ArrayList<>(); int partitionSize = -1; int partitionParallelism = -1; String partitionFunction = null; @@ -186,7 +186,7 @@ private static ColocationGroup toReducibleGroup(List members, Map members, Map parents, int fragmentId) { static class ColocationGroup { /// The number of partition classes, and of workers before reduction, i.e. the hinted `partition_size`. final int _partitionSize; - /// The members whose data decides which classes survive. - final List _partitionedLeafFragmentIds; + /// The members whose data decides which classes survive. The whole fragment rather than its id because deciding + /// survival reads each member's own filter off its leaf stage tree, see `WorkerManager#assignPartitionClasses`. + final List _partitionedLeafFragments; - ColocationGroup(int partitionSize, List partitionedLeafFragmentIds) { + ColocationGroup(int partitionSize, List partitionedLeafFragments) { _partitionSize = partitionSize; - _partitionedLeafFragmentIds = partitionedLeafFragmentIds; + _partitionedLeafFragments = partitionedLeafFragments; } } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java index d7b32f3e4b37..0b6a4ff65c3f 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java @@ -85,6 +85,27 @@ public static PinotQuery createPinotQueryForRouting(String tableName, PlanNode l return pinotQuery; } + /// Whether [#createPinotQueryForRouting] can fold the given leaf stage tree, i.e. whether it holds no multi-input + /// node. Lets a caller skip the attempt rather than pay a thrown-and-caught exception per query, which a colocated + /// semi-join's probe leaf -- the one holding the join -- would otherwise do on every query. It lives here so that it + /// walks the tree the same way [#accumulateBottomToTop] does; the two disagreeing would either bring the exception + /// back or, worse, silently refuse shapes that fold perfectly well. + /// + /// A `true` result is not a promise that the fold succeeds: the tree may still be missing a table scan, which only + /// the fold itself detects. + public static boolean canBuildRoutingQuery(PlanNode leafStageRoot) { + List inputs = leafStageRoot.getInputs(); + if (inputs.size() > 1) { + return false; + } + for (PlanNode input : inputs) { + if (!canBuildRoutingQuery(input)) { + return false; + } + } + return true; + } + private static void accumulateBottomToTop(PlanNode root, List parentNodes) { Preconditions.checkState(root.getInputs().size() <= 1, "Leaf stage nodes should have at most one input, found: %s", root.getInputs().size()); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java index a6f9e3e3303d..1cab2a37a28c 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java @@ -136,12 +136,24 @@ public void assignWorkers(PlanFragment rootFragment, DispatchablePlanContext con /// publishes the decision on each partitioned leaf of the group (see /// [DispatchablePlanMetadata#getPartitionClassIds()] and [DispatchablePlanMetadata#getPaddedClassCandidates()]). /// - /// A class survives when *any* member holds a segment in it: the union, not the intersection, because a class that - /// holds data for one member must keep its worker on every member or the members stop agreeing on what a worker id - /// stands for. A member holding no data in a surviving class gets a worker with no segments (see - /// [#assignPaddedWorker]). Emptiness is computed in class space (`0..partitionSize-1`) rather than over raw partition - /// ids because members may declare different partition counts; a member carrying no per-class visibility (replicated, - /// non-partitioned, or deriving its worker map from a peer) contributes nothing to the union. + /// A class survives when *any* member holds a segment in it that the query's filter does not provably exclude: the + /// union, not the intersection, because a class that holds matching data for one member must keep its worker on + /// every member or the members stop agreeing on what a worker id stands for. Dropping a class the group as a whole + /// has no matching row in is what turns broker pruning into fewer workers, and therefore fewer dispatched servers, + /// for a colocated join. The union direction is also what keeps this layer free of relational semantics: no input + /// row anywhere in a class means no output row attributable to it, whatever operator sits above, whereas dropping a + /// class only one side filtered away would be wrong for a RIGHT or FULL join, a union, or an anti-join. + /// + /// A member holding no data in a surviving class gets a worker with no segments (see [#assignPaddedWorker]). A + /// member that holds data the filter excludes does *not*: it keeps every segment of the class. Survival is decided + /// per class rather than per partition on purpose -- a class's worker is placed on the servers shared by all of its + /// populated partitions, so dropping some of them on one member and not on another would move that member's worker + /// off its peer's server and turn an in-process exchange into a network one. Filtering inside a surviving class is a + /// possible follow-up. + /// + /// Emptiness is computed in class space (`0..partitionSize-1`) rather than over raw partition ids because members + /// may declare different partition counts; a member carrying no per-class visibility (replicated, non-partitioned, + /// or deriving its worker map from a peer) contributes nothing to the union. /// /// Marking no group keeps the assignment as it is without one: every class gets a worker, so a class holding no /// segment fails the assignment instead of being dropped or padded. @@ -151,67 +163,108 @@ private void assignPartitionClasses(PlanFragment rootFragment, DispatchablePlanC for (ColocationGroupAnalyzer.ColocationGroup group : ColocationGroupAnalyzer.findReducibleGroups(rootFragment, metadataMap)) { int numWorkers = group._partitionSize; - List memberFragmentIds = group._partitionedLeafFragmentIds; - // The servers each member can scan each class on, in the same order as the member fragment ids. - List>> memberClassServers = new ArrayList<>(memberFragmentIds.size()); + List memberFragments = group._partitionedLeafFragments; + List memberFragmentIds = new ArrayList<>(memberFragments.size()); + // The servers each member can scan each class on, in the same order as the member fragments. + List>> memberClassServers = new ArrayList<>(memberFragments.size()); + // The partition layouts, same order again, kept only to count what a filter-dropped class cost each member. + List memberPartitionInfoMaps = new ArrayList<>(memberFragments.size()); // Allocated lazily, once the first member has checked the hint against its table: numWorkers is the raw hinted // partition size, so sizing anything from it before that check would let a bogus hint allocate unboundedly. The // check also bounds it by the table's partition count. - boolean[] survivingClasses = null; + // populatedClasses ignores the filter and is what decides padding and the all-pruned fallback; matchingClasses + // is the same union taken over the segments the filter leaves, and is a subset of it. + boolean[] populatedClasses = null; + boolean[] matchingClasses = null; boolean reducible = true; - for (Integer fragmentId : memberFragmentIds) { - DispatchablePlanMetadata metadata = metadataMap.get(fragmentId); + for (PlanFragment fragment : memberFragments) { + DispatchablePlanMetadata metadata = metadataMap.get(fragment.getFragmentId()); String tableName = metadata.getScannedTables().get(0); // NOTE: A failure here is the same one the leaf assignment would hit for this table, only raised earlier. PartitionTableInfo partitionTableInfo = partitionTableInfoCache.computeIfAbsent(tableName, this::calculatePartitionTableInfo); - int numPartitions = partitionTableInfo._partitionInfoMap.length; + PartitionInfo[] partitionInfoMap = partitionTableInfo._partitionInfoMap; + int numPartitions = partitionInfoMap.length; if (numPartitions == 0 || numPartitions % numWorkers != 0) { // The table does not match the hinted partition size. Leave the group alone so that checkPartitionInfoMap // reports it during the leaf assignment. reducible = false; break; } - if (survivingClasses == null) { - survivingClasses = new boolean[numWorkers]; + if (populatedClasses == null) { + populatedClasses = new boolean[numWorkers]; + matchingClasses = new boolean[numWorkers]; } - List> classServers = collectClassServers(partitionTableInfo._partitionInfoMap, numWorkers); + List> classServers = collectClassServers(partitionInfoMap, numWorkers); + Set prunedSegments = getPrunedSegments(fragment, tableName, context); boolean anyPopulated = false; for (int classId = 0; classId < numWorkers; classId++) { if (classServers.get(classId) != null) { - survivingClasses[classId] = true; + populatedClasses[classId] = true; anyPopulated = true; + if (prunedSegments == null) { + // No filter to prune this member with, so it contributes every class it holds data in. + matchingClasses[classId] = true; + } } } // A member holding no data at all leaves nothing to assign: no class to place its single empty worker in, and // no server known to host the table to place it on. Check the deferred cause first though -- a table whose // every partition is deferred also has no populated class, and reports far more actionably. That is the - // pre-pass' only deferred check: a group it marks gets no broker pruning, so the leaf assignment covers the - // rest. + // pre-pass' only deferred check: a group it marks gets no broker pruning at the leaf, so the leaf assignment + // covers the rest. + // + // NOTE: This reads the unfiltered population on purpose. A member whose filter matches nothing holds data + // all the same, and failing it here would turn a correct empty result into a query error. if (!anyPopulated) { checkNoPartitionsWithOnlyDeferredSegments(partitionTableInfo, tableName); } Preconditions.checkState(anyPopulated, "Failed to find any segment in any partition for table: %s, which is required for a partitioned worker " + "assignment", tableName); + if (prunedSegments != null) { + markClassesWithMatchingData(partitionInfoMap, numWorkers, prunedSegments, matchingClasses); + } memberClassServers.add(classServers); + memberPartitionInfoMaps.add(partitionInfoMap); + memberFragmentIds.add(fragment.getFragmentId()); } if (!reducible) { continue; } - // The member list is never empty (see ColocationGroupAnalyzer#toReducibleGroup), so the loop allocated this, and - // the class list is never empty either: every member holds data in at least one class, and the union keeps it. - assert survivingClasses != null; + // The member list is never empty (see ColocationGroupAnalyzer#toReducibleGroup), so the loop allocated these, + // and the class list is never empty either: every member holds data in at least one class, and the union keeps + // it. + assert populatedClasses != null && matchingClasses != null; + // A group the filter empties keeps all of its populated classes, mirroring the leaf-level fallback in + // computePartitionsToKeep: a zero-worker leaf has no handling on a 1-to-1 exchange, and the server-side filter + // still returns the correct empty result from an unreduced plan. + boolean[] survivingClasses = anyTrue(matchingClasses) ? matchingClasses : populatedClasses; int[] partitionClassIds = toClassIds(survivingClasses); Map>> padding = computePadding(memberFragmentIds, memberClassServers, partitionClassIds); if (padding.isEmpty() && partitionClassIds.length == numWorkers) { // Worker k already stands for class k on every member: nothing to reduce, nothing to pad. Leaving the group - // unmarked also keeps broker pruning on for its leaves (see computePartitionsToKeep). A group that needs + // unmarked also keeps leaf-level broker pruning on for the members that are eligible for it (see + // computePartitionsToKeep), which prunes at partition rather than class granularity. A group that needs // padding is marked even when it keeps every class, because a padded worker's id is its index in the class // list. continue; } + // Report what the filter cost, not what the class reduction did: a class no member holds data in is empty + // rather than pruned, and it is already dropped above without being counted. Derived from the decision actually + // taken, so the all-pruned fallback above reports nothing. + long numPrunedSegments = 0; + for (int classId = 0; classId < numWorkers; classId++) { + if (populatedClasses[classId] && !survivingClasses[classId]) { + for (PartitionInfo[] partitionInfoMap : memberPartitionInfoMaps) { + numPrunedSegments += countSegmentsInClass(partitionInfoMap, numWorkers, classId); + } + } + } + if (numPrunedSegments > 0) { + context.addNumSegmentsPrunedByBroker(numPrunedSegments); + } // One shared array instance, so that the agreement check in MailboxAssignmentVisitor compares one list rather // than copies of it. The padding goes on the same metadata: a padded worker's id only means something within the // list. @@ -223,6 +276,48 @@ private void assignPartitionClasses(PlanFragment rootFragment, DispatchablePlanC } } + private static long countSegmentsInClass(PartitionInfo[] partitionInfoMap, int numWorkers, int classId) { + long numSegments = 0; + for (int partitionId = classId; partitionId < partitionInfoMap.length; partitionId += numWorkers) { + PartitionInfo partitionInfo = partitionInfoMap[partitionId]; + if (partitionInfo != null) { + numSegments += CollectionUtils.size(partitionInfo._offlineSegments) + + CollectionUtils.size(partitionInfo._realtimeSegments); + } + } + return numSegments; + } + + /// Sets, for every class holding at least one segment the given pruned set does not cover, the corresponding entry + /// of `matchingClasses`. Only presence in the pruned set is a proof (see [RoutingManager#getPrunedSegments]), so a + /// partition with no segment listed at all -- which the hybrid layout allows for one of the two table types -- is + /// read as matching rather than as empty. + private static void markClassesWithMatchingData(PartitionInfo[] partitionInfoMap, int numWorkers, + Set prunedSegments, boolean[] matchingClasses) { + int numPartitions = partitionInfoMap.length; + for (int classId = 0; classId < numWorkers; classId++) { + if (matchingClasses[classId]) { + continue; + } + for (int partitionId = classId; partitionId < numPartitions; partitionId += numWorkers) { + PartitionInfo partitionInfo = partitionInfoMap[partitionId]; + if (partitionInfo != null && !allSegmentsPruned(partitionInfo, prunedSegments)) { + matchingClasses[classId] = true; + break; + } + } + } + } + + private static boolean anyTrue(boolean[] flags) { + for (boolean flag : flags) { + if (flag) { + return true; + } + } + return false; + } + /// Returns the servers that can scan each partition class of the given layout as a whole, in class-id order, or /// `null` for a class that holds no segment at all. This is the intersection of the fully replicated servers of the /// class's populated partitions, i.e. the candidate set its worker is picked from (see @@ -687,13 +782,14 @@ private void assignWorkersToLeafFragment(PlanFragment fragment, DispatchablePlan metadata.setPartitionParallelism(partitionHints.getPartitionParallelism()); if (partitionHints.getPartitionKey() != null) { - // Broker pruning: build a filter-bearing routing query (null when disabled/unsupported) so the partitioned - // assignment can drop partitions with no matching segments. Reuses the same gate as the non-partitioned path. - // Skip pre-partitioned leaves and leaves of a reduced colocated group up front: pruning is disabled for them - // (see computePartitionsToKeep), so don't spend planning time building the routing query. - PinotQuery routingPinotQuery = metadata.isPrePartitioned() || metadata.getPartitionClassIds() != null ? null - : extractRoutingQuery(fragment.getFragmentRoot(), metadata.getScannedTables().get(0), context); - assignWorkersToPartitionedLeafFragment(metadata, context, partitionHints, routingPinotQuery); + // Broker pruning: the segments the pruners provably eliminated (empty when disabled/unsupported) so the + // partitioned assignment can drop partitions holding none of them. Skip the lookup for a pre-partitioned leaf + // and for a leaf of a marked colocated group: leaf-level pruning is disabled for both (see + // computePartitionsToKeep) because the group's shared class list already carries their verdict, so asking + // would only cost planning time. + Set prunedSegments = metadata.isPrePartitioned() || metadata.getPartitionClassIds() != null ? null + : getPrunedSegments(fragment, metadata.getScannedTables().get(0), context); + assignWorkersToPartitionedLeafFragment(metadata, context, partitionHints, prunedSegments); updateContextForLeafStage(metadata, context); return; } @@ -867,6 +963,86 @@ private Map getRoutingTable(PinotQuery pinotQuery, long re } } + /// Returns the segments the broker's pruners provably eliminated for the given leaf fragment, or `null` when there + /// was no filter to prune with at all: broker pruning off, an unsupported leaf shape, a filterless leaf, or a + /// routing failure (pruning is best-effort and never fails a query that would otherwise route). + /// + /// The empty set and `null` mean different things and callers depend on it. An empty set is "a filter ran and + /// proved nothing", which still lets the partitioned assignment skip a partition holding no segment at all -- + /// behaviour that predates this and that a query with an empty partition relies on to plan. + /// + /// Memoised per fragment for the query, because the colocation pre-pass and the leaf assignment both ask for it and + /// each answer costs a routing call. Keyed by fragment rather than by table: the two sides of a self-join scan one + /// table under two different filters. + @Nullable + private Set getPrunedSegments(PlanFragment fragment, String tableName, DispatchablePlanContext context) { + // Not computeIfAbsent: null is a meaningful answer here and would be recomputed on every call. + Map> prunedSegmentsCache = context.getPrunedSegmentsCache(); + Integer fragmentId = fragment.getFragmentId(); + if (prunedSegmentsCache.containsKey(fragmentId)) { + return prunedSegmentsCache.get(fragmentId); + } + Set prunedSegments = computePrunedSegments(fragment, tableName, context); + prunedSegmentsCache.put(fragmentId, prunedSegments); + return prunedSegments; + } + + @Nullable + private Set computePrunedSegments(PlanFragment fragment, String tableName, + DispatchablePlanContext context) { + PinotQuery routingPinotQuery = extractRoutingQuery(fragment.getFragmentRoot(), tableName, context); + if (routingPinotQuery == null || routingPinotQuery.getFilterExpression() == null) { + return null; + } + try { + TableType tableType = TableNameBuilder.getTableTypeFromTableName(routingPinotQuery.getDataSource() + .getTableName()); + if (tableType != null) { + return getPrunedSegmentsHelper(routingPinotQuery); + } + // A raw table name may resolve to either or both physical tables. Segment names are unique across them, and a + // segment is pruned by the table that holds it, so the two verdicts simply add up. Only merge when both prove + // something: a table of one type alone is the common case, and copying its verdict to union it with an empty + // set would allocate a second set over every segment name for nothing. + Set offlinePrunedSegments = getPrunedSegmentsHelper(routingPinotQuery, TableType.OFFLINE); + Set realtimePrunedSegments = getPrunedSegmentsHelper(routingPinotQuery, TableType.REALTIME); + if (offlinePrunedSegments.isEmpty()) { + return realtimePrunedSegments; + } + if (realtimePrunedSegments.isEmpty()) { + return offlinePrunedSegments; + } + Set prunedSegments = new HashSet<>(offlinePrunedSegments); + prunedSegments.addAll(realtimePrunedSegments); + return prunedSegments; + } catch (RuntimeException e) { + // Pruning is best-effort: never fail a query that would otherwise route successfully unpruned. + LOGGER.warn("Broker pruning skipped for table {} due to routing failure", tableName, e); + return null; + } + } + + /// A table the routing manager does not have is reported as `null` there; here it is simply a table that proves + /// nothing, which is the same thing this path does with a table whose pruners eliminated no segment. + private Set getPrunedSegmentsHelper(PinotQuery pinotQuery) { + Set prunedSegments = + _routingManager.getPrunedSegments(CalciteSqlCompiler.convertToBrokerRequest(pinotQuery)); + return prunedSegments != null ? prunedSegments : Set.of(); + } + + private Set getPrunedSegmentsHelper(PinotQuery pinotQuery, TableType tableType) { + return getPrunedSegmentsHelper(withTableType(pinotQuery, tableType)); + } + + /// Returns a copy of the given routing query aimed at one physical table, so that a query written against a raw + /// table name can be routed against each type in turn. + private static PinotQuery withTableType(PinotQuery pinotQuery, TableType tableType) { + PinotQuery copy = pinotQuery.deepCopy(); + copy.getDataSource().setTableName(TableNameBuilder.forType(tableType).tableNameWithType( + TableNameBuilder.extractRawTableName(pinotQuery.getDataSource().getTableName()))); + return copy; + } + /// Builds a [PinotQuery] from the leaf stage tree for broker-side segment pruning on the logical planner path. /// Returns `null` if broker pruning is disabled or the leaf stage shape is unsupported. @Nullable @@ -878,6 +1054,9 @@ private PinotQuery extractRoutingQuery(PlanNode leafStageRoot, String tableName, if (!useBrokerPruning) { return null; } + if (!PlanNodeRoutingQueryBuilder.canBuildRoutingQuery(leafStageRoot)) { + return null; + } try { PinotQuery pinotQuery = PlanNodeRoutingQueryBuilder.createPinotQueryForRouting(tableName, leafStageRoot, false); Map queryOptions = context.getPlannerContext().getOptions(); @@ -911,10 +1090,7 @@ private RoutingTable getRoutingTableHelper(PinotQuery pinotQuery, long requestId @Nullable private RoutingTable getRoutingTableHelper(PinotQuery pinotQuery, long requestId, TableType tableType) { - PinotQuery copy = pinotQuery.deepCopy(); - copy.getDataSource().setTableName(TableNameBuilder.forType(tableType).tableNameWithType( - TableNameBuilder.extractRawTableName(pinotQuery.getDataSource().getTableName()))); - return getRoutingTableHelper(copy, requestId); + return getRoutingTableHelper(withTableType(pinotQuery, tableType), requestId); } // -------------------------------------------------------------------------- @@ -1137,8 +1313,7 @@ private static void transferToServerInstanceLogicalSegmentsMap(String physicalTa /// Assigns one worker per partition class of a leaf that scans a partitioned table. private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata metadata, - DispatchablePlanContext context, LeafPartitionHints partitionHints, - @Nullable PinotQuery routingPinotQuery) { + DispatchablePlanContext context, LeafPartitionHints partitionHints, @Nullable Set prunedSegments) { // when partition key exist, we assign workers for leaf-stage in partitioned fashion. String partitionKey = partitionHints.getPartitionKey(); assert partitionKey != null; @@ -1172,8 +1347,7 @@ private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata met collectHostingServers(partitionInfoMap)) : null; // Broker pruning: the partitions to keep (null means keep all). Partitions absent from the set are skipped below. - Set partitionsToKeep = - computePartitionsToKeep(routingPinotQuery, metadata, context.getRequestId(), partitionInfoMap); + Set partitionsToKeep = computePartitionsToKeep(prunedSegments, metadata, partitionInfoMap); if (partitionsToKeep != null) { long numSegmentsPrunedByBroker = countPrunedSegments(partitionInfoMap, partitionsToKeep); if (numSegmentsPrunedByBroker > 0) { @@ -1202,35 +1376,42 @@ private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata met partitionClassIds, partitionsToKeep, paddingInfo, _routingManager.getEnabledServerInstanceMap(), workerIdToServerInstanceMap, workerIdToSegmentsMap); } - checkLeafWorkerAssignment(tableName, workerIdToServerInstanceMap, workerIdToSegmentsMap); + checkLeafWorkerAssignment(tableName, partitionClassIds, workerIdToServerInstanceMap, workerIdToSegmentsMap); metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap); metadata.setWorkerIdToSegmentsMap(workerIdToSegmentsMap); metadata.setTimeBoundaryInfo(partitionTableInfo._timeBoundaryInfo); metadata.setPartitionFunction(partitionFunction); } - /// Broker pruning for the partitioned leaf path. Returns the set of partition ids that still have at least one - /// segment matching the query filter, or `null` to keep all partitions. + /// Broker pruning for the partitioned leaf path. Returns the set of partition ids that are not provably empty for + /// this query, or `null` to keep all partitions. + /// + /// Note that an *empty* pruned set is not the same as an absent one and does not return `null` here: a filter that + /// proved nothing still leaves this the job of skipping a partition holding no segment at all, which is what lets a + /// table with an empty partition plan rather than fail on a worker it cannot place. /// /// Returns `null` (no pruning) when any of the following hold: /// - /// - broker pruning is disabled or the leaf shape is unsupported (the routing query is `null`), or there is - /// no filter to prune with; + /// - there was no filter to prune with at all (see [#getPrunedSegments]) -- broker pruning is disabled, the leaf + /// shape is unsupported, the leaf carries no filter, or routing failed (pruning is best-effort); /// - the leaf feeds a pre-partitioned (1-to-1 direct) exchange, or it belongs to a colocated group that agreed on a - /// partition class list -- dropping/compacting workers would misalign sender/receiver worker ids in - /// `MailboxAssignmentVisitor`. A non-pre-partitioned leaf is shuffled via `connectWorkers`, which re-hashes across - /// any worker count, so pruning is safe there; - /// - routing fails (pruning is best-effort); + /// partition class list. Both get their verdict from the group instead, in `assignPartitionClasses`, because it is + /// the only place that sees every member before any of them is assigned: a leaf deciding on its own would drop a + /// class its peer keeps, and the two would stop agreeing on what a worker id stands for. A non-pre-partitioned, + /// unmarked leaf is shuffled via `connectWorkers`, which re-hashes across any worker count, so it can decide alone + /// -- and at partition rather than class granularity; /// - every partition would be pruned -- an empty worker map would break exchanges in a multi-leaf plan (the /// all-leaves-empty short-circuit does not fire for a partially-empty plan), and the server-side filter still /// yields the correct empty result unpruned. /// - /// Partition survival is decided by routing the filter-bearing query through the [RoutingManager] (the same - /// mechanism the non-partitioned path uses), so the segment-level pruners judge survival using each segment's own - /// partition metadata. This is correct for every partition function and configuration, unlike recomputing the - /// partition id from the table-level function name (which lacks the per-segment function config). A partition is - /// dropped only when every one of its segments was pruned; a segment that merely became unavailable keeps its - /// partition alive so matching data is never silently dropped. + /// A partition is dropped only when every one of its segments is in the *provably pruned* set (see + /// [RoutingManager#getPrunedSegments]), never because a segment failed to appear somewhere. That direction is the + /// whole point: absence from a routing result has innocent causes -- a segment classified as optional by instance + /// selection, one whose server left the enabled server map, one that entered the partition metadata before it became + /// selectable -- and each would otherwise be read as "this partition is empty" and silently drop matching rows. It + /// also makes the verdict independent of the request id, so two leaves scanning one table under one filter cannot + /// disagree. Judging by pruner verdict rather than by recomputing the partition id from the table-level function + /// name is also what keeps it correct for every partition function and per-segment function config. /// /// Note that pruning here is partition-level, not segment-level: a surviving partition dispatches all of its /// segments, including ones the pruners eliminated (the server-side pruners drop those again cheaply). This keeps @@ -1238,37 +1419,15 @@ private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata met /// workers -- at the cost of a lower pruning ceiling than the non-partitioned path for partitions with mixed-match /// segments. Segment-level pruning within surviving partitions is a possible follow-up. @Nullable - private Set computePartitionsToKeep(@Nullable PinotQuery routingPinotQuery, - DispatchablePlanMetadata metadata, long requestId, PartitionInfo[] partitionInfoMap) { - if (routingPinotQuery == null || routingPinotQuery.getFilterExpression() == null || metadata.isPrePartitioned() - || metadata.getPartitionClassIds() != null) { - return null; - } - Map routingTableMap; - try { - routingTableMap = getRoutingTable(routingPinotQuery, requestId); - } catch (RuntimeException e) { - // Pruning is best-effort: never fail a query that would otherwise route successfully unpruned. - LOGGER.warn("Broker pruning skipped for partitioned table {} due to routing failure", - routingPinotQuery.getDataSource().getTableName(), e); + private static Set computePartitionsToKeep(@Nullable Set prunedSegments, + DispatchablePlanMetadata metadata, PartitionInfo[] partitionInfoMap) { + if (prunedSegments == null || metadata.isPrePartitioned() || metadata.getPartitionClassIds() != null) { return null; } - if (routingTableMap.isEmpty()) { - return null; - } - Set matchedSegments = new HashSet<>(); - for (RoutingTable routingTable : routingTableMap.values()) { - for (SegmentsToQuery segmentsToQuery : routingTable.getServerInstanceToSegmentsMap().values()) { - matchedSegments.addAll(segmentsToQuery.getSegments()); - } - // Keep a partition alive if any of its segments is merely unavailable (rather than pruned) so we never drop data. - matchedSegments.addAll(routingTable.getUnavailableSegments()); - } Set partitionsToKeep = new HashSet<>(); for (int i = 0; i < partitionInfoMap.length; i++) { PartitionInfo partitionInfo = partitionInfoMap[i]; - if (partitionInfo != null && (containsAny(partitionInfo._offlineSegments, matchedSegments) || containsAny( - partitionInfo._realtimeSegments, matchedSegments))) { + if (partitionInfo != null && !allSegmentsPruned(partitionInfo, prunedSegments)) { partitionsToKeep.add(i); } } @@ -1276,18 +1435,28 @@ private Set computePartitionsToKeep(@Nullable PinotQuery routingPinotQu return partitionsToKeep.isEmpty() ? null : partitionsToKeep; } - private static boolean containsAny(@Nullable List segments, Set matchedSegments) { + /// Returns whether every segment of the given partition is provably pruned. Vacuously true for a partition listing + /// no segment at all, which has no rows to contribute either way -- note that a partition holding data the broker + /// cannot route yet has no entry in the map rather than an empty one, so it is not this case (see + /// [#checkNoPartitionsWithOnlyDeferredSegments]). + private static boolean allSegmentsPruned(PartitionInfo partitionInfo, Set prunedSegments) { + return allPruned(partitionInfo._offlineSegments, prunedSegments) + && allPruned(partitionInfo._realtimeSegments, prunedSegments); + } + + private static boolean allPruned(@Nullable List segments, Set prunedSegments) { if (segments != null) { for (String segment : segments) { - if (matchedSegments.contains(segment)) { - return true; + if (!prunedSegments.contains(segment)) { + return false; } } } - return false; + return true; } - /// Counts the segments in partitions dropped by broker pruning (those absent from `partitionsToKeep`). + /// Counts the segments in partitions dropped by broker pruning (those absent from `partitionsToKeep`). A partition + /// dropped for holding no segment rather than for being pruned contributes nothing, so it is not miscounted. private static long countPrunedSegments(PartitionInfo[] partitionInfoMap, Set partitionsToKeep) { long numPrunedSegments = 0; for (int i = 0; i < partitionInfoMap.length; i++) { @@ -1508,12 +1677,20 @@ private static Map collectHostingServers(PartitionInfo[] partiti /// server map by worker id, where a gap leaves a null entry; /// - every worker must have a segments map keyed by 1 or 2 [TableType] names, with non-null lists, because the server /// splits the request on the number of entries and resolves one table data manager per key: an unexpected key - /// becomes an opaque server-side failure. + /// becomes an opaque server-side failure; + /// - a leaf of a colocated group must produce exactly one worker per class of the group's shared list, because a + /// worker id *is* an index into that list. Nothing downstream can catch a leaf that skipped one: + /// `MailboxAssignmentVisitor#checkPartitionClassAgreement` compares the shared array against itself, so a member + /// that quietly assigned fewer workers than the array claims still agrees with its peers on the array while + /// disagreeing with them on what every worker id after the gap means. @VisibleForTesting - static void checkLeafWorkerAssignment(String tableName, + static void checkLeafWorkerAssignment(String tableName, @Nullable int[] partitionClassIds, Map workerIdToServerInstanceMap, Map>> workerIdToSegmentsMap) { int numWorkers = workerIdToServerInstanceMap.size(); + Preconditions.checkState(partitionClassIds == null || partitionClassIds.length == numWorkers, + "Got %s workers for partition classes: %s of table: %s", numWorkers, + partitionClassIds != null ? Arrays.toString(partitionClassIds) : null, tableName); Preconditions.checkState(workerIdToSegmentsMap.size() == numWorkers, "Got %s workers but %s worker segment entries for table: %s", numWorkers, workerIdToSegmentsMap.size(), tableName); diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java index 3a779db87312..5908900e33c4 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/ColocationGroupAnalyzerTest.java @@ -22,6 +22,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.stream.Collectors; import javax.annotation.Nullable; import org.apache.calcite.rel.RelDistribution; import org.apache.pinot.calcite.rel.hint.PinotHintOptions; @@ -54,7 +55,7 @@ public void testGroupWithOnlyPrePartitionedSendsIsReducible() { assertEquals(groups.size(), 1); assertEquals(groups.get(0)._partitionSize, 4); - assertEquals(Set.copyOf(groups.get(0)._partitionedLeafFragmentIds), Set.of(2, 3)); + assertEquals(Set.copyOf(fragmentIds(groups.get(0))), Set.of(2, 3)); } /// A member that also receives a shuffled send must keep today's worker count, or that sender's rows land on @@ -78,7 +79,7 @@ public void testSingletonSendFormsAnEdgeWithoutPrePartitioning() { ColocationGroupAnalyzer.findReducibleGroups(twoLeafPlan(RelDistribution.Type.SINGLETON), metadataMap); assertEquals(groups.size(), 1); - assertEquals(Set.copyOf(groups.get(0)._partitionedLeafFragmentIds), Set.of(2, 3)); + assertEquals(Set.copyOf(fragmentIds(groups.get(0))), Set.of(2, 3)); } /// Reducing the worker count must not turn mismatched counts into a match for a pre-partitioned BROADCAST send, which @@ -227,7 +228,7 @@ public void testReplicatedLeafDoesNotBlockTheGroup() { assertEquals(groups.size(), 1); // Only the partitioned leaf decides which classes survive. - assertEquals(groups.get(0)._partitionedLeafFragmentIds, List.of(2)); + assertEquals(fragmentIds(groups.get(0)), List.of(2)); } /// A lookup join's workers come from its single local exchange child, so its own hints (a different partition size @@ -249,7 +250,7 @@ public void testLookupJoinMemberIsIgnored() { assertEquals(groups.size(), 1); assertEquals(groups.get(0)._partitionSize, 4); - assertEquals(groups.get(0)._partitionedLeafFragmentIds, List.of(2)); + assertEquals(fragmentIds(groups.get(0)), List.of(2)); } /// A group of intermediate stages only has nothing to reduce: only a partitioned leaf's data decides the classes. @@ -287,7 +288,11 @@ public void testSpooledFragmentTiesEveryReceiverIntoOneGroup() { // One group, and the spooled leaf is listed once rather than once per receiver. assertEquals(groups.size(), 1); - assertEquals(groups.get(0)._partitionedLeafFragmentIds, List.of(3)); + assertEquals(fragmentIds(groups.get(0)), List.of(3)); + } + + private static List fragmentIds(ColocationGroupAnalyzer.ColocationGroup group) { + return group._partitionedLeafFragments.stream().map(PlanFragment::getFragmentId).collect(Collectors.toList()); } /// Builds a 4 stage plan: 2 partitioned leaves (stages 2 and 3) sending to a join stage (stage 1), which sends diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java index 84ce1a4e108e..49802eba16f7 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java @@ -720,10 +720,14 @@ public void testBuildLogicalTableRoutingRequestWithoutFilterUsesSelectStar() { } @Test - public void testBrokerPruningPartitionedLeafSkippedForColocatedJoin() { - // A pre-partitioned leaf (here both sides of a colocated self-join) feeds a 1-to-1 direct exchange wired by worker - // id. Compacting a side's workers for pruned partitions could pair mismatched partitions across the exchange, so - // pruning is skipped for any pre-partitioned leaf and all partitions stay assigned on every scan. + public void testBrokerPruningColocatedJoinDropsClassesBothSidesPrune() { + // A pre-partitioned leaf feeds a 1-to-1 direct exchange wired by worker id, so no leaf may decide on its own what + // to drop. The colocation pre-pass decides for the whole group instead: a class survives when ANY member still + // holds a segment its own filter leaves, so dropping one is safe without knowing which operator sits above. + // + // This is the shape the reduction is worth having for. The filter is on the partition key, so it transfers across + // the join equality to both sides, both prune every class but 2, and the group's class list shrinks to [2]: one + // worker per leaf, and the query is dispatched to the single server holding partition 2 instead of all four. QueryEnvironment queryEnvironment = newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4, List.of("seg2"), 3); try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( @@ -733,13 +737,76 @@ public void testBrokerPruningPartitionedLeafSkippedForColocatedJoin() { + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ t2 " + "ON t1.col1 = t2.col1 WHERE t1.col1 = 'foo'")) { DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); - assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); List leafFragments = leafFragments(dispatchableSubPlan); - assertFalse(leafFragments.isEmpty()); + assertEquals(leafFragments.size(), 2); + for (DispatchablePlanFragment leaf : leafFragments) { + assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 1); + // Worker 0 stands for class 2 on both members, and carries that class's whole segment list. + assertEquals(assignedSegments(leaf, 0), List.of("seg2")); + } + // Three classes dropped, one segment each, counted once per member of the group rather than from the routing + // table's own self-reported count (the fixture reports 3). + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 6); + // The point of the exercise: partition 2 lives only on server 3, so that is the whole dispatched set. + assertEquals(dispatchedServers(dispatchableSubPlan), Set.of(getServerInstance("localhost", 3).getInstanceId())); + } + } + + @Test + public void testBrokerPruningSelfJoinKeepsTheTwoSidesVerdictsApart() { + // Two leaves scanning ONE table under two different filters. The pruning verdict is memoised per leaf fragment, + // and this is what says so: keyed by table instead, one side would inherit the other's verdict and the union would + // be computed from one filter applied twice. + // + // Left keeps only class 1, right only class 2, so the union is {1, 2} -- a result neither side's verdict produces + // on its own, and neither does either verdict applied to both sides ({1} or {2}). + QueryEnvironment queryEnvironment = + newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4, 1, List.of(), List.of(), 0, false, Set.of(), Set.of(), + Map.of("left", List.of("seg1"), "right", List.of("seg2"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; SELECT t1.col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ t1 " + + "JOIN testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ t2 " + + "ON t1.col1 = t2.col1 WHERE t1.col2 = 'left' AND t2.col2 = 'right'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + List leafFragments = leafFragments(dispatchableSubPlan); + assertEquals(leafFragments.size(), 2); + Map workerIdToClass = new HashMap<>(); for (DispatchablePlanFragment leaf : leafFragments) { - assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 4, "Expected all partitions assigned for a pruned-gated " - + "colocated join leaf"); + // Both classes survive on both sides, and a class the group keeps dispatches all of its segments even on the + // member whose own filter excluded them. + assertEquals(workerIdToPartitions(leaf, "seg"), Map.of(0, Set.of(1), 1, Set.of(2))); + mergeWorkerIdToClass(workerIdToClass, leaf, "seg", 4); } + assertEquals(workerIdToClass, Map.of(0, 1, 1, 2)); + // Classes 0 and 3 dropped, one segment each, on each of the two leaves. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 4); + } + } + + @Test + public void testBrokerPruningColocatedJoinKeepsClassOnlyOneSidePrunes() { + // The union rule, and the reason it is not an intersection: only t1 is filtered, t2 still holds every class, so + // every class keeps its worker and the fan-out is unchanged. Dropping the classes t1's filter empties would be + // wrong the moment the operator above is a RIGHT or FULL join, a union, or an anti-join -- and the worker + // assignment deliberately knows nothing about which one it is. A one-sided filter buying nothing is the price. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).survivingSegments(List.of("a_seg1")), + new ColocatedTableSpec(4, false).survivingSegments(List.of("b_seg0", "b_seg1", "b_seg2", "b_seg3"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + assertEquals(leafA.getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafB.getWorkerIdToSegmentsMap().size(), 4); + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + // A kept class dispatches all of its segments on every member, including the ones the member's own filter + // excluded, so worker 0 of the filtered side still scans a_seg0. + assertEquals(workerIdToPartitions(leafA, "a_seg"), + Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2), 3, Set.of(3))); + assertEquals(workerIdToServer(leafA), workerIdToServer(leafB)); } } @@ -914,6 +981,29 @@ public void testPlainPartitionedLeafWithPartitionWithOnlyDeferredSegmentsStillPr } } + @Test + public void testPlainPartitionedLeafWithEmptyPartitionPlansWhenNothingIsPruned() { + // Partition 3 holds no segment at all and the filter prunes nothing, so the pruning verdict is an EMPTY set rather + // than an absent one. It still has to be acted on: without a verdict every partition gets a worker, and the one + // with no segment has no server to place it on, which fails a query that plans perfectly well with 3 workers. + // Hence the verdict distinguishes "a filter ran and proved nothing" from "there was no filter". + QueryEnvironment queryEnvironment = + newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4, 1, List.of("seg0", "seg1", "seg2"), List.of(), 0, + false, Set.of(3), Set.of()); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; SELECT col2 FROM testTable " + + "/*+ tableOptions(partition_function='hashcode', partition_key='col1', partition_size='4') */ " + + "WHERE col1 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan); + assertNotNull(leaf); + assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 3); + assertEquals(assignedSegments(leaf), List.of("seg0", "seg1", "seg2")); + // Nothing was pruned, only skipped for holding nothing, so nothing is reported as pruned either. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + } + } + @Test public void testPartitionedLeafRejectsTableWithoutAnyPartition() { // An empty partition info map passes the "partitions must be a multiple of the hinted partition size" check @@ -1320,36 +1410,176 @@ public void testColocatedJoinReducesFanOutToPopulatedClasses() { } @Test - public void testColocatedJoinReducedGroupIgnoresBrokerPruning() { - // A reduced group's worker id is a position in the group's surviving class list, not a running counter over what a - // filter leaves behind, so broker pruning has to be off for its leaves. This shape is the only one that reaches - // that gate: a join written with is_colocated_by_join_keys marks its leaves pre-partitioned and is gated one step - // earlier (see testBrokerPruningPartitionedLeafSkippedForColocatedJoin), while a fact table joined with a - // replicated dimension table over an explicit local exchange is not marked pre-partitioned. + public void testBrokerPruningColocatedJoinDropsClassWhosePartitionsListNoSegment() { + // A partition whose entry lists no segment holds no rows, so its class is dropped even though the pruners proved + // nothing about any segment -- emptiness the planner can see for itself, not a pruning verdict. It is reported as + // such: numSegmentsPrunedByBroker stays 0 because no segment was pruned. + // + // The broker does not publish this shape today (a partition's entry is created together with its first segment), + // so this pins the behaviour rather than describing something reachable. It is what the golden physical plans + // record for a leaf whose fixture builds partitions this way. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).partitionsWithoutSegments(Set.of(2, 3)) + .survivingSegments(List.of("a_seg0", "a_seg1")), + new ColocatedTableSpec(4, false).partitionsWithoutSegments(Set.of(2, 3)) + .survivingSegments(List.of("b_seg0", "b_seg1"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo' AND " + + COLOCATED_TABLE_B + ".col2 = 'bar'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0), 1, Set.of(1))); + assertEquals(leafA.getWorkerIdToSegmentsMap().keySet(), Set.of(0, 1)); + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + } + } + + @Test + public void testBrokerPruningColocatedJoinAllPrunedKeepsEveryPopulatedClass() { + // Both sides prune everything, so the filtered union is empty and the group falls back to its populated classes. + // A group reduced to zero workers would leave a 1-to-1 exchange with no worker to wire on either side, and the + // server-side filter returns the same empty result from the unreduced plan anyway. So the most selective query + // gets the least reduction, which is the same trade the leaf-level fallback makes. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).survivingSegments(List.of()), + new ColocatedTableSpec(4, false).survivingSegments(List.of())); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + } + } + + @Test + public void testBrokerPruningColocatedJoinDisabledByQueryOption() { + // The existing useBrokerPruning switch is the kill switch for this too: with it off no routing query is built, so + // nothing is provably pruned and every populated class keeps its worker. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).survivingSegments(List.of("a_seg1")), + new ColocatedTableSpec(4, false).survivingSegments(List.of("b_seg1"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=false; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + } + } + + @Test + public void testBrokerPruningColocatedJoinFallsBackOnRoutingFailure() { + // Pruning is best-effort on this path too: a routing call that throws must leave the group with every populated + // class rather than fail a query that would otherwise plan. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).survivingSegments(List.of("a_seg1")), + new ColocatedTableSpec(4, false).survivingSegments(List.of("b_seg1")), TableType.OFFLINE, true); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B).getWorkerIdToSegmentsMap().size(), 4); + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + } + } + + @Test + public void testBrokerPruningColocatedJoinKeepsClassWithUnavailableSegment() { + // An unavailable segment was selected and not pruned, so the broker cannot prove its class empty. Class 2 keeps + // its worker on the strength of a_seg2 being unavailable rather than eliminated; only class 3, which both sides + // prune, is dropped. Reading "absent from the routing table" as "pruned" instead would silently drop the rows a + // transient outage hid. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).survivingSegments(List.of("a_seg0", "a_seg1")) + .unavailableSegments(List.of("a_seg2")), + new ColocatedTableSpec(4, false).survivingSegments(List.of("b_seg0", "b_seg1"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo' AND " + + COLOCATED_TABLE_B + ".col2 = 'bar'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2))); + assertEquals(workerIdToPartitions(leafB, "b_seg"), Map.of(0, Set.of(0), 1, Set.of(1), 2, Set.of(2))); + // Only class 3 dropped: one segment on each side. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 2); + } + } + + @Test + public void testBrokerPruningColocatedJoinStillPadsAMemberWithNoDataInASurvivingClass() { + // Emptiness padding and filter reduction have to compose. Table A holds nothing in class 3 while B does, and B's + // filter keeps class 3, so the class survives and A gets an empty worker for it -- placed on the server B picks + // for that class, so the exchange stays in process. Classes 1 and 2 are dropped because BOTH sides prune them. + // + // Padding is decided from unfiltered presence on purpose: a member that holds data the filter excludes dispatches + // it rather than being padded, so a surviving class always has the same segment list it would have unpruned. + QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( + new ColocatedTableSpec(4, false).emptyPartitions(Set.of(3)).survivingSegments(List.of("a_seg0")), + new ColocatedTableSpec(4, false).survivingSegments(List.of("b_seg0", "b_seg3"))); + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile( + "SET useBrokerPruning=true; " + colocatedJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo' AND " + + COLOCATED_TABLE_B + ".col2 = 'bar'")) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); + DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); + // Class list [0, 3]: worker 1 of A is the padded one, so it has an entry with no segment in it. + assertEquals(leafA.getWorkerIdToSegmentsMap().keySet(), Set.of(0, 1)); + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0))); + assertEquals(workerIdToPartitions(leafB, "b_seg"), Map.of(0, Set.of(0), 1, Set.of(3))); + Map workerIdToClass = new HashMap<>(); + mergeWorkerIdToClass(workerIdToClass, leafA, "a_seg", 4); + mergeWorkerIdToClass(workerIdToClass, leafB, "b_seg", 4); + assertEquals(workerIdToClass, Map.of(0, 0, 1, 3)); + // Worker 0 stands on the same server on both sides, which is what the 1-to-1 exchange is for. + assertEquals(workerIdToServer(leafA).get(0), workerIdToServer(leafB).get(0)); + // Worker 1 is A's padded one. It would rather land on B's server for class 3 to keep the exchange in process, + // but here A holds no data on that server at all -- partition 3 lives only on server 4 and A is empty there -- + // and a server with no data manager for the table fails the query outright, so it falls back to a server that + // provably hosts A and accepts one cross-server send. + assertEquals(workerIdToServer(leafB).get(1), getServerInstance("localhost", 4).getInstanceId()); + assertTrue(Set.of(getServerInstance("localhost", 1).getInstanceId(), + getServerInstance("localhost", 2).getInstanceId(), getServerInstance("localhost", 3).getInstanceId()) + .contains(workerIdToServer(leafA).get(1)), workerIdToServer(leafA).toString()); + // Classes 1 and 2 dropped, one segment each on each side. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 4); + } + } + + @Test + public void testColocatedJoinReducedGroupPrunesWholeClasses() { + // Emptiness reduction and filter reduction compose, on the one shape whose partitioned leaf is not marked + // pre-partitioned: a fact table joined with a replicated dimension table over an explicit local exchange. + // + // 8 partitions over 4 classes, so class c holds partitions c and c+4. The fact table's class 3 is empty, and the + // filter leaves only a_seg0 and a_seg4, which are both class 0. Emptiness drops class 3, the filter drops classes + // 1 and 2, and the class list ends up [0] -- a single worker holding the whole of class 0. // - // The fact table's class 3 is empty, so the group is reduced to [0, 1, 2], and the filter leaves only class 0. Were - // pruning left on, assignMultiplePartitionsPerWorker would find no segment for classes 1 and 2 and skip them - // WITHOUT consuming a worker id, leaving the leaf one worker while its class list still claimed three. + // This is the shape that used to be gated: reducing to [0] by dropping classes from the SHARED list keeps the + // worker id equal to its index in that list. Compacting at the leaf instead -- which is what + // assignMultiplePartitionsPerWorker would do if a per-leaf verdict ever reached it -- would leave the leaf one + // worker while its class list still claimed three. QueryEnvironment queryEnvironment = newColocatedJoinQueryEnvironment( new ColocatedTableSpec(8, false).emptyPartitions(Set.of(3, 7)).survivingSegments(List.of("a_seg0", "a_seg4")), new ColocatedTableSpec(8, false)); try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile("SET useBrokerPruning=true; " + replicatedDimensionJoinQuery(4) + " WHERE " + COLOCATED_TABLE_A + ".col2 = 'foo'")) { DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); - // Nothing was pruned, even though the filtered routing query would have dropped 2 of the 3 surviving classes. - assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0); + // Classes 1 and 2, two segments each. Class 3 is empty, not pruned, so it is not counted. + assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 4); DispatchablePlanFragment leafA = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_A); DispatchablePlanFragment leafB = leafFragmentForTable(dispatchableSubPlan, COLOCATED_TABLE_B); - // One worker per surviving class, holding both partitions of that class, and no worker dropped or padded. - assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0, 4), 1, Set.of(1, 5), 2, Set.of(2, 6))); - assertEquals(leafA.getWorkerIdToSegmentsMap().keySet(), Set.of(0, 1, 2)); + // One worker for the one surviving class, holding both of its partitions, and no worker dropped or padded. + assertEquals(workerIdToPartitions(leafA, "a_seg"), Map.of(0, Set.of(0, 4))); + assertEquals(leafA.getWorkerIdToSegmentsMap().keySet(), Set.of(0)); Map workerIdToClass = new HashMap<>(); mergeWorkerIdToClass(workerIdToClass, leafA, "a_seg", 4); - assertEquals(workerIdToClass, Map.of(0, 0, 1, 1, 2, 2)); + assertEquals(workerIdToClass, Map.of(0, 0)); // The replicated leaf and the join derive their workers from the fact leaf, so they follow it class for class. assertEquals(leafB.getWorkerIdToSegmentsMap().keySet(), leafA.getWorkerIdToSegmentsMap().keySet()); assertEquals(workerIdToServer(leafB), workerIdToServer(leafA)); - assertEquals(joinFragment(dispatchableSubPlan).getWorkerMetadataList().size(), 3); + assertEquals(joinFragment(dispatchableSubPlan).getWorkerMetadataList().size(), 1); } } @@ -1516,7 +1746,7 @@ public void testCheckLeafWorkerAssignmentRejectsSparseWorkerIds() { Map>> segmentsMap = Map.of(0, offlineSegments("seg0"), 2, offlineSegments("seg2")); IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", serverMap, segmentsMap)); + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, serverMap, segmentsMap)); assertTrue(e.getMessage().contains("Missing server instance for worker: 1"), e.getMessage()); } @@ -1526,7 +1756,7 @@ public void testCheckLeafWorkerAssignmentRejectsKeySetMismatch() { Map>> segmentsMap = Map.of(0, offlineSegments("seg0"), 5, offlineSegments("seg5")); IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", serverMap, segmentsMap)); + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, serverMap, segmentsMap)); assertTrue(e.getMessage().contains("Missing segments for worker: 1"), e.getMessage()); } @@ -1535,7 +1765,7 @@ public void testCheckLeafWorkerAssignmentRejectsNullSegmentList() { Map> nullList = new HashMap<>(); nullList.put(TableType.OFFLINE.name(), null); IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, Map.of(0, queryServerInstance(1)), Map.of(0, nullList))); assertTrue(e.getMessage().contains("Null segment list for table type: OFFLINE"), e.getMessage()); } @@ -1545,7 +1775,7 @@ public void testCheckLeafWorkerAssignmentRejectsEmptyTableTypeMap() { // The server splits the request on the number of entries in this map, so a worker with no table type at all would // produce no server request. IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, Map.of(0, queryServerInstance(1)), Map.of(0, Map.of()))); assertTrue(e.getMessage().contains("Expected 1 or 2 table types for worker: 0, got: 0"), e.getMessage()); } @@ -1557,7 +1787,7 @@ public void testCheckLeafWorkerAssignmentRejectsThreeTableTypeMap() { threeTypes.put(TableType.REALTIME.name(), List.of()); threeTypes.put("HYBRID", List.of()); IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, Map.of(0, queryServerInstance(1)), Map.of(0, threeTypes))); assertTrue(e.getMessage().contains("Expected 1 or 2 table types for worker: 0, got: 3"), e.getMessage()); } @@ -1566,11 +1796,30 @@ public void testCheckLeafWorkerAssignmentRejectsThreeTableTypeMap() { public void testCheckLeafWorkerAssignmentRejectsUnknownTableType() { // The server resolves one table data manager per key in this map, and reports a missing table for an unknown one. IllegalStateException e = expectThrows(IllegalStateException.class, - () -> WorkerManager.checkLeafWorkerAssignment("testTable", Map.of(0, queryServerInstance(1)), + () -> WorkerManager.checkLeafWorkerAssignment("testTable", null, Map.of(0, queryServerInstance(1)), Map.of(0, Map.of("HYBRID", List.of())))); assertTrue(e.getMessage().contains("Unexpected table type: HYBRID for worker: 0"), e.getMessage()); } + @Test + public void testCheckLeafWorkerAssignmentRejectsFewerWorkersThanPartitionClasses() { + // A worker id of a colocated leaf IS an index into the group's shared class list, so a member that assigned fewer + // workers than the list has renumbered every class after the gap. checkPartitionClassAgreement cannot see it: it + // compares the shared array against itself, and both sides still hold the same instance. + IllegalStateException e = expectThrows(IllegalStateException.class, + () -> WorkerManager.checkLeafWorkerAssignment("testTable", new int[]{0, 2, 5}, + Map.of(0, queryServerInstance(1), 1, queryServerInstance(2)), + Map.of(0, offlineSegments("seg0"), 1, offlineSegments("seg2")))); + assertTrue(e.getMessage().contains("Got 2 workers for partition classes: [0, 2, 5]"), e.getMessage()); + } + + @Test + public void testCheckLeafWorkerAssignmentAcceptsOneWorkerPerPartitionClass() { + WorkerManager.checkLeafWorkerAssignment("testTable", new int[]{0, 2}, + Map.of(0, queryServerInstance(1), 1, queryServerInstance(2)), + Map.of(0, offlineSegments("seg0"), 1, offlineSegments("seg2"))); + } + @Test public void testCheckLeafWorkerAssignmentAcceptsHybridAndEmptySegmentWorkers() { // The two shapes the partitioned assignment produces: a hybrid worker with both table types, and one with a single @@ -1578,7 +1827,7 @@ public void testCheckLeafWorkerAssignmentAcceptsHybridAndEmptySegmentWorkers() { Map> hybridSegments = new HashMap<>(); hybridSegments.put(TableType.OFFLINE.name(), List.of("segO0")); hybridSegments.put(TableType.REALTIME.name(), List.of("segR0")); - WorkerManager.checkLeafWorkerAssignment("testTable", + WorkerManager.checkLeafWorkerAssignment("testTable", null, Map.of(0, queryServerInstance(1), 1, queryServerInstance(2)), Map.of(0, hybridSegments, 1, Map.of(TableType.OFFLINE.name(), new ArrayList<>()))); } @@ -1694,6 +1943,13 @@ private static QueryEnvironment newColocatedJoinQueryEnvironment(ColocatedTableS /// tables are registered under, so that the realtime-only shape can be covered too. private static QueryEnvironment newColocatedJoinQueryEnvironment(ColocatedTableSpec specA, ColocatedTableSpec specB, TableType tableType) { + return newColocatedJoinQueryEnvironment(specA, specB, tableType, false); + } + + /// Same as [#newColocatedJoinQueryEnvironment(ColocatedTableSpec, ColocatedTableSpec, TableType)], with a routing + /// manager that throws on every routing call, to exercise the best-effort fallback. + private static QueryEnvironment newColocatedJoinQueryEnvironment(ColocatedTableSpec specA, ColocatedTableSpec specB, + TableType tableType, boolean throwOnRouting) { int numServers = 4; ServerInstance[] servers = new ServerInstance[numServers]; Map enabledServers = new HashMap<>(); @@ -1710,13 +1966,15 @@ private static QueryEnvironment newColocatedJoinQueryEnvironment(ColocatedTableS colocatedTablePartitionInfo(tableBWithType, "b_seg", servers, specB)); Map routingTableByTable = new HashMap<>(); if (specA._survivingSegments != null) { - routingTableByTable.put(tableAWithType, colocatedRoutingTable(servers, "a_seg", specA._survivingSegments)); + routingTableByTable.put(tableAWithType, + colocatedRoutingTable(servers, "a_seg", specA._survivingSegments, specA._unavailableSegments)); } if (specB._survivingSegments != null) { - routingTableByTable.put(tableBWithType, colocatedRoutingTable(servers, "b_seg", specB._survivingSegments)); + routingTableByTable.put(tableBWithType, + colocatedRoutingTable(servers, "b_seg", specB._survivingSegments, specB._unavailableSegments)); } PartitionedRoutingManager routingManager = - new PartitionedRoutingManager(enabledServers, partitionInfoByTable, routingTableByTable, false); + new PartitionedRoutingManager(enabledServers, partitionInfoByTable, routingTableByTable, throwOnRouting); Map tableNameMap = new HashMap<>(); tableNameMap.put(tableAWithType, tableAWithType); @@ -1764,8 +2022,9 @@ private static TablePartitionReplicatedServersInfo colocatedTablePartitionInfo(S : Set.of(servers[p % servers.length].getInstanceId()); } // Mutable, like the lists the broker publishes: the assignment must hand out a copy rather than this instance. - partitionInfoMap[p] = new TablePartitionReplicatedServersInfo.PartitionInfo(partitionServers, - new ArrayList<>(List.of(segmentPrefix + p))); + List segments = spec._partitionsWithoutSegments.contains(p) ? new ArrayList<>() + : new ArrayList<>(List.of(segmentPrefix + p)); + partitionInfoMap[p] = new TablePartitionReplicatedServersInfo.PartitionInfo(partitionServers, segments); } } return new TablePartitionReplicatedServersInfo(tableNameWithType, "col1", "Hashcode", numPartitions, @@ -1775,7 +2034,7 @@ private static TablePartitionReplicatedServersInfo colocatedTablePartitionInfo(S /// Buckets the given surviving segments onto the server hosting their partition (partition `p` lives on server /// `p % 4`), i.e. builds what the routing manager returns for one colocated table's filtered routing query. private static RoutingTable colocatedRoutingTable(ServerInstance[] servers, String segmentPrefix, - List survivingSegments) { + List survivingSegments, List unavailableSegments) { Map> serverToSegmentList = new HashMap<>(); for (String segment : survivingSegments) { int partition = Integer.parseInt(segment.substring(segmentPrefix.length())); @@ -1784,7 +2043,7 @@ private static RoutingTable colocatedRoutingTable(ServerInstance[] servers, Stri Map serverToSegments = new HashMap<>(); serverToSegmentList.forEach((server, segments) -> serverToSegments.put(server, new SegmentsToQuery(segments, List.of()))); - return new RoutingTable(serverToSegments, List.of(), 0); + return new RoutingTable(serverToSegments, new ArrayList<>(unavailableSegments), 0); } /// How one side of a colocated join is laid out, for [#newColocatedJoinQueryEnvironment(ColocatedTableSpec, @@ -1809,6 +2068,8 @@ private static class ColocatedTableSpec { /// null rather than a silently empty answer. @Nullable List _survivingSegments; + List _unavailableSegments = List.of(); + Set _partitionsWithoutSegments = Set.of(); ColocatedTableSpec(int numPartitions, boolean everyServerHostsEveryPartition) { _numPartitions = numPartitions; @@ -1844,6 +2105,21 @@ ColocatedTableSpec survivingSegments(List survivingSegments) { _survivingSegments = survivingSegments; return this; } + + /// Partitions that get an entry listing no segment, as opposed to no entry at all. The broker never publishes + /// this shape today -- a partition's entry is created together with its first segment -- so it exists only to pin + /// what the assignment does if one ever appears. + ColocatedTableSpec partitionsWithoutSegments(Set partitionsWithoutSegments) { + _partitionsWithoutSegments = partitionsWithoutSegments; + return this; + } + + /// Segments the routing table reports as unavailable. They were selected and not pruned, so the broker cannot + /// prove them empty and their partition class has to survive. + ColocatedTableSpec unavailableSegments(List unavailableSegments) { + _unavailableSegments = unavailableSegments; + return this; + } } /// Returns the only fragment below the reduce stage with neither segments nor children, i.e. the join stage. @@ -2054,6 +2330,18 @@ private static QueryEnvironment newPartitionedQueryEnvironment(int[] serverIdxPe int replicasPerPartition, List survivingSegments, List unavailableSegments, int reportedPrunedByRouting, boolean throwOnRouting, Set emptyPartitions, Set partitionsWithOnlyDeferredSegments) { + return newPartitionedQueryEnvironment(serverIdxPerPartition, numServers, replicasPerPartition, survivingSegments, + unavailableSegments, reportedPrunedByRouting, throwOnRouting, emptyPartitions, + partitionsWithOnlyDeferredSegments, Map.of()); + } + + /// Same again, with the surviving segments a leaf sees when its own filter carries a given string literal. Lets two + /// leaves scanning the SAME table be given different verdicts, which is the only way to tell a per-leaf pruning + /// verdict from a per-table one. + private static QueryEnvironment newPartitionedQueryEnvironment(int[] serverIdxPerPartition, int numServers, + int replicasPerPartition, List survivingSegments, List unavailableSegments, + int reportedPrunedByRouting, boolean throwOnRouting, Set emptyPartitions, + Set partitionsWithOnlyDeferredSegments, Map> survivingSegmentsByFilterLiteral) { int numPartitions = serverIdxPerPartition.length; ServerInstance[] servers = new ServerInstance[numServers]; Map enabledServers = new HashMap<>(); @@ -2095,6 +2383,7 @@ private static QueryEnvironment newPartitionedQueryEnvironment(int[] serverIdxPe PartitionedRoutingManager routingManager = new PartitionedRoutingManager(enabledServers, Map.of(PARTITIONED_TABLE_OFFLINE, tablePartitionInfo), Map.of(PARTITIONED_TABLE_OFFLINE, prunedRoutingTable), throwOnRouting); + survivingSegmentsByFilterLiteral.forEach(routingManager::survivingSegmentsForFilterLiteral); Map tableNameMap = new HashMap<>(); tableNameMap.put(PARTITIONED_TABLE_OFFLINE, PARTITIONED_TABLE_OFFLINE); @@ -2494,6 +2783,9 @@ private static class PartitionedRoutingManager implements RoutingManager { private final Map _enabledServers; private final Map _partitionInfoByTable; private final Map _routingTableByTable; + /// Surviving segments for a leaf whose filter carries the given string literal, which is how a test gives two + /// leaves scanning ONE table two different verdicts. Falls back to the per-table routing table when absent. + private final Map> _survivingSegmentsByFilterLiteral = new HashMap<>(); private final boolean _throwOnRouting; @Nullable private final TimeBoundaryInfo _timeBoundaryInfo; @@ -2541,8 +2833,67 @@ public RoutingTable getRoutingTable(BrokerRequest brokerRequest, String tableNam @Nullable @Override public List getSegments(BrokerRequest brokerRequest) { - TablePartitionReplicatedServersInfo partitionInfo = - _partitionInfoByTable.get(brokerRequest.getQuerySource().getTableName()); + return allSegments(brokerRequest.getQuerySource().getTableName()); + } + + /// Derives the pruned set the way `BaseBrokerRoutingManager` does -- everything selection offered minus what the + /// pruners kept -- from the same configured routing table. Unavailable segments were selected and not pruned, so + /// they stay out of it and keep their partition alive. A table with no routing table registered proves nothing, + /// which is what an unexpected routing call should look like. + @Override + public Set getPrunedSegments(BrokerRequest brokerRequest) { + if (_throwOnRouting) { + throw new RuntimeException("Simulated routing failure"); + } + validatePrunableFilter(brokerRequest.getPinotQuery().getFilterExpression()); + String tableNameWithType = brokerRequest.getQuerySource().getTableName(); + Set prunedSegments = new HashSet<>(allSegments(tableNameWithType)); + String filterLiteral = firstStringLiteral(brokerRequest.getPinotQuery().getFilterExpression()); + List survivingSegments = + filterLiteral != null ? _survivingSegmentsByFilterLiteral.get(filterLiteral) : null; + if (survivingSegments != null) { + prunedSegments.removeAll(survivingSegments); + return prunedSegments; + } + RoutingTable routingTable = _routingTableByTable.get(tableNameWithType); + if (routingTable == null) { + return Set.of(); + } + for (SegmentsToQuery segmentsToQuery : routingTable.getServerInstanceToSegmentsMap().values()) { + prunedSegments.removeAll(segmentsToQuery.getSegments()); + } + prunedSegments.removeAll(routingTable.getUnavailableSegments()); + return prunedSegments; + } + + PartitionedRoutingManager survivingSegmentsForFilterLiteral(String filterLiteral, List segments) { + _survivingSegmentsByFilterLiteral.put(filterLiteral, segments); + return this; + } + + /// The first string literal in the filter, which the tests use as a stand-in for "which filter is this". + @Nullable + private static String firstStringLiteral(@Nullable Expression expression) { + if (expression == null) { + return null; + } + if (expression.getLiteral() != null && expression.getLiteral().isSetStringValue()) { + return expression.getLiteral().getStringValue(); + } + Function function = expression.getFunctionCall(); + if (function != null) { + for (Expression operand : function.getOperands()) { + String literal = firstStringLiteral(operand); + if (literal != null) { + return literal; + } + } + } + return null; + } + + private List allSegments(String tableNameWithType) { + TablePartitionReplicatedServersInfo partitionInfo = _partitionInfoByTable.get(tableNameWithType); if (partitionInfo == null) { return List.of(); } diff --git a/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json b/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json index 20ef9aed09ee..5f6c4b84e9d5 100644 --- a/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json +++ b/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json @@ -421,21 +421,15 @@ "\n ├── [2]@localhost:1|[1] PROJECT", "\n │ └── [2]@localhost:1|[1] TABLE SCAN (a) null", "\n └── [2]@localhost:1|[1] MAIL_RECEIVE(BROADCAST_DISTRIBUTED)", - "\n ├── [3]@localhost:2|[2] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2, 3]} (Subtree Omitted)", - "\n ├── [3]@localhost:2|[3] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2, 3]} (Subtree Omitted)", - "\n ├── [3]@localhost:1|[0] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2, 3]} (Subtree Omitted)", - "\n └── [3]@localhost:1|[1] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2, 3]}", - "\n └── [3]@localhost:1|[1] PROJECT", - "\n └── [3]@localhost:1|[1] FILTER", - "\n └── [3]@localhost:1|[1] AGGREGATE_FINAL", - "\n └── [3]@localhost:1|[1] MAIL_RECEIVE(HASH_DISTRIBUTED)", - "\n ├── [4]@localhost:2|[2] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[3]@localhost:2|[2]} (Subtree Omitted)", - "\n ├── [4]@localhost:2|[3] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[3]@localhost:2|[3]} (Subtree Omitted)", - "\n ├── [4]@localhost:1|[0] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[3]@localhost:1|[0]} (Subtree Omitted)", - "\n └── [4]@localhost:1|[1] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[3]@localhost:1|[1]}", - "\n └── [4]@localhost:1|[1] AGGREGATE_LEAF", - "\n └── [4]@localhost:1|[1] FILTER", - "\n └── [4]@localhost:1|[1] TABLE SCAN (b) null", + "\n └── [3]@localhost:1|[0] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2, 3]}", + "\n └── [3]@localhost:1|[0] PROJECT", + "\n └── [3]@localhost:1|[0] FILTER", + "\n └── [3]@localhost:1|[0] AGGREGATE_FINAL", + "\n └── [3]@localhost:1|[0] MAIL_RECEIVE(HASH_DISTRIBUTED)", + "\n └── [4]@localhost:1|[0] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[3]@localhost:1|[0]}", + "\n └── [4]@localhost:1|[0] AGGREGATE_LEAF", + "\n └── [4]@localhost:1|[0] FILTER", + "\n └── [4]@localhost:1|[0] TABLE SCAN (b) null", "\n" ] }, diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 7357c7c4be5e..24b0816b2ae6 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -636,6 +636,11 @@ public static class Broker { /// Separated from [#CONFIG_OF_USE_BROKER_PRUNING] so the two paths can be rolled out independently; both /// default to enabled now that all logical-planner leaf paths (non-partitioned, partitioned, logical tables) /// support broker pruning. Actual pruning still requires segment pruners to be configured on the table. + /// + /// On a colocated join this governs more than which segments are dispatched: a partition class that every member + /// of the colocated group prunes away is dropped from the group's shared class list, so the leaves and the stages + /// derived from them run fewer workers and the query is dispatched to fewer servers. Turning it off restores one + /// worker per populated class. public static final String CONFIG_OF_LOGICAL_PLANNER_USE_BROKER_PRUNING = "pinot.broker.multistage.logical.planner.use.broker.pruning"; public static final boolean DEFAULT_LOGICAL_PLANNER_USE_BROKER_PRUNING = true;