Jackie-Jiang commented on code in PR #20035:
URL: https://github.com/apache/pinot/pull/20035#discussion_r4406571988
##########
pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java:
##########
@@ -353,6 +355,44 @@ 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<>();
+ } else {
+ // FILTER-only changes can mutate keys even when no aggregation
function is recompiled.
+ _filteredAggregationsIndexMap.clear();
+ }
+ for (int i = 0; i < numAggregations; i++) {
+ Pair<FunctionContext, FilterContext> context =
_aggregationFunctionContexts.get(i);
+ if (initialPreparation ||
changedFunctions.contains(context.getLeft())) {
+ AggregationFunction function =
+
AggregationFunctionFactory.getAggregationFunction(context.getLeft(),
_nullHandlingEnabled);
+ if (initialPreparation) {
+ _filteredAggregationFunctions.add(Pair.of(function,
context.getRight()));
+ _hasFilteredAggregations |= context.getRight() != null;
+ } else if (function != _aggregationFunctions[i]) {
+ _filteredAggregationFunctions.set(i, Pair.of(function,
context.getRight()));
+ }
+ _aggregationFunctions[i] = function;
+ }
+ _filteredAggregationsIndexMap.put(context, i);
+ }
+ }
+ Builder.extractColumns(this);
+ if (_groupByExpressions != null) {
+ _isUnsafeTrim = isGroupingSets() || _havingFilter != null ||
!isSameOrderAndGroupByColumns(this);
Review Comment:
The isSameOrderAndGroupByColumns() helper called here still constructs two
HashSets when GROUP BY and ORDER BY each contain exactly one expression (helper
lines 185–198). Can we compare the two expressions directly when both list
sizes are one, then use the existing set comparison for other cases? That
preserves the current semantics while avoiding the two sets and structural
hashing during initial preparation and again after rewriting. This is an
optimization to adjacent existing code invoked by the new preparation path, not
a correctness regression.
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java:
##########
@@ -657,86 +661,187 @@ 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));
-
- 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);
+ /// Keeps availability checks local to this segment set without changing the
broker-provided hints.
+ private static Map<ExpressionContext, ExpressionContext> getApplicableHints(
+ Map<ExpressionContext, ExpressionContext> hints, List<SegmentContext>
segmentContexts) {
+ if (hints.isEmpty()) {
+ return Map.of();
+ }
+ if (hints.size() == 1) {
+ Map.Entry<ExpressionContext, ExpressionContext> hint =
hints.entrySet().iterator().next();
+ String column = getOverrideColumn(hint.getValue());
+ for (SegmentContext segmentContext : segmentContexts) {
+ if
(!segmentContext.getIndexSegment().getColumnNames().contains(column)) {
+ return Map.of();
}
}
+ // The single-entry map uses equality directly, avoiding recursive
hashing of unrelated expression trees.
+ return Map.of(hint.getKey(), hint.getValue());
+ }
+ Set<String> availableColumns = new HashSet<>();
+ for (ExpressionContext target : hints.values()) {
+ availableColumns.add(getOverrideColumn(target));
+ }
+ for (SegmentContext segmentContext : segmentContexts) {
+ Set<String> columns = segmentContext.getIndexSegment().getColumnNames();
+ availableColumns.removeIf(column -> !columns.contains(column));
+ if (availableColumns.isEmpty()) {
+ return Map.of();
+ }
}
+ Map<ExpressionContext, ExpressionContext> applicableHints = new
HashMap<>();
Review Comment:
For multiple hints, can we return the original hints map when every target
column survives? Save the initial distinct-target count before scanning
segments; since availableColumns only shrinks, an unchanged count proves that
every hint is still applicable. The subsequent matcher only reads the map. This
avoids allocating another HashMap/table and entry nodes, structurally hashing
all H keys again, and performing this filtering pass for fully indexed
multi-hint queries. Keep the current filtering and singleton normalization when
targets were removed.
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java:
##########
@@ -657,86 +661,187 @@ 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));
-
- 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);
+ /// Keeps availability checks local to this segment set without changing the
broker-provided hints.
+ private static Map<ExpressionContext, ExpressionContext> getApplicableHints(
+ Map<ExpressionContext, ExpressionContext> hints, List<SegmentContext>
segmentContexts) {
+ if (hints.isEmpty()) {
+ return Map.of();
+ }
+ if (hints.size() == 1) {
+ Map.Entry<ExpressionContext, ExpressionContext> hint =
hints.entrySet().iterator().next();
+ String column = getOverrideColumn(hint.getValue());
+ for (SegmentContext segmentContext : segmentContexts) {
+ if
(!segmentContext.getIndexSegment().getColumnNames().contains(column)) {
+ return Map.of();
}
}
+ // The single-entry map uses equality directly, avoiding recursive
hashing of unrelated expression trees.
+ return Map.of(hint.getKey(), hint.getValue());
+ }
+ Set<String> availableColumns = new HashSet<>();
+ for (ExpressionContext target : hints.values()) {
+ availableColumns.add(getOverrideColumn(target));
+ }
+ for (SegmentContext segmentContext : segmentContexts) {
+ Set<String> columns = segmentContext.getIndexSegment().getColumnNames();
+ availableColumns.removeIf(column -> !columns.contains(column));
+ if (availableColumns.isEmpty()) {
+ return Map.of();
+ }
}
+ Map<ExpressionContext, ExpressionContext> applicableHints = new
HashMap<>();
+ for (Map.Entry<ExpressionContext, ExpressionContext> hint :
hints.entrySet()) {
+ if (availableColumns.contains(hint.getValue().getIdentifier())) {
+ applicableHints.put(hint.getKey(), hint.getValue());
+ }
+ }
+ if (applicableHints.size() == 1) {
+ Map.Entry<ExpressionContext, ExpressionContext> hint =
applicableHints.entrySet().iterator().next();
+ return Map.of(hint.getKey(), hint.getValue());
+ }
+ return applicableHints;
+ }
+ private static String getOverrideColumn(ExpressionContext target) {
+ Preconditions.checkArgument(target.getType() ==
ExpressionContext.Type.IDENTIFIER,
+ "Expression override target must be an identifier");
+ return target.getIdentifier();
+ }
+
+ /// Gates structural hashing by function name and arity; expression hashes
cannot be cached across mutations.
+ @Nullable
+ private static Map<String, Set<Integer>>
getHintFunctionArities(Map<ExpressionContext, ExpressionContext> hints) {
+ if (hints.size() <= 1) {
+ return null;
+ }
+ Map<String, Set<Integer>> arities = new HashMap<>();
+ for (ExpressionContext expression : hints.keySet()) {
+ FunctionContext function = expression.getFunction();
+ if (function != null) {
+ arities.computeIfAbsent(function.getFunctionName(), key -> new
HashSet<>()).add(function.getArguments().size());
+ }
+ }
+ return arities;
+ }
+
+ private static void rewriteQueryContextWithHints(QueryContext queryContext,
+ Map<ExpressionContext, ExpressionContext> hints) {
+ Set<FunctionContext> changedFunctions = Collections.newSetFromMap(new
IdentityHashMap<>());
+ Map<String, Set<Integer>> hintFunctionArities =
getHintFunctionArities(hints);
+ // FILTER predicates have separate contexts; aggregation inputs are
visited through SELECT/HAVING/ORDER-BY.
+ List<Pair<AggregationFunction, FilterContext>> filteredAggregations =
+ queryContext.getFilteredAggregationFunctions();
+ if (filteredAggregations != null) {
Review Comment:
Can we guard this separate FILTER-context pass with
queryContext.hasFilteredAggregations()? The list contains every aggregation
slot, so ordinary hinted aggregation queries currently inspect all A slots
merely to find null filters. The flag is already initialized during query
construction across SELECT, HAVING, and ORDER BY. Using it skips a guaranteed
no-op O(A) pass without new state. Retain the per-entry null checks when the
flag is true, and retain the separate FILTER traversal itself because those
predicate containers are distinct from the original expression-tree slots.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java:
##########
@@ -876,70 +910,35 @@ 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);
+ getAggregations(selectExpression, contexts);
}
- 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);
- }
-
- // Add aggregation functions in the HAVING and ORDER-BY clause but not
in the SELECT clause
- filteredAggregations.clear();
+ List<Pair<FunctionContext, FilterContext>> otherAggregations = new
ArrayList<>();
Review Comment:
Small allocation cleanup: create otherAggregations only when HAVING or ORDER
BY exists. It currently creates a throwaway empty ArrayList on every
conversion, including ordinary queries without hints or either clause. Guard
the collection/deduplication block while still assigning the captured SELECT
slots to _aggregationFunctionContexts. The existing nonempty check before
creating the deduplication HashSet should remain.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/request/context/utils/QueryContextConverterUtils.java:
##########
@@ -186,9 +186,12 @@ public static QueryContext getQueryContext(PinotQuery
pinotQuery) {
}
// EXPRESSION OVERRIDE HINTS
- Map<ExpressionContext, ExpressionContext> expressionContextOverrideHints =
new HashMap<>();
+ Map<ExpressionContext, ExpressionContext> expressionContextOverrideHints;
Map<Expression, Expression> expressionOverrideHints =
pinotQuery.getExpressionOverrideHints();
- if (expressionOverrideHints != null) {
+ if (expressionOverrideHints == null || expressionOverrideHints.isEmpty()) {
+ expressionContextOverrideHints = Map.of();
+ } else {
+ expressionContextOverrideHints = new HashMap<>();
Review Comment:
Can we also specialize singleton conversion here? Empty hints now use
Map.of(), but one hint still allocates a HashMap/backing table/entry and
structurally hashes its key. Instance planning then normalizes that entry to
Map.of(key, value) for matching. Constructing a singleton map directly during
conversion avoids the initial hash-table allocation and key hashing on the
common single-hint path. The planner can retain normalization for arbitrary
maps supplied directly through QueryContext.Builder; the inspected production
consumers only read converted hints.
--
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]
