Jackie-Jiang commented on code in PR #20035:
URL: https://github.com/apache/pinot/pull/20035#discussion_r4404338490
##########
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:
Can we avoid constructing this set unless HAVING/ORDER BY actually
contributes aggregation candidates? Every SELECT aggregation is now
structurally hashed into `existing`, then hashed again into
`_filteredAggregationsIndexMap` during initial preparation. The base used the
index map directly for deduplication. This adds recursive key hashing and
transient set entries to broker and server contexts even when no rewrite hints
exist.
A small fix is to collect `otherAggregations` first and create `existing`
only when that list is nonempty. A stronger alternative is to construct and
reuse the initial slot-index map for deduplication, rebuilding it only after a
rewrite. SELECT duplicates must continue retaining their original result slots.
##########
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:
Can preparation reuse the filtered-aggregation list and unchanged Pair
entries? Every actual rewrite currently allocates a new list, one Pair per
aggregation, and a new index map, including when only WHERE changes. FILTER
contexts already mutate in place, so only pairs whose compiled aggregation
function is replaced need new wrappers. The presence of aggregation FILTER
clauses is also unchanged by this rewrite, so `_hasFilteredAggregations` need
not be recomputed each time.
A WHERE-only rewrite could additionally skip aggregation-map rebuilding if
change tracking distinguishes affected state, while still refreshing required
columns. Do not skip rehashing merely because no aggregation was recompiled: a
FILTER-only rewrite can change a map key's hash. Selective aggregation
recompilation and stable slot ordering should remain.
##########
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:
Can we stop the outer loop as soon as `applicableHints` becomes empty?
Rejected hints are already removed, but all remaining segments still incur
`getColumnNames()` and empty-map iteration. With one hint missing on the first
of N segments, there is only one membership check but still N segment visits.
This is a realistic MSE case: `ServerPlanRequestVisitor.applyTimestampIndex()`
generates hints without checking index availability.
We can also validate each distinct replacement identifier once. The
two-argument DATETRUNC form and the explicit MILLISECONDS form can both map to
the same `$ts$DAY`, so the current implementation repeats identical checks. A
local availability cache reduces worst-case membership checks from N × H hint
entries to N × U distinct target columns, stopping at the first missing segment
for each target. Optionally resolve availability when the AST visitor actually
matches a hint, eliminating the upfront copied-and-filtered hint map. Keep any
cache scoped to this invocation/segment set; a positive result must still cover
every selected segment.
##########
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:
Small short-circuit improvement: evaluate `_havingFilter != null` before
`!isSameOrderAndGroupByColumns(this)`. HAVING already forces unsafe trimming,
so there is no need to build and compare the two temporary expression sets in
that case. This avoids the work during both initial preparation and
post-rewrite preparation.
--
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]
