xiangfu0 commented on code in PR #20035:
URL: https://github.com/apache/pinot/pull/20035#discussion_r4384033804
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java:
##########
@@ -657,86 +660,130 @@ private static boolean
fitsByColumnUnionCardinalities(List<IndexSegment> segment
return true;
}
- /// In-place rewrite QueryContext based on the information from local
IndexSegment.
- ///
- /// @param queryContext
- /// @param indexSegment
+ /// Rewrites the shared context before fetch planning, using only columns
available on every selected segment.
+ /// A partially reloaded timestamp index must leave the original expression
executable on all segments.
+ @VisibleForTesting
+ static void rewriteQueryContextWithHints(QueryContext queryContext,
List<SegmentContext> segmentContexts) {
+ Map<ExpressionContext, ExpressionContext> hints =
queryContext.getExpressionOverrideHints();
+ if (MapUtils.isEmpty(hints) || segmentContexts.isEmpty()) {
+ return;
+ }
+ queryContext.getOrComputeSharedValue(Boolean.class,
EXPRESSION_HINTS_REWRITTEN_KEY, ignored -> {
+ Set<String> columns = getOverrideColumns(hints);
+ for (SegmentContext segmentContext : segmentContexts) {
+ columns.retainAll(segmentContext.getIndexSegment().getColumnNames());
+ }
+ rewriteQueryContextWithHints(queryContext, columns);
+ return true;
+ });
+ }
+
+ /// Rewrites a standalone segment plan once; instance plans first constrain
the rewrite across all their segments.
@VisibleForTesting
public static void rewriteQueryContextWithHints(QueryContext queryContext,
IndexSegment indexSegment) {
- Map<ExpressionContext, ExpressionContext> expressionOverrideHints =
queryContext.getExpressionOverrideHints();
- if (MapUtils.isEmpty(expressionOverrideHints)) {
+ if (MapUtils.isEmpty(queryContext.getExpressionOverrideHints())) {
return;
}
+ queryContext.getOrComputeSharedValue(Boolean.class,
EXPRESSION_HINTS_REWRITTEN_KEY, ignored -> {
+ Set<String> columns =
getOverrideColumns(queryContext.getExpressionOverrideHints());
+ columns.retainAll(indexSegment.getColumnNames());
+ rewriteQueryContextWithHints(queryContext, columns);
+ return true;
+ });
+ }
+
+ private static Set<String> getOverrideColumns(Map<ExpressionContext,
ExpressionContext> hints) {
+ Set<String> columns = new HashSet<>();
+ for (ExpressionContext expression : hints.values()) {
+ if (expression.getIdentifier() != null) {
+ columns.add(expression.getIdentifier());
+ }
+ }
+ return columns;
+ }
+ private static void rewriteQueryContextWithHints(QueryContext queryContext,
Set<String> columns) {
+ if (columns.isEmpty()) {
+ return;
+ }
+ Map<ExpressionContext, ExpressionContext> expressionOverrideHints =
queryContext.getExpressionOverrideHints();
List<ExpressionContext> selectExpressions =
queryContext.getSelectExpressions();
selectExpressions.replaceAll(
- expression -> overrideWithExpressionHints(expression, indexSegment,
expressionOverrideHints));
+ expression -> overrideWithExpressionHints(expression, columns,
expressionOverrideHints));
List<Pair<AggregationFunction, FilterContext>> filtAggrFuns =
queryContext.getFilteredAggregationFunctions();
if (filtAggrFuns != null) {
for (Pair<AggregationFunction, FilterContext>
filteredAggregationFunction : filtAggrFuns) {
FilterContext right = filteredAggregationFunction.getRight();
if (right != null) {
- overrideWithExpressionHints(right, indexSegment,
expressionOverrideHints);
+ overrideWithExpressionHints(right, columns, expressionOverrideHints);
}
}
}
List<ExpressionContext> groupByExpressions =
queryContext.getGroupByExpressions();
if (CollectionUtils.isNotEmpty(groupByExpressions)) {
groupByExpressions.replaceAll(
- expression -> overrideWithExpressionHints(expression, indexSegment,
expressionOverrideHints));
+ expression -> overrideWithExpressionHints(expression, columns,
expressionOverrideHints));
}
List<OrderByExpressionContext> orderByExpressions =
queryContext.getOrderByExpressions();
if (CollectionUtils.isNotEmpty(orderByExpressions)) {
orderByExpressions.replaceAll(expression -> new OrderByExpressionContext(
- overrideWithExpressionHints(expression.getExpression(),
indexSegment, expressionOverrideHints),
+ overrideWithExpressionHints(expression.getExpression(), columns,
expressionOverrideHints),
expression.isAsc()));
}
- // In-place override
FilterContext filter = queryContext.getFilter();
if (filter != null) {
- overrideWithExpressionHints(filter, indexSegment,
expressionOverrideHints);
+ overrideWithExpressionHints(filter, columns, expressionOverrideHints);
}
-
- // In-place override
FilterContext havingFilter = queryContext.getHavingFilter();
if (havingFilter != null) {
- overrideWithExpressionHints(havingFilter, indexSegment,
expressionOverrideHints);
+ overrideWithExpressionHints(havingFilter, columns,
expressionOverrideHints);
}
+ // Aggregation constructors capture their input expressions; AST
replacement alone leaves those inputs stale.
+ // Rebuilding also repairs hash-map keys whose FunctionContext arguments
were changed in place.
+ queryContext.refreshAfterExpressionRewrite();
}
@VisibleForTesting
public static void overrideWithExpressionHints(FilterContext filter,
IndexSegment indexSegment,
Map<ExpressionContext, ExpressionContext> expressionOverrideHints) {
+ overrideWithExpressionHints(filter, indexSegment.getColumnNames(),
expressionOverrideHints);
+ }
+
+ private static void overrideWithExpressionHints(FilterContext filter,
Set<String> columns,
+ Map<ExpressionContext, ExpressionContext> expressionOverrideHints) {
if (filter.getChildren() != null) {
- // AND, OR, NOT
for (FilterContext child : filter.getChildren()) {
- overrideWithExpressionHints(child, indexSegment,
expressionOverrideHints);
+ overrideWithExpressionHints(child, columns, expressionOverrideHints);
}
} else {
- // PREDICATE
Predicate predicate = filter.getPredicate();
- predicate.setLhs(overrideWithExpressionHints(predicate.getLhs(),
indexSegment, expressionOverrideHints));
+ predicate.setLhs(overrideWithExpressionHints(predicate.getLhs(),
columns, expressionOverrideHints));
}
}
@VisibleForTesting
public static ExpressionContext
overrideWithExpressionHints(ExpressionContext expression, IndexSegment
indexSegment,
Map<ExpressionContext, ExpressionContext> expressionOverrideHints) {
+ return overrideWithExpressionHints(expression,
indexSegment.getColumnNames(), expressionOverrideHints);
+ }
+
+ private static ExpressionContext
overrideWithExpressionHints(ExpressionContext expression, Set<String> columns,
+ Map<ExpressionContext, ExpressionContext> expressionOverrideHints) {
if (expression.getType() != ExpressionContext.Type.FUNCTION) {
return expression;
}
ExpressionContext overrideExpression =
expressionOverrideHints.get(expression);
- if (overrideExpression != null && overrideExpression.getIdentifier() !=
null && indexSegment.getColumnNames()
- .contains(overrideExpression.getIdentifier())) {
+ if (overrideExpression != null && overrideExpression.getIdentifier() !=
null
+ && columns.contains(overrideExpression.getIdentifier())) {
return overrideExpression;
}
expression.getFunction()
.getArguments()
- .replaceAll(argument -> overrideWithExpressionHints(argument,
indexSegment, expressionOverrideHints));
+ .replaceAll(argument -> overrideWithExpressionHints(argument, columns,
expressionOverrideHints));
Review Comment:
Yes, the old rewrite mutates the shared query expressions in place. If an
indexed segment replaces an expression with `$ts$DAY`, a subsequent segment
without that column can inherit the replacement; compiled aggregation inputs
can also remain stale.
This PR performs the rewrite once at the instance-plan level, before fetch
planning, using derived columns present on every selected segment. Both normal
and streaming planning use that path. Fully indexed sets use the index, while
partially reloaded sets retain the original expression throughout; the
subsequent segment-level calls are no-ops.
[af7d522ed8](https://github.com/apache/pinot/commit/af7d522ed8d0e208911f4d1ed325f2c2ee192c68)
now documents the shared-context contract explicitly.
Grouping segments by available override columns would preserve the
optimization for indexed subsets during a partial reload. Each group would need
isolated mutable expression/compiled-aggregation state and compatible combine
layouts. I am keeping the common-column fallback in this focused correctness
fix and leaving grouped planning as a separate optimization.
The rebased branch passed all 28 focused tests (`QueryOverrideWithHintsTest`
and `InstancePlanMakerImplV2Test`) on JDK 25, including fully
indexed/mixed-index normal and streaming plans, plus formatting, license, and
Checkstyle checks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
