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 849c6206237 Enable broker segment pruning by default for the MSE
logical planner (and fix routing-filter construction bugs it surfaces) (#18941)
849c6206237 is described below
commit 849c6206237fd7ce7151229db87638bac9ff54c0
Author: Yash Mayya <[email protected]>
AuthorDate: Thu Jul 9 11:58:49 2026 -0700
Enable broker segment pruning by default for the MSE logical planner (and
fix routing-filter construction bugs it surfaces) (#18941)
---
...PartitionLLCRealtimeClusterIntegrationTest.java | 12 +-
.../BaseLogicalTableIntegrationTest.java | 6 +-
.../planner/logical/LeafStageToPinotQuery.java | 69 ++++-
.../query/routing/PlanNodeRoutingQueryBuilder.java | 27 +-
.../apache/pinot/query/routing/WorkerManager.java | 54 +++-
.../planner/logical/LeafStageToPinotQueryTest.java | 143 +++++++++++
.../routing/PlanNodeRoutingQueryBuilderTest.java | 60 +++++
.../pinot/query/routing/WorkerManagerTest.java | 286 ++++++++++++++++++++-
.../resources/queries/ExplainPhysicalPlans.json | 33 +--
.../apache/pinot/spi/utils/CommonConstants.java | 10 +-
10 files changed, 639 insertions(+), 61 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 75e1ad0e166..6846d06918d 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
@@ -215,20 +215,22 @@ public class
SegmentPartitionLLCRealtimeClusterIntegrationTest extends BaseClust
// '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);
+ // Baseline: with broker pruning explicitly disabled, the query scans
all segments across both partitions.
+ JsonNode unpruned = postQuery("SET useBrokerPruning=false; " +
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);
+ // Broker pruning is on by default: without any SET, the non-matching
partition's segments are pruned at the
+ // broker before dispatch.
+ JsonNode pruned = postQuery(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.
+ // Broker pruning actually engaged by default: 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,
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 db52283c71d..c9c61871d14 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
@@ -470,10 +470,12 @@ public abstract class BaseLogicalTableIntegrationTest
extends BaseClusterIntegra
String logicalTableName = getLogicalTableName();
String filteredQuery = "SELECT COUNT(*) FROM " + logicalTableName + "
WHERE DaysSinceEpoch > 16312";
- JsonNode withoutPruning = postQuery(filteredQuery);
+ // Broker pruning is on by default on the MSE logical planner path;
disable it explicitly for the baseline so the
+ // two runs exercise different routing paths.
+ JsonNode withoutPruning = postQuery("SET useBrokerPruning=false; " +
filteredQuery);
assertTrue(withoutPruning.get("exceptions").isEmpty(), "Unexpected
exceptions without broker pruning");
- JsonNode withPruning = postQuery("SET useBrokerPruning=true; " +
filteredQuery);
+ JsonNode withPruning = postQuery(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.
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 99d2235f99a..f6777b9dcc3 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
@@ -22,10 +22,12 @@ import com.google.common.base.Preconditions;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
+import javax.annotation.Nullable;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Filter;
import org.apache.calcite.rel.core.Project;
import org.apache.calcite.rel.core.TableScan;
+import org.apache.commons.lang3.EnumUtils;
import org.apache.pinot.common.request.DataSource;
import org.apache.pinot.common.request.Expression;
import org.apache.pinot.common.request.ExpressionType;
@@ -38,6 +40,11 @@ import org.apache.pinot.sql.FilterKind;
/**
* Utility to convert a leaf stage to a {@link PinotQuery}.
+ *
+ * <p>{@link #createPinotQueryForRouting} folds a Calcite {@link RelNode} leaf
stage into a routing query for the
+ * physical optimizer; {@code PlanNodeRoutingQueryBuilder} is the {@code
PlanNode} (logical-planner) counterpart that
+ * must be kept in sync (same leaf-boundary stop and filter-combining
behavior). The shared filter-normalization
+ * helpers ({@link #ensureFilterIsFunctionExpression} and {@link
#addFilterExpression}) live here and are used by both.
*/
public class LeafStageToPinotQuery {
private LeafStageToPinotQuery() {
@@ -59,13 +66,20 @@ public class LeafStageToPinotQuery {
"Could not find table scan");
TableScan tableScan = (TableScan) bottomToTopNodes.get(0);
PinotQuery pinotQuery = initializePinotQueryForTableScan(tableName,
tableScan);
- for (RelNode parentNode : bottomToTopNodes) {
+ for (int i = 1; i < bottomToTopNodes.size(); i++) {
+ RelNode parentNode = bottomToTopNodes.get(i);
if (parentNode instanceof Filter) {
if (!skipFilter) {
handleFilter((Filter) parentNode, pinotQuery);
}
} else if (parentNode instanceof Project) {
handleProject((Project) parentNode, pinotQuery);
+ } else {
+ // Leaf boundary: the first node that is neither Filter nor Project
(e.g. an un-split DIRECT aggregate)
+ // changes the row space -- InputRefs in nodes above it index that
node's output, not the scan/project
+ // columns -- so folding anything above it (e.g. a HAVING filter)
would mis-resolve columns and corrupt the
+ // routing filter. Everything below this boundary is a genuine
row-level condition; ignore everything above.
+ break;
}
}
return pinotQuery;
@@ -103,8 +117,31 @@ public class LeafStageToPinotQuery {
RexExpression rexExpression =
RexExpressionUtils.fromRexNode(filter.getCondition());
Expression filterExpression =
CalciteRexExpressionParser.toExpression(rexExpression,
pinotQuery.getSelectList());
-
pinotQuery.setFilterExpression(ensureFilterIsFunctionExpression(filterExpression));
+ addFilterExpression(pinotQuery,
ensureFilterIsFunctionExpression(filterExpression));
+ }
+ }
+
+ /**
+ * Adds the given filter to the query, AND-ing it with any previously set
filter (rather than overwriting it, which
+ * would silently discard a genuine row-level condition when a leaf stage
contains multiple filter nodes).
+ * <p>
+ * Note: this mutates {@code pinotQuery} (and, when combining, reuses the
existing filter expression in-place). It
+ * assumes the query is freshly constructed for routing and not shared
across concurrent callers.
+ */
+ public static void addFilterExpression(PinotQuery pinotQuery, @Nullable
Expression filterExpression) {
+ if (filterExpression == null) {
+ return;
}
+ Expression existingFilter = pinotQuery.getFilterExpression();
+ if (existingFilter == null) {
+ pinotQuery.setFilterExpression(filterExpression);
+ return;
+ }
+ Function andFunction = new Function(FilterKind.AND.name());
+ andFunction.setOperands(new ArrayList<>(List.of(existingFilter,
filterExpression)));
+ Expression combined = new Expression(ExpressionType.FUNCTION);
+ combined.setFunctionCall(andFunction);
+ pinotQuery.setFilterExpression(combined);
}
/**
@@ -117,10 +154,18 @@ public class LeafStageToPinotQuery {
* always-false predicate EQUALS(0, 1), and drops LITERAL true expressions
(null = no filter).
* For AND/OR/NOT nodes, operands are recursively fixed.
* <p>
+ * Boolean scalar functions used directly as predicates (e.g. {@code WHERE
contains(col, 'foo')}) are wrapped as
+ * {@code EQUALS(fn(...), true)}, mirroring the single-stage engine's {@code
PredicateComparisonRewriter}: segment
+ * pruners resolve filter operators via {@code FilterKind.valueOf} and would
otherwise throw on them.
+ * <p>
* 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.
+ *
+ * @return the pruner-safe filter expression, or {@code null} when the input
carries no filter constraint (e.g.
+ * LITERAL true, or an AND/OR/NOT that reduces away entirely).
*/
- public static Expression ensureFilterIsFunctionExpression(Expression
expression) {
+ @Nullable
+ public static Expression ensureFilterIsFunctionExpression(@Nullable
Expression expression) {
if (expression == null) {
return null;
}
@@ -157,17 +202,16 @@ public class LeafStageToPinotQuery {
return null; // NOT(constant) — drop entire filter
}
operands.set(0, fixed);
+ } else if (!EnumUtils.isValidEnum(FilterKind.class, operator)) {
+ // Boolean scalar function used directly as a predicate (e.g.
contains(col, 'foo')).
+ return wrapAsEqualsTrue(expression);
}
return expression;
}
if (expression.getIdentifier() != null) {
// Bare boolean column reference (e.g., "is_active" after REINTERPRET
stripped).
// Wrap as EQUALS(col, true) so pruners see a standard predicate.
- Function equalsFunction = new Function(FilterKind.EQUALS.name());
- equalsFunction.setOperands(new ArrayList<>(List.of(expression,
RequestUtils.getLiteralExpression(true))));
- Expression wrapped = new Expression(ExpressionType.FUNCTION);
- wrapped.setFunctionCall(equalsFunction);
- return wrapped;
+ return wrapAsEqualsTrue(expression);
}
// LITERAL expression (constant-folded predicate, e.g., TRUE/FALSE).
// Treat LITERAL false as an always-false predicate that pruners can
process so they
@@ -185,4 +229,13 @@ public class LeafStageToPinotQuery {
// LITERAL true or non-boolean literals have no filter constraint so we
skip pruning
return null;
}
+
+ /** Wraps a boolean-valued expression as the standard predicate {@code
EQUALS(expression, true)}. */
+ private static Expression wrapAsEqualsTrue(Expression expression) {
+ Function equalsFunction = new Function(FilterKind.EQUALS.name());
+ equalsFunction.setOperands(new ArrayList<>(List.of(expression,
RequestUtils.getLiteralExpression(true))));
+ Expression wrapped = new Expression(ExpressionType.FUNCTION);
+ wrapped.setFunctionCall(equalsFunction);
+ return wrapped;
+ }
}
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java
index 3171ec0bef9..aecd0590c2d 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilder.java
@@ -38,15 +38,24 @@ 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.
+ *
+ * <p>This is the {@link PlanNode} (logical-planner) counterpart of {@link
LeafStageToPinotQuery}, which performs the
+ * same leaf-stage → routing-query fold over Calcite {@link
org.apache.calcite.rel.RelNode}s for the physical
+ * optimizer. The two traversals are intentionally kept as separate builders
because they walk different node
+ * hierarchies; the shared filter-normalization logic ({@code
ensureFilterIsFunctionExpression} /
+ * {@code addFilterExpression}) lives in {@link LeafStageToPinotQuery}. Keep
the leaf-boundary and filter-combining
+ * behavior of the two in sync.
*/
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.
+ * Converts a PlanNode leaf stage root to a {@link PinotQuery} for routing
purposes. Folds Filter and Project nodes
+ * bottom-up starting from the TableScan, stopping at the first node of any
other type (a leaf boundary such as an
+ * un-split aggregate): nodes above the boundary operate on a different row
space and must not contribute to the
+ * routing filter. 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.
@@ -61,13 +70,20 @@ public class PlanNodeRoutingQueryBuilder {
"Could not find table scan");
TableScanNode tableScan = (TableScanNode) bottomToTopNodes.get(0);
PinotQuery pinotQuery = initializePinotQueryForTableScan(tableName,
tableScan);
- for (PlanNode parentNode : bottomToTopNodes) {
+ for (int i = 1; i < bottomToTopNodes.size(); i++) {
+ PlanNode parentNode = bottomToTopNodes.get(i);
if (parentNode instanceof FilterNode) {
if (!skipFilter) {
handleFilter((FilterNode) parentNode, pinotQuery);
}
} else if (parentNode instanceof ProjectNode) {
handleProject((ProjectNode) parentNode, pinotQuery);
+ } else {
+ // Leaf boundary: the first node that is neither Filter nor Project
(e.g. an un-split DIRECT aggregate)
+ // changes the row space -- InputRefs in nodes above it index that
node's output, not the scan/project
+ // columns -- so folding anything above it (e.g. a HAVING filter)
would mis-resolve columns and corrupt the
+ // routing filter. Everything below this boundary is a genuine
row-level condition; ignore everything above.
+ break;
}
}
return pinotQuery;
@@ -100,6 +116,7 @@ public class PlanNodeRoutingQueryBuilder {
private static void handleFilter(FilterNode filter, PinotQuery pinotQuery) {
Expression filterExpression =
CalciteRexExpressionParser.toExpression(filter.getCondition(),
pinotQuery.getSelectList());
-
pinotQuery.setFilterExpression(LeafStageToPinotQuery.ensureFilterIsFunctionExpression(filterExpression));
+ LeafStageToPinotQuery.addFilterExpression(pinotQuery,
+
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 a0edc857847..3f84a2c4881 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
@@ -538,9 +538,19 @@ public class WorkerManager {
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());
+ Map<String, RoutingTable> routingTableMap = null;
+ if (routingPinotQuery != null) {
+ try {
+ routingTableMap = getRoutingTable(routingPinotQuery,
context.getRequestId());
+ } catch (RuntimeException e) {
+ // Pruning is best-effort: never fail a query that would otherwise
route successfully unpruned.
+ LOGGER.warn("Broker pruning skipped for table {} due to routing
failure", tableName, e);
+ routingPinotQuery = null;
+ }
+ }
+ if (routingTableMap == null) {
+ routingTableMap = 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.
@@ -807,18 +817,40 @@ public class WorkerManager {
: logicalTableRouteInfo.getRealtimeTableName());
PinotQuery routingPinotQuery =
extractRoutingQuery(fragment.getFragmentRoot(), rawTableName, context);
+ boolean routed = false;
+ if (routingPinotQuery != null) {
+ try {
+ calculateLogicalTableRoutes(tableRouteProvider, logicalTableRouteInfo,
routingPinotQuery, queryOptions,
+ context);
+
context.addNumSegmentsPrunedByBroker(logicalTableRouteInfo.getNumPrunedSegmentsTotal());
+ routed = true;
+ } catch (RuntimeException e) {
+ // Pruning is best-effort: never fail a query that would otherwise
route successfully unpruned. Re-running
+ // unfiltered below is safe because calculateRoutes assigns (rather
than accumulates) the per-table routing
+ // state, fully overwriting anything the failed attempt wrote.
+ LOGGER.warn("Broker pruning skipped for logical table {} due to
routing failure", rawTableName, e);
+ }
+ }
+ if (!routed) {
+ calculateLogicalTableRoutes(tableRouteProvider, logicalTableRouteInfo,
null, queryOptions, context);
+ }
+
+ assignTableSegmentsToWorkers(logicalTableRouteInfo, metadata);
+ }
+
+ /**
+ * Builds the per-table-type routing {@link BrokerRequest}s for a logical
table (filter-bearing when
+ * {@code routingPinotQuery} is non-null, bare {@code SELECT *} otherwise)
and calculates the routes.
+ */
+ private void calculateLogicalTableRoutes(LogicalTableRouteProvider
tableRouteProvider,
+ LogicalTableRouteInfo logicalTableRouteInfo, @Nullable PinotQuery
routingPinotQuery,
+ Map<String, String> queryOptions, DispatchablePlanContext context) {
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);
}
/**
@@ -1008,8 +1040,8 @@ public class WorkerManager {
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());
+ LOGGER.warn("Broker pruning skipped for partitioned table {} due to
routing failure",
+ routingPinotQuery.getDataSource().getTableName(), e);
return null;
}
if (routingTableMap.isEmpty()) {
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/LeafStageToPinotQueryTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/LeafStageToPinotQueryTest.java
index b80ca4b4d48..76ac02898af 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/LeafStageToPinotQueryTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/LeafStageToPinotQueryTest.java
@@ -18,22 +18,82 @@
*/
package org.apache.pinot.query.planner.logical;
+import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.List;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.plan.volcano.VolcanoPlanner;
+import org.apache.calcite.rel.logical.LogicalAggregate;
+import org.apache.calcite.rel.logical.LogicalFilter;
+import org.apache.calcite.rel.logical.LogicalTableScan;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.util.ImmutableBitSet;
import org.apache.pinot.common.request.Expression;
import org.apache.pinot.common.request.ExpressionType;
import org.apache.pinot.common.request.Function;
+import org.apache.pinot.common.request.PinotQuery;
import org.apache.pinot.common.utils.request.RequestUtils;
+import org.apache.pinot.query.type.TypeFactory;
+import org.mockito.Mockito;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertTrue;
public class LeafStageToPinotQueryTest {
+ // --- RelNode leaf-boundary handling (createPinotQueryForRouting, the
physical-optimizer path) ---
+
+ @Test
+ public void testCreatePinotQueryForRoutingStopsAtLeafBoundary() {
+ // Build a RelNode leaf tree: TableScan(col1 INT, col2 INT) -> Filter(col1
= 5) -> Aggregate -> Filter($1 > 10).
+ // The Aggregate is a leaf boundary: InputRefs above it index the
aggregate output, not the scan columns, so its
+ // HAVING filter must NOT be folded into the routing query. The routing
filter must be exactly the WHERE (col1 = 5).
+ // Removing the leaf-boundary break in createPinotQueryForRouting makes
this fail (HAVING folds in against the wrong
+ // row space). This is the RelNode-path mirror of
PlanNodeRoutingQueryBuilderTest#...StopsAtLeafBoundary.
+ TypeFactory typeFactory = new TypeFactory();
+ RexBuilder rexBuilder = new RexBuilder(typeFactory);
+ RelOptCluster cluster = RelOptCluster.create(new VolcanoPlanner(),
rexBuilder);
+ RelDataType intType = typeFactory.createSqlType(SqlTypeName.INTEGER);
+ RelDataType rowType = typeFactory.builder().add("col1",
intType).add("col2", intType).build();
+
+ RelOptTable table = Mockito.mock(RelOptTable.class);
+ Mockito.when(table.getRowType()).thenReturn(rowType);
+ Mockito.when(table.getRelOptSchema()).thenReturn(null);
+ LogicalTableScan tableScan = LogicalTableScan.create(cluster, table,
List.of());
+
+ RexNode whereCondition = rexBuilder.makeCall(SqlStdOperatorTable.EQUALS,
+ rexBuilder.makeInputRef(intType, 0),
rexBuilder.makeExactLiteral(BigDecimal.valueOf(5), intType));
+ LogicalFilter whereFilter = LogicalFilter.create(tableScan,
whereCondition);
+
+ LogicalAggregate aggregate =
+ LogicalAggregate.create(whereFilter, List.of(), ImmutableBitSet.of(0),
null, List.of());
+
+ // HAVING references the aggregate's output (index 0 here = the group
key). Its InputRef space is the aggregate
+ // output, so folding it into the routing query would mis-resolve against
the scan columns.
+ RexNode havingCondition =
rexBuilder.makeCall(SqlStdOperatorTable.GREATER_THAN,
+
rexBuilder.makeInputRef(aggregate.getRowType().getFieldList().get(0).getType(),
0),
+ rexBuilder.makeExactLiteral(BigDecimal.valueOf(10), intType));
+ LogicalFilter havingFilter = LogicalFilter.create(aggregate,
havingCondition);
+
+ PinotQuery pinotQuery =
LeafStageToPinotQuery.createPinotQueryForRouting("testTable", havingFilter,
false);
+
+ Expression filter = pinotQuery.getFilterExpression();
+ assertNotNull(filter);
+ assertEquals(filter.getFunctionCall().getOperator(), "EQUALS");
+
assertEquals(filter.getFunctionCall().getOperands().get(0).getIdentifier().getName(),
"col1");
+ }
+
// --- Basic expression type handling ---
@Test
@@ -94,6 +154,89 @@ public class LeafStageToPinotQueryTest {
assertEquals(result.getFunctionCall().getOperator(), "EQUALS");
}
+ @DataProvider(name = "validFilterKindOperators")
+ public static Object[][] validFilterKindOperators() {
+ // Function-style predicates whose operator IS a FilterKind must NOT be
wrapped as EQUALS(fn, true) -- only
+ // non-FilterKind scalar functions (e.g. contains) are wrapped. This
guards the !isValidEnum boundary.
+ return new Object[][]{
+ {"IN"}, {"NOT_IN"}, {"RANGE"}, {"LIKE"}, {"REGEXP_LIKE"},
{"TEXT_MATCH"}, {"JSON_MATCH"},
+ {"IS_NULL"}, {"IS_NOT_NULL"}, {"NOT_EQUALS"}, {"BETWEEN"}
+ };
+ }
+
+ @Test(dataProvider = "validFilterKindOperators")
+ public void testValidFilterKindFunctionPassedThroughUnchanged(String
operator) {
+ Function function = new Function(operator);
+ function.setOperands(new
ArrayList<>(List.of(RequestUtils.getIdentifierExpression("col"))));
+ Expression funcExpr = new Expression(ExpressionType.FUNCTION);
+ funcExpr.setFunctionCall(function);
+
+ Expression result =
LeafStageToPinotQuery.ensureFilterIsFunctionExpression(funcExpr);
+
+ assertSame(result, funcExpr, operator + " (a valid FilterKind) must not be
wrapped");
+ assertEquals(result.getFunctionCall().getOperator(), operator);
+ }
+
+ @Test
+ public void testBooleanScalarFunctionWrappedAsEqualsTrue() {
+ // A boolean scalar function used directly as a predicate (WHERE
contains(col, 'foo')) is not a FilterKind;
+ // segment pruners resolve operators via FilterKind.valueOf and would
throw on it. It must be wrapped as
+ // EQUALS(contains(col, 'foo'), true), mirroring the single-stage engine's
PredicateComparisonRewriter.
+ Expression containsExpr = makeCompound("contains",
+ RequestUtils.getIdentifierExpression("col"),
RequestUtils.getLiteralExpression("foo"));
+
+ Expression result =
LeafStageToPinotQuery.ensureFilterIsFunctionExpression(containsExpr);
+
+ assertNotNull(result);
+ assertEquals(result.getFunctionCall().getOperator(), "EQUALS");
+ List<Expression> operands = result.getFunctionCall().getOperands();
+ assertEquals(operands.size(), 2);
+ assertEquals(operands.get(0).getFunctionCall().getOperator(), "contains");
+ assertTrue(operands.get(1).getLiteral().getBoolValue());
+ }
+
+ @Test
+ public void testAndWithBooleanScalarFunctionWrapsIt() {
+ assertCompoundWithBooleanScalarFunctionWraps("AND");
+ }
+
+ @Test
+ public void testOrWithBooleanScalarFunctionWrapsIt() {
+ assertCompoundWithBooleanScalarFunctionWraps("OR");
+ }
+
+ private void assertCompoundWithBooleanScalarFunctionWraps(String op) {
+ // OP(contains(col, 'foo'), EQUALS(col2, 'val')) → OP(EQUALS(contains(col,
'foo'), true), EQUALS(col2, 'val'))
+ Expression containsExpr = makeCompound("contains",
+ RequestUtils.getIdentifierExpression("col"),
RequestUtils.getLiteralExpression("foo"));
+ Expression compoundExpr = makeCompound(op, containsExpr,
makeEquals("col2", "val"));
+
+ Expression result =
LeafStageToPinotQuery.ensureFilterIsFunctionExpression(compoundExpr);
+
+ assertNotNull(result);
+ assertEquals(result.getFunctionCall().getOperator(), op);
+ List<Expression> operands = result.getFunctionCall().getOperands();
+ assertEquals(operands.size(), 2);
+ assertEquals(operands.get(0).getFunctionCall().getOperator(), "EQUALS");
+
assertEquals(operands.get(0).getFunctionCall().getOperands().get(0).getFunctionCall().getOperator(),
"contains");
+ assertEquals(operands.get(1).getFunctionCall().getOperator(), "EQUALS");
+ }
+
+ @Test
+ public void testNotWithBooleanScalarFunctionWrapsOperand() {
+ Expression containsExpr = makeCompound("contains",
+ RequestUtils.getIdentifierExpression("col"),
RequestUtils.getLiteralExpression("foo"));
+ Expression notExpr = makeCompound("NOT", containsExpr);
+
+ Expression result =
LeafStageToPinotQuery.ensureFilterIsFunctionExpression(notExpr);
+
+ assertNotNull(result);
+ assertEquals(result.getFunctionCall().getOperator(), "NOT");
+ Expression operand = result.getFunctionCall().getOperands().get(0);
+ assertEquals(operand.getFunctionCall().getOperator(), "EQUALS");
+
assertEquals(operand.getFunctionCall().getOperands().get(0).getFunctionCall().getOperator(),
"contains");
+ }
+
// --- AND/OR handling (shared logic, tested symmetrically) ---
@Test
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilderTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilderTest.java
index a1532eff487..125249b64ea 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilderTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/PlanNodeRoutingQueryBuilderTest.java
@@ -19,9 +19,11 @@
package org.apache.pinot.query.routing;
import java.util.List;
+import org.apache.pinot.common.request.Expression;
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.AggregateNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
import org.apache.pinot.query.planner.plannode.PlanNode;
import org.apache.pinot.query.planner.plannode.ProjectNode;
@@ -58,4 +60,62 @@ public class PlanNodeRoutingQueryBuilderTest {
assertEquals(pinotQuery.getFilterExpression().getFunctionCall().getOperands().get(0).getIdentifier().getName(),
"col1");
}
+
+ @Test
+ public void testCreatePinotQueryForRoutingCombinesStackedLeafFilters() {
+ // Two filter nodes below the leaf boundary (both operate on the scan row
stream) must be AND-combined into the
+ // routing filter, not overwritten -- otherwise a genuine row-level
condition is silently dropped and pruning
+ // becomes incorrect. Reverting the builder to overwrite-the-filter would
make this fail.
+ TableScanNode tableScanNode =
+ new TableScanNode(1, TEST_SCHEMA, PlanNode.NodeHint.EMPTY, List.of(),
"testTable", List.of("col1", "col2"));
+ FilterNode lowerFilter = 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"))));
+ FilterNode upperFilter = new FilterNode(1, TEST_SCHEMA,
PlanNode.NodeHint.EMPTY, List.of(lowerFilter),
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.BOOLEAN,
"EQUALS",
+ List.of(new RexExpression.InputRef(1),
+ new RexExpression.Literal(DataSchema.ColumnDataType.STRING,
"bar"))));
+
+ PinotQuery pinotQuery =
PlanNodeRoutingQueryBuilder.createPinotQueryForRouting("testTable",
upperFilter, false);
+
+ Expression filter = pinotQuery.getFilterExpression();
+ assertNotNull(filter);
+ assertEquals(filter.getFunctionCall().getOperator(), "AND");
+ List<Expression> operands = filter.getFunctionCall().getOperands();
+ assertEquals(operands.size(), 2);
+ // Combined as AND(lower, upper): EQUALS(col1,'foo') AND
EQUALS(col2,'bar').
+
assertEquals(operands.get(0).getFunctionCall().getOperands().get(0).getIdentifier().getName(),
"col1");
+
assertEquals(operands.get(1).getFunctionCall().getOperands().get(0).getIdentifier().getName(),
"col2");
+ }
+
+ @Test
+ public void testCreatePinotQueryForRoutingStopsAtLeafBoundary() {
+ // A node that is neither Filter nor Project (here an AggregateNode) is a
leaf boundary: anything above it operates
+ // on a different (post-aggregate) row space, so its filters/projects must
NOT be folded into the routing query.
+ // The routing query should carry only the WHERE filter below the
aggregate.
+ TableScanNode tableScanNode =
+ new TableScanNode(1, TEST_SCHEMA, PlanNode.NodeHint.EMPTY, List.of(),
"testTable", List.of("col1", "col2"));
+ FilterNode whereFilter = 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"))));
+ DataSchema aggSchema = new DataSchema(new String[]{"col1", "EXPR$1"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING,
DataSchema.ColumnDataType.LONG});
+ AggregateNode aggregateNode = new AggregateNode(1, aggSchema,
PlanNode.NodeHint.EMPTY, List.of(whereFilter),
+ List.of(new RexExpression.FunctionCall(DataSchema.ColumnDataType.LONG,
"COUNT", List.of())),
+ List.of(), List.of(0), AggregateNode.AggType.DIRECT, false, List.of(),
-1);
+ // A HAVING filter above the aggregate; its InputRef(1) indexes the
aggregate output, not the scan columns.
+ FilterNode havingFilter = new FilterNode(1, aggSchema,
PlanNode.NodeHint.EMPTY, List.of(aggregateNode),
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.BOOLEAN,
"GREATER_THAN",
+ List.of(new RexExpression.InputRef(1),
+ new RexExpression.Literal(DataSchema.ColumnDataType.LONG,
10))));
+
+ PinotQuery pinotQuery =
PlanNodeRoutingQueryBuilder.createPinotQueryForRouting("testTable",
havingFilter, false);
+
+ Expression filter = pinotQuery.getFilterExpression();
+ assertNotNull(filter);
+ assertEquals(filter.getFunctionCall().getOperator(), "EQUALS");
+
assertEquals(filter.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 0c5ac12cfeb..4d99eeb537e 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.Function;
import org.apache.pinot.common.request.PinotQuery;
import org.apache.pinot.core.routing.RoutingManager;
import org.apache.pinot.core.routing.RoutingTable;
@@ -46,6 +47,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.FilterKind;
import org.apache.pinot.sql.parsers.CalciteSqlCompiler;
import org.testng.annotations.Test;
@@ -257,7 +259,7 @@ public class WorkerManagerTest {
}
@Test
- public void testBrokerPruningDefaultsToUnfilteredRoutingOnThisPath() {
+ public void testBrokerPruningOnByDefaultAndExplicitlyDisabledOnThisPath() {
Schema schema = getSchemaBuilder("testTable").build();
ServerInstance server = getServerInstance("localhost", 1);
Map<String, ServerInstance> serverInstanceMap =
Map.of(server.getInstanceId(), server);
@@ -280,19 +282,255 @@ public class WorkerManagerTest {
QueryEnvironment queryEnvironment = new
QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
workerManager);
- // Without SET useBrokerPruning=true, this path should fall back to
unfiltered SELECT * routing.
+ // Broker pruning is on by default: without any SET, the routing query
should carry the leaf filter.
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);
+ assertNotNull(brokerRequest.getPinotQuery().getFilterExpression());
+
+ // Explicitly disabling falls back to unfiltered SELECT * routing (segment
lookup only).
+ try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(
+ "SET useBrokerPruning=false; SELECT col2 FROM testTable WHERE col1 =
'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan =
compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan);
+ }
+ brokerRequest =
routingManager.getCapturedRoutingRequest("testTable_OFFLINE");
+ assertNotNull(brokerRequest);
+ assertNull(brokerRequest.getPinotQuery().getFilterExpression());
+ }
+
+ @Test
+ public void testBrokerPruningWrapsBooleanScalarFunctionPredicate() {
+ // Regression: a boolean scalar function used directly as a predicate
(WHERE contains(...)) is not a FilterKind,
+ // and segment pruners resolve filter operators via FilterKind.valueOf.
The routing query must carry it wrapped
+ // as EQUALS(contains(...), true); before this was handled, such queries
failed to plan with
+ // "No enum constant FilterKind.contains" when broker pruning was enabled.
+ 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));
+
+ QueryEnvironment queryEnvironment = newQueryEnvironment(schema,
routingManager);
+ try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(
+ "SELECT col2 FROM testTable WHERE contains(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");
+ Expression wrappedFunction =
filterExpression.getFunctionCall().getOperands().get(0);
+ assertNotNull(wrappedFunction.getFunctionCall());
+ assertEquals(wrappedFunction.getFunctionCall().getOperator(), "contains");
+ }
+
+ @Test
+ public void testBrokerPruningIgnoresFilterAboveLeafAggregate() {
+ // is_partitioned_by_group_by_keys produces a DIRECT (un-split) aggregate
with no exchange under it, so the
+ // HAVING filter lands in the SAME leaf fragment, above the aggregate. Its
InputRefs index the aggregate's
+ // OUTPUT row space ([col1, SUM(col3)]), not the scan columns: folding it
into the routing query would
+ // mis-resolve the refs against scan columns AND overwrite the genuine
WHERE filter, causing incorrect
+ // pruning. The routing query must carry exactly the WHERE filter and
nothing above the aggregate boundary.
+ 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));
+
+ QueryEnvironment queryEnvironment = newQueryEnvironment(schema,
routingManager);
+ try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(
+ "SELECT /*+ aggOptions(is_partitioned_by_group_by_keys='true') */
col1, SUM(col3) FROM testTable "
+ + "WHERE col2 = 'x' GROUP BY col1 HAVING SUM(col3) > 10")) {
+ 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(),
"col2");
+ }
+
+ @Test
+ public void testBrokerPruningPhysicalOptimizerRoutingFilterExcludesHaving() {
+ // The physical optimizer (usePhysicalOptimizer=true, where broker pruning
is already on by default) builds its
+ // routing query via the shared LeafStageToPinotQuery. This exercises that
path end-to-end: the captured routing
+ // filter must carry only the WHERE predicate and must NOT contain the
HAVING predicate. (On the v2 path the
+ // aggregate is split across an exchange, so HAVING stays out of the leaf;
this guards that the shared builder
+ // produces a correct WHERE-only routing filter on the v2 path -- the
un-split-aggregate boundary case that the
+ // leaf-boundary break specifically protects is covered on the logical
path by
+ // testBrokerPruningIgnoresFilterAboveLeafAggregate and
PlanNodeRoutingQueryBuilderTest.)
+ 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));
+
+ QueryEnvironment queryEnvironment = newQueryEnvironment(schema,
routingManager);
+ try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(
+ "SET usePhysicalOptimizer=true; SELECT col1, SUM(col3) FROM testTable "
+ + "WHERE col2 = 'x' GROUP BY col1 HAVING SUM(col3) > 10")) {
+ DispatchableSubPlan dispatchableSubPlan =
compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan);
+ }
+
+ BrokerRequest brokerRequest =
routingManager.getCapturedRoutingRequest("testTable_OFFLINE");
+ assertNotNull(brokerRequest, "Physical optimizer should route through the
capturing routing manager");
+ Expression filterExpression =
brokerRequest.getPinotQuery().getFilterExpression();
+ assertNotNull(filterExpression);
+ // The routing filter is the WHERE predicate only; the HAVING
(GREATER_THAN on SUM) must not appear.
+ assertFalse(containsOperatorOnColumn(filterExpression, "GREATER_THAN"),
+ "HAVING predicate leaked into the physical-optimizer routing filter: "
+ filterExpression);
+ assertTrue(containsIdentifier(filterExpression, "col2"),
+ "WHERE predicate missing from the physical-optimizer routing filter: "
+ filterExpression);
+ }
+
+ private static boolean containsOperatorOnColumn(Expression expression,
String operator) {
+ if (expression == null || expression.getFunctionCall() == null) {
+ return false;
+ }
+ Function function = expression.getFunctionCall();
+ if (function.getOperator().equals(operator)) {
+ return true;
+ }
+ for (Expression operand : function.getOperands()) {
+ if (containsOperatorOnColumn(operand, operator)) {
+ return true;
+ }
+ }
+ return false;
+ }
+ private static boolean containsIdentifier(Expression expression, String
identifier) {
+ if (expression == null) {
+ return false;
+ }
+ if (expression.getIdentifier() != null) {
+ return expression.getIdentifier().getName().equals(identifier);
+ }
+ if (expression.getFunctionCall() != null) {
+ for (Expression operand : expression.getFunctionCall().getOperands()) {
+ if (containsIdentifier(operand, identifier)) {
+ return true;
+ }
+ }
+ }
+ return false;
+ }
+
+ @Test
+ public void testBrokerPruningAllPrunedLeafPlansAcrossExchangeShapes() {
+ // When the filter prunes every segment, the leaf gets zero workers.
Planning (including mailbox assignment,
+ // which runs before the all-leaves-empty short-circuit rewrite) must
still succeed for every exchange shape a
+ // leaf can feed: plain select, global sort/limit (singleton receiver),
aggregations, empty OVER() windows and
+ // set-ops. A planning exception here is a regression: the same query
planned fine with pruning off.
+ List<String> queries = List.of(
+ "SELECT col2 FROM testTable WHERE col1 = 'foo'",
+ "SELECT col2 FROM testTable WHERE col1 = 'foo' ORDER BY col2 LIMIT 5",
+ "SELECT COUNT(*) FROM testTable WHERE col1 = 'foo'",
+ "SELECT col1, COUNT(*) FROM testTable WHERE col1 = 'foo' GROUP BY col1
ORDER BY COUNT(*) LIMIT 3",
+ "SELECT SUM(col3) OVER () FROM testTable WHERE col1 = 'foo'",
+ "SELECT col2 FROM testTable WHERE col1 = 'foo' UNION ALL SELECT col2
FROM testTable",
+ "SELECT DISTINCT col2 FROM testTable WHERE col1 = 'foo' LIMIT 4",
+ // Dynamic-broadcast semi-join: the build side (subquery) is a
separate prunable leaf feeding a
+ // PIPELINE_BREAKER exchange into the main-scan leaf.
+ "SELECT /*+ joinOptions(join_strategy='dynamic_broadcast') */ col2
FROM testTable "
+ + "WHERE col2 IN (SELECT col2 FROM testTable WHERE col1 = 'foo')");
+ for (String query : queries) {
+ 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));
+ routingManager.setEmptyOnFilteredRouting(true);
+
+ QueryEnvironment queryEnvironment = newQueryEnvironment(schema,
routingManager);
+ try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(query)) {
+ DispatchableSubPlan dispatchableSubPlan =
compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan, "Planning failed for all-pruned
query: " + query);
+ } catch (RuntimeException e) {
+ throw new AssertionError("All-pruned leaf broke planning for query: "
+ query + " -- " + e, e);
+ }
+ }
+ }
+
+ @Test
+ public void testBrokerPruningFallsBackToUnfilteredRoutingOnRoutingFailure() {
+ // Pruning is best-effort: if routing the filtered query throws (e.g. a
segment pruner failing on an exotic
+ // filter shape), the query must still plan via the unfiltered SELECT *
fallback rather than fail.
+ 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), true);
+
+ QueryEnvironment queryEnvironment = newQueryEnvironment(schema,
routingManager);
+ try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(
+ "SELECT col2 FROM testTable WHERE col1 = 'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan =
compiledQuery.planQuery(0).getQueryPlan();
+ assertNotNull(dispatchableSubPlan);
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0);
+ }
+
+ // The captured request is the fallback: unfiltered SELECT * (the filtered
attempt threw and was not recorded).
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());
}
+ /**
+ * Mimics how segment pruners consume a routing filter: operators are
resolved via {@code FilterKind.valueOf}
+ * (which throws on non-FilterKind operators, e.g. bare boolean scalar
functions like {@code contains}) and
+ * AND/OR/NOT operands are walked recursively. Keeps the mock routing
managers honest: a routing query that would
+ * crash the real segment pruners also fails the unit tests.
+ */
+ private static void validatePrunableFilter(@Nullable Expression expression) {
+ if (expression == null || expression.getFunctionCall() == null) {
+ return;
+ }
+ Function function = expression.getFunctionCall();
+ FilterKind filterKind = FilterKind.valueOf(function.getOperator());
+ if (filterKind == FilterKind.AND || filterKind == FilterKind.OR ||
filterKind == FilterKind.NOT) {
+ for (Expression operand : function.getOperands()) {
+ validatePrunableFilter(operand);
+ }
+ }
+ }
+
+ /** Builds a QueryEnvironment over a single offline table "testTable" backed
by the given routing manager. */
+ private static QueryEnvironment newQueryEnvironment(Schema schema,
RoutingManager routingManager) {
+ 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);
+ return new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache,
workerManager);
+ }
+
@Test
public void testBrokerPruningPreservesQueryOptionsOnRoutingRequest() {
Schema schema = getSchemaBuilder("testTable").build();
@@ -398,8 +636,8 @@ public class WorkerManagerTest {
}
@Test
- public void testBrokerPruningPartitionedLeafDisabledKeepsAllPartitions() {
- // Without SET useBrokerPruning=true the logical planner default (false)
applies, so no pruning occurs.
+ public void testBrokerPruningPartitionedLeafOnByDefault() {
+ // Broker pruning is on by default: without any SET, the partitioned leaf
prunes non-matching partitions.
QueryEnvironment queryEnvironment =
newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4,
List.of("seg2"), 3);
try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(
@@ -407,6 +645,23 @@ public class WorkerManagerTest {
+ "/*+ tableOptions(partition_function='hashcode',
partition_key='col1', partition_size='4') */ "
+ "WHERE col1 = 'foo'")) {
DispatchableSubPlan dispatchableSubPlan =
compiledQuery.planQuery(0).getQueryPlan();
+ assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 3);
+ DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan);
+ assertNotNull(leaf);
+ assertEquals(leaf.getWorkerIdToSegmentsMap().size(), 1);
+ }
+ }
+
+ @Test
+ public void testBrokerPruningPartitionedLeafDisabledKeepsAllPartitions() {
+ // Explicitly disabling broker pruning keeps all partitions assigned.
+ QueryEnvironment queryEnvironment =
+ newPartitionedQueryEnvironment(new int[]{0, 1, 2, 3}, 4,
List.of("seg2"), 3);
+ try (QueryEnvironment.CompiledQuery compiledQuery =
queryEnvironment.compile(
+ "SET useBrokerPruning=false; SELECT col2 FROM testTable "
+ + "/*+ tableOptions(partition_function='hashcode',
partition_key='col1', partition_size='4') */ "
+ + "WHERE col1 = 'foo'")) {
+ DispatchableSubPlan dispatchableSubPlan =
compiledQuery.planQuery(0).getQueryPlan();
assertEquals(dispatchableSubPlan.getNumSegmentsPrunedByBroker(), 0);
DispatchablePlanFragment leaf = leafFragment(dispatchableSubPlan);
assertNotNull(leaf);
@@ -1009,12 +1264,25 @@ public class WorkerManagerTest {
private static class CapturingRoutingManager implements RoutingManager {
private final Map<String, ServerInstance> _serverInstanceMap;
private final Map<String, RoutingTable> _routingTableByName;
+ private final boolean _throwOnFilteredRouting;
+ private boolean _emptyOnFilteredRouting;
private final Map<String, BrokerRequest> _capturedRoutingRequests = new
LinkedHashMap<>();
CapturingRoutingManager(Map<String, ServerInstance> serverInstanceMap,
Map<String, RoutingTable> routingTableByName) {
+ this(serverInstanceMap, routingTableByName, false);
+ }
+
+ CapturingRoutingManager(Map<String, ServerInstance> serverInstanceMap,
+ Map<String, RoutingTable> routingTableByName, boolean
throwOnFilteredRouting) {
_serverInstanceMap = serverInstanceMap;
_routingTableByName = routingTableByName;
+ _throwOnFilteredRouting = throwOnFilteredRouting;
+ }
+
+ /** When set, filter-bearing routing requests return an all-pruned (empty)
routing table. */
+ void setEmptyOnFilteredRouting(boolean emptyOnFilteredRouting) {
+ _emptyOnFilteredRouting = emptyOnFilteredRouting;
}
@Nullable
@@ -1030,8 +1298,15 @@ public class WorkerManagerTest {
@Nullable
@Override
public RoutingTable getRoutingTable(BrokerRequest brokerRequest, long
requestId) {
+ if (_throwOnFilteredRouting &&
brokerRequest.getPinotQuery().getFilterExpression() != null) {
+ throw new RuntimeException("Simulated routing failure for filtered
request");
+ }
+
validatePrunableFilter(brokerRequest.getPinotQuery().getFilterExpression());
String tableNameWithType = brokerRequest.getQuerySource().getTableName();
_capturedRoutingRequests.put(tableNameWithType, brokerRequest);
+ if (_emptyOnFilteredRouting &&
brokerRequest.getPinotQuery().getFilterExpression() != null) {
+ return new RoutingTable(Map.of(), List.of(), 1);
+ }
return _routingTableByName.get(tableNameWithType);
}
@@ -1124,6 +1399,7 @@ public class WorkerManagerTest {
if (_throwOnRouting) {
throw new RuntimeException("Simulated routing failure");
}
+
validatePrunableFilter(brokerRequest.getPinotQuery().getFilterExpression());
return
_routingTableByTable.get(brokerRequest.getQuerySource().getTableName());
}
diff --git
a/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
b/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
index 55535941b97..641603e9468 100644
--- a/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
+++ b/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
@@ -70,13 +70,10 @@
"\n │ └── [2]@localhost:1|[1] PROJECT",
"\n │ └── [2]@localhost:1|[1] TABLE SCAN (a)
null",
"\n └── [1]@localhost:1|[0]
MAIL_RECEIVE(HASH_DISTRIBUTED)",
- "\n ├── [3]@localhost:2|[2]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]} (Subtree
Omitted)",
- "\n ├── [3]@localhost:2|[3]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]} (Subtree
Omitted)",
- "\n ├── [3]@localhost:1|[0]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]} (Subtree
Omitted)",
- "\n └── [3]@localhost:1|[1]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]}",
- "\n └── [3]@localhost:1|[1] PROJECT",
- "\n └── [3]@localhost:1|[1] FILTER",
- "\n └── [3]@localhost:1|[1] TABLE SCAN
(b) null",
+ "\n └── [3]@localhost:1|[0]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]}",
+ "\n └── [3]@localhost:1|[0] PROJECT",
+ "\n └── [3]@localhost:1|[0] FILTER",
+ "\n └── [3]@localhost:1|[0] TABLE SCAN
(b) null",
"\n"
]
},
@@ -184,13 +181,10 @@
"\n │ └── [2]@localhost:1|[1] FILTER",
"\n │ └── [2]@localhost:1|[1] TABLE
SCAN (a) null",
"\n └── [2]@localhost:1|[1]
MAIL_RECEIVE(BROADCAST_DISTRIBUTED)",
- "\n ├── [3]@localhost:2|[2]
MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2,
3]} (Subtree Omitted)",
- "\n ├── [3]@localhost:2|[3]
MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2,
3]} (Subtree Omitted)",
- "\n ├── [3]@localhost:1|[0]
MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2,
3]} (Subtree Omitted)",
- "\n └── [3]@localhost:1|[1]
MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2,
3]}",
- "\n └── [3]@localhost:1|[1] PROJECT",
- "\n └── [3]@localhost:1|[1]
FILTER",
- "\n └── [3]@localhost:1|[1]
TABLE SCAN (b) null",
+ "\n └── [3]@localhost:1|[0]
MAIL_SEND(BROADCAST_DISTRIBUTED)->{[2]@localhost:1|[0, 1],[2]@localhost:2|[2,
3]}",
+ "\n └── [3]@localhost:1|[0] PROJECT",
+ "\n └── [3]@localhost:1|[0]
FILTER",
+ "\n └── [3]@localhost:1|[0]
TABLE SCAN (b) null",
"\n"
]
},
@@ -402,13 +396,10 @@
"\n └── [3]@localhost:1|[0]
FILTER",
"\n └── [3]@localhost:1|[0]
AGGREGATE_FINAL",
"\n └──
[3]@localhost:1|[0] MAIL_RECEIVE(HASH_DISTRIBUTED)",
- "\n ├──
[4]@localhost:2|[2]
MAIL_SEND(HASH_DISTRIBUTED)->{[3]@localhost:1|[0],[3]@localhost:2|[1]} (Subtree
Omitted)",
- "\n ├──
[4]@localhost:2|[3]
MAIL_SEND(HASH_DISTRIBUTED)->{[3]@localhost:1|[0],[3]@localhost:2|[1]} (Subtree
Omitted)",
- "\n ├──
[4]@localhost:1|[0]
MAIL_SEND(HASH_DISTRIBUTED)->{[3]@localhost:1|[0],[3]@localhost:2|[1]} (Subtree
Omitted)",
- "\n └──
[4]@localhost:1|[1]
MAIL_SEND(HASH_DISTRIBUTED)->{[3]@localhost:1|[0],[3]@localhost:2|[1]}",
- "\n └──
[4]@localhost:1|[1] AGGREGATE_LEAF",
- "\n └──
[4]@localhost:1|[1] FILTER",
- "\n └──
[4]@localhost:1|[1] TABLE SCAN (b) null",
+ "\n └──
[4]@localhost:1|[0]
MAIL_SEND(HASH_DISTRIBUTED)->{[3]@localhost:1|[0],[3]@localhost:2|[1]}",
+ "\n └──
[4]@localhost:1|[0] AGGREGATE_LEAF",
+ "\n └──
[4]@localhost:1|[0] FILTER",
+ "\n └──
[4]@localhost:1|[0] TABLE SCAN (b) null",
"\n"
]
},
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
index ee25435988a..f617e91b74b 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
@@ -640,11 +640,13 @@ public class CommonConstants {
/**
* 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.
+ * Separated from {@link #CONFIG_OF_USE_BROKER_PRUNING} so the two paths
can be rolled out independently; both
+ * default to enabled now that all logical-planner leaf paths
(non-partitioned, partitioned, logical tables)
+ * support broker pruning. Actual pruning still requires segment pruners
to be configured on the table.
*/
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;
+ public static final boolean DEFAULT_LOGICAL_PLANNER_USE_BROKER_PRUNING =
true;
/**
* Default server stage limit for lite mode queries.
@@ -1000,8 +1002,8 @@ public class CommonConstants {
// Not user-settable; used by servers to detect silent truncation at
execution time.
public static final String LITE_MODE_IMPLICIT_LEAF_STAGE_LIMIT =
"liteModeImplicitLeafStageLimit";
// 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.
+ // path defaults to DEFAULT_USE_BROKER_PRUNING and the logical planner
path defaults to
+ // DEFAULT_LOGICAL_PLANNER_USE_BROKER_PRUNING (both true). 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: [email protected]
For additional commands, e-mail: [email protected]
