xiangfu0 commented on code in PR #20035:
URL: https://github.com/apache/pinot/pull/20035#discussion_r4406908487
##########
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:
Addressed in faa46f1388887d7dcb750727c90c6a1c97e9a575. Singleton conversion
now directly creates Map.of() from the converted key/value, avoiding the
temporary hash table and key hashing. Empty hints and larger maps retain their
respective fast/fallback paths, and the planner still normalizes arbitrary maps
supplied through Builder. The existing singleton regression verifies the
converted mapping.
##########
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:
Addressed in faa46f1388887d7dcb750727c90c6a1c97e9a575:
hasFilteredAggregations() now guards the separate FILTER-context pass. The
per-entry null checks and traversal of distinct FILTER predicate containers
remain. Existing filtered, constant-filter, and WHERE-only regressions pass.
##########
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:
Addressed in faa46f1388887d7dcb750727c90c6a1c97e9a575. The initial
distinct-target count is saved before availability only shrinks; if the count
is unchanged after checking every selected segment, the original hint map is
returned for read-only matching. Removed-target cases retain filtering and
singleton normalization. The existing fixture now covers both fully applicable
and partially unavailable hint sets.
##########
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:
Addressed in faa46f1388887d7dcb750727c90c6a1c97e9a575. After the existing
empty-list guards, two singleton lists compare their expressions directly.
Other sizes retain the existing set comparison, including reordered and
duplicate-key semantics. The existing preparation regression covers
equal/different singleton keys and both fallback cases.
##########
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:
Addressed in faa46f1388887d7dcb750727c90c6a1c97e9a575: the extra-candidate
collection/deduplication block now runs only when HAVING or ORDER BY exists.
Captured SELECT slots are still assigned outside that block, and the
deduplication set remains conditional on a nonempty candidate list.
--
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]
