Jackie-Jiang commented on code in PR #20035:
URL: https://github.com/apache/pinot/pull/20035#discussion_r4404338521
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java:
##########
@@ -657,86 +661,125 @@ 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.
+ /// Rewriting per segment would mutate the expressions shared by already
constructed segment plans.
@VisibleForTesting
- public static void rewriteQueryContextWithHints(QueryContext queryContext,
IndexSegment indexSegment) {
- Map<ExpressionContext, ExpressionContext> expressionOverrideHints =
queryContext.getExpressionOverrideHints();
- if (MapUtils.isEmpty(expressionOverrideHints)) {
+ static void rewriteQueryContextWithHints(QueryContext queryContext,
List<SegmentContext> segmentContexts) {
+ Map<ExpressionContext, ExpressionContext> hints =
queryContext.getExpressionOverrideHints();
+ if (MapUtils.isEmpty(hints) || segmentContexts.isEmpty()) {
return;
}
+ Map<ExpressionContext, ExpressionContext> applicableHints =
getApplicableHints(hints, segmentContexts);
+ if (!applicableHints.isEmpty()) {
+ rewriteQueryContextWithHints(queryContext, applicableHints);
+ }
+ }
- List<ExpressionContext> selectExpressions =
queryContext.getSelectExpressions();
- selectExpressions.replaceAll(
- expression -> overrideWithExpressionHints(expression, indexSegment,
expressionOverrideHints));
+ /// Makes a private copy so segment availability checks never mutate the
broker-provided hint map.
+ private static Map<ExpressionContext, ExpressionContext> getApplicableHints(
+ Map<ExpressionContext, ExpressionContext> hints, List<SegmentContext>
segmentContexts) {
+ Map<ExpressionContext, ExpressionContext> applicableHints = new
HashMap<>(hints);
+ applicableHints.values().removeIf(expression -> expression.getType() !=
ExpressionContext.Type.IDENTIFIER);
+ for (SegmentContext segmentContext : segmentContexts) {
+ Set<String> columns = segmentContext.getIndexSegment().getColumnNames();
+ applicableHints.values().removeIf(expression ->
!columns.contains(expression.getIdentifier()));
+ }
+ return applicableHints;
+ }
- 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);
+ private static void rewriteQueryContextWithHints(QueryContext queryContext,
+ Map<ExpressionContext, ExpressionContext> hints) {
+ Set<FunctionContext> changedFunctions = Collections.newSetFromMap(new
IdentityHashMap<>());
+ // FILTER predicates have separate contexts; aggregation inputs are
visited through SELECT/HAVING/ORDER-BY.
+ List<Pair<AggregationFunction, FilterContext>> filteredAggregations =
+ queryContext.getFilteredAggregationFunctions();
+ if (filteredAggregations != null) {
+ for (Pair<AggregationFunction, FilterContext> aggregation :
filteredAggregations) {
+ if (aggregation.getRight() != null) {
+ overrideWithExpressionHints(aggregation.getRight(), hints,
changedFunctions);
}
}
}
-
+ queryContext.getSelectExpressions().replaceAll(
+ expression -> overrideWithExpressionHints(expression, hints,
changedFunctions));
List<ExpressionContext> groupByExpressions =
queryContext.getGroupByExpressions();
if (CollectionUtils.isNotEmpty(groupByExpressions)) {
groupByExpressions.replaceAll(
- expression -> overrideWithExpressionHints(expression, indexSegment,
expressionOverrideHints));
+ expression -> overrideWithExpressionHints(expression, hints,
changedFunctions));
}
-
List<OrderByExpressionContext> orderByExpressions =
queryContext.getOrderByExpressions();
if (CollectionUtils.isNotEmpty(orderByExpressions)) {
- orderByExpressions.replaceAll(expression -> new OrderByExpressionContext(
- overrideWithExpressionHints(expression.getExpression(),
indexSegment, expressionOverrideHints),
- expression.isAsc()));
+ orderByExpressions.replaceAll(orderBy -> {
+ ExpressionContext expression = orderBy.getExpression();
+ ExpressionContext rewritten = overrideWithExpressionHints(expression,
hints, changedFunctions);
+ return rewritten == expression ? orderBy
+ : new OrderByExpressionContext(rewritten, orderBy.isAsc(),
orderBy.isNullsLast());
+ });
}
-
- // In-place override
- FilterContext filter = queryContext.getFilter();
- if (filter != null) {
- overrideWithExpressionHints(filter, indexSegment,
expressionOverrideHints);
+ if (queryContext.getFilter() != null) {
+ overrideWithExpressionHints(queryContext.getFilter(), hints,
changedFunctions);
}
-
- // In-place override
- FilterContext havingFilter = queryContext.getHavingFilter();
- if (havingFilter != null) {
- overrideWithExpressionHints(havingFilter, indexSegment,
expressionOverrideHints);
+ if (queryContext.getHavingFilter() != null) {
+ overrideWithExpressionHints(queryContext.getHavingFilter(), hints,
changedFunctions);
+ }
+ if (!changedFunctions.isEmpty()) {
+ queryContext.prepareForExecution(changedFunctions);
}
}
@VisibleForTesting
public static void overrideWithExpressionHints(FilterContext filter,
IndexSegment indexSegment,
Map<ExpressionContext, ExpressionContext> expressionOverrideHints) {
+ overrideWithExpressionHints(filter,
+ getApplicableHints(expressionOverrideHints, List.of(new
SegmentContext(indexSegment))),
+ Collections.newSetFromMap(new IdentityHashMap<>()));
+ }
+
+ private static void overrideWithExpressionHints(FilterContext filter,
+ Map<ExpressionContext, ExpressionContext> expressionOverrideHints,
Set<FunctionContext> changedFunctions) {
+ if (filter.isConstant()) {
+ return;
+ }
if (filter.getChildren() != null) {
- // AND, OR, NOT
for (FilterContext child : filter.getChildren()) {
- overrideWithExpressionHints(child, indexSegment,
expressionOverrideHints);
+ overrideWithExpressionHints(child, expressionOverrideHints,
changedFunctions);
}
} else {
- // PREDICATE
Predicate predicate = filter.getPredicate();
- predicate.setLhs(overrideWithExpressionHints(predicate.getLhs(),
indexSegment, expressionOverrideHints));
+ predicate.setLhs(
+ overrideWithExpressionHints(predicate.getLhs(),
expressionOverrideHints, changedFunctions));
}
}
@VisibleForTesting
public static ExpressionContext
overrideWithExpressionHints(ExpressionContext expression, IndexSegment
indexSegment,
Map<ExpressionContext, ExpressionContext> expressionOverrideHints) {
+ return overrideWithExpressionHints(expression,
+ getApplicableHints(expressionOverrideHints, List.of(new
SegmentContext(indexSegment))),
+ Collections.newSetFromMap(new IdentityHashMap<>()));
+ }
+
+ private static ExpressionContext
overrideWithExpressionHints(ExpressionContext expression,
+ Map<ExpressionContext, ExpressionContext> expressionOverrideHints,
Set<FunctionContext> changedFunctions) {
if (expression.getType() != ExpressionContext.Type.FUNCTION) {
return expression;
}
ExpressionContext overrideExpression =
expressionOverrideHints.get(expression);
Review Comment:
A further optimization is to gate this structural lookup by candidate
function name/arity derived from the hint keys. `FunctionContext.hashCode()`
recursively hashes the argument tree, and the visitor then descends and hashes
each child subtree again. For a deep function chain this can produce quadratic
hashing work; ordinary enclosing SUM/ADD expressions cannot match a DATETRUNC
hint but still pay the lookup cost.
This pattern predates the PR, and moving rewriting to instance planning
reduces its frequency. A cheap candidate-signature check would avoid hashing
unrelated subtrees while preserving recursive traversal. Do not cache
structural hashes on these mutable expression contexts.
--
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]
