This is an automated email from the ASF dual-hosted git repository.
yashmayya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new da15d13df5e Support broker segment pruning for colocated joins (#19174)
da15d13df5e is described below
commit da15d13df5e7640dc183e9f12c30c099f3ea7f3c
Author: Yash Mayya <[email protected]>
AuthorDate: Tue Aug 25 14:54:15 2026 -0400
Support broker segment pruning for colocated joins (#19174)
---
.../routing/manager/BaseBrokerRoutingManager.java | 86 +++-
.../manager/MultiClusterRoutingManager.java | 56 +++
.../routing/manager/BrokerRoutingManagerTest.java | 168 +++++++-
.../manager/MultiClusterRoutingManagerTest.java | 86 ++++
.../apache/pinot/core/routing/RoutingManager.java | 32 ++
.../custom/ColocatedJoinEmptyPartitionTest.java | 176 ++++++++-
.../planner/physical/DispatchablePlanContext.java | 9 +
.../query/routing/ColocationGroupAnalyzer.java | 19 +-
.../query/routing/PlanNodeRoutingQueryBuilder.java | 21 +
.../apache/pinot/query/routing/WorkerManager.java | 356 +++++++++++++----
.../query/routing/ColocationGroupAnalyzerTest.java | 15 +-
.../routing/PlanNodeRoutingQueryBuilderTest.java | 56 +++
.../pinot/query/routing/WorkerManagerTest.java | 431 +++++++++++++++++++--
.../resources/queries/ExplainPhysicalPlans.json | 24 +-
.../apache/pinot/spi/utils/CommonConstants.java | 5 +
15 files changed, 1349 insertions(+), 191 deletions(-)
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java
index 631b67f63cc..269243cb5b6 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 @@ package org.apache.pinot.broker.routing.manager;
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 abstract class BaseBrokerRoutingManager
implements RoutingManager, Cluste
return routingEntry.getSegments(brokerRequest, samplerName);
}
+ @Nullable
+ @Override
+ public Set<String> 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 @@ public abstract class BaseBrokerRoutingManager
implements RoutingManager, Cluste
}
}
- 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<String> selectedSegments = segmentSelector.select(brokerRequest);
- int numTotalSelectedSegments = selectedSegments.size();
+ Set<String> 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<String> 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 @@ public abstract class BaseBrokerRoutingManager
implements RoutingManager, Cluste
}
List<String> getSegments(BrokerRequest brokerRequest, @Nullable String
samplerName) {
- SamplerInfo samplerInfo = getSamplerInfo(samplerName);
- SegmentSelector segmentSelector = samplerInfo != null ?
samplerInfo._segmentSelector : _segmentSelector;
- Set<String> 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<String> 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<String> 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<String> _selected;
+ final Set<String> _surviving;
+
+ SelectedSegments(Set<String> selected, Set<String> 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 312d7b0b7d1..9f04a4ca3d4 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 class MultiClusterRoutingManager implements
RoutingManager {
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<String> getPrunedSegments(BrokerRequest brokerRequest) {
+ String tableNameWithType = brokerRequest.getQuerySource().getTableName();
+ Set<String> 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<String> intersectPrunedSegments(@Nullable Set<String> 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<String> 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 32be5f024cf..75c2f3570ba 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.tablesampler.TableSampler;
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.doThrow;
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 class BrokerRoutingManagerTest {
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<Set<String>> 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<String> 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<String> 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<String> droppedSegments = Set.of(segments);
+ SegmentPruner segmentPruner = mock(SegmentPruner.class);
+ when(segmentPruner.prune(any(), any())).thenAnswer(invocation -> {
+ Set<String> candidateSegments = invocation.getArgument(1);
+ if (droppedSegments.stream().noneMatch(candidateSegments::contains)) {
+ return candidateSegments;
+ }
+ Set<String> survivingSegments = new HashSet<>(candidateSegments);
+ survivingSegments.removeAll(droppedSegments);
+ return survivingSegments;
+ });
+ return segmentPruner;
+ }
+
private static Object createRoutingEntry(String tableNameWithType,
TimeBoundaryManager timeBoundaryManager,
SegmentPartitionMetadataManager partitionMetadataManager, Map<String, ?>
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<String, ?>
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<SegmentPruner> 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<SegmentPruner> segmentPruners, InstanceSelector instanceSelector,
TimeBoundaryManager timeBoundaryManager,
+ SegmentPartitionMetadataManager partitionMetadataManager, Map<String, ?>
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 @@ public class BrokerRoutingManagerTest {
Map.class, boolean.class);
constructor.setAccessible(true);
return constructor.newInstance(tableNameWithType, "/IDEALSTATES/" +
tableNameWithType,
- "/EXTERNALVIEW/" + tableNameWithType, mock(SegmentPreSelector.class),
mock(SegmentSelector.class),
- Collections.<SegmentPruner>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 8d8cdb34c26..0ea105a56e1 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.never;
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 class MultiClusterRoutingManagerTest {
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<String> 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<String>
segments) {
Map<ServerInstance, SegmentsToQuery> 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 9f33b868b51..53198ede031 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 @@ public interface RoutingManager {
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<String> 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 797ff14d35f..cfc96bbc721 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.List;
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.avro.generic.GenericDatumWriter;
import org.apache.commons.io.FileUtils;
import org.apache.pinot.integration.tests.ClusterIntegrationTestUtils;
import org.apache.pinot.spi.config.table.ColumnPartitionConfig;
+import org.apache.pinot.spi.config.table.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 @@ import static org.testng.Assert.assertTrue;
/// 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<Integer> FILTERED_KEYS = List.of(1, 2);
+ /// A key in a partition neither table populates, so the filter matches
nothing anywhere.
+ private static final List<Integer> UNMATCHED_KEYS = List.of(5);
+
private static final int LEFT_METRIC_MULTIPLIER = 10;
private static final int RIGHT_METRIC_MULTIPLIER = 100;
@@ -235,31 +249,143 @@ public class ColocatedJoinEmptyPartitionTest extends
CustomDataQueryClusterInteg
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<JsonNode> 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<Integer> 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<Integer> 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<Integer> keptClassesFor(List<Integer> 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<List<Long>> expectedFilteredRows(List<Integer> keys) {
+ List<List<Long>> 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) {
@@ -315,6 +441,22 @@ public class ColocatedJoinEmptyPartitionTest extends
CustomDataQueryClusterInteg
}
}
+ /// 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<JsonNode> 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<JsonNode>
leafStageSends) {
JsonNode children = node.get("children");
if ("MAILBOX_SEND".equals(node.path("type").asText()) && children != null
&& children.size() == 1 && "LEAF".equals(
@@ -406,6 +548,10 @@ public class ColocatedJoinEmptyPartitionTest extends
CustomDataQueryClusterInteg
.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 f50a3db9e1c..585576dfa99 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<Integer, DispatchablePlanMetadata>
_dispatchablePlanMetadataMap = new HashMap<>();
private final Map<Integer, PlanNode> _dispatchablePlanStageRootMap = new
HashMap<>();
private final Map<String, WorkerManager.PartitionTableInfo>
_partitionTableInfoCache = new HashMap<>();
+ private final Map<Integer, Set<String>> _prunedSegmentsCache = new
HashMap<>();
private long _numSegmentsPrunedByBroker;
private int _leafStagesAssigned;
private int _leafStagesEmpty;
@@ -141,6 +142,14 @@ public class DispatchablePlanContext {
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<Integer, Set<String>> 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 88b43fb61e7..0d5508f6fc9 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 @@ class ColocationGroupAnalyzer {
if (members.size() < 2) {
return null;
}
- List<Integer> partitionedLeafFragmentIds = new ArrayList<>();
+ List<PlanFragment> partitionedLeafFragments = new ArrayList<>();
int partitionSize = -1;
int partitionParallelism = -1;
String partitionFunction = null;
@@ -186,7 +186,7 @@ class ColocationGroupAnalyzer {
return null;
}
String leafPartitionFunction = hints.getHintedPartitionFunction();
- if (partitionedLeafFragmentIds.isEmpty()) {
+ if (partitionedLeafFragments.isEmpty()) {
partitionSize = hints.getPartitionSize();
partitionParallelism = hints.getPartitionParallelism();
partitionFunction = leafPartitionFunction;
@@ -195,12 +195,12 @@ class ColocationGroupAnalyzer {
|| !isSamePartitionFunction(partitionFunction,
leafPartitionFunction)) {
return null;
}
- partitionedLeafFragmentIds.add(fragmentId);
+ partitionedLeafFragments.add(fragment);
}
- if (partitionedLeafFragmentIds.isEmpty()) {
+ if (partitionedLeafFragments.isEmpty()) {
return null;
}
- return new ColocationGroup(partitionSize, partitionedLeafFragmentIds);
+ return new ColocationGroup(partitionSize, partitionedLeafFragments);
}
/// Compares two `partition_function` hints the way the rest of the engine
compares partition function names:
@@ -242,12 +242,13 @@ class ColocationGroupAnalyzer {
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<Integer> _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<PlanFragment> _partitionedLeafFragments;
- ColocationGroup(int partitionSize, List<Integer>
partitionedLeafFragmentIds) {
+ ColocationGroup(int partitionSize, List<PlanFragment>
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 d7b32f3e4b3..e946e4b68b0 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 class PlanNodeRoutingQueryBuilder {
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.
+ static boolean canBuildRoutingQuery(PlanNode leafStageRoot) {
+ List<PlanNode> 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<PlanNode>
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 a6f9e3e3303..917cbab89fe 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 class WorkerManager {
/// 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 @@ public class WorkerManager {
for (ColocationGroupAnalyzer.ColocationGroup group :
ColocationGroupAnalyzer.findReducibleGroups(rootFragment,
metadataMap)) {
int numWorkers = group._partitionSize;
- List<Integer> memberFragmentIds = group._partitionedLeafFragmentIds;
- // The servers each member can scan each class on, in the same order as
the member fragment ids.
- List<List<Set<String>>> memberClassServers = new
ArrayList<>(memberFragmentIds.size());
+ List<PlanFragment> memberFragments = group._partitionedLeafFragments;
+ List<Integer> memberFragmentIds = new
ArrayList<>(memberFragments.size());
+ // The servers each member can scan each class on, in the same order as
the member fragments.
+ List<List<Set<String>>> memberClassServers = new
ArrayList<>(memberFragments.size());
+ // The partition layouts, same order again, kept only to count what a
filter-dropped class cost each member.
+ List<PartitionInfo[]> 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<Set<String>> classServers =
collectClassServers(partitionTableInfo._partitionInfoMap, numWorkers);
+ List<Set<String>> classServers = collectClassServers(partitionInfoMap,
numWorkers);
+ Set<String> 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<Integer, Map<Integer, Set<String>>> 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,55 @@ public class WorkerManager {
}
}
+ 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
+ /// segment missing from it counts as matching. A hybrid partition that
lists segments for one table type and none
+ /// for the other is therefore decided by the type that has them.
+ ///
+ /// This reads "holds data" more strictly than [#collectClassServers], which
counts every partition that has an entry
+ /// at all. A partition whose entry lists no segment is empty here and its
class goes unmarked, while
+ /// `collectClassServers` still reports it as populated. The two can only
disagree on a shape the broker does not
+ /// publish -- a partition's entry is created together with its first
segment -- so the difference shows up in test
+ /// fixtures rather than on a live table, and it is emptiness the planner
sees for itself rather than a pruning
+ /// verdict, so nothing is counted as pruned for it.
+ private static void markClassesWithMatchingData(PartitionInfo[]
partitionInfoMap, int numWorkers,
+ Set<String> 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 +789,14 @@ public class WorkerManager {
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<String> 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 +970,86 @@ public class WorkerManager {
}
}
+ /// 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<String> getPrunedSegments(PlanFragment fragment, String
tableName, DispatchablePlanContext context) {
+ // Not computeIfAbsent: null is a meaningful answer here and would be
recomputed on every call.
+ Map<Integer, Set<String>> prunedSegmentsCache =
context.getPrunedSegmentsCache();
+ Integer fragmentId = fragment.getFragmentId();
+ if (prunedSegmentsCache.containsKey(fragmentId)) {
+ return prunedSegmentsCache.get(fragmentId);
+ }
+ Set<String> prunedSegments = computePrunedSegments(fragment, tableName,
context);
+ prunedSegmentsCache.put(fragmentId, prunedSegments);
+ return prunedSegments;
+ }
+
+ @Nullable
+ private Set<String> 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<String> offlinePrunedSegments =
getPrunedSegmentsHelper(routingPinotQuery, TableType.OFFLINE);
+ Set<String> realtimePrunedSegments =
getPrunedSegmentsHelper(routingPinotQuery, TableType.REALTIME);
+ if (offlinePrunedSegments.isEmpty()) {
+ return realtimePrunedSegments;
+ }
+ if (realtimePrunedSegments.isEmpty()) {
+ return offlinePrunedSegments;
+ }
+ Set<String> 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<String> getPrunedSegmentsHelper(PinotQuery pinotQuery) {
+ Set<String> prunedSegments =
+
_routingManager.getPrunedSegments(CalciteSqlCompiler.convertToBrokerRequest(pinotQuery));
+ return prunedSegments != null ? prunedSegments : Set.of();
+ }
+
+ private Set<String> 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 +1061,9 @@ public class WorkerManager {
if (!useBrokerPruning) {
return null;
}
+ if (!PlanNodeRoutingQueryBuilder.canBuildRoutingQuery(leafStageRoot)) {
+ return null;
+ }
try {
PinotQuery pinotQuery =
PlanNodeRoutingQueryBuilder.createPinotQueryForRouting(tableName,
leafStageRoot, false);
Map<String, String> queryOptions =
context.getPlannerContext().getOptions();
@@ -911,10 +1097,7 @@ public class WorkerManager {
@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 +1320,7 @@ public class WorkerManager {
/// 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<String> prunedSegments) {
// when partition key exist, we assign workers for leaf-stage in
partitioned fashion.
String partitionKey = partitionHints.getPartitionKey();
assert partitionKey != null;
@@ -1172,8 +1354,7 @@ public class WorkerManager {
collectHostingServers(partitionInfoMap)) : null;
// Broker pruning: the partitions to keep (null means keep all).
Partitions absent from the set are skipped below.
- Set<Integer> partitionsToKeep =
- computePartitionsToKeep(routingPinotQuery, metadata,
context.getRequestId(), partitionInfoMap);
+ Set<Integer> partitionsToKeep = computePartitionsToKeep(prunedSegments,
metadata, partitionInfoMap);
if (partitionsToKeep != null) {
long numSegmentsPrunedByBroker = countPrunedSegments(partitionInfoMap,
partitionsToKeep);
if (numSegmentsPrunedByBroker > 0) {
@@ -1202,35 +1383,42 @@ public class WorkerManager {
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 +1426,15 @@ public class WorkerManager {
/// 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<Integer> computePartitionsToKeep(@Nullable PinotQuery
routingPinotQuery,
- DispatchablePlanMetadata metadata, long requestId, PartitionInfo[]
partitionInfoMap) {
- if (routingPinotQuery == null || routingPinotQuery.getFilterExpression()
== null || metadata.isPrePartitioned()
- || metadata.getPartitionClassIds() != null) {
- return null;
- }
- Map<String, RoutingTable> 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<Integer> computePartitionsToKeep(@Nullable Set<String>
prunedSegments,
+ DispatchablePlanMetadata metadata, PartitionInfo[] partitionInfoMap) {
+ if (prunedSegments == null || metadata.isPrePartitioned() ||
metadata.getPartitionClassIds() != null) {
return null;
}
- if (routingTableMap.isEmpty()) {
- return null;
- }
- Set<String> 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<Integer> 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 +1442,28 @@ public class WorkerManager {
return partitionsToKeep.isEmpty() ? null : partitionsToKeep;
}
- private static boolean containsAny(@Nullable List<String> segments,
Set<String> 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<String> prunedSegments) {
+ return allPruned(partitionInfo._offlineSegments, prunedSegments)
+ && allPruned(partitionInfo._realtimeSegments, prunedSegments);
+ }
+
+ private static boolean allPruned(@Nullable List<String> segments,
Set<String> 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<Integer> partitionsToKeep) {
long numPrunedSegments = 0;
for (int i = 0; i < partitionInfoMap.length; i++) {
@@ -1508,12 +1684,20 @@ public class WorkerManager {
/// 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<Integer, QueryServerInstance> workerIdToServerInstanceMap,
Map<Integer, Map<String, List<String>>> 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 3a779db8731..5908900e33c 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.HashMap;
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 class ColocationGroupAnalyzerTest {
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 class ColocationGroupAnalyzerTest {
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 class ColocationGroupAnalyzerTest {
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 class ColocationGroupAnalyzerTest {
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 class ColocationGroupAnalyzerTest {
// 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<Integer>
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/PlanNodeRoutingQueryBuilderTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilderTest.java
index 125249b64ea..0b0e0fc1491 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilderTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilderTest.java
@@ -27,11 +27,15 @@ import
org.apache.pinot.query.planner.plannode.AggregateNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
import org.apache.pinot.query.planner.plannode.PlanNode;
import org.apache.pinot.query.planner.plannode.ProjectNode;
+import org.apache.pinot.query.planner.plannode.SetOpNode;
import org.apache.pinot.query.planner.plannode.TableScanNode;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
public class PlanNodeRoutingQueryBuilderTest {
@@ -118,4 +122,56 @@ public class PlanNodeRoutingQueryBuilderTest {
assertEquals(filter.getFunctionCall().getOperator(), "EQUALS");
assertEquals(filter.getFunctionCall().getOperands().get(0).getIdentifier().getName(),
"col1");
}
+
+ @Test
+ public void testCanBuildRoutingQueryAcceptsASingleInputChain() {
+ TableScanNode tableScanNode =
+ new TableScanNode(1, TEST_SCHEMA, PlanNode.NodeHint.EMPTY, List.of(),
"testTable", List.of("col1", "col2"));
+ FilterNode filterNode = new FilterNode(1, TEST_SCHEMA,
PlanNode.NodeHint.EMPTY, List.of(tableScanNode),
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.BOOLEAN,
"EQUALS",
+ List.of(new RexExpression.InputRef(0),
+ new RexExpression.Literal(DataSchema.ColumnDataType.STRING,
"foo"))));
+
+ assertTrue(PlanNodeRoutingQueryBuilder.canBuildRoutingQuery(filterNode));
+ }
+
+ /// The pre-check has to recurse exactly as the fold does. Checking only the
root would call this tree foldable and
+ /// hand the caller back the thrown-and-caught exception the pre-check
exists to avoid.
+ @Test
+ public void testCanBuildRoutingQueryRejectsAMultiInputNodeBelowTheRoot() {
+ PlanNode multiInputNode = setOpOverTwoScans();
+ FilterNode filterAbove = new FilterNode(1, TEST_SCHEMA,
PlanNode.NodeHint.EMPTY, List.of(multiInputNode),
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.BOOLEAN,
"EQUALS",
+ List.of(new RexExpression.InputRef(0),
+ new RexExpression.Literal(DataSchema.ColumnDataType.STRING,
"foo"))));
+
+
assertFalse(PlanNodeRoutingQueryBuilder.canBuildRoutingQuery(multiInputNode));
+ assertFalse(PlanNodeRoutingQueryBuilder.canBuildRoutingQuery(filterAbove));
+ }
+
+ /// The two must agree: whatever the pre-check refuses is exactly what the
fold would have thrown on, and whatever it
+ /// accepts the fold attempts. Disagreement either restores the per-query
exception or silently disables pruning for
+ /// shapes that fold perfectly well.
+ @Test
+ public void testCanBuildRoutingQueryAgreesWithTheFold() {
+ PlanNode foldable =
+ new TableScanNode(1, TEST_SCHEMA, PlanNode.NodeHint.EMPTY, List.of(),
"testTable", List.of("col1", "col2"));
+ PlanNode notFoldable = setOpOverTwoScans();
+
+ assertTrue(PlanNodeRoutingQueryBuilder.canBuildRoutingQuery(foldable));
+
assertNotNull(PlanNodeRoutingQueryBuilder.createPinotQueryForRouting("testTable",
foldable, false));
+
+ assertFalse(PlanNodeRoutingQueryBuilder.canBuildRoutingQuery(notFoldable));
+ assertThrows(IllegalStateException.class,
+ () ->
PlanNodeRoutingQueryBuilder.createPinotQueryForRouting("testTable",
notFoldable, false));
+ }
+
+ private static PlanNode setOpOverTwoScans() {
+ TableScanNode first =
+ new TableScanNode(1, TEST_SCHEMA, PlanNode.NodeHint.EMPTY, List.of(),
"testTable", List.of("col1", "col2"));
+ TableScanNode second =
+ new TableScanNode(1, TEST_SCHEMA, PlanNode.NodeHint.EMPTY, List.of(),
"testTable", List.of("col1", "col2"));
+ return new SetOpNode(1, TEST_SCHEMA, PlanNode.NodeHint.EMPTY,
List.of(first, second), SetOpNode.SetOpType.UNION,
+ true);
+ }
}
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 84ce1a4e108..49802eba16f 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 class WorkerManagerTest {
}
@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 class WorkerManagerTest {
+ "/*+ 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<DispatchablePlanFragment> 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<DispatchablePlanFragment> leafFragments =
leafFragments(dispatchableSubPlan);
+ assertEquals(leafFragments.size(), 2);
+ Map<Integer, Integer> 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 class WorkerManagerTest {
}
}
+ @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 class WorkerManagerTest {
}
@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<Integer, Integer> 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<Integer, Integer> 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 class WorkerManagerTest {
Map<Integer, Map<String, List<String>>> 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 class WorkerManagerTest {
Map<Integer, Map<String, List<String>>> 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 class WorkerManagerTest {
Map<String, List<String>> 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 class WorkerManagerTest {
// 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 class WorkerManagerTest {
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 class WorkerManagerTest {
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 class WorkerManagerTest {
Map<String, List<String>> 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 @@ public class WorkerManagerTest {
/// 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<String, ServerInstance> enabledServers = new HashMap<>();
@@ -1710,13 +1966,15 @@ public class WorkerManagerTest {
colocatedTablePartitionInfo(tableBWithType, "b_seg", servers, specB));
Map<String, RoutingTable> 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<String, String> tableNameMap = new HashMap<>();
tableNameMap.put(tableAWithType, tableAWithType);
@@ -1764,8 +2022,9 @@ public class WorkerManagerTest {
: 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<String> 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 @@ public class WorkerManagerTest {
/// 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<String> survivingSegments) {
+ List<String> survivingSegments, List<String> unavailableSegments) {
Map<ServerInstance, List<String>> serverToSegmentList = new HashMap<>();
for (String segment : survivingSegments) {
int partition =
Integer.parseInt(segment.substring(segmentPrefix.length()));
@@ -1784,7 +2043,7 @@ public class WorkerManagerTest {
Map<ServerInstance, SegmentsToQuery> 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 @@ public class WorkerManagerTest {
/// null rather than a silently empty answer.
@Nullable
List<String> _survivingSegments;
+ List<String> _unavailableSegments = List.of();
+ Set<Integer> _partitionsWithoutSegments = Set.of();
ColocatedTableSpec(int numPartitions, boolean
everyServerHostsEveryPartition) {
_numPartitions = numPartitions;
@@ -1844,6 +2105,21 @@ public class WorkerManagerTest {
_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<Integer>
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<String> 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 @@ public class WorkerManagerTest {
int replicasPerPartition, List<String> survivingSegments, List<String>
unavailableSegments,
int reportedPrunedByRouting, boolean throwOnRouting, Set<Integer>
emptyPartitions,
Set<Integer> 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<String> survivingSegments, List<String>
unavailableSegments,
+ int reportedPrunedByRouting, boolean throwOnRouting, Set<Integer>
emptyPartitions,
+ Set<Integer> partitionsWithOnlyDeferredSegments, Map<String,
List<String>> survivingSegmentsByFilterLiteral) {
int numPartitions = serverIdxPerPartition.length;
ServerInstance[] servers = new ServerInstance[numServers];
Map<String, ServerInstance> enabledServers = new HashMap<>();
@@ -2095,6 +2383,7 @@ public class WorkerManagerTest {
PartitionedRoutingManager routingManager = new
PartitionedRoutingManager(enabledServers,
Map.of(PARTITIONED_TABLE_OFFLINE, tablePartitionInfo),
Map.of(PARTITIONED_TABLE_OFFLINE, prunedRoutingTable), throwOnRouting);
+
survivingSegmentsByFilterLiteral.forEach(routingManager::survivingSegmentsForFilterLiteral);
Map<String, String> tableNameMap = new HashMap<>();
tableNameMap.put(PARTITIONED_TABLE_OFFLINE, PARTITIONED_TABLE_OFFLINE);
@@ -2494,6 +2783,9 @@ public class WorkerManagerTest {
private final Map<String, ServerInstance> _enabledServers;
private final Map<String, TablePartitionReplicatedServersInfo>
_partitionInfoByTable;
private final Map<String, RoutingTable> _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<String, List<String>> _survivingSegmentsByFilterLiteral
= new HashMap<>();
private final boolean _throwOnRouting;
@Nullable
private final TimeBoundaryInfo _timeBoundaryInfo;
@@ -2541,8 +2833,67 @@ public class WorkerManagerTest {
@Nullable
@Override
public List<String> 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<String> getPrunedSegments(BrokerRequest brokerRequest) {
+ if (_throwOnRouting) {
+ throw new RuntimeException("Simulated routing failure");
+ }
+
validatePrunableFilter(brokerRequest.getPinotQuery().getFilterExpression());
+ String tableNameWithType = brokerRequest.getQuerySource().getTableName();
+ Set<String> prunedSegments = new
HashSet<>(allSegments(tableNameWithType));
+ String filterLiteral =
firstStringLiteral(brokerRequest.getPinotQuery().getFilterExpression());
+ List<String> 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<String> 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<String> 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 20ef9aed09e..5f6c4b84e9d 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 a60b43ddd58..7cdf8a56398 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 class CommonConstants {
/// 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;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
