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 637b6b3babd Add broker pruning support to partitioned leaf path and
logical tables in MSE (#18924)
637b6b3babd is described below
commit 637b6b3babd59b1c9ccf46e40b656ca53e0f7039
Author: Yash Mayya <[email protected]>
AuthorDate: Wed Jul 8 13:01:07 2026 -0700
Add broker pruning support to partitioned leaf path and logical tables in
MSE (#18924)
---
...PartitionLLCRealtimeClusterIntegrationTest.java | 31 ++
.../BaseLogicalTableIntegrationTest.java | 29 ++
.../apache/pinot/query/routing/WorkerManager.java | 242 ++++++++--
.../pinot/query/routing/WorkerManagerTest.java | 527 +++++++++++++++++++++
4 files changed, 785 insertions(+), 44 deletions(-)
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java
index 201fd136e09..75e1ad0e166 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java
@@ -207,6 +207,37 @@ public class
SegmentPartitionLLCRealtimeClusterIntegrationTest extends BaseClust
}
}
+ @Test(dependsOnMethods = "testPartitionRouting")
+ public void testMultiStageBrokerPruningOnPartitionedTable()
+ throws Exception {
+ setUseMultiStageQueryEngine(true);
+ try {
+ // 'CA' hashes to partition 0, so only that partition's segments should
be routed when broker pruning is enabled.
+ String filteredQuery = "SELECT COUNT(*) FROM mytable WHERE DestState =
'CA'";
+
+ // Baseline: the multi-stage engine without broker pruning scans all
segments across both partitions.
+ JsonNode unpruned = postQuery(filteredQuery);
+ assertTrue(unpruned.get("exceptions").isEmpty(), "Unexpected exceptions
without broker pruning");
+ int unprunedSegmentsQueried =
unpruned.get(MetadataKey.NUM_SEGMENTS_QUERIED.getName()).asInt();
+
+ // With broker pruning enabled, the non-matching partition's segments
are pruned at the broker before dispatch.
+ JsonNode pruned = postQuery("SET useBrokerPruning=true; " +
filteredQuery);
+ assertTrue(pruned.get("exceptions").isEmpty(), "Unexpected exceptions
with broker pruning");
+
+ // Broker pruning is a routing optimization only; the result must be
identical to the unpruned run.
+
assertEquals(pruned.get("resultTable").get("rows").get(0).get(0).asLong(),
+ unpruned.get("resultTable").get("rows").get(0).get(0).asLong(),
+ "Broker pruning changed the query result");
+ // Broker pruning actually engaged: the non-matching partition's
segments were dropped, so fewer were queried.
+ assertTrue(pruned.get("numSegmentsPrunedByBroker").asInt() > 0,
+ "Expected the broker to prune the non-matching partition's
segments");
+
assertTrue(pruned.get(MetadataKey.NUM_SEGMENTS_QUERIED.getName()).asInt() <
unprunedSegmentsQueried,
+ "Expected fewer segments queried with broker pruning");
+ } finally {
+ setUseMultiStageQueryEngine(false);
+ }
+ }
+
@Test(dependsOnMethods = "testPartitionRouting")
public void testPartitionIdVirtualColumn()
throws Exception {
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/logicaltable/BaseLogicalTableIntegrationTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/logicaltable/BaseLogicalTableIntegrationTest.java
index cefa8f2c7f9..db52283c71d 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/logicaltable/BaseLogicalTableIntegrationTest.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/logicaltable/BaseLogicalTableIntegrationTest.java
@@ -453,6 +453,35 @@ public abstract class BaseLogicalTableIntegrationTest
extends BaseClusterIntegra
super.testGeneratedQueries(true, useMultiStageQueryEngine);
}
+ /**
+ * End-to-end smoke test that enabling broker segment pruning on the MSE
logical-planner path executes cleanly for a
+ * logical table and does not change results. It exercises the logical-table
routing path where the leaf-stage filter
+ * is forwarded (via {@code
WorkerManager.buildLogicalTableRoutingBrokerRequest}) into each physical
table's routing
+ * request so the {@code LogicalTableRouteProvider} runs segment pruners
over a filter-bearing request rather than a
+ * bare {@code SELECT *}. A malformed forwarded request would surface here
as an exception or a changed result.
+ *
+ * <p>Note: these physical tables carry no segment-pruner config, so this
asserts correctness, not that a segment was
+ * actually pruned; the exact pruning behavior is unit-tested in {@code
WorkerManagerTest}.
+ */
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testBrokerPruningPreservesLogicalTableResults(boolean
useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+ String logicalTableName = getLogicalTableName();
+ String filteredQuery = "SELECT COUNT(*) FROM " + logicalTableName + "
WHERE DaysSinceEpoch > 16312";
+
+ JsonNode withoutPruning = postQuery(filteredQuery);
+ assertTrue(withoutPruning.get("exceptions").isEmpty(), "Unexpected
exceptions without broker pruning");
+
+ JsonNode withPruning = postQuery("SET useBrokerPruning=true; " +
filteredQuery);
+ assertTrue(withPruning.get("exceptions").isEmpty(), "Unexpected exceptions
with broker pruning");
+
+ // Broker pruning is a routing optimization only; the result must be
identical to the unpruned run.
+
assertEquals(withPruning.get("resultTable").get("rows").get(0).get(0).asLong(),
+ withoutPruning.get("resultTable").get("rows").get(0).get(0).asLong(),
+ "Broker pruning changed the result of a logical table query");
+ }
+
@Test
public void testDisableGroovyQueryTableConfigOverride()
throws Exception {
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 973078771c3..a0edc857847 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
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.query.routing;
+import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.Maps;
import java.util.ArrayList;
@@ -492,14 +493,20 @@ public class WorkerManager {
String partitionKey =
tableOptions.get(PinotHintOptions.TableHintOptions.PARTITION_KEY);
if (partitionKey != null) {
- assignWorkersToPartitionedLeafFragment(metadata, context,
partitionKey, tableOptions);
+ // Broker pruning: build a filter-bearing routing query (null when
disabled/unsupported) so the partitioned
+ // assignment can drop partitions with no matching segments. Reuses
the same gate as the non-partitioned path.
+ // Skip pre-partitioned leaves up front: pruning is disabled for them
(see computePartitionsToKeep), so don't
+ // spend planning time building the routing query, e.g. for
colocated-join leaves.
+ PinotQuery routingPinotQuery = metadata.isPrePartitioned() ? null
+ : extractRoutingQuery(fragment.getFragmentRoot(),
metadata.getScannedTables().get(0), context);
+ assignWorkersToPartitionedLeafFragment(metadata, context,
partitionKey, tableOptions, routingPinotQuery);
updateContextForLeafStage(metadata, context);
return;
}
}
if (metadata.getLogicalTableRouteInfo() != null) {
- assignWorkersToNonPartitionedLeafFragmentForLogicalTable(metadata,
context);
+ assignWorkersToNonPartitionedLeafFragmentForLogicalTable(fragment,
metadata, context);
} else {
assignWorkersToNonPartitionedLeafFragment(fragment, metadata, context);
}
@@ -780,8 +787,8 @@ public class WorkerManager {
CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM \"" +
tableNameWithType + "\""), samplerName);
}
- private void
assignWorkersToNonPartitionedLeafFragmentForLogicalTable(DispatchablePlanMetadata
metadata,
- DispatchablePlanContext context) {
+ private void
assignWorkersToNonPartitionedLeafFragmentForLogicalTable(PlanFragment fragment,
+ DispatchablePlanMetadata metadata, DispatchablePlanContext context) {
LogicalTableRouteInfo logicalTableRouteInfo =
metadata.getLogicalTableRouteInfo();
Preconditions.checkNotNull(logicalTableRouteInfo);
LogicalTableRouteProvider tableRouteProvider = new
LogicalTableRouteProvider();
@@ -789,32 +796,57 @@ public class WorkerManager {
if (logicalTableRouteInfo.getTimeBoundaryInfo() != null) {
metadata.setTimeBoundaryInfo(logicalTableRouteInfo.getTimeBoundaryInfo());
}
- BrokerRequest offlineBrokerRequest = null;
- BrokerRequest realtimeBrokerRequest = null;
Map<String, String> queryOptions =
context.getPlannerContext().getOptions();
- if (logicalTableRouteInfo.hasOffline()) {
- offlineBrokerRequest = CalciteSqlCompiler.compileToBrokerRequest(
- "SELECT * FROM \"" + logicalTableRouteInfo.getOfflineTableName() +
"\"");
- if (MapUtils.isNotEmpty(queryOptions) &&
offlineBrokerRequest.isSetPinotQuery()) {
- offlineBrokerRequest.getPinotQuery().setQueryOptions(new
HashMap<>(queryOptions));
- }
- }
+ // Broker pruning: build a filter-bearing routing query (null when
disabled/unsupported). When non-null, each
+ // physical table's route below carries the filter so segment pruners can
eliminate segments; when null, fall back
+ // to an unfiltered SELECT * per table type. The physical route resolution
ignores the data source table name (it
+ // uses the physical table names), so we set the typed logical name here
purely to build a valid broker request.
+ String rawTableName = TableNameBuilder.extractRawTableName(
+ logicalTableRouteInfo.hasOffline() ?
logicalTableRouteInfo.getOfflineTableName()
+ : logicalTableRouteInfo.getRealtimeTableName());
+ PinotQuery routingPinotQuery =
extractRoutingQuery(fragment.getFragmentRoot(), rawTableName, context);
- if (logicalTableRouteInfo.hasRealtime()) {
- realtimeBrokerRequest = CalciteSqlCompiler.compileToBrokerRequest(
- "SELECT * FROM \"" + logicalTableRouteInfo.getRealtimeTableName() +
"\"");
- if (MapUtils.isNotEmpty(queryOptions) &&
realtimeBrokerRequest.isSetPinotQuery()) {
- realtimeBrokerRequest.getPinotQuery().setQueryOptions(new
HashMap<>(queryOptions));
- }
- }
+ BrokerRequest offlineBrokerRequest = logicalTableRouteInfo.hasOffline() ?
buildLogicalTableRoutingBrokerRequest(
+ logicalTableRouteInfo.getOfflineTableName(), routingPinotQuery,
queryOptions) : null;
+ BrokerRequest realtimeBrokerRequest = logicalTableRouteInfo.hasRealtime()
? buildLogicalTableRoutingBrokerRequest(
+ logicalTableRouteInfo.getRealtimeTableName(), routingPinotQuery,
queryOptions) : null;
tableRouteProvider.calculateRoutes(logicalTableRouteInfo, _routingManager,
offlineBrokerRequest,
realtimeBrokerRequest, context.getRequestId());
+ if (routingPinotQuery != null) {
+
context.addNumSegmentsPrunedByBroker(logicalTableRouteInfo.getNumPrunedSegmentsTotal());
+ }
assignTableSegmentsToWorkers(logicalTableRouteInfo, metadata);
}
+ /**
+ * Builds the routing {@link BrokerRequest} for one physical table type of a
logical table. When {@code
+ * routingPinotQuery} is non-null it carries the leaf-stage filter so
segment pruners can eliminate segments;
+ * otherwise a bare {@code SELECT *} is used (no pruning). The data source
is set to {@code logicalTableNameWithType}
+ * (the typed logical table name; physical table names are resolved later by
the route provider).
+ *
+ * <p>The given {@code routingPinotQuery} is not modified: it is deep-copied
before the table name is rewritten, so
+ * the same instance can be reused to build both the offline and realtime
requests.
+ */
+ @VisibleForTesting
+ static BrokerRequest buildLogicalTableRoutingBrokerRequest(String
logicalTableNameWithType,
+ @Nullable PinotQuery routingPinotQuery, Map<String, String>
queryOptions) {
+ BrokerRequest brokerRequest;
+ if (routingPinotQuery != null) {
+ PinotQuery pinotQuery = routingPinotQuery.deepCopy();
+ pinotQuery.getDataSource().setTableName(logicalTableNameWithType);
+ brokerRequest = CalciteSqlCompiler.convertToBrokerRequest(pinotQuery);
+ } else {
+ brokerRequest = CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM
\"" + logicalTableNameWithType + "\"");
+ }
+ if (MapUtils.isNotEmpty(queryOptions) && brokerRequest.isSetPinotQuery()) {
+ brokerRequest.getPinotQuery().setQueryOptions(new
HashMap<>(queryOptions));
+ }
+ return brokerRequest;
+ }
+
private static void assignTableSegmentsToWorkers(LogicalTableRouteInfo
logicalTableRouteInfo,
DispatchablePlanMetadata metadata) {
Map<ServerInstance, Map<String, List<String>>>
serverInstanceToLogicalSegmentsMap =
@@ -884,7 +916,8 @@ public class WorkerManager {
// Partitioned leaf stage assignment
// --------------------------------------------------------------------------
private void assignWorkersToPartitionedLeafFragment(DispatchablePlanMetadata
metadata,
- DispatchablePlanContext context, String partitionKey, Map<String,
String> tableOptions) {
+ DispatchablePlanContext context, String partitionKey, Map<String,
String> tableOptions,
+ @Nullable PinotQuery routingPinotQuery) {
// when partition key exist, we assign workers for leaf-stage in
partitioned fashion.
String numPartitionsStr =
tableOptions.get(PinotHintOptions.TableHintOptions.PARTITION_SIZE);
@@ -907,41 +940,152 @@ public class WorkerManager {
int numPartitions = partitionInfoMap.length;
assert numPartitions % numWorkers == 0;
int numPartitionsPerWorker = numPartitions / numWorkers;
- Map<Integer, QueryServerInstance> workedIdToServerInstanceMap = new
HashMap<>();
+
+ // 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);
+ if (partitionsToKeep != null) {
+ long numSegmentsPrunedByBroker = countPrunedSegments(partitionInfoMap,
partitionsToKeep);
+ if (numSegmentsPrunedByBroker > 0) {
+ context.addNumSegmentsPrunedByBroker(numSegmentsPrunedByBroker);
+ }
+ }
+
+ Map<Integer, QueryServerInstance> workerIdToServerInstanceMap = new
HashMap<>();
Map<Integer, Map<String, List<String>>> workerIdToSegmentsMap = new
HashMap<>();
if (numPartitionsPerWorker == 1) {
- assignOnePartitionPerWorker(tableName, context.getRequestId(),
partitionInfoMap,
- _routingManager.getEnabledServerInstanceMap(),
workedIdToServerInstanceMap, workerIdToSegmentsMap);
+ assignOnePartitionPerWorker(tableName, context.getRequestId(),
partitionInfoMap, partitionsToKeep,
+ _routingManager.getEnabledServerInstanceMap(),
workerIdToServerInstanceMap, workerIdToSegmentsMap);
} else {
assignMultiplePartitionsPerWorker(tableName, context.getRequestId(),
numPartitionsPerWorker, partitionInfoMap,
- _routingManager.getEnabledServerInstanceMap(),
workedIdToServerInstanceMap, workerIdToSegmentsMap);
+ partitionsToKeep, _routingManager.getEnabledServerInstanceMap(),
workerIdToServerInstanceMap,
+ workerIdToSegmentsMap);
}
- metadata.setWorkerIdToServerInstanceMap(workedIdToServerInstanceMap);
+ metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap);
metadata.setWorkerIdToSegmentsMap(workerIdToSegmentsMap);
metadata.setTimeBoundaryInfo(partitionTableInfo._timeBoundaryInfo);
metadata.setPartitionFunction(partitionFunction);
}
- /// Pick one worker per partition for partitioned leaf stage.
+ /**
+ * 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 {@code null} to keep all partitions.
+ *
+ * <p>Returns {@code null} (no pruning) when any of the following hold:
+ * <ul>
+ * <li>broker pruning is disabled or the leaf shape is unsupported (the
routing query is {@code null}), or there is
+ * no filter to prune with;</li>
+ * <li>the leaf feeds a pre-partitioned (1-to-1 direct) exchange --
dropping/compacting workers would misalign
+ * sender/receiver worker ids in {@code MailboxAssignmentVisitor}. A
non-pre-partitioned leaf is shuffled via
+ * {@code connectWorkers}, which re-hashes across any worker count, so
pruning is safe there;</li>
+ * <li>routing fails (pruning is best-effort);</li>
+ * <li>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.</li>
+ * </ul>
+ *
+ * <p>Partition survival is decided by routing the filter-bearing query
through the {@link 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.
+ *
+ * <p>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
+ * surviving partitions' assignments identical to the unpruned path -- the
only behavioral delta is dropped
+ * 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()) {
+ 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.getMessage());
+ 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))) {
+ partitionsToKeep.add(i);
+ }
+ }
+ // If everything would be pruned, keep all partitions to avoid an empty
worker map (see the javadoc above).
+ return partitionsToKeep.isEmpty() ? null : partitionsToKeep;
+ }
+
+ private static boolean containsAny(@Nullable List<String> segments,
Set<String> matchedSegments) {
+ if (segments != null) {
+ for (String segment : segments) {
+ if (matchedSegments.contains(segment)) {
+ return true;
+ }
+ }
+ }
+ return false;
+ }
+
+ /// Counts the segments in partitions dropped by broker pruning (those
absent from {@code partitionsToKeep}).
+ private static long countPrunedSegments(PartitionInfo[] partitionInfoMap,
Set<Integer> partitionsToKeep) {
+ long numPrunedSegments = 0;
+ for (int i = 0; i < partitionInfoMap.length; i++) {
+ PartitionInfo partitionInfo = partitionInfoMap[i];
+ if (partitionInfo != null && !partitionsToKeep.contains(i)) {
+ numPrunedSegments +=
CollectionUtils.size(partitionInfo._offlineSegments)
+ + CollectionUtils.size(partitionInfo._realtimeSegments);
+ }
+ }
+ return numPrunedSegments;
+ }
+
+ /// Pick one worker per partition for partitioned leaf stage. When {@code
partitionsToKeep} is non-null (broker
+ /// pruning is active), partitions absent from the set are pruned by the
query filter and skipped.
private void assignOnePartitionPerWorker(String tableName, long requestId,
PartitionInfo[] partitionInfoMap,
- Map<String, ServerInstance> enabledServerInstanceMap,
- Map<Integer, QueryServerInstance> workedIdToServerInstanceMap,
+ @Nullable Set<Integer> partitionsToKeep, Map<String, ServerInstance>
enabledServerInstanceMap,
+ Map<Integer, QueryServerInstance> workerIdToServerInstanceMap,
Map<Integer, Map<String, List<String>>> workerIdToSegmentsMap) {
int numPartitions = partitionInfoMap.length;
int workerId = 0;
for (int i = 0; i < numPartitions; i++) {
+ // Skip partitions pruned by the broker filter. Empty partitions are
never in partitionsToKeep, so under pruning
+ // they are skipped here too; the precondition below only fires when
pruning is inactive (partitionsToKeep null).
+ if (partitionsToKeep != null && !partitionsToKeep.contains(i)) {
+ continue;
+ }
PartitionInfo partitionInfo = partitionInfoMap[i];
// TODO: Currently we don't support the case when a partition doesn't
contain any segment. The reason is that
// the leaf stage won't be able to directly return empty response.
Preconditions.checkState(partitionInfo != null, "Failed to find any
segment for table: %s, partition: %s",
tableName, i);
- // NOTE: Pick worker based on the request id so that the same worker is
picked across different table scan when
- // the segments for the same partition is colocated
+ // NOTE: Pick worker based on the request id plus the partition id (not
a running counter) so that the same worker
+ // is picked across different table scans when the segments for
the same partition are colocated, and so
+ // that skipping pruned partitions does not shift the server
assignment of the surviving ones.
ServerInstance serverInstance =
- pickEnabledServer(partitionInfo._fullyReplicatedServers,
enabledServerInstanceMap, requestId++);
+ pickEnabledServer(partitionInfo._fullyReplicatedServers,
enabledServerInstanceMap, requestId + i);
Preconditions.checkState(serverInstance != null,
"Failed to find enabled fully replicated server for table: %s,
partition: %s", tableName, i);
- workedIdToServerInstanceMap.put(workerId, new
QueryServerInstance(serverInstance));
+ workerIdToServerInstanceMap.put(workerId, new
QueryServerInstance(serverInstance));
workerIdToSegmentsMap.put(workerId,
getSegmentsMap(partitionInfo._offlineSegments,
partitionInfo._realtimeSegments));
workerId++;
@@ -954,8 +1098,9 @@ public class WorkerManager {
/// E.g. when there are 16 partitions for table A and 4 partitions for table
B, we may assign 16 partitions for table
/// A to 4 workers, where partition 0, 4, 8, 12 goes to worker 0, partition
1, 5, 9, 13 goes to worker 1, etc.
private void assignMultiplePartitionsPerWorker(String tableName, long
requestId, int numPartitionsPerWorker,
- PartitionInfo[] partitionInfoMap, Map<String, ServerInstance>
enabledServerInstanceMap,
- Map<Integer, QueryServerInstance> workedIdToServerInstanceMap,
+ PartitionInfo[] partitionInfoMap, @Nullable Set<Integer>
partitionsToKeep,
+ Map<String, ServerInstance> enabledServerInstanceMap,
+ Map<Integer, QueryServerInstance> workerIdToServerInstanceMap,
Map<Integer, Map<String, List<String>>> workerIdToSegmentsMap) {
int numPartitions = partitionInfoMap.length;
assert numPartitions % numPartitionsPerWorker == 0;
@@ -966,6 +1111,10 @@ public class WorkerManager {
List<String> offlineSegments = null;
List<String> realtimeSegments = null;
for (int j = i; j < numPartitions; j += numWorkers) {
+ if (partitionsToKeep != null && !partitionsToKeep.contains(j)) {
+ // Partition pruned by the broker filter.
+ continue;
+ }
PartitionInfo partitionInfo = partitionInfoMap[j];
if (partitionInfo == null) {
continue;
@@ -990,18 +1139,23 @@ public class WorkerManager {
}
}
}
- // TODO: Currently we don't support the case when all partitions for a
worker don't contain any segment. The
- // reason is that the leaf stage won't be able to directly return
empty response.
- Preconditions.checkState(fullyReplicatedServers != null,
- "Failed to find any segment for table: %s, worker: %s, partitions
per worker: %s", tableName, i,
- numPartitionsPerWorker);
- // NOTE: Pick worker based on the request id so that the same worker is
picked across different table scan when
- // the segments for the same partition is colocated
- ServerInstance serverInstance =
pickEnabledServer(fullyReplicatedServers, enabledServerInstanceMap,
requestId++);
+ // Without broker pruning we don't support a worker whose partitions all
lack segments, because the leaf stage
+ // can't directly return an empty response. With pruning active a
fully-pruned worker is legitimate and skipped.
+ if (fullyReplicatedServers == null) {
+ Preconditions.checkState(partitionsToKeep != null,
+ "Failed to find any segment for table: %s, worker: %s, partitions
per worker: %s", tableName, i,
+ numPartitionsPerWorker);
+ continue;
+ }
+ // NOTE: Pick worker based on the request id plus the worker index (not
a running counter) so that the same worker
+ // is picked across different table scans when the segments for
the same partition are colocated, and so
+ // that skipping fully-pruned workers does not shift the server
assignment of the surviving ones.
+ ServerInstance serverInstance =
+ pickEnabledServer(fullyReplicatedServers, enabledServerInstanceMap,
requestId + i);
Preconditions.checkState(serverInstance != null,
"Failed to find enabled fully replicated server for table: %s,
worker: %s, partitions per worker: %s",
tableName, i, numPartitionsPerWorker);
- workedIdToServerInstanceMap.put(workerId, new
QueryServerInstance(serverInstance));
+ workerIdToServerInstanceMap.put(workerId, new
QueryServerInstance(serverInstance));
workerIdToSegmentsMap.put(workerId, getSegmentsMap(offlineSegments,
realtimeSegments));
workerId++;
}
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 e12f60c4693..0c5ac12cfeb 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
@@ -31,6 +31,7 @@ import org.apache.helix.zookeeper.datamodel.ZNRecord;
import org.apache.pinot.common.config.provider.TableCache;
import org.apache.pinot.common.request.BrokerRequest;
import org.apache.pinot.common.request.Expression;
+import org.apache.pinot.common.request.PinotQuery;
import org.apache.pinot.core.routing.RoutingManager;
import org.apache.pinot.core.routing.RoutingTable;
import org.apache.pinot.core.routing.SegmentsToQuery;
@@ -45,6 +46,7 @@ import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.data.FieldSpec;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.sql.parsers.CalciteSqlCompiler;
import org.testng.annotations.Test;
import static org.mockito.ArgumentMatchers.anyString;
@@ -331,6 +333,439 @@ public class WorkerManagerTest {
assertEquals(queryOptions.get("useLeafServerForIntermediateStage"),
"true");
}
+ //
---------------------------------------------------------------------------
+ // Broker pruning: partitioned leaf path
+ //
---------------------------------------------------------------------------
+
+ private static final String PARTITIONED_TABLE = "testTable";
+ private static final String PARTITIONED_TABLE_OFFLINE = "testTable_OFFLINE";
+
+ @Test
+ public void testBrokerPruningPartitionedLeafPrunesNonMatchingPartitions() {
+ // 4 partitions, one segment + one server each. The routing manager
reports only partition 2's segment survives.
+ QueryEnvironment queryEnvironment =
+ newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4,
List.of("seg2"), /*reportedPrunedByRouting=*/999);
+ 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();
+ // seg0, seg1, seg3 pruned (3). The count is computed from dropped
partitions, not the routing table's own
+ // numPrunedSegments (999), proving the partitioned path computes it
independently.
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 3);
+ DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan);
+ assertNotNull(leaf);
+ assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 1);
+ assertEquals(assignedSegments(leaf), List.of("seg2"));
+ }
+ }
+
+ @Test
+ public void
testBrokerPruningPartitionedLeafKeepsMultipleMatchingPartitions() {
+ QueryEnvironment queryEnvironment =
+ newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4,
List.of("seg1", "seg3"), 0);
+ 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();
+ // seg0, seg2 pruned.
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 2);
+ DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan);
+ assertNotNull(leaf);
+ assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 2);
+ assertEquals(new HashSet<>(assignedSegments(leaf)), Set.of("seg1",
"seg3"));
+ }
+ }
+
+ @Test
+ public void testBrokerPruningPartitionedLeafAllPrunedFallsBackToUnpruned() {
+ // The routing manager reports that no segment survives. Rather than
produce an empty worker map (which would break
+ // exchanges in a multi-leaf plan), the partitioned path falls back to an
unpruned assignment: the server-side
+ // filter still yields the correct empty result.
+ QueryEnvironment queryEnvironment = newPartitionedQueryEnvironment(new
int[]{0, 1, 2, 3}, 4, List.of(), 4);
+ 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();
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0);
+ DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan);
+ assertNotNull(leaf);
+ assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 4);
+ assertEquals(new HashSet<>(assignedSegments(leaf)), Set.of("seg0",
"seg1", "seg2", "seg3"));
+ }
+ }
+
+ @Test
+ public void testBrokerPruningPartitionedLeafDisabledKeepsAllPartitions() {
+ // Without SET useBrokerPruning=true the logical planner default (false)
applies, so no pruning occurs.
+ QueryEnvironment queryEnvironment =
+ newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4,
List.of("seg2"), 3);
+ try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(
+ "SELECT col2 FROM testTable "
+ + "/*+ tableOptions(partition_function='hashcode',
partition_key='col1', partition_size='4') */ "
+ + "WHERE col1 = 'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan =
compiledQuery.planQuery(0).getQueryPlan();
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0);
+ DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan);
+ assertNotNull(leaf);
+ assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 4);
+ }
+ }
+
+ @Test
+ public void testBrokerPruningPartitionedLeafMultiplePartitionsPerWorker() {
+ // 4 partitions across 2 workers (partition_size=2 => 2 partitions per
worker). Worker 0 owns partitions {0, 2}
+ // (colocated on server 0), worker 1 owns {1, 3} (colocated on server 1).
Only worker 0's partitions survive.
+ QueryEnvironment queryEnvironment =
+ newPartitionedQueryEnvironment(new int[]{0, 1, 0, 1}, 2,
List.of("seg0", "seg2"), 0);
+ try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(
+ "SET useBrokerPruning=true; SELECT col2 FROM testTable "
+ + "/*+ tableOptions(partition_function='hashcode',
partition_key='col1', partition_size='2') */ "
+ + "WHERE col1 = 'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan =
compiledQuery.planQuery(0).getQueryPlan();
+ // Worker 1's partitions (seg1, seg3) are fully pruned.
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 2);
+ DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan);
+ assertNotNull(leaf);
+ assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 1);
+ assertEquals(new HashSet<>(assignedSegments(leaf)), Set.of("seg0",
"seg2"));
+ }
+ }
+
+ //
---------------------------------------------------------------------------
+ // Broker pruning: logical-table leaf path (filter forwarding into
per-physical-table routing requests)
+ //
---------------------------------------------------------------------------
+
+ @Test
+ public void testBuildLogicalTableRoutingRequestForwardsFilter() {
+ // With a routing query present, the physical table's routing request must
carry the leaf filter (so segment
+ // pruners can run), the typed logical table name, and the query options.
+ PinotQuery routingPinotQuery =
+ CalciteSqlCompiler.compileToBrokerRequest("SELECT col2 FROM someTable
WHERE col1 = 'foo'").getPinotQuery();
+ BrokerRequest brokerRequest =
WorkerManager.buildLogicalTableRoutingBrokerRequest("logicalTable_OFFLINE",
+ routingPinotQuery, Map.of("useLeafServerForIntermediateStage",
"true"));
+
+ assertEquals(brokerRequest.getPinotQuery().getDataSource().getTableName(),
"logicalTable_OFFLINE");
+ Expression filterExpression =
brokerRequest.getPinotQuery().getFilterExpression();
+ assertNotNull(filterExpression);
+ assertEquals(filterExpression.getFunctionCall().getOperator(), "EQUALS");
+
assertEquals(filterExpression.getFunctionCall().getOperands().get(0).getIdentifier().getName(),
"col1");
+
assertEquals(brokerRequest.getPinotQuery().getQueryOptions().get("useLeafServerForIntermediateStage"),
"true");
+ // The source routing query must not be mutated (the method deep-copies
before rewriting the table name).
+ assertEquals(routingPinotQuery.getDataSource().getTableName(),
"someTable");
+ }
+
+ @Test
+ public void testBuildLogicalTableRoutingRequestWithoutFilterUsesSelectStar()
{
+ // With no routing query (pruning disabled/unsupported), a bare SELECT *
request is used so no pruning occurs.
+ BrokerRequest brokerRequest =
+
WorkerManager.buildLogicalTableRoutingBrokerRequest("logicalTable_REALTIME",
null, Map.of());
+ assertEquals(brokerRequest.getPinotQuery().getDataSource().getTableName(),
"logicalTable_REALTIME");
+ assertNull(brokerRequest.getPinotQuery().getFilterExpression());
+ }
+
+ @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.
+ QueryEnvironment queryEnvironment =
+ newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4,
List.of("seg2"), 3);
+ 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.col1 = 'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan =
compiledQuery.planQuery(0).getQueryPlan();
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0);
+ List<DispatchablePlanFragment> leafFragments =
leafFragments(dispatchableSubPlan);
+ assertFalse(leafFragments.isEmpty());
+ for (DispatchablePlanFragment leaf : leafFragments) {
+ assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 4, "Expected all
partitions assigned for a pruned-gated "
+ + "colocated join leaf");
+ }
+ }
+ }
+
+ @Test
+ public void
testBrokerPruningPartitionedLeafKeepsPartitionWithUnavailableSegment() {
+ // Partition 2's segment survives; partition 1's segment is merely
unavailable (not pruned). A partition must be
+ // kept when its only segment is unavailable, so we never drop data that a
transient outage hid -- only partitions
+ // 0 and 3 (whose segments were actually pruned) are dropped.
+ QueryEnvironment queryEnvironment = newPartitionedQueryEnvironment(new
int[]{0, 1, 2, 3}, 4, 1, List.of("seg2"),
+ List.of("seg1"), 0, false);
+ 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();
+ // Only seg0 and seg3 pruned; seg1 is unavailable (kept), seg2 survives.
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 2);
+ DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan);
+ assertNotNull(leaf);
+ assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 2);
+ assertEquals(new HashSet<>(assignedSegments(leaf)), Set.of("seg1",
"seg2"));
+ }
+ }
+
+ @Test
+ public void testBrokerPruningPartitionedLeafFallsBackWhenRoutingFails() {
+ // If the routing call used to compute partition survival fails, pruning
is best-effort and must fall back to the
+ // unpruned assignment rather than failing the query.
+ QueryEnvironment queryEnvironment =
+ newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4, 1, List.of(),
List.of(), 0, /*throwOnRouting=*/true);
+ 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();
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0);
+ DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan);
+ assertNotNull(leaf);
+ assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 4);
+ }
+ }
+
+ @Test
+ public void testBrokerPruningPartitionedLeafKeepsServerPlacementStable() {
+ // Every partition is fully replicated on the same two servers, so
pickEnabledServer's choice depends on the
+ // per-partition seed. Pruning partition 0 must NOT shift the surviving
partitions' server assignments -- the seed
+ // is requestId + partitionId, not a running counter that would shift when
earlier partitions are skipped. Assert
+ // each surviving segment lands on the same server whether or not
partition 0 was pruned.
+ QueryEnvironment unprunedEnv = newPartitionedQueryEnvironment(new int[]{0,
0, 0, 0}, 2, 2,
+ List.of("seg0", "seg1", "seg2", "seg3"), List.of(), 0, false);
+ Map<String, String> unprunedPlacement;
+ try (QueryEnvironment.CompiledQuery compiledQuery = unprunedEnv.compile(
+ "SET useBrokerPruning=true; SELECT col2 FROM testTable "
+ + "/*+ tableOptions(partition_function='hashcode',
partition_key='col1', partition_size='4') */ "
+ + "WHERE col1 = 'foo'")) {
+ unprunedPlacement =
segmentToServer(leafFragment(compiledQuery.planQuery(0).getQueryPlan()));
+ }
+ assertEquals(unprunedPlacement.size(), 4);
+
+ // Prune partition 0 (only seg1/seg2/seg3 survive).
+ QueryEnvironment prunedEnv = newPartitionedQueryEnvironment(new int[]{0,
0, 0, 0}, 2, 2,
+ List.of("seg1", "seg2", "seg3"), List.of(), 0, false);
+ Map<String, String> prunedPlacement;
+ try (QueryEnvironment.CompiledQuery compiledQuery = prunedEnv.compile(
+ "SET useBrokerPruning=true; SELECT col2 FROM testTable "
+ + "/*+ tableOptions(partition_function='hashcode',
partition_key='col1', partition_size='4') */ "
+ + "WHERE col1 = 'foo'")) {
+ prunedPlacement =
segmentToServer(leafFragment(compiledQuery.planQuery(0).getQueryPlan()));
+ }
+
+ assertEquals(prunedPlacement.size(), 3);
+ for (String segment : List.of("seg1", "seg2", "seg3")) {
+ assertEquals(prunedPlacement.get(segment),
unprunedPlacement.get(segment),
+ "Pruning partition 0 shifted the server assignment of surviving " +
segment);
+ }
+ }
+
+ @Test
+ public void testBrokerPruningPartitionedLeafHybridTable() {
+ // Hybrid table: partition p holds offline segment segO{p} and realtime
segment segR{p}, both colocated on
+ // server p. The routing manager reports segO2 surviving on the offline
side and segR1 on the realtime side, so
+ // partitions {1, 2} are kept (a partition survives if a segment of EITHER
type matches) and {0, 3} are pruned.
+ QueryEnvironment queryEnvironment =
newHybridPartitionedQueryEnvironment(List.of("segO2"), List.of("segR1"));
+ 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();
+ // Pruned = both segments of partitions 0 and 3.
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 4);
+ DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan);
+ assertNotNull(leaf);
+ assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 2);
+ // Pruning is partition-level: surviving partitions dispatch ALL their
segments, so the non-matching segO1 and
+ // segR2 are still included alongside the matching segR1 and segO2.
+ assertEquals(new HashSet<>(assignedSegments(leaf)), Set.of("segO1",
"segR1", "segO2", "segR2"));
+ }
+ }
+
+ /**
+ * Builds a QueryEnvironment for a hybrid partitioned table "testTable"
(function Hashcode on col1, 4 partitions).
+ * Partition {@code p} holds offline segment {@code "segO{p}"} and realtime
segment {@code "segR{p}"}, both fully
+ * replicated on server {@code p}. The given surviving segment lists are
what the {@link RoutingManager} returns for
+ * the filtered routing query of each table type.
+ */
+ private static QueryEnvironment
newHybridPartitionedQueryEnvironment(List<String> survivingOfflineSegments,
+ List<String> survivingRealtimeSegments) {
+ int numPartitions = 4;
+ ServerInstance[] servers = new ServerInstance[numPartitions];
+ Map<String, ServerInstance> enabledServers = new HashMap<>();
+ for (int i = 0; i < numPartitions; i++) {
+ servers[i] = getServerInstance("localhost", i + 1);
+ enabledServers.put(servers[i].getInstanceId(), servers[i]);
+ }
+ TablePartitionReplicatedServersInfo.PartitionInfo[] offlinePartitions =
+ new TablePartitionReplicatedServersInfo.PartitionInfo[numPartitions];
+ TablePartitionReplicatedServersInfo.PartitionInfo[] realtimePartitions =
+ new TablePartitionReplicatedServersInfo.PartitionInfo[numPartitions];
+ for (int p = 0; p < numPartitions; p++) {
+ Set<String> partitionServers = Set.of(servers[p].getInstanceId());
+ offlinePartitions[p] = new
TablePartitionReplicatedServersInfo.PartitionInfo(partitionServers,
+ List.of("segO" + p));
+ realtimePartitions[p] = new
TablePartitionReplicatedServersInfo.PartitionInfo(partitionServers,
+ List.of("segR" + p));
+ }
+ String realtimeTableName = PARTITIONED_TABLE + "_REALTIME";
+ TablePartitionReplicatedServersInfo offlineInfo = new
TablePartitionReplicatedServersInfo(
+ PARTITIONED_TABLE_OFFLINE, "col1", "Hashcode", numPartitions,
offlinePartitions, List.of());
+ TablePartitionReplicatedServersInfo realtimeInfo = new
TablePartitionReplicatedServersInfo(
+ realtimeTableName, "col1", "Hashcode", numPartitions,
realtimePartitions, List.of());
+
+ PartitionedRoutingManager routingManager = new
PartitionedRoutingManager(enabledServers,
+ Map.of(PARTITIONED_TABLE_OFFLINE, offlineInfo, realtimeTableName,
realtimeInfo),
+ Map.of(PARTITIONED_TABLE_OFFLINE, hybridRoutingTable(servers,
survivingOfflineSegments),
+ realtimeTableName, hybridRoutingTable(servers,
survivingRealtimeSegments)),
+ false, new TimeBoundaryInfo("col3", "100"));
+
+ Map<String, String> tableNameMap = new HashMap<>();
+ tableNameMap.put(PARTITIONED_TABLE_OFFLINE, PARTITIONED_TABLE_OFFLINE);
+ tableNameMap.put(realtimeTableName, realtimeTableName);
+ tableNameMap.put(PARTITIONED_TABLE, PARTITIONED_TABLE);
+ TableCache tableCache = mock(TableCache.class);
+ when(tableCache.getTableNameMap()).thenReturn(tableNameMap);
+ when(tableCache.getActualTableName(anyString())).thenAnswer(inv ->
tableNameMap.get(inv.getArgument(0)));
+
when(tableCache.getSchema(anyString())).thenReturn(getSchemaBuilder(PARTITIONED_TABLE).build());
+
when(tableCache.getTableConfig(PARTITIONED_TABLE_OFFLINE)).thenReturn(mock(TableConfig.class));
+
when(tableCache.getTableConfig(realtimeTableName)).thenReturn(mock(TableConfig.class));
+
+ WorkerManager workerManager = new WorkerManager("Broker_localhost",
"localhost", 5, routingManager);
+ return new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
workerManager);
+ }
+
+ /** Buckets the given surviving segments onto their owning server (partition
p's segment lives on server p). */
+ private static RoutingTable hybridRoutingTable(ServerInstance[] servers,
List<String> survivingSegments) {
+ Map<ServerInstance, List<String>> serverToSegmentList = new HashMap<>();
+ for (String segment : survivingSegments) {
+ int partition = Integer.parseInt(segment.substring("segO".length()));
+ serverToSegmentList.computeIfAbsent(servers[partition], k -> new
ArrayList<>()).add(segment);
+ }
+ 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);
+ }
+
+ private static QueryEnvironment newPartitionedQueryEnvironment(int[]
serverIdxPerPartition, int numServers,
+ List<String> survivingSegments, int reportedPrunedByRouting) {
+ return newPartitionedQueryEnvironment(serverIdxPerPartition, numServers,
1, survivingSegments, List.of(),
+ reportedPrunedByRouting, false);
+ }
+
+ /**
+ * Builds a QueryEnvironment for an offline partitioned table "testTable"
(function Hashcode on col1). Partition
+ * {@code p} holds one segment {@code "seg{p}"} fully replicated on the
{@code replicasPerPartition} servers starting
+ * at {@code serverIdxPerPartition[p]}. {@code survivingSegments} is what
the {@link RoutingManager} returns from
+ * getRoutingTable for the filtered routing query -- i.e. the segments that
survive broker pruning; the corresponding
+ * partitions are kept. {@code unavailableSegments} are reported by the
routing table as unavailable (they must keep
+ * their partition alive). When {@code throwOnRouting} is true,
getRoutingTable throws to exercise the fail-open path.
+ */
+ private static QueryEnvironment newPartitionedQueryEnvironment(int[]
serverIdxPerPartition, int numServers,
+ int replicasPerPartition, List<String> survivingSegments, List<String>
unavailableSegments,
+ int reportedPrunedByRouting, boolean throwOnRouting) {
+ int numPartitions = serverIdxPerPartition.length;
+ ServerInstance[] servers = new ServerInstance[numServers];
+ Map<String, ServerInstance> enabledServers = new HashMap<>();
+ for (int i = 0; i < numServers; i++) {
+ servers[i] = getServerInstance("localhost", i + 1);
+ enabledServers.put(servers[i].getInstanceId(), servers[i]);
+ }
+ TablePartitionReplicatedServersInfo.PartitionInfo[] partitionInfoMap =
+ new TablePartitionReplicatedServersInfo.PartitionInfo[numPartitions];
+ for (int p = 0; p < numPartitions; p++) {
+ Set<String> fullyReplicatedServers = new HashSet<>();
+ for (int r = 0; r < replicasPerPartition; r++) {
+ fullyReplicatedServers.add(servers[serverIdxPerPartition[p] +
r].getInstanceId());
+ }
+ partitionInfoMap[p] =
+ new
TablePartitionReplicatedServersInfo.PartitionInfo(fullyReplicatedServers,
List.of("seg" + p));
+ }
+ TablePartitionReplicatedServersInfo tablePartitionInfo = new
TablePartitionReplicatedServersInfo(
+ PARTITIONED_TABLE_OFFLINE, "col1", "Hashcode", numPartitions,
partitionInfoMap, List.of());
+
+ // Model the pruned routing table: surviving segments bucketed onto their
owning server.
+ Map<ServerInstance, List<String>> serverToSegmentList = new HashMap<>();
+ for (String segment : survivingSegments) {
+ int partition = Integer.parseInt(segment.substring("seg".length()));
+
serverToSegmentList.computeIfAbsent(servers[serverIdxPerPartition[partition]],
k -> new ArrayList<>())
+ .add(segment);
+ }
+ Map<ServerInstance, SegmentsToQuery> serverToSegments = new HashMap<>();
+ serverToSegmentList.forEach((server, segments) ->
serverToSegments.put(server,
+ new SegmentsToQuery(segments, List.of())));
+ RoutingTable prunedRoutingTable =
+ new RoutingTable(serverToSegments, new
ArrayList<>(unavailableSegments), reportedPrunedByRouting);
+
+ PartitionedRoutingManager routingManager = new
PartitionedRoutingManager(enabledServers,
+ Map.of(PARTITIONED_TABLE_OFFLINE, tablePartitionInfo),
+ Map.of(PARTITIONED_TABLE_OFFLINE, prunedRoutingTable), throwOnRouting);
+
+ Map<String, String> tableNameMap = new HashMap<>();
+ tableNameMap.put(PARTITIONED_TABLE_OFFLINE, PARTITIONED_TABLE_OFFLINE);
+ tableNameMap.put(PARTITIONED_TABLE, PARTITIONED_TABLE);
+ TableCache tableCache = mock(TableCache.class);
+ when(tableCache.getTableNameMap()).thenReturn(tableNameMap);
+ when(tableCache.getActualTableName(anyString())).thenAnswer(inv ->
tableNameMap.get(inv.getArgument(0)));
+
when(tableCache.getSchema(anyString())).thenReturn(getSchemaBuilder(PARTITIONED_TABLE).build());
+
when(tableCache.getTableConfig(PARTITIONED_TABLE_OFFLINE)).thenReturn(mock(TableConfig.class));
+
+ WorkerManager workerManager = new WorkerManager("Broker_localhost",
"localhost", 5, routingManager);
+ return new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
workerManager);
+ }
+
+ /** Returns the leaf table-scan fragment: the only fragment with segment
assignments. */
+ @Nullable
+ private static DispatchablePlanFragment leafFragment(DispatchableSubPlan
dispatchableSubPlan) {
+ List<DispatchablePlanFragment> leafFragments =
leafFragments(dispatchableSubPlan);
+ return leafFragments.isEmpty() ? null : leafFragments.get(0);
+ }
+
+ /** Returns all fragments with segment assignments (the table-scan leaves).
*/
+ private static List<DispatchablePlanFragment>
leafFragments(DispatchableSubPlan dispatchableSubPlan) {
+ List<DispatchablePlanFragment> leafFragments = new ArrayList<>();
+ for (DispatchablePlanFragment fragment :
dispatchableSubPlan.getQueryStageMap().values()) {
+ if (!fragment.getWorkerIdToSegmentsMap().isEmpty()) {
+ leafFragments.add(fragment);
+ }
+ }
+ return leafFragments;
+ }
+
+ /** Maps each assigned segment to the instance id of the server its worker
was placed on. */
+ private static Map<String, String> segmentToServer(DispatchablePlanFragment
leafFragment) {
+ Map<Integer, String> workerIdToServer = new HashMap<>();
+ for (Map.Entry<QueryServerInstance, List<Integer>> entry
+ : leafFragment.getServerInstanceToWorkerIdMap().entrySet()) {
+ for (Integer workerId : entry.getValue()) {
+ workerIdToServer.put(workerId, entry.getKey().getInstanceId());
+ }
+ }
+ Map<String, String> segmentToServer = new HashMap<>();
+ for (Map.Entry<Integer, Map<String, List<String>>> entry :
leafFragment.getWorkerIdToSegmentsMap().entrySet()) {
+ String server = workerIdToServer.get(entry.getKey());
+ for (List<String> segments : entry.getValue().values()) {
+ for (String segment : segments) {
+ segmentToServer.put(segment, server);
+ }
+ }
+ }
+ return segmentToServer;
+ }
+
+ private static List<String> assignedSegments(DispatchablePlanFragment
leafFragment) {
+ List<String> segments = new ArrayList<>();
+ for (Map<String, List<String>> segmentsByType :
leafFragment.getWorkerIdToSegmentsMap().values()) {
+ segmentsByType.values().forEach(segments::addAll);
+ }
+ return segments;
+ }
+
/**
* Tests that literal-only stages (e.g. UNION ALL of constant values) are
assigned to the same
* servers as the table-scanning stages, not to all enabled servers across
all tenants.
@@ -645,4 +1080,96 @@ public class WorkerManagerTest {
return false;
}
}
+
+ /**
+ * A RoutingManager for the partitioned leaf path. It exposes a {@link
TablePartitionReplicatedServersInfo} per typed
+ * table (driving {@code calculatePartitionTableInfo}) and returns a
pre-configured {@link RoutingTable} of surviving
+ * segments from getRoutingTable (simulating what the real segment pruners
would return for the query filter). This
+ * lets the test drive which partitions survive without depending on the
pruner internals (which are tested
+ * separately).
+ */
+ private static class PartitionedRoutingManager implements RoutingManager {
+ private final Map<String, ServerInstance> _enabledServers;
+ private final Map<String, TablePartitionReplicatedServersInfo>
_partitionInfoByTable;
+ private final Map<String, RoutingTable> _routingTableByTable;
+ private final boolean _throwOnRouting;
+ @Nullable
+ private final TimeBoundaryInfo _timeBoundaryInfo;
+
+ PartitionedRoutingManager(Map<String, ServerInstance> enabledServers,
+ Map<String, TablePartitionReplicatedServersInfo> partitionInfoByTable,
+ Map<String, RoutingTable> routingTableByTable, boolean throwOnRouting)
{
+ this(enabledServers, partitionInfoByTable, routingTableByTable,
throwOnRouting, null);
+ }
+
+ PartitionedRoutingManager(Map<String, ServerInstance> enabledServers,
+ Map<String, TablePartitionReplicatedServersInfo> partitionInfoByTable,
+ Map<String, RoutingTable> routingTableByTable, boolean throwOnRouting,
+ @Nullable TimeBoundaryInfo timeBoundaryInfo) {
+ _enabledServers = enabledServers;
+ _partitionInfoByTable = partitionInfoByTable;
+ _routingTableByTable = routingTableByTable;
+ _throwOnRouting = throwOnRouting;
+ _timeBoundaryInfo = timeBoundaryInfo;
+ }
+
+ @Override
+ public Map<String, ServerInstance> getEnabledServerInstanceMap() {
+ return _enabledServers;
+ }
+
+ @Nullable
+ @Override
+ public RoutingTable getRoutingTable(BrokerRequest brokerRequest, long
requestId) {
+ if (_throwOnRouting) {
+ throw new RuntimeException("Simulated routing failure");
+ }
+ return
_routingTableByTable.get(brokerRequest.getQuerySource().getTableName());
+ }
+
+ @Nullable
+ @Override
+ public RoutingTable getRoutingTable(BrokerRequest brokerRequest, String
tableNameWithType, long requestId) {
+ return getRoutingTable(brokerRequest, requestId);
+ }
+
+ @Nullable
+ @Override
+ public List<String> getSegments(BrokerRequest brokerRequest) {
+ return List.of();
+ }
+
+ @Override
+ public boolean routingExists(String tableNameWithType) {
+ return _partitionInfoByTable.containsKey(tableNameWithType);
+ }
+
+ @Nullable
+ @Override
+ public TimeBoundaryInfo getTimeBoundaryInfo(String offlineTableName) {
+ return _timeBoundaryInfo;
+ }
+
+ @Nullable
+ @Override
+ public TablePartitionInfo getTablePartitionInfo(String tableNameWithType) {
+ return null;
+ }
+
+ @Nullable
+ @Override
+ public TablePartitionReplicatedServersInfo
getTablePartitionReplicatedServersInfo(String tableNameWithType) {
+ return _partitionInfoByTable.get(tableNameWithType);
+ }
+
+ @Override
+ public Set<String> getServingInstances(String tableNameWithType) {
+ return new HashSet<>(_enabledServers.keySet());
+ }
+
+ @Override
+ public boolean isTableDisabled(String tableNameWithType) {
+ return false;
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
