You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2026/05/01 06:45:21 UTC
(pinot) branch master updated: [multistage] add broker pruning support to non-partitioned leaf path (#18237)
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 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 a4bef7bc7c3 [multistage] add broker pruning support to non-partitioned leaf path (#18237)
a4bef7bc7c3 is described below
commit a4bef7bc7c31dda01773e9255deb6f9145d6e303
Author: dang-stripe <50...@users.noreply.github.com>
AuthorDate: Thu Apr 30 23:45:14 2026 -0700
[multistage] add broker pruning support to non-partitioned leaf path (#18237)
* add broker pruning support to non-partitioned leaf path
* fix config option
* address feedback
* address feedback
* add null guard and fix test
* empty commit to trigger CI
---
.../MultiStageBrokerRequestHandler.java | 15 ++
.../response/broker/BrokerResponseNativeV2.java | 3 +-
.../org/apache/pinot/query/QueryEnvironment.java | 14 +-
.../planner/logical/LeafStageToPinotQuery.java | 2 +-
.../planner/physical/DispatchablePlanContext.java | 9 +
.../planner/physical/DispatchableSubPlan.java | 9 +-
.../planner/physical/PinotDispatchPlanner.java | 3 +-
.../query/routing/PlanNodeRoutingQueryBuilder.java | 105 ++++++++
.../apache/pinot/query/routing/WorkerManager.java | 87 +++++-
.../routing/PlanNodeRoutingQueryBuilderTest.java | 61 +++++
.../pinot/query/routing/WorkerManagerTest.java | 292 ++++++++++++++++++++-
.../apache/pinot/spi/utils/CommonConstants.java | 18 +-
12 files changed, 596 insertions(+), 22 deletions(-)
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
index 9ccfad0a612..846f6acb3c8 100644
--- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
@@ -53,6 +53,7 @@ import org.apache.pinot.broker.querylog.QueryLogger;
import org.apache.pinot.broker.queryquota.QueryQuotaManager;
import org.apache.pinot.common.config.TlsConfig;
import org.apache.pinot.common.config.provider.TableCache;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.failuredetector.FailureDetector;
import org.apache.pinot.common.metrics.BrokerGauge;
import org.apache.pinot.common.metrics.BrokerMeter;
@@ -494,6 +495,9 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
boolean defaultUseBrokerPruning = _config.getProperty(
CommonConstants.Broker.CONFIG_OF_USE_BROKER_PRUNING,
CommonConstants.Broker.DEFAULT_USE_BROKER_PRUNING);
+ boolean defaultLogicalPlannerUseBrokerPruning = _config.getProperty(
+ CommonConstants.Broker.CONFIG_OF_LOGICAL_PLANNER_USE_BROKER_PRUNING,
+ CommonConstants.Broker.DEFAULT_LOGICAL_PLANNER_USE_BROKER_PRUNING);
int defaultLiteModeLeafStageLimit = _config.getProperty(
CommonConstants.Broker.CONFIG_OF_LITE_MODE_LEAF_STAGE_LIMIT,
CommonConstants.Broker.DEFAULT_LITE_MODE_LEAF_STAGE_LIMIT);
@@ -531,6 +535,7 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
.defaultUseLiteMode(defaultUseLiteMode)
.defaultRunInBroker(defaultRunInBroker)
.defaultUseBrokerPruning(defaultUseBrokerPruning)
+ .defaultLogicalPlannerUseBrokerPruning(defaultLogicalPlannerUseBrokerPruning)
.defaultLiteModeLeafStageLimit(defaultLiteModeLeafStageLimit)
.defaultLiteModeLeafStageFanOutAdjustedLimit(defaultLiteModeFanoutAdjustedLimit)
.defaultLiteModeEnableJoins(defaultLiteModeEnableJoins)
@@ -824,6 +829,16 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
stageStats.forEach((type, stats) -> type.mergeInto(brokerResponse, stats));
}
}
+ // Broker-pruned segments are computed during routing in WorkerManager, not reported by servers.
+ // Inject the count directly since it doesn't flow through the LeafOperator stat pipeline.
+ long numSegmentsPrunedByBroker = dispatchableSubPlan.getNumSegmentsPrunedByBroker();
+ if (numSegmentsPrunedByBroker > 0) {
+ StatMap<BrokerResponseNativeV2.StatKey> brokerPruningStats =
+ new StatMap<>(BrokerResponseNativeV2.StatKey.class);
+ brokerPruningStats.merge(BrokerResponseNativeV2.StatKey.NUM_SEGMENTS_PRUNED_BY_BROKER,
+ (int) numSegmentsPrunedByBroker);
+ brokerResponse.addBrokerStats(brokerPruningStats);
+ }
} catch (Exception e) {
LOGGER.warn("Error encountered while collecting multi-stage stats", e);
brokerResponse.setStageStats(JsonNodeFactory.instance.objectNode()
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
index e485912feea..69ae90b675a 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
@@ -333,7 +333,7 @@ public class BrokerResponseNativeV2 implements BrokerResponse {
@Override
public long getNumSegmentsPrunedByBroker() {
- return 0;
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PRUNED_BY_BROKER);
}
@Override
@@ -474,6 +474,7 @@ public class BrokerResponseNativeV2 implements BrokerResponse {
return StatMap.Key.minPositive(value1, value2);
}
},
+ NUM_SEGMENTS_PRUNED_BY_BROKER(StatMap.Type.INT),
NUM_SEGMENTS_PRUNED_BY_SERVER(StatMap.Type.INT),
NUM_SEGMENTS_PRUNED_INVALID(StatMap.Type.INT),
NUM_SEGMENTS_PRUNED_BY_LIMIT(StatMap.Type.INT),
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
index 7535cdcbca0..e55566d1066 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
@@ -786,7 +786,7 @@ public class QueryEnvironment {
}
/**
- * Whether to use broker pruning by default.
+ * Whether to use broker pruning by default on the physical optimizer path.
*
* This is treated as the default value for the broker and it is expected to be obtained from a Pinot configuration.
* This default value can be always overridden at query level by the query option
@@ -797,6 +797,18 @@ public class QueryEnvironment {
return CommonConstants.Broker.DEFAULT_USE_BROKER_PRUNING;
}
+ /**
+ * Whether to use broker pruning by default on the logical planner (non-physical-optimizer) path.
+ *
+ * Separated from {@link #defaultUseBrokerPruning()} so the two planner paths can be rolled out independently.
+ * This default value can be always overridden at query level by the query option
+ * {@link CommonConstants.Broker.Request.QueryOptionKey#USE_BROKER_PRUNING}.
+ */
+ @Value.Default
+ default boolean defaultLogicalPlannerUseBrokerPruning() {
+ return CommonConstants.Broker.DEFAULT_LOGICAL_PLANNER_USE_BROKER_PRUNING;
+ }
+
/**
* Default server stage limit for lite mode queries.
*
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/LeafStageToPinotQuery.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/LeafStageToPinotQuery.java
index bc19313b55f..99d2235f99a 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/LeafStageToPinotQuery.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/LeafStageToPinotQuery.java
@@ -120,7 +120,7 @@ public class LeafStageToPinotQuery {
* Note: This method mutates the input expression's operand lists in-place for AND/OR/NOT nodes.
* It assumes the expression tree is freshly constructed and not shared across concurrent callers.
*/
- static Expression ensureFilterIsFunctionExpression(Expression expression) {
+ public static Expression ensureFilterIsFunctionExpression(Expression expression) {
if (expression == null) {
return null;
}
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 fe89e51a151..75e82dac475 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java
@@ -60,6 +60,7 @@ public class DispatchablePlanContext {
private final Map<Integer, DispatchablePlanMetadata> _dispatchablePlanMetadataMap = new HashMap<>();
private final Map<Integer, PlanNode> _dispatchablePlanStageRootMap = new HashMap<>();
+ private long _numSegmentsPrunedByBroker;
public DispatchablePlanContext(WorkerManager workerManager, long requestId, PlannerContext plannerContext,
@@ -130,6 +131,14 @@ public class DispatchablePlanContext {
return _dispatchablePlanStageRootMap;
}
+ public long getNumSegmentsPrunedByBroker() {
+ return _numSegmentsPrunedByBroker;
+ }
+
+ public void addNumSegmentsPrunedByBroker(long count) {
+ _numSegmentsPrunedByBroker += count;
+ }
+
public Map<Integer, DispatchablePlanFragment> constructDispatchablePlanFragmentMap(PlanFragment subPlanRoot) {
Map<Integer, DispatchablePlanFragment> dispatchablePlanFragmentMap = createDispatchablePlanFragmentMap(subPlanRoot);
for (Map.Entry<Integer, DispatchablePlanMetadata> planMetadataEntry : _dispatchablePlanMetadataMap.entrySet()) {
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchableSubPlan.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchableSubPlan.java
index 83458f6cc7e..660e7c02a06 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchableSubPlan.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchableSubPlan.java
@@ -50,14 +50,17 @@ public class DispatchableSubPlan {
private final Map<Integer, DispatchablePlanFragment> _queryStageMap;
private final Set<String> _tableNames;
private final Map<String, Set<String>> _tableToUnavailableSegmentsMap;
+ private final long _numSegmentsPrunedByBroker;
public DispatchableSubPlan(PairList<Integer, String> fields,
Map<Integer, DispatchablePlanFragment> queryStageMap,
- Set<String> tableNames, Map<String, Set<String>> tableToUnavailableSegmentsMap) {
+ Set<String> tableNames, Map<String, Set<String>> tableToUnavailableSegmentsMap,
+ long numSegmentsPrunedByBroker) {
_queryResultFields = fields;
_queryStageMap = queryStageMap;
_tableNames = tableNames;
_tableToUnavailableSegmentsMap = tableToUnavailableSegmentsMap;
+ _numSegmentsPrunedByBroker = numSegmentsPrunedByBroker;
}
/**
@@ -122,6 +125,10 @@ public class DispatchableSubPlan {
return _tableToUnavailableSegmentsMap;
}
+ public long getNumSegmentsPrunedByBroker() {
+ return _numSegmentsPrunedByBroker;
+ }
+
/**
* Get the estimated total number of threads that will be spawned for this query (across all stages and servers).
*/
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java
index 6cb3e51bf6c..a1414d48d64 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java
@@ -123,7 +123,8 @@ public class PinotDispatchPlanner {
return new DispatchableSubPlan(dispatchablePlanContext.getResultFields(),
dispatchablePlanContext.constructDispatchablePlanFragmentMap(subPlanRoot),
dispatchablePlanContext.getTableNames(),
- populateTableUnavailableSegments(dispatchablePlanContext.getDispatchablePlanMetadataMap()));
+ populateTableUnavailableSegments(dispatchablePlanContext.getDispatchablePlanMetadataMap()),
+ dispatchablePlanContext.getNumSegmentsPrunedByBroker());
}
private static Map<String, Set<String>> populateTableUnavailableSegments(
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
new file mode 100644
index 00000000000..3171ec0bef9
--- /dev/null
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java
@@ -0,0 +1,105 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.routing;
+
+import com.google.common.base.Preconditions;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.pinot.common.request.DataSource;
+import org.apache.pinot.common.request.Expression;
+import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.common.utils.request.RequestUtils;
+import org.apache.pinot.query.parser.CalciteRexExpressionParser;
+import org.apache.pinot.query.planner.logical.LeafStageToPinotQuery;
+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.TableScanNode;
+
+
+/**
+ * Converts an MSE leaf stage {@link PlanNode} tree to a {@link PinotQuery} for routing purposes. Used by the broker
+ * pruning path in {@link WorkerManager} to build a filter-bearing routing query that enables segment pruning at the
+ * broker.
+ */
+public class PlanNodeRoutingQueryBuilder {
+ private PlanNodeRoutingQueryBuilder() {
+ }
+
+ /**
+ * Converts a PlanNode leaf stage root to a {@link PinotQuery} for routing purposes. Only handles Project, Filter
+ * and TableScan nodes — other node types in the chain are silently skipped. Callers should expect this method
+ * to throw on malformed trees (e.g., missing TableScanNode, multi-input nodes) and handle failures gracefully.
+ *
+ * @param tableName the table name (with or without type suffix). Passed explicitly because PlanNode trees
+ * don't carry the resolved table name.
+ * @param leafStageRoot the root of the leaf stage
+ * @param skipFilter whether to skip the filter in the query
+ * @return a {@link PinotQuery} representing the leaf stage
+ */
+ public static PinotQuery createPinotQueryForRouting(String tableName, PlanNode leafStageRoot, boolean skipFilter) {
+ List<PlanNode> bottomToTopNodes = new ArrayList<>();
+ accumulateBottomToTop(leafStageRoot, bottomToTopNodes);
+ Preconditions.checkState(!bottomToTopNodes.isEmpty() && bottomToTopNodes.get(0) instanceof TableScanNode,
+ "Could not find table scan");
+ TableScanNode tableScan = (TableScanNode) bottomToTopNodes.get(0);
+ PinotQuery pinotQuery = initializePinotQueryForTableScan(tableName, tableScan);
+ for (PlanNode parentNode : bottomToTopNodes) {
+ if (parentNode instanceof FilterNode) {
+ if (!skipFilter) {
+ handleFilter((FilterNode) parentNode, pinotQuery);
+ }
+ } else if (parentNode instanceof ProjectNode) {
+ handleProject((ProjectNode) parentNode, pinotQuery);
+ }
+ }
+ return pinotQuery;
+ }
+
+ 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());
+ for (PlanNode input : root.getInputs()) {
+ accumulateBottomToTop(input, parentNodes);
+ }
+ parentNodes.add(root);
+ }
+
+ private static PinotQuery initializePinotQueryForTableScan(String tableName, TableScanNode tableScan) {
+ PinotQuery pinotQuery = new PinotQuery();
+ pinotQuery.setDataSource(new DataSource());
+ pinotQuery.getDataSource().setTableName(tableName);
+ pinotQuery.setSelectList(tableScan.getColumns().stream().map(
+ RequestUtils::getIdentifierExpression).collect(Collectors.toList()));
+ return pinotQuery;
+ }
+
+ private static void handleProject(ProjectNode project, PinotQuery pinotQuery) {
+ List<Expression> selectList = CalciteRexExpressionParser.convertRexNodes(project.getProjects(),
+ pinotQuery.getSelectList());
+ pinotQuery.setSelectList(selectList);
+ }
+
+ private static void handleFilter(FilterNode filter, PinotQuery pinotQuery) {
+ Expression filterExpression = CalciteRexExpressionParser.toExpression(filter.getCondition(),
+ pinotQuery.getSelectList());
+ pinotQuery.setFilterExpression(LeafStageToPinotQuery.ensureFilterIsFunctionExpression(filterExpression));
+ }
+}
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 9c9fb6d7901..6c10f2af2b5 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
@@ -40,6 +40,7 @@ import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType;
import org.apache.pinot.calcite.rel.rules.ImmutableTableOptions;
import org.apache.pinot.calcite.rel.rules.TableOptions;
import org.apache.pinot.common.request.BrokerRequest;
+import org.apache.pinot.common.request.PinotQuery;
import org.apache.pinot.common.utils.config.QueryOptionsUtils;
import org.apache.pinot.core.routing.LogicalTableRouteInfo;
import org.apache.pinot.core.routing.LogicalTableRouteProvider;
@@ -501,7 +502,7 @@ public class WorkerManager {
if (metadata.getLogicalTableRouteInfo() != null) {
assignWorkersToNonPartitionedLeafFragmentForLogicalTable(metadata, context);
} else {
- assignWorkersToNonPartitionedLeafFragment(metadata, context);
+ assignWorkersToNonPartitionedLeafFragment(fragment, metadata, context);
}
addLeafServersToContext(metadata, context);
}
@@ -517,11 +518,15 @@ public class WorkerManager {
// --------------------------------------------------------------------------
// Non-partitioned leaf stage assignment
// --------------------------------------------------------------------------
- private void assignWorkersToNonPartitionedLeafFragment(DispatchablePlanMetadata metadata,
+ private void assignWorkersToNonPartitionedLeafFragment(PlanFragment fragment, DispatchablePlanMetadata metadata,
DispatchablePlanContext context) {
String tableName = metadata.getScannedTables().get(0);
- Map<String, RoutingTable> routingTableMap =
- getRoutingTable(tableName, context.getRequestId(), context.getPlannerContext().getOptions());
+ PinotQuery routingPinotQuery = extractRoutingQuery(fragment.getFragmentRoot(), tableName, context);
+ // When broker pruning is enabled, routingPinotQuery carries the leaf stage filter so that segment pruners can
+ // eliminate segments. When disabled (null), fall back to an unfiltered SELECT * routing request.
+ Map<String, RoutingTable> routingTableMap = routingPinotQuery != null
+ ? getRoutingTable(routingPinotQuery, context.getRequestId())
+ : getRoutingTable(tableName, context.getRequestId(), context.getPlannerContext().getOptions());
Preconditions.checkState(!routingTableMap.isEmpty(), "Unable to find routing entries for table: %s", tableName);
// acquire time boundary info if it is a hybrid table.
@@ -555,6 +560,9 @@ public class WorkerManager {
if (!routingTable.getUnavailableSegments().isEmpty()) {
metadata.addUnavailableSegments(tableName, routingTable.getUnavailableSegments());
}
+ if (routingPinotQuery != null) {
+ context.addNumSegmentsPrunedByBroker(routingTable.getNumPrunedSegments());
+ }
}
// Sort server instances to ensure deterministic worker ID assignment.
// This is critical for pre-partitioned exchanges where worker ID N on one stage
@@ -583,7 +591,8 @@ public class WorkerManager {
}
/**
- * Acquire routing table for items listed in {@link TableScanNode}.
+ * Acquire routing table for items listed in {@link TableScanNode}. Creates a bare {@code SELECT *} broker request
+ * with no filter, so no broker-side segment pruning occurs.
*
* @param tableName table name with or without type suffix.
* @return keyed-map from table type(s) to routing table(s).
@@ -616,6 +625,61 @@ public class WorkerManager {
}
}
+ /**
+ * Acquire routing table using a pre-built {@link PinotQuery} that carries filter expressions for segment pruning.
+ * Unlike {@link #getRoutingTable(String, long)} which creates a bare {@code SELECT *} broker request,
+ * this overload forwards the filter to the routing manager so broker-side segment pruners can eliminate
+ * segments before dispatching to servers.
+ *
+ * @param pinotQuery the routing query with filter expressions for pruning. Table name may be raw or typed.
+ * @return keyed-map from table type(s) to routing table(s).
+ */
+ private Map<String, RoutingTable> getRoutingTable(PinotQuery pinotQuery, long requestId) {
+ TableType tableType = TableNameBuilder.getTableTypeFromTableName(pinotQuery.getDataSource().getTableName());
+ if (tableType == null) {
+ Map<String, RoutingTable> routingTableMap = new HashMap<>(4);
+ RoutingTable offlineRoutingTable = getRoutingTableHelper(pinotQuery, requestId, TableType.OFFLINE);
+ if (offlineRoutingTable != null) {
+ routingTableMap.put(TableType.OFFLINE.name(), offlineRoutingTable);
+ }
+ RoutingTable realtimeRoutingTable = getRoutingTableHelper(pinotQuery, requestId, TableType.REALTIME);
+ if (realtimeRoutingTable != null) {
+ routingTableMap.put(TableType.REALTIME.name(), realtimeRoutingTable);
+ }
+ return routingTableMap;
+ } else {
+ RoutingTable routingTable = getRoutingTableHelper(pinotQuery, requestId);
+ return routingTable != null ? Map.of(tableType.name(), routingTable) : Map.of();
+ }
+ }
+
+ /**
+ * Builds a {@link PinotQuery} from the leaf stage tree for broker-side segment pruning on the logical planner path.
+ * Returns {@code null} if broker pruning is disabled or the leaf stage shape is unsupported.
+ */
+ @Nullable
+ private PinotQuery extractRoutingQuery(PlanNode leafStageRoot, String tableName, DispatchablePlanContext context) {
+ boolean defaultLogicalPlannerUseBrokerPruning =
+ context.getPlannerContext().getEnvConfig().defaultLogicalPlannerUseBrokerPruning();
+ boolean useBrokerPruning = QueryOptionsUtils.isUseBrokerPruning(
+ context.getPlannerContext().getOptions(), defaultLogicalPlannerUseBrokerPruning);
+ if (!useBrokerPruning) {
+ return null;
+ }
+ try {
+ PinotQuery pinotQuery = PlanNodeRoutingQueryBuilder.createPinotQueryForRouting(tableName, leafStageRoot, false);
+ Map<String, String> queryOptions = context.getPlannerContext().getOptions();
+ if (MapUtils.isNotEmpty(queryOptions)) {
+ pinotQuery.setQueryOptions(new HashMap<>(queryOptions));
+ }
+ return pinotQuery;
+ } catch (RuntimeException e) {
+ LOGGER.warn("Broker pruning skipped for table {} due to unsupported leaf stage shape: {}",
+ tableName, e.getMessage());
+ return null;
+ }
+ }
+
@Nullable
private RoutingTable getRoutingTableHelper(String tableNameWithType, long requestId,
Map<String, String> queryOptions) {
@@ -628,6 +692,19 @@ public class WorkerManager {
return _routingManager.getRoutingTable(brokerRequest, requestId);
}
+ @Nullable
+ private RoutingTable getRoutingTableHelper(PinotQuery pinotQuery, long requestId) {
+ return _routingManager.getRoutingTable(CalciteSqlCompiler.convertToBrokerRequest(pinotQuery), requestId);
+ }
+
+ @Nullable
+ private RoutingTable getRoutingTableHelper(PinotQuery pinotQuery, long requestId, TableType tableType) {
+ PinotQuery copy = pinotQuery.deepCopy();
+ copy.getDataSource().setTableName(TableNameBuilder.forType(tableType).tableNameWithType(
+ TableNameBuilder.extractRawTableName(pinotQuery.getDataSource().getTableName())));
+ return getRoutingTableHelper(copy, requestId);
+ }
+
// --------------------------------------------------------------------------
// Replicated non-partitioned leaf stage assignment
// --------------------------------------------------------------------------
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
new file mode 100644
index 00000000000..a1532eff487
--- /dev/null
+++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilderTest.java
@@ -0,0 +1,61 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.routing;
+
+import java.util.List;
+import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.query.planner.logical.RexExpression;
+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.TableScanNode;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+
+
+public class PlanNodeRoutingQueryBuilderTest {
+ private static final DataSchema TEST_SCHEMA = new DataSchema(new String[]{"col1", "col2"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING, DataSchema.ColumnDataType.STRING});
+
+ @Test
+ public void testCreatePinotQueryForRoutingPlanNodeTree() {
+ 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"))));
+ ProjectNode projectNode = new ProjectNode(1, new DataSchema(new String[]{"col2"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING}), PlanNode.NodeHint.EMPTY,
+ List.of(filterNode), List.of(new RexExpression.InputRef(1)));
+
+ PinotQuery pinotQuery = PlanNodeRoutingQueryBuilder.createPinotQueryForRouting("testTable", projectNode, false);
+
+ assertEquals(pinotQuery.getDataSource().getTableName(), "testTable");
+ assertEquals(pinotQuery.getSelectList().size(), 1);
+ assertEquals(pinotQuery.getSelectList().get(0).getIdentifier().getName(), "col2");
+ assertNotNull(pinotQuery.getFilterExpression());
+ assertEquals(pinotQuery.getFilterExpression().getFunctionCall().getOperator(), "EQUALS");
+ assertEquals(pinotQuery.getFilterExpression().getFunctionCall().getOperands().get(0).getIdentifier().getName(),
+ "col1");
+ }
+}
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 0643901760d..07b614c9ef5 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
@@ -22,6 +22,7 @@ import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
+import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -30,6 +31,7 @@ import org.apache.helix.model.InstanceConfig;
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.core.routing.RoutingManager;
import org.apache.pinot.core.routing.RoutingTable;
import org.apache.pinot.core.routing.SegmentsToQuery;
@@ -40,6 +42,7 @@ import org.apache.pinot.core.transport.ServerInstance;
import org.apache.pinot.query.QueryEnvironment;
import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
import org.apache.pinot.query.planner.physical.DispatchableSubPlan;
+import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.data.FieldSpec;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.utils.CommonConstants;
@@ -48,8 +51,10 @@ import org.testng.annotations.Test;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
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;
@@ -111,8 +116,7 @@ public class WorkerManagerTest {
when(tableCache.getTableNameMap()).thenReturn(tableNameMap);
when(tableCache.getActualTableName(anyString())).thenAnswer(inv -> tableNameMap.get(inv.getArgument(0)));
when(tableCache.getSchema(anyString())).thenReturn(emptyTableSchema);
- when(tableCache.getTableConfig("emptyTable_OFFLINE"))
- .thenReturn(mock(org.apache.pinot.spi.config.table.TableConfig.class));
+ when(tableCache.getTableConfig("emptyTable_OFFLINE")).thenReturn(mock(TableConfig.class));
WorkerManager workerManager = new WorkerManager("Broker_localhost", "localhost", 3, routingManager);
QueryEnvironment queryEnvironment = new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
@@ -123,9 +127,209 @@ public class WorkerManagerTest {
String query = "SET useLeafServerForIntermediateStage=true; SELECT * FROM emptyTable LIMIT 10";
// This should not throw "bound must be positive" error anymore
- @SuppressWarnings("deprecation")
- DispatchableSubPlan dispatchableSubPlan = queryEnvironment.planQuery(query);
- assertNotNull(dispatchableSubPlan);
+ try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(query)) {
+ DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan);
+ }
+ }
+
+ @Test
+ public void testBrokerPruningUsesFilteredRoutingQueryOnThisPath() {
+ Schema schema = getSchemaBuilder("testTable").build();
+ ServerInstance server = getServerInstance("localhost", 1);
+ Map<String, ServerInstance> serverInstanceMap = Map.of(server.getInstanceId(), server);
+ RoutingTable routingTable = new RoutingTable(Map.of(server, new SegmentsToQuery(List.of("segment1"), List.of())),
+ List.of(), 0);
+ CapturingRoutingManager routingManager = new CapturingRoutingManager(serverInstanceMap,
+ Map.of("testTable_OFFLINE", routingTable));
+
+ Map<String, String> tableNameMap = new HashMap<>();
+ tableNameMap.put("testTable_OFFLINE", "testTable_OFFLINE");
+ tableNameMap.put("testTable", "testTable");
+
+ 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(schema);
+ when(tableCache.getTableConfig("testTable_OFFLINE")).thenReturn(mock(TableConfig.class));
+
+ WorkerManager workerManager = new WorkerManager("Broker_localhost", "localhost", 3, routingManager);
+ QueryEnvironment queryEnvironment = new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
+ workerManager);
+
+ try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(
+ "SET useBrokerPruning=true; SELECT col2 FROM testTable WHERE col1 = 'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan);
+ }
+
+ BrokerRequest brokerRequest = routingManager.getCapturedRoutingRequest("testTable_OFFLINE");
+ assertNotNull(brokerRequest);
+ Expression filterExpression = brokerRequest.getPinotQuery().getFilterExpression();
+ assertNotNull(filterExpression);
+ assertEquals(filterExpression.getFunctionCall().getOperator(), "EQUALS");
+ assertEquals(filterExpression.getFunctionCall().getOperands().get(0).getIdentifier().getName(), "col1");
+ assertEquals(brokerRequest.getPinotQuery().getSelectList().size(), 1);
+ assertEquals(brokerRequest.getPinotQuery().getSelectList().get(0).getIdentifier().getName(), "col2");
+ }
+
+ @Test
+ public void testBrokerPruningRoutesFilterToBothHybridTableTypesOnThisPath() {
+ Schema schema = getSchemaBuilder("testTable").build();
+ ServerInstance server1 = getServerInstance("localhost", 1);
+ ServerInstance server2 = getServerInstance("localhost", 2);
+ Map<String, ServerInstance> serverInstanceMap = Map.of(
+ server1.getInstanceId(), server1, server2.getInstanceId(), server2);
+ RoutingTable offlineRoutingTable = new RoutingTable(
+ Map.of(server1, new SegmentsToQuery(List.of("offline_seg1"), List.of())), List.of(), 0);
+ RoutingTable realtimeRoutingTable = new RoutingTable(
+ Map.of(server2, new SegmentsToQuery(List.of("realtime_seg1"), List.of())), List.of(), 0);
+ CapturingRoutingManager routingManager = new CapturingRoutingManager(serverInstanceMap,
+ Map.of("testTable_OFFLINE", offlineRoutingTable, "testTable_REALTIME", realtimeRoutingTable));
+
+ Map<String, String> tableNameMap = new HashMap<>();
+ tableNameMap.put("testTable_OFFLINE", "testTable_OFFLINE");
+ tableNameMap.put("testTable_REALTIME", "testTable_REALTIME");
+ tableNameMap.put("testTable", "testTable");
+
+ 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(schema);
+ when(tableCache.getTableConfig("testTable_OFFLINE")).thenReturn(mock(TableConfig.class));
+ when(tableCache.getTableConfig("testTable_REALTIME")).thenReturn(mock(TableConfig.class));
+
+ WorkerManager workerManager = new WorkerManager("Broker_localhost", "localhost", 3, routingManager);
+ QueryEnvironment queryEnvironment = new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
+ workerManager);
+
+ try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(
+ "SET useBrokerPruning=true; SELECT col2 FROM testTable WHERE col1 = 'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan);
+ }
+
+ // Verify both table types received the filter from the query
+ for (String tableType : List.of("testTable_OFFLINE", "testTable_REALTIME")) {
+ BrokerRequest brokerRequest = routingManager.getCapturedRoutingRequest(tableType);
+ assertNotNull(brokerRequest, "Missing routing request for " + tableType);
+ assertEquals(brokerRequest.getPinotQuery().getDataSource().getTableName(), tableType);
+ Expression filterExpression = brokerRequest.getPinotQuery().getFilterExpression();
+ assertNotNull(filterExpression, "Missing filter for " + tableType);
+ assertEquals(filterExpression.getFunctionCall().getOperator(), "EQUALS");
+ assertEquals(filterExpression.getFunctionCall().getOperands().get(0).getIdentifier().getName(), "col1");
+ }
+ }
+
+ @Test
+ public void testBrokerPruningCountPropagatedToDispatchableSubPlan() {
+ Schema schema = getSchemaBuilder("testTable").build();
+ ServerInstance server = getServerInstance("localhost", 1);
+ Map<String, ServerInstance> serverInstanceMap = Map.of(server.getInstanceId(), server);
+ // RoutingTable with 42 pruned segments
+ RoutingTable routingTable = new RoutingTable(Map.of(server, new SegmentsToQuery(List.of("segment1"), List.of())),
+ List.of(), 42);
+ CapturingRoutingManager routingManager = new CapturingRoutingManager(serverInstanceMap,
+ Map.of("testTable_OFFLINE", routingTable));
+
+ Map<String, String> tableNameMap = new HashMap<>();
+ tableNameMap.put("testTable_OFFLINE", "testTable_OFFLINE");
+ tableNameMap.put("testTable", "testTable");
+
+ 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(schema);
+ when(tableCache.getTableConfig("testTable_OFFLINE")).thenReturn(mock(TableConfig.class));
+
+ WorkerManager workerManager = new WorkerManager("Broker_localhost", "localhost", 3, routingManager);
+ QueryEnvironment queryEnvironment = new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
+ workerManager);
+
+ try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(
+ "SET useBrokerPruning=true; SELECT col2 FROM testTable WHERE col1 = 'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan);
+ // Pruned count should propagate from RoutingTable through DispatchablePlanContext to DispatchableSubPlan
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 42);
+ }
+ }
+
+ @Test
+ public void testBrokerPruningDefaultsToUnfilteredRoutingOnThisPath() {
+ Schema schema = getSchemaBuilder("testTable").build();
+ ServerInstance server = getServerInstance("localhost", 1);
+ Map<String, ServerInstance> serverInstanceMap = Map.of(server.getInstanceId(), server);
+ RoutingTable routingTable = new RoutingTable(Map.of(server, new SegmentsToQuery(List.of("segment1"), List.of())),
+ List.of(), 0);
+ CapturingRoutingManager routingManager = new CapturingRoutingManager(serverInstanceMap,
+ Map.of("testTable_OFFLINE", routingTable));
+
+ Map<String, String> tableNameMap = new HashMap<>();
+ tableNameMap.put("testTable_OFFLINE", "testTable_OFFLINE");
+ tableNameMap.put("testTable", "testTable");
+
+ 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(schema);
+ when(tableCache.getTableConfig("testTable_OFFLINE")).thenReturn(mock(TableConfig.class));
+
+ WorkerManager workerManager = new WorkerManager("Broker_localhost", "localhost", 3, routingManager);
+ QueryEnvironment queryEnvironment = new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
+ workerManager);
+
+ // Without SET useBrokerPruning=true, this path should fall back to unfiltered SELECT * routing.
+ try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(
+ "SELECT col2 FROM testTable WHERE col1 = 'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan);
+ }
+
+ BrokerRequest brokerRequest = routingManager.getCapturedRoutingRequest("testTable_OFFLINE");
+ assertNotNull(brokerRequest);
+ // No filter should be present — the routing query is a plain SELECT * used for segment lookup only.
+ assertNull(brokerRequest.getPinotQuery().getFilterExpression());
+ }
+
+ @Test
+ public void testBrokerPruningPreservesQueryOptionsOnRoutingRequest() {
+ Schema schema = getSchemaBuilder("testTable").build();
+ ServerInstance server = getServerInstance("localhost", 1);
+ Map<String, ServerInstance> serverInstanceMap = Map.of(server.getInstanceId(), server);
+ RoutingTable routingTable = new RoutingTable(Map.of(server, new SegmentsToQuery(List.of("segment1"), List.of())),
+ List.of(), 0);
+ CapturingRoutingManager routingManager = new CapturingRoutingManager(serverInstanceMap,
+ Map.of("testTable_OFFLINE", routingTable));
+
+ Map<String, String> tableNameMap = new HashMap<>();
+ tableNameMap.put("testTable_OFFLINE", "testTable_OFFLINE");
+ tableNameMap.put("testTable", "testTable");
+
+ 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(schema);
+ when(tableCache.getTableConfig("testTable_OFFLINE")).thenReturn(mock(TableConfig.class));
+
+ WorkerManager workerManager = new WorkerManager("Broker_localhost", "localhost", 3, routingManager);
+ QueryEnvironment queryEnvironment = new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
+ workerManager);
+
+ try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(
+ "SET useBrokerPruning=true; SET useLeafServerForIntermediateStage=true;"
+ + " SELECT col2 FROM testTable WHERE col1 = 'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan);
+ }
+
+ BrokerRequest brokerRequest = routingManager.getCapturedRoutingRequest("testTable_OFFLINE");
+ assertNotNull(brokerRequest);
+ // Query options must be preserved on the broker-pruning routing path so that
+ // routing-affecting options are visible to the routing manager.
+ Map<String, String> queryOptions = brokerRequest.getPinotQuery().getQueryOptions();
+ assertNotNull(queryOptions);
+ assertEquals(queryOptions.get("useLeafServerForIntermediateStage"), "true");
}
/**
@@ -175,8 +379,7 @@ public class WorkerManagerTest {
when(tableCache.getTableNameMap()).thenReturn(tableNameMap);
when(tableCache.getActualTableName(anyString())).thenAnswer(inv -> tableNameMap.get(inv.getArgument(0)));
when(tableCache.getSchema(anyString())).thenReturn(tableSchema);
- when(tableCache.getTableConfig("testTable_OFFLINE"))
- .thenReturn(mock(org.apache.pinot.spi.config.table.TableConfig.class));
+ when(tableCache.getTableConfig("testTable_OFFLINE")).thenReturn(mock(TableConfig.class));
WorkerManager workerManager = new WorkerManager("Broker_localhost", "localhost", 5, routingManager);
QueryEnvironment queryEnvironment = new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
@@ -368,4 +571,79 @@ public class WorkerManagerTest {
return false;
}
}
+
+ private static class CapturingRoutingManager implements RoutingManager {
+ private final Map<String, ServerInstance> _serverInstanceMap;
+ private final Map<String, RoutingTable> _routingTableByName;
+ private final Map<String, BrokerRequest> _capturedRoutingRequests = new LinkedHashMap<>();
+
+ CapturingRoutingManager(Map<String, ServerInstance> serverInstanceMap,
+ Map<String, RoutingTable> routingTableByName) {
+ _serverInstanceMap = serverInstanceMap;
+ _routingTableByName = routingTableByName;
+ }
+
+ @Nullable
+ BrokerRequest getCapturedRoutingRequest(String tableNameWithType) {
+ return _capturedRoutingRequests.get(tableNameWithType);
+ }
+
+ @Override
+ public Map<String, ServerInstance> getEnabledServerInstanceMap() {
+ return _serverInstanceMap;
+ }
+
+ @Nullable
+ @Override
+ public RoutingTable getRoutingTable(BrokerRequest brokerRequest, long requestId) {
+ String tableNameWithType = brokerRequest.getQuerySource().getTableName();
+ _capturedRoutingRequests.put(tableNameWithType, brokerRequest);
+ return _routingTableByName.get(tableNameWithType);
+ }
+
+ @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 _routingTableByName.containsKey(tableNameWithType);
+ }
+
+ @Nullable
+ @Override
+ public TimeBoundaryInfo getTimeBoundaryInfo(String offlineTableName) {
+ return null;
+ }
+
+ @Nullable
+ @Override
+ public TablePartitionInfo getTablePartitionInfo(String tableNameWithType) {
+ return null;
+ }
+
+ @Nullable
+ @Override
+ public TablePartitionReplicatedServersInfo getTablePartitionReplicatedServersInfo(String tableNameWithType) {
+ return null;
+ }
+
+ @Override
+ public Set<String> getServingInstances(String tableNameWithType) {
+ return new HashSet<>(_serverInstanceMap.keySet());
+ }
+
+ @Override
+ public boolean isTableDisabled(String tableNameWithType) {
+ return false;
+ }
+ }
}
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 916ca208a4a..30c074a422c 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
@@ -574,12 +574,21 @@ public class CommonConstants {
public static final boolean DEFAULT_RUN_IN_BROKER = true;
/**
- * Whether to use broker pruning by default.
+ * Whether to use broker pruning by default on the physical optimizer path.
* This value can always be overridden by {@link Request.QueryOptionKey#USE_BROKER_PRUNING} query option
*/
public static final String CONFIG_OF_USE_BROKER_PRUNING = "pinot.broker.multistage.use.broker.pruning";
public static final boolean DEFAULT_USE_BROKER_PRUNING = true;
+ /**
+ * Whether to use broker pruning by default on the logical planner (non-physical-optimizer) path.
+ * This value can always be overridden by {@link Request.QueryOptionKey#USE_BROKER_PRUNING} query option.
+ * Separated from {@link #CONFIG_OF_USE_BROKER_PRUNING} so the two paths can be rolled out independently.
+ */
+ 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 = false;
+
/**
* Default server stage limit for lite mode queries.
* This value can always be overridden by {@link Request.QueryOptionKey#LITE_MODE_LEAF_STAGE_LIMIT} query option
@@ -884,10 +893,9 @@ public class CommonConstants {
// Server stage limit for lite mode queries.
public static final String LITE_MODE_LEAF_STAGE_LIMIT = "liteModeLeafStageLimit";
public static final String LITE_MODE_LEAF_STAGE_FANOUT_ADJUSTED_LIMIT = "liteModeLeafStageFanOutAdjustedLimit";
- // Used by the MSE Engine to determine whether to use the broker pruning logic. Only supported by the
- // new MSE query optimizer.
- // TODO(mse-physical): Consider removing this query option and making this the default, since there's already
- // a table config to enable broker pruning (it is disabled by default).
+ // Used by the MSE engine to enable broker-side segment pruning during routing. The physical optimizer
+ // path defaults to DEFAULT_USE_BROKER_PRUNING (true); the logical planner path defaults to
+ // DEFAULT_LOGICAL_PLANNER_USE_BROKER_PRUNING (false). Both can be overridden per-query.
public static final String USE_BROKER_PRUNING = "useBrokerPruning";
// When lite mode is enabled, if this flag is set, we will run all the non-leaf stage operators within the
// broker itself. That way, the MSE queries will model the scatter gather pattern used by the V1 Engine.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org