xiangfu0 commented on code in PR #20035:
URL: https://github.com/apache/pinot/pull/20035#discussion_r4406022973
##########
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) {
Review Comment:
Addressed in df653c2d17e2cd833b110ff50d570892c4f5b26c. Empty hints return
immediately. A single hint validates its identifier target once, checks that
column across segments, and returns at the first missing column without a
copied map or availability cache. Multiple hints deduplicate target columns and
stop scanning when none remain. Successful eligibility still covers every
selected segment, and the original hint map is unchanged. Unsupported
non-identifier targets now fail the requested precondition before scanning.
Focused regressions verify the early exit, duplicate-target checks, and
positive cross-segment case.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java:
##########
@@ -876,70 +903,33 @@ public QueryContext build() {
}
}
+ queryContext.prepareForExecution(Set.of());
return queryContext;
}
- /// Helper method to generate the aggregation functions for the query.
- private void generateAggregationFunctions(QueryContext queryContext) {
- List<Pair<AggregationFunction, FilterContext>>
filteredAggregationFunctions = new ArrayList<>();
- Map<Pair<FunctionContext, FilterContext>, Integer>
filteredAggregationsIndexMap = new HashMap<>();
-
- // Add aggregation functions in the SELECT clause
- // NOTE: DO NOT deduplicate the aggregation functions in the SELECT
clause because that involves protocol change.
- List<Pair<FunctionContext, FilterContext>> filteredAggregations = new
ArrayList<>();
+ /// Captures logical aggregation slots before rewriting; execution
preparation must not deduplicate them again.
+ private void extractAggregationFunctionContexts(QueryContext queryContext)
{
+ List<Pair<FunctionContext, FilterContext>> contexts = new ArrayList<>();
+ // SELECT duplicates are protocol-visible and must retain independent
result slots.
for (ExpressionContext selectExpression :
queryContext._selectExpressions) {
- getAggregations(selectExpression, filteredAggregations);
- }
- for (Pair<FunctionContext, FilterContext> pair : filteredAggregations) {
- FunctionContext aggregation = pair.getLeft();
- FilterContext filter = pair.getRight();
- int functionIndex = filteredAggregationFunctions.size();
- AggregationFunction aggregationFunction =
- AggregationFunctionFactory.getAggregationFunction(aggregation,
queryContext._nullHandlingEnabled);
- filteredAggregationFunctions.add(Pair.of(aggregationFunction, filter));
- filteredAggregationsIndexMap.put(Pair.of(aggregation, filter),
functionIndex);
+ getAggregations(selectExpression, contexts);
}
-
- // Add aggregation functions in the HAVING and ORDER-BY clause but not
in the SELECT clause
- filteredAggregations.clear();
+ Set<Pair<FunctionContext, FilterContext>> existing = new
HashSet<>(contexts);
Review Comment:
Addressed in df653c2d17e2cd833b110ff50d570892c4f5b26c: HAVING/ORDER BY
candidates are collected before creating the deduplication set, which is now
allocated only when that list is nonempty. SELECT duplicates keep their
separate original slots; the preparation regression checks those slots and
unchanged aggregation reuse.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java:
##########
@@ -353,6 +355,37 @@ public Map<ExpressionContext, ExpressionContext>
getExpressionOverrideHints() {
return _expressionOverrideHints;
}
+ /// Prepares derived execution state after logical expressions change,
before fetch or segment planning.
+ /// Reuses compiled aggregations whose input trees did not change. The
supplied set uses function identity because
+ /// expression rewrites can mutate function arguments, including their hash
codes, in place.
+ public void prepareForExecution(Set<FunctionContext> changedFunctions) {
+ int numAggregations = _aggregationFunctionContexts.size();
+ if (numAggregations > 0) {
+ boolean initialPreparation = _aggregationFunctions == null;
+ if (initialPreparation) {
+ _aggregationFunctions = new AggregationFunction[numAggregations];
+ }
+ _filteredAggregationFunctions = new ArrayList<>(numAggregations);
Review Comment:
Addressed in df653c2d17e2cd833b110ff50d570892c4f5b26c. Preparation retains
the existing filtered-aggregation list and index map, preserves unchanged Pair
instances, and replaces a Pair only when its compiled function instance
changes. Filter-clause presence is computed during initial preparation only.
Actual rewrites still clear and rehash the map, including FILTER-only changes;
this keeps the change-tracking API small. Regressions verify list/pair reuse
and a successful lookup after a FILTER-only key mutation, plus refreshed
WHERE-only required columns.
##########
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:
Addressed in df653c2d17e2cd833b110ff50d570892c4f5b26c. Multiple-hint
traversal checks candidate function name/arity before structural lookup while
continuing recursive traversal. Singleton matching uses the JDK 25 single-entry
Map.of lookup, which compares equality directly. Expression hashes are not
cached, and function-identity tracking remains in place for shared mutable
descendants.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java:
##########
@@ -353,6 +355,37 @@ public Map<ExpressionContext, ExpressionContext>
getExpressionOverrideHints() {
return _expressionOverrideHints;
}
+ /// Prepares derived execution state after logical expressions change,
before fetch or segment planning.
+ /// Reuses compiled aggregations whose input trees did not change. The
supplied set uses function identity because
+ /// expression rewrites can mutate function arguments, including their hash
codes, in place.
+ public void prepareForExecution(Set<FunctionContext> changedFunctions) {
+ int numAggregations = _aggregationFunctionContexts.size();
+ if (numAggregations > 0) {
+ boolean initialPreparation = _aggregationFunctions == null;
+ if (initialPreparation) {
+ _aggregationFunctions = new AggregationFunction[numAggregations];
+ }
+ _filteredAggregationFunctions = new ArrayList<>(numAggregations);
+ _filteredAggregationsIndexMap = new HashMap<>();
+ _hasFilteredAggregations = false;
+ for (int i = 0; i < numAggregations; i++) {
+ Pair<FunctionContext, FilterContext> context =
_aggregationFunctionContexts.get(i);
+ if (initialPreparation ||
changedFunctions.contains(context.getLeft())) {
+ _aggregationFunctions[i] =
+
AggregationFunctionFactory.getAggregationFunction(context.getLeft(),
_nullHandlingEnabled);
+ }
+ _filteredAggregationFunctions.add(Pair.of(_aggregationFunctions[i],
context.getRight()));
+ _filteredAggregationsIndexMap.put(context, i);
+ _hasFilteredAggregations |= context.getRight() != null;
+ }
+ }
+ Builder.extractColumns(this);
+ if (_groupByExpressions != null) {
+ _isUnsafeTrim = isGroupingSets() || !isSameOrderAndGroupByColumns(this)
|| _havingFilter != null;
Review Comment:
Addressed in df653c2d17e2cd833b110ff50d570892c4f5b26c: the non-null HAVING
check now precedes the order/group-column comparison, during both initial and
post-rewrite preparation. All 76 focused tests and scoped checks passed.
--
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]
