This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new fcd05cca42d Fix compiled aggregation inputs after timestamp index
rewrites (#20035)
fcd05cca42d is described below
commit fcd05cca42d1ce4350a3246ba7e5cac203b534a9
Author: Xiang Fu <[email protected]>
AuthorDate: Wed Oct 28 16:58:24 2026 -0700
Fix compiled aggregation inputs after timestamp index rewrites (#20035)
* Fix compiled aggregation inputs after timestamp index rewrites
* Cover hinted timestamp aggregation with GROUP BY
* Clarify timestamp rewrite planning contract
* Prepare rewritten query execution state once per instance
* Preserve null semantics during timestamp index rewrites
* Reduce planning work for expression override hints
* Avoid redundant common query planning allocations
---
.../core/plan/maker/InstancePlanMakerImplV2.java | 199 +++++++++---
.../core/query/request/context/QueryContext.java | 128 ++++----
.../context/utils/QueryContextConverterUtils.java | 11 +-
.../plan/maker/QueryOverrideWithHintsTest.java | 349 ++++++++++++++++++++-
.../apache/pinot/queries/TimestampQueriesTest.java | 70 +++++
5 files changed, 639 insertions(+), 118 deletions(-)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
index 68a141afb15..305ad2f2b32 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
@@ -21,7 +21,10 @@ package org.apache.pinot.core.plan.maker;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
import java.util.HashSet;
+import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -33,6 +36,7 @@ import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.FilterContext;
+import org.apache.pinot.common.request.context.FunctionContext;
import org.apache.pinot.common.request.context.OrderByExpressionContext;
import org.apache.pinot.common.request.context.predicate.Predicate;
import org.apache.pinot.common.utils.config.QueryOptionsUtils;
@@ -221,6 +225,7 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
@Override
public Plan makeInstancePlan(List<SegmentContext> segmentContexts,
QueryContext queryContext,
ExecutorService executorService) {
+ rewriteQueryContextWithHints(queryContext, segmentContexts);
applyQueryOptions(queryContext);
boundGroupingSetsBaseAggregation(segmentContexts, queryContext);
@@ -406,7 +411,6 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
@Override
public PlanNode makeSegmentPlanNode(SegmentContext segmentContext,
QueryContext queryContext) {
- rewriteQueryContextWithHints(queryContext,
segmentContext.getIndexSegment());
if (QueryContextUtils.isAggregationQuery(queryContext)) {
List<ExpressionContext> groupByExpressions =
queryContext.getGroupByExpressions();
if (groupByExpressions != null) {
@@ -427,6 +431,7 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
@Override
public Plan makeStreamingInstancePlan(List<SegmentContext> segmentContexts,
QueryContext queryContext,
ExecutorService executorService, ResultsBlockStreamer streamer) {
+ rewriteQueryContextWithHints(queryContext, segmentContexts);
applyQueryOptions(queryContext);
boundGroupingSetsBaseAggregation(segmentContexts, queryContext);
@@ -459,7 +464,6 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
public PlanNode makeStreamingSegmentPlanNode(SegmentContext segmentContext,
QueryContext queryContext) {
if (QueryContextUtils.isSelectionOnlyQuery(queryContext) &&
queryContext.getLimit() != 0) {
// Use streaming operator only for non-empty selection-only query
- rewriteQueryContextWithHints(queryContext,
segmentContext.getIndexSegment());
return new StreamingSelectionPlanNode(segmentContext, queryContext);
} else {
return makeSegmentPlanNode(segmentContext, queryContext);
@@ -657,86 +661,191 @@ public class InstancePlanMakerImplV2 implements
PlanMaker {
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));
+ }
+ int numTargets = availableColumns.size();
+ for (SegmentContext segmentContext : segmentContexts) {
+ Set<String> columns = segmentContext.getIndexSegment().getColumnNames();
+ availableColumns.removeIf(column -> !columns.contains(column));
+ if (availableColumns.isEmpty()) {
+ return Map.of();
+ }
+ }
+ if (availableColumns.size() == numTargets) {
+ return hints;
+ }
+ 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.
+ if (queryContext.hasFilteredAggregations()) {
+ List<Pair<AggregationFunction, FilterContext>> filteredAggregations =
+ queryContext.getFilteredAggregationFunctions();
+ for (Pair<AggregationFunction, FilterContext> aggregation :
filteredAggregations) {
+ if (aggregation.getRight() != null) {
+ overrideWithExpressionHints(aggregation.getRight(), hints,
hintFunctionArities, changedFunctions);
+ }
+ }
+ }
+ queryContext.getSelectExpressions().replaceAll(
+ expression -> overrideWithExpressionHints(expression, hints,
hintFunctionArities, changedFunctions));
List<ExpressionContext> groupByExpressions =
queryContext.getGroupByExpressions();
if (CollectionUtils.isNotEmpty(groupByExpressions)) {
groupByExpressions.replaceAll(
- expression -> overrideWithExpressionHints(expression, indexSegment,
expressionOverrideHints));
+ expression -> overrideWithExpressionHints(expression, hints,
hintFunctionArities, 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,
hintFunctionArities, 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,
hintFunctionArities, changedFunctions);
}
-
- // In-place override
- FilterContext havingFilter = queryContext.getHavingFilter();
- if (havingFilter != null) {
- overrideWithExpressionHints(havingFilter, indexSegment,
expressionOverrideHints);
+ if (queryContext.getHavingFilter() != null) {
+ overrideWithExpressionHints(queryContext.getHavingFilter(), hints,
hintFunctionArities, changedFunctions);
+ }
+ if (!changedFunctions.isEmpty()) {
+ queryContext.prepareForExecution(changedFunctions);
}
}
@VisibleForTesting
public static void overrideWithExpressionHints(FilterContext filter,
IndexSegment indexSegment,
Map<ExpressionContext, ExpressionContext> expressionOverrideHints) {
+ Map<ExpressionContext, ExpressionContext> hints =
+ getApplicableHints(expressionOverrideHints, List.of(new
SegmentContext(indexSegment)));
+ overrideWithExpressionHints(filter, hints, getHintFunctionArities(hints),
+ Collections.newSetFromMap(new IdentityHashMap<>()));
+ }
+
+ private static void overrideWithExpressionHints(FilterContext filter,
+ Map<ExpressionContext, ExpressionContext> expressionOverrideHints,
+ @Nullable Map<String, Set<Integer>> hintFunctionArities,
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,
hintFunctionArities, changedFunctions);
}
} else {
- // PREDICATE
Predicate predicate = filter.getPredicate();
- predicate.setLhs(overrideWithExpressionHints(predicate.getLhs(),
indexSegment, expressionOverrideHints));
+ ExpressionContext rewritten =
overrideWithExpressionHints(predicate.getLhs(), expressionOverrideHints,
+ hintFunctionArities, changedFunctions);
+ predicate.setLhs(rewritten);
}
}
@VisibleForTesting
public static ExpressionContext
overrideWithExpressionHints(ExpressionContext expression, IndexSegment
indexSegment,
Map<ExpressionContext, ExpressionContext> expressionOverrideHints) {
+ Map<ExpressionContext, ExpressionContext> hints =
+ getApplicableHints(expressionOverrideHints, List.of(new
SegmentContext(indexSegment)));
+ return overrideWithExpressionHints(expression, hints,
getHintFunctionArities(hints),
+ Collections.newSetFromMap(new IdentityHashMap<>()));
+ }
+
+ private static ExpressionContext
overrideWithExpressionHints(ExpressionContext expression,
+ Map<ExpressionContext, ExpressionContext> expressionOverrideHints,
+ @Nullable Map<String, Set<Integer>> hintFunctionArities,
Set<FunctionContext> changedFunctions) {
if (expression.getType() != ExpressionContext.Type.FUNCTION) {
return expression;
}
- ExpressionContext overrideExpression =
expressionOverrideHints.get(expression);
- if (overrideExpression != null && overrideExpression.getIdentifier() !=
null && indexSegment.getColumnNames()
- .contains(overrideExpression.getIdentifier())) {
- return overrideExpression;
+ FunctionContext function = expression.getFunction();
+ List<ExpressionContext> arguments = function.getArguments();
+ Set<Integer> arities = hintFunctionArities != null ?
hintFunctionArities.get(function.getFunctionName()) : null;
+ if (hintFunctionArities == null || arities != null &&
arities.contains(arguments.size())) {
+ ExpressionContext overrideExpression =
expressionOverrideHints.get(expression);
+ if (overrideExpression != null) {
+ changedFunctions.add(function);
+ return overrideExpression;
+ }
+ }
+ for (int i = 0; i < arguments.size(); i++) {
+ ExpressionContext argument = arguments.get(i);
+ ExpressionContext rewritten =
+ overrideWithExpressionHints(argument, expressionOverrideHints,
hintFunctionArities, changedFunctions);
+ if (rewritten != argument ||
changedFunctions.contains(argument.getFunction())) {
+ changedFunctions.add(function);
+ }
+ arguments.set(i, rewritten);
}
- expression.getFunction()
- .getArguments()
- .replaceAll(argument -> overrideWithExpressionHints(argument,
indexSegment, expressionOverrideHints));
return expression;
}
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
index 24548c551dd..857f11e5b2f 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
@@ -89,6 +89,8 @@ public class QueryContext implements QueryScopedCache {
private AggregationFunction[] _aggregationFunctions;
private List<Pair<AggregationFunction, FilterContext>>
_filteredAggregationFunctions;
private Map<Pair<FunctionContext, FilterContext>, Integer>
_filteredAggregationsIndexMap;
+ // Logical aggregation slots are fixed before expression rewrites to
preserve the broker/server result layout.
+ private List<Pair<FunctionContext, FilterContext>>
_aggregationFunctionContexts;
private boolean _hasFilteredAggregations;
// Set by the instance planner after bounding all local BASE and derived
grouping-set keys.
private boolean _groupingSetsBaseAggregationAllowed = true;
@@ -190,6 +192,9 @@ public class QueryContext implements QueryScopedCache {
return false;
}
+ if (groupByKeys.size() == 1 && orderByKeys.size() == 1) {
+ return groupByKeys.get(0).equals(orderByKeys.get(0).getExpression());
+ }
Set<ExpressionContext> groupByKeyIdentSet = new HashSet<>(groupByKeys);
Set<ExpressionContext> orderByKeyIdentSet = new HashSet<>();
orderByKeys.forEach(key -> orderByKeyIdentSet.add(key.getExpression()));
@@ -353,6 +358,44 @@ public class QueryContext implements QueryScopedCache {
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);
+ _effectiveSegmentGroupTrimSize =
calculateEffectiveSegmentGroupTrimSize();
+ }
+ }
+
/// Returns `true` if the query is an EXPLAIN query, `false` otherwise.
public boolean isExplain() {
return _explain != ExplainMode.NONE;
@@ -850,20 +893,14 @@ public class QueryContext implements QueryScopedCache {
}
// Pre-calculate the aggregation functions and columns for the query
- generateAggregationFunctions(queryContext);
- extractColumns(queryContext);
+ extractAggregationFunctionContexts(queryContext);
// Pre-calculate group-by configs
if (queryContext.getGroupByExpressions() != null) {
- /// Grouping-set queries expand each row into multiple groups; force
the unsafe-trim path so the
- /// sort-aggregate fast path (which assumes a single flat grouping) is
not taken.
- queryContext._isUnsafeTrim = queryContext.isGroupingSets()
- || !queryContext.isSameOrderAndGroupByColumns(queryContext) ||
queryContext.getHavingFilter() != null;
Integer sortAggregateLimitThreshold =
QueryOptionsUtils.getSortAggregateLimitThreshold(_queryOptions);
if (sortAggregateLimitThreshold != null) {
queryContext.setSortAggregateLimitThreshold(sortAggregateLimitThreshold);
}
-
queryContext.setEffectiveSegmentGroupTrimSize(queryContext.calculateEffectiveSegmentGroupTrimSize());
// sortAggregateSequentialCombineNumSegmentsThreshold is defaulted to
hardware concurrency
// if not specified. this allows one more parallel thread (the main
thread) to do only combine
@@ -876,70 +913,37 @@ public class QueryContext implements QueryScopedCache {
}
}
+ 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);
- }
-
- // Add aggregation functions in the HAVING and ORDER-BY clause but not
in the SELECT clause
- filteredAggregations.clear();
- if (queryContext._havingFilter != null) {
- getAggregations(queryContext._havingFilter, filteredAggregations);
+ getAggregations(selectExpression, contexts);
}
- if (queryContext._orderByExpressions != null) {
- for (OrderByExpressionContext orderByExpression :
queryContext._orderByExpressions) {
- getAggregations(orderByExpression.getExpression(),
filteredAggregations);
+ if (queryContext._havingFilter != null ||
queryContext._orderByExpressions != null) {
+ List<Pair<FunctionContext, FilterContext>> otherAggregations = new
ArrayList<>();
+ if (queryContext._havingFilter != null) {
+ getAggregations(queryContext._havingFilter, otherAggregations);
}
- }
- for (Pair<FunctionContext, FilterContext> pair : filteredAggregations) {
- if (!filteredAggregationsIndexMap.containsKey(pair)) {
- 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);
+ if (queryContext._orderByExpressions != null) {
+ for (OrderByExpressionContext orderByExpression :
queryContext._orderByExpressions) {
+ getAggregations(orderByExpression.getExpression(),
otherAggregations);
+ }
}
- }
-
- if (!filteredAggregationFunctions.isEmpty()) {
- int numAggregations = filteredAggregationFunctions.size();
- AggregationFunction[] aggregationFunctions = new
AggregationFunction[numAggregations];
- for (int i = 0; i < numAggregations; i++) {
- Pair<AggregationFunction, FilterContext> pair =
filteredAggregationFunctions.get(i);
- aggregationFunctions[i] = pair.getLeft();
- // NOTE: Compute the flag over all the collected aggregations
(SELECT, HAVING and ORDER-BY) instead of only
- // the SELECT ones. A FILTER clause on an aggregation that is
referenced only in HAVING or ORDER-BY
- // must still be honored, else the plan nodes pick the
non-filtered operator and silently evaluate the
- // aggregation over the main filter alone.
- if (pair.getRight() != null) {
- queryContext._hasFilteredAggregations = true;
+ if (!otherAggregations.isEmpty()) {
+ Set<Pair<FunctionContext, FilterContext>> existing = new
HashSet<>(contexts);
+ for (Pair<FunctionContext, FilterContext> context :
otherAggregations) {
+ if (existing.add(context)) {
+ contexts.add(context);
+ }
}
}
- queryContext._aggregationFunctions = aggregationFunctions;
- queryContext._filteredAggregationFunctions =
filteredAggregationFunctions;
- queryContext._filteredAggregationsIndexMap =
filteredAggregationsIndexMap;
}
+ queryContext._aggregationFunctionContexts = contexts;
}
/// Helper method to extract AGGREGATION FunctionContexts and FILTER
FilterContexts from the given expression.
@@ -986,7 +990,7 @@ public class QueryContext implements QueryScopedCache {
}
/// Helper method to extract the columns (IDENTIFIER expressions) for the
query.
- private void extractColumns(QueryContext query) {
+ private static void extractColumns(QueryContext query) {
Set<String> columns = new HashSet<>();
for (ExpressionContext expression : query._selectExpressions) {
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/utils/QueryContextConverterUtils.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/utils/QueryContextConverterUtils.java
index a5cf6b84e02..6d88a294579 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/utils/QueryContextConverterUtils.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/utils/QueryContextConverterUtils.java
@@ -186,9 +186,16 @@ public class QueryContextConverterUtils {
}
// 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 if (expressionOverrideHints.size() == 1) {
+ Map.Entry<Expression, Expression> entry =
expressionOverrideHints.entrySet().iterator().next();
+ expressionContextOverrideHints =
Map.of(RequestContextUtils.getExpression(entry.getKey()),
+ RequestContextUtils.getExpression(entry.getValue()));
+ } else {
+ expressionContextOverrideHints = new HashMap<>();
for (Map.Entry<Expression, Expression> entry :
expressionOverrideHints.entrySet()) {
expressionContextOverrideHints.put(RequestContextUtils.getExpression(entry.getKey()),
RequestContextUtils.getExpression(entry.getValue()));
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/QueryOverrideWithHintsTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/QueryOverrideWithHintsTest.java
index 274b891b129..f9e3ba24dcb 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/QueryOverrideWithHintsTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/QueryOverrideWithHintsTest.java
@@ -19,29 +19,50 @@
package org.apache.pinot.core.plan.maker;
import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import org.apache.commons.lang3.tuple.Pair;
import org.apache.pinot.common.request.Expression;
import org.apache.pinot.common.request.PinotQuery;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.FilterContext;
import org.apache.pinot.common.request.context.FunctionContext;
+import org.apache.pinot.common.request.context.OrderByExpressionContext;
import org.apache.pinot.common.request.context.RequestContextUtils;
import org.apache.pinot.common.request.context.predicate.EqPredicate;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.common.utils.request.RequestUtils;
+import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+import org.apache.pinot.core.query.executor.ResultsBlockStreamer;
+import org.apache.pinot.core.query.reduce.ReducerDataSchemaUtils;
import org.apache.pinot.core.query.request.context.QueryContext;
import
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.SegmentContext;
import org.apache.pinot.spi.data.FieldSpec;
import org.apache.pinot.sql.parsers.CalciteSqlParser;
import org.mockito.Mockito;
import org.testng.annotations.BeforeClass;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNotEquals;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertSame;
public class QueryOverrideWithHintsTest {
@@ -160,19 +181,52 @@ public class QueryOverrideWithHintsTest {
ExpressionContext.forLiteral(FieldSpec.DataType.INT, 1000));
}
- @Test
- public void testRewriteExpressionsWithHints() {
+ @DataProvider
+ public Object[][] multiHintCases() {
+ return new Object[][]{{false}, {true}};
+ }
+
+ @Test(dataProvider = "multiHintCases")
+ public void testRewriteExpressionsWithHints(boolean includeUnavailableHint) {
PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(
- "SELECT datetrunc('MONTH', ts), count(*), sum(abc) from myTable group
by datetrunc('MONTH', ts) ");
+ "SELECT datetrunc('MONTH', ts), count(*), sum(datetrunc('MONTH', ts))
from myTable"
+ + " WHERE datetrunc('MONTH', ts) > 0 group by datetrunc('MONTH',
ts)");
Expression dateTruncFunctionExpr =
RequestUtils.getFunctionExpression("datetrunc",
RequestUtils.getLiteralExpression("MONTH"),
RequestUtils.getIdentifierExpression("ts"));
Expression timestampIndexColumn =
RequestUtils.getIdentifierExpression("$ts$MONTH");
- pinotQuery.setExpressionOverrideHints(Map.of(dateTruncFunctionExpr,
timestampIndexColumn));
+ Map<Expression, Expression> hints = new
HashMap<>(Map.of(dateTruncFunctionExpr, timestampIndexColumn,
+ CalciteSqlParser.compileToExpression("DATETRUNC('MONTH', ts,
'MILLISECONDS')"), timestampIndexColumn));
+ if (includeUnavailableHint) {
+ hints.put(CalciteSqlParser.compileToExpression("DATETRUNC('DAY', ts)"),
+ RequestUtils.getIdentifierExpression("$ts$DAY"));
+ }
+ pinotQuery.setExpressionOverrideHints(hints);
QueryContext queryContext =
QueryContextConverterUtils.getQueryContext(pinotQuery);
- InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext,
_indexSegment);
+ AggregationFunction count = queryContext.getAggregationFunctions()[0];
+ var filteredAggregations = queryContext.getFilteredAggregationFunctions();
+ var countPair = filteredAggregations.get(0);
+ var sumPair = filteredAggregations.get(1);
+ IndexSegment first = mock(IndexSegment.class);
+ Set<String> columns = spy(new HashSet<>(Set.of("$ts$MONTH")));
+ when(first.getColumnNames()).thenReturn(columns);
+ IndexSegment second = mock(IndexSegment.class);
+ when(second.getColumnNames()).thenReturn(Set.of("$ts$MONTH"));
+ InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext,
+ List.of(new SegmentContext(first), new SegmentContext(second)));
+ verify(columns, times(1)).contains("$ts$MONTH");
+ verify(second, times(1)).getColumnNames();
+ assertSame(queryContext.getAggregationFunctions()[0], count);
+ assertSame(queryContext.getFilteredAggregationFunctions(),
filteredAggregations);
+ assertSame(filteredAggregations.get(0), countPair);
+ assertNotSame(filteredAggregations.get(1), sumPair);
+ assertEquals(queryContext.getFilter().getPredicate().getLhs(),
ExpressionContext.forIdentifier("$ts$MONTH"));
+ assertEquals(queryContext.getColumns(), Set.of("$ts$MONTH"));
assertEquals(queryContext.getSelectExpressions().get(0).getIdentifier(),
"$ts$MONTH");
assertEquals(queryContext.getGroupByExpressions().get(0).getIdentifier(),
"$ts$MONTH");
+
assertEquals(queryContext.getAggregationFunctions()[1].getInputExpressions(),
+ List.of(ExpressionContext.forIdentifier("$ts$MONTH")));
+ assertEquals(queryContext.getExpressionOverrideHints().size(),
includeUnavailableHint ? 3 : 2);
}
@Test
@@ -183,10 +237,287 @@ public class QueryOverrideWithHintsTest {
RequestUtils.getFunctionExpression("datetrunc",
RequestUtils.getLiteralExpression("DAY"),
RequestUtils.getIdentifierExpression("ts"));
Expression timestampIndexColumn =
RequestUtils.getIdentifierExpression("$ts$DAY");
- pinotQuery.setExpressionOverrideHints(Map.of(dateTruncFunctionExpr,
timestampIndexColumn));
+ pinotQuery.setExpressionOverrideHints(Map.of(dateTruncFunctionExpr,
timestampIndexColumn,
+ CalciteSqlParser.compileToExpression("DATETRUNC('DAY', ts,
'MILLISECONDS')"), timestampIndexColumn));
QueryContext queryContext =
QueryContextConverterUtils.getQueryContext(pinotQuery);
- InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext,
_indexSegment);
- assertEquals(queryContext.getSelectExpressions().get(0).getFunction(),
-
queryContext.getExpressionOverrideHints().keySet().iterator().next().getFunction());
+ IndexSegment first = mock(IndexSegment.class);
+ Set<String> columns = spy(new HashSet<>(Set.of("ts")));
+ when(first.getColumnNames()).thenReturn(columns);
+ IndexSegment later = mock(IndexSegment.class);
+ Map<ExpressionContext, ExpressionContext> originalHints =
Map.copyOf(queryContext.getExpressionOverrideHints());
+ InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext,
+ List.of(new SegmentContext(first), new SegmentContext(later)));
+ verify(columns, times(1)).contains("$ts$DAY");
+ verify(later, never()).getColumnNames();
+ assertEquals(queryContext.getExpressionOverrideHints(), originalHints);
+
assertEquals(queryContext.getSelectExpressions().get(0).getFunction().getFunctionName(),
"datetrunc");
+ }
+
+ @DataProvider
+ public Object[][] aggregationRewriteCases() {
+ List<Object[]> cases = new ArrayList<>();
+ for (boolean nullHandling : new boolean[]{false, true}) {
+ for (boolean nested : new boolean[]{false, true}) {
+ for (String predicate : List.of("", "DATETRUNC('MONTH', ts,
'MILLISECONDS') > 0", "TRUE", "FALSE")) {
+ cases.add(new Object[]{nullHandling, predicate, nested});
+ }
+ }
+ }
+ return cases.toArray(new Object[0][]);
+ }
+
+ @Test(dataProvider = "aggregationRewriteCases")
+ public void testRewriteCompiledAggregation(boolean nullHandlingEnabled,
String predicate, boolean nested) {
+ boolean filtered = !predicate.isEmpty();
+ String filter = filtered ? " FILTER(WHERE " + predicate + ")" : "";
+ String input = "DATETRUNC('MONTH', ts, 'MILLISECONDS')";
+ if (nested) {
+ input = "ADD(" + input + ", 1)";
+ }
+ QueryContext queryContext = aggregationQuery(
+ "SELECT DISTINCTCOUNT(" + input + ")" + filter + " AS active_days FROM
myTable", "MONTH", nullHandlingEnabled);
+ if (predicate.equals("TRUE") || predicate.equals("FALSE")) {
+
assertEquals(queryContext.getFilteredAggregationFunctions().get(0).getRight().isConstant(),
true);
+ }
+ InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext,
List.of(new SegmentContext(_indexSegment)));
+
+ ExpressionContext derived = ExpressionContext.forIdentifier("$ts$MONTH");
+ ExpressionContext expectedInput = nested
+ ?
RequestContextUtils.getExpression(CalciteSqlParser.compileToExpression("ADD(\"$ts$MONTH\",
1)")) : derived;
+
assertEquals(queryContext.getAggregationFunctions()[0].getInputExpressions(),
List.of(expectedInput));
+ assertEquals(queryContext.getColumns(), Set.of("$ts$MONTH"));
+ assertEquals(queryContext.getAliasList(), List.of("active_days"));
+ assertEquals(queryContext.isNullHandlingEnabled(), nullHandlingEnabled);
+ assertEquals(queryContext.hasFilteredAggregations(), filtered);
+ var pair = queryContext.getFilteredAggregationFunctions().get(0);
+ if (filtered) {
+ if (predicate.equals("TRUE") || predicate.equals("FALSE")) {
+ assertEquals(pair.getRight().isConstant(), true);
+ assertEquals(pair.getRight().isConstantTrue(),
Boolean.parseBoolean(predicate));
+ assertEquals(pair.getRight().isConstantFalse(),
!Boolean.parseBoolean(predicate));
+ } else {
+ assertEquals(pair.getRight().getPredicate().getLhs(), derived);
+ }
+ }
+ FunctionContext function =
queryContext.getSelectExpressions().get(0).getFunction();
+ if (filtered) {
+ function = function.getArguments().get(0).getFunction();
+ }
+
assertEquals(queryContext.getFilteredAggregationsIndexMap().get(Pair.of(function,
pair.getRight())),
+ Integer.valueOf(0));
+
+ AggregationFunction compiled = queryContext.getAggregationFunctions()[0];
+ InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext,
List.of(new SegmentContext(_indexSegment)));
+ assertSame(queryContext.getAggregationFunctions()[0], compiled);
+ }
+
+ @DataProvider
+ public Object[][] orderByRewriteCases() {
+ List<Object[]> cases = new ArrayList<>();
+ for (String expression : List.of("ts", "DATETRUNC('MONTH', ts,
'MILLISECONDS')",
+ "ADD(DATETRUNC('MONTH', ts, 'MILLISECONDS'), 1)")) {
+ for (String ordering : List.of("ASC NULLS FIRST", "ASC NULLS LAST",
"DESC NULLS FIRST", "DESC NULLS LAST")) {
+ cases.add(new Object[]{expression, ordering});
+ }
+ }
+ return cases.toArray(new Object[0][]);
+ }
+
+ @Test(dataProvider = "orderByRewriteCases")
+ public void testRewritePreservesNullOrdering(String expression, String
ordering) {
+ QueryContext context = aggregationQuery(
+ "SELECT DATETRUNC('MONTH', ts, 'MILLISECONDS') FROM myTable ORDER BY "
+ expression + " " + ordering,
+ "MONTH", true);
+ OrderByExpressionContext original = context.getOrderByExpressions().get(0);
+ InstancePlanMakerImplV2.rewriteQueryContextWithHints(context, List.of(new
SegmentContext(_indexSegment)));
+ OrderByExpressionContext rewritten =
context.getOrderByExpressions().get(0);
+ assertEquals(rewritten.isAsc(), original.isAsc());
+ assertEquals(rewritten.isNullsLast(), original.isNullsLast());
+ if (expression.startsWith("DATETRUNC")) {
+ assertEquals(rewritten.getExpression(),
ExpressionContext.forIdentifier("$ts$MONTH"));
+ assertNotSame(rewritten, original);
+ } else {
+ assertSame(rewritten, original);
+ if (expression.startsWith("ADD")) {
+
assertEquals(rewritten.getExpression().getFunction().getArguments().get(0),
+ ExpressionContext.forIdentifier("$ts$MONTH"));
+ }
+ }
+ }
+
+ @Test
+ public void testPreparationWithoutHints() {
+ QueryContext queryContext =
+ QueryContextConverterUtils.getQueryContext("SELECT COUNT(*), SUM(ts),
SUM(ts) FROM myTable");
+ assertEquals(queryContext.getAggregationFunctions().length, 3);
+ var filteredAggregations = queryContext.getFilteredAggregationFunctions();
+ var countPair = filteredAggregations.get(0);
+ var sumPair = filteredAggregations.get(1);
+ var duplicatePair = filteredAggregations.get(2);
+ var indexMap = queryContext.getFilteredAggregationsIndexMap();
+ AggregationFunction count = queryContext.getAggregationFunctions()[0];
+ AggregationFunction sum = queryContext.getAggregationFunctions()[1];
+ FunctionContext sumContext =
queryContext.getSelectExpressions().get(1).getFunction();
+ sumContext.getArguments().set(0, ExpressionContext.forIdentifier("other"));
+ Set<FunctionContext> changed = Collections.newSetFromMap(new
IdentityHashMap<>());
+ changed.add(sumContext);
+ queryContext.prepareForExecution(changed);
+ assertSame(queryContext.getAggregationFunctions()[0], count);
+ assertSame(queryContext.getFilteredAggregationFunctions(),
filteredAggregations);
+ assertSame(queryContext.getFilteredAggregationsIndexMap(), indexMap);
+ assertSame(filteredAggregations.get(0), countPair);
+ assertNotSame(filteredAggregations.get(1), sumPair);
+ assertSame(filteredAggregations.get(2), duplicatePair);
+ assertNotSame(queryContext.getAggregationFunctions()[1], sum);
+
assertEquals(queryContext.getAggregationFunctions()[1].getInputExpressions(),
+ List.of(ExpressionContext.forIdentifier("other")));
+ assertEquals(queryContext.getColumns(), Set.of("ts", "other"));
+
+ ExpressionContext first = ExpressionContext.forIdentifier("first");
+ ExpressionContext second = ExpressionContext.forIdentifier("second");
+ for (Pair<List<ExpressionContext>, Boolean> testCase :
List.of(Pair.of(List.of(first), false),
+ Pair.of(List.of(second), true), Pair.of(List.of(first, second), false),
+ Pair.of(List.of(first, first, second), false))) {
+ List<ExpressionContext> groupKeys = testCase.getLeft();
+ List<OrderByExpressionContext> orderKeys = groupKeys.size() == 1
+ ? List.of(new OrderByExpressionContext(first, true))
+ : List.of(new OrderByExpressionContext(second, true), new
OrderByExpressionContext(first, true));
+ QueryContext grouped = new
QueryContext.Builder().setSelectExpressions(List.of(first, second))
+
.setGroupByExpressions(groupKeys).setOrderByExpressions(orderKeys).build();
+ assertEquals(grouped.isUnsafeTrim(), testCase.getRight().booleanValue());
+ }
+ }
+
+ @Test
+ public void testFilterOnlyRewriteRetainsAggregationEntries() {
+ for (boolean aggregationFilter : new boolean[]{false, true}) {
+ String predicate = "DATETRUNC('MONTH', ts, 'MILLISECONDS') > 0";
+ String sql = aggregationFilter ? "SELECT SUM(ts) FILTER(WHERE " +
predicate + ") FROM myTable"
+ : "SELECT SUM(ts) FROM myTable WHERE " + predicate;
+ QueryContext context = aggregationQuery(sql, "MONTH", false);
+ var aggregations = context.getFilteredAggregationFunctions();
+ var pair = aggregations.get(0);
+ var indexMap = context.getFilteredAggregationsIndexMap();
+ AggregationFunction function = context.getAggregationFunctions()[0];
+ InstancePlanMakerImplV2.rewriteQueryContextWithHints(context,
List.of(new SegmentContext(_indexSegment)));
+ assertSame(context.getFilteredAggregationFunctions(), aggregations);
+ assertSame(aggregations.get(0), pair);
+ assertSame(context.getAggregationFunctions()[0], function);
+ assertSame(context.getFilteredAggregationsIndexMap(), indexMap);
+ assertEquals(context.hasFilteredAggregations(), aggregationFilter);
+ FilterContext filter = aggregationFilter ? pair.getRight() :
context.getFilter();
+ assertEquals(filter.getPredicate().getLhs(),
ExpressionContext.forIdentifier("$ts$MONTH"));
+ FunctionContext aggregation =
context.getSelectExpressions().get(0).getFunction();
+ if (aggregationFilter) {
+ aggregation = aggregation.getArguments().get(0).getFunction();
+ }
+ assertEquals(indexMap.get(Pair.of(aggregation, pair.getRight())),
Integer.valueOf(0));
+ assertEquals(context.getColumns(), Set.of("ts", "$ts$MONTH"));
+ }
+ }
+
+ @Test(expectedExceptions = IllegalArgumentException.class,
+ expectedExceptionsMessageRegExp = "Expression override target must be an
identifier")
+ public void testRejectNonIdentifierOverrideTarget() {
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery("SELECT
DATETRUNC('MONTH', ts) FROM myTable");
+
query.setExpressionOverrideHints(Map.of(CalciteSqlParser.compileToExpression("DATETRUNC('MONTH',
ts)"),
+ RequestUtils.getLiteralExpression(0)));
+
InstancePlanMakerImplV2.rewriteQueryContextWithHints(QueryContextConverterUtils.getQueryContext(query),
+ List.of(new SegmentContext(_indexSegment)));
+ }
+
+ @Test
+ public void testRewritePreservesAggregationSlots() {
+ String sql = "SELECT DISTINCTCOUNT(DATETRUNC('MONTH', ts, 'MILLISECONDS'))
AS days, "
+ + "DISTINCTCOUNT(DATETRUNC('MONTH', ts, 'MILLISECONDS')) AS
duplicate_days FROM myTable "
+ + "HAVING DISTINCTCOUNT(\"$ts$MONTH\") > 0 ORDER BY
DISTINCTCOUNT(\"$ts$MONTH\")";
+ QueryContext queryContext = aggregationQuery(sql, "MONTH", false);
+ QueryContext brokerContext = aggregationQuery(sql, "MONTH", false);
+ assertEquals(queryContext.getAggregationFunctions().length, 3);
+ InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext,
List.of(new SegmentContext(_indexSegment)));
+ // Converging SELECT and HAVING/ORDER-BY expressions must keep their
original wire slots.
+ assertEquals(queryContext.getAggregationFunctions().length, 3);
+ assertEquals(queryContext.getFilteredAggregationFunctions().size(), 3);
+ for (AggregationFunction function :
queryContext.getAggregationFunctions()) {
+ assertEquals(function.getInputExpressions(),
List.of(ExpressionContext.forIdentifier("$ts$MONTH")));
+ }
+ FunctionContext having =
queryContext.getHavingFilter().getPredicate().getLhs().getFunction();
+
assertEquals(queryContext.getFilteredAggregationsIndexMap().get(Pair.of(having,
null)), Integer.valueOf(2));
+ assertEquals(queryContext.getAliasList(), List.of("days",
"duplicate_days"));
+ String[] physicalNames = new String[3];
+ String[] logicalNames = new String[3];
+ ColumnDataType[] types = new ColumnDataType[3];
+ for (int i = 0; i < 3; i++) {
+ physicalNames[i] =
queryContext.getAggregationFunctions()[i].getResultColumnName();
+ logicalNames[i] =
brokerContext.getAggregationFunctions()[i].getResultColumnName();
+ types[i] =
queryContext.getAggregationFunctions()[i].getIntermediateResultColumnType();
+ }
+ DataSchema canonical =
ReducerDataSchemaUtils.canonicalizeDataSchemaForAggregation(brokerContext,
+ new DataSchema(physicalNames, types));
+ assertEquals(canonical.getColumnNames(), logicalNames);
+ assertEquals(canonical.getColumnDataTypes(), types);
+ }
+
+ @Test
+ public void testMissingIndexRetainsCompiledAggregation() {
+ QueryContext queryContext = aggregationQuery(
+ "SELECT DISTINCTCOUNT(DATETRUNC('DAY', ts, 'MILLISECONDS')) AS
active_days FROM myTable", "DAY", false);
+ assertEquals(queryContext.getExpressionOverrideHints().size(), 1);
+
assertEquals(queryContext.getExpressionOverrideHints().get(RequestContextUtils.getExpression(
+ CalciteSqlParser.compileToExpression("DATETRUNC('DAY', ts,
'MILLISECONDS')"))),
+ ExpressionContext.forIdentifier("$ts$DAY"));
+ AggregationFunction function = queryContext.getAggregationFunctions()[0];
+ List<ExpressionContext> inputs = function.getInputExpressions();
+ IndexSegment later = mock(IndexSegment.class);
+ InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext,
+ List.of(new SegmentContext(_indexSegment), new SegmentContext(later)));
+ verify(later, never()).getColumnNames();
+ assertSame(queryContext.getAggregationFunctions()[0], function);
+
assertEquals(queryContext.getAggregationFunctions()[0].getInputExpressions(),
inputs);
+ assertEquals(queryContext.getColumns(), Set.of("ts"));
+ assertEquals(queryContext.getAliasList(), List.of("active_days"));
+ }
+
+ @DataProvider
+ public Object[][] instanceRewriteCases() {
+ return new Object[][]{{false, false}, {true, false}, {false, true}, {true,
true}};
+ }
+
+ @Test(dataProvider = "instanceRewriteCases")
+ public void testInstanceRewriteRequiresIndexOnEverySegment(boolean
streaming, boolean allIndexed) {
+ QueryContext queryContext = aggregationQuery(
+ "SELECT DISTINCTCOUNT(DATETRUNC('MONTH', ts, 'MILLISECONDS')) AS
active_days FROM myTable", "MONTH", false);
+ IndexSegment second = mock(IndexSegment.class);
+ when(second.getColumnNames()).thenReturn(allIndexed ? Set.of("ts",
"$ts$MONTH") : Set.of("ts"));
+ List<SegmentContext> segments = List.of(new SegmentContext(_indexSegment),
new SegmentContext(second));
+ Map<ExpressionContext, ExpressionContext> hints =
queryContext.getExpressionOverrideHints();
+ Map<ExpressionContext, ExpressionContext> originalHints =
Map.copyOf(hints);
+ InstancePlanMakerImplV2 planMaker = new InstancePlanMakerImplV2();
+ ExecutorService executor = mock(ExecutorService.class);
+ if (streaming) {
+ planMaker.makeStreamingInstancePlan(segments, queryContext, executor,
mock(ResultsBlockStreamer.class));
+ } else {
+ planMaker.makeInstancePlan(segments, queryContext, executor);
+ }
+ assertSame(queryContext.getExpressionOverrideHints(), hints);
+ assertEquals(hints, originalHints);
+ AggregationFunction<?, ?> aggregation =
queryContext.getAggregationFunctions()[0];
+ ExpressionContext input = aggregation.getInputExpressions().get(0);
+ if (allIndexed) {
+ assertEquals(input, ExpressionContext.forIdentifier("$ts$MONTH"));
+ assertEquals(queryContext.getColumns(), Set.of("$ts$MONTH"));
+ } else {
+ assertEquals(input.getFunction().getFunctionName(), "datetrunc");
+ assertEquals(queryContext.getColumns(), Set.of("ts"));
+ }
+ }
+
+ private static QueryContext aggregationQuery(String sql, String granularity,
boolean nullHandlingEnabled) {
+ PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql);
+ pinotQuery.setQueryOptions(Map.of("enableNullHandling",
Boolean.toString(nullHandlingEnabled)));
+ pinotQuery.setExpressionOverrideHints(Map.of(
+ CalciteSqlParser.compileToExpression("DATETRUNC('" + granularity + "',
ts, 'MILLISECONDS')"),
+ RequestUtils.getIdentifierExpression("$ts$" + granularity)));
+ return QueryContextConverterUtils.getQueryContext(pinotQuery);
}
}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/queries/TimestampQueriesTest.java
b/pinot-core/src/test/java/org/apache/pinot/queries/TimestampQueriesTest.java
index 824c762ea6e..72a1cda43ef 100644
---
a/pinot-core/src/test/java/org/apache/pinot/queries/TimestampQueriesTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/queries/TimestampQueriesTest.java
@@ -24,30 +24,46 @@ import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
+import java.util.Map;
+import java.util.Set;
import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.request.PinotQuery;
import org.apache.pinot.common.response.broker.BrokerResponseNative;
import org.apache.pinot.common.response.broker.ResultTable;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.common.utils.request.RequestUtils;
+import org.apache.pinot.core.operator.blocks.InstanceResponseBlock;
+import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
import
org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
import
org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
import org.apache.pinot.segment.spi.ImmutableSegment;
import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.SegmentContext;
import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.spi.config.table.FieldConfig;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.TimestampConfig;
+import org.apache.pinot.spi.config.table.TimestampIndexGranularity;
import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.data.readers.GenericRow;
import org.apache.pinot.spi.exception.BadQueryRequestException;
+import org.apache.pinot.spi.query.QueryThreadContext;
import org.apache.pinot.spi.utils.ReadMode;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.apache.pinot.sql.parsers.CalciteSqlParser;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
/// Queries test for TIMESTAMP data type.
@@ -66,6 +82,7 @@ public class TimestampQueriesTest extends BaseQueriesTest {
new
TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build();
private IndexSegment _indexSegment;
+ private IndexSegment _timestampIndexedSegment;
private List<IndexSegment> _indexSegments;
@Override
@@ -115,6 +132,58 @@ public class TimestampQueriesTest extends BaseQueriesTest {
ImmutableSegment immutableSegment = ImmutableSegmentLoader.load(new
File(INDEX_DIR, SEGMENT_NAME), ReadMode.mmap);
_indexSegment = immutableSegment;
_indexSegments = Arrays.asList(immutableSegment, immutableSegment);
+
+ TableConfig indexedConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME)
+ .setFieldConfigList(List.of(new
FieldConfig.Builder(TIMESTAMP_COLUMN).withTimestampConfig(
+ new
TimestampConfig(List.of(TimestampIndexGranularity.DAY))).build())).build();
+ Schema indexedSchema = new Schema.SchemaBuilder()
+ .addDateTime(TIMESTAMP_COLUMN, DataType.TIMESTAMP,
"1:MILLISECONDS:EPOCH", "1:MILLISECONDS").build();
+ SegmentGeneratorConfig indexedGenerator = new
SegmentGeneratorConfig(indexedConfig, indexedSchema);
+ indexedGenerator.setSegmentName(SEGMENT_NAME + "_indexed");
+ indexedGenerator.setOutDir(INDEX_DIR.getPath());
+ List<GenericRow> indexedRecords = new ArrayList<>(NUM_RECORDS);
+ for (int i = 0; i < NUM_RECORDS; i++) {
+ GenericRow indexedRecord = new GenericRow();
+ indexedRecord.putValue(TIMESTAMP_COLUMN, BASE_TIMESTAMP + i);
+ indexedRecords.add(indexedRecord);
+ }
+ SegmentIndexCreationDriverImpl indexedDriver = new
SegmentIndexCreationDriverImpl();
+ indexedDriver.init(indexedGenerator, new
GenericRowRecordReader(indexedRecords));
+ indexedDriver.build();
+ _timestampIndexedSegment =
+ ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME +
"_indexed"), ReadMode.mmap);
+ }
+
+ @DataProvider
+ public Object[][] timestampIndexSegmentOrders() {
+ return new Object[][]{{false, false}, {true, false}, {false, true}, {true,
true}};
+ }
+
+ @Test(dataProvider = "timestampIndexSegmentOrders")
+ public void testTimestampAggregationWithPrefetch(boolean reversed, boolean
allIndexed)
+ throws Exception {
+ String expression = "DATETRUNC('DAY', " + TIMESTAMP_COLUMN + ",
'MILLISECONDS')";
+ String derived = "$" + TIMESTAMP_COLUMN + "$DAY";
+ assertTrue(_timestampIndexedSegment.getColumnNames().contains(derived));
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery(
+ "SELECT DISTINCTCOUNT(" + expression + "), COUNT(*) FROM " +
RAW_TABLE_NAME);
+
query.setExpressionOverrideHints(Map.of(CalciteSqlParser.compileToExpression(expression),
+ RequestUtils.getIdentifierExpression(derived)));
+ QueryContext context = QueryContextConverterUtils.getQueryContext(query);
+ context.setEnablePrefetch(true);
+ context.setServerReturnFinalResult(true);
+ context.setEndTimeMs(System.currentTimeMillis() + 30_000);
+ IndexSegment second = allIndexed ? _timestampIndexedSegment :
_indexSegment;
+ List<SegmentContext> segments = reversed
+ ? List.of(new SegmentContext(second), new
SegmentContext(_timestampIndexedSegment))
+ : List.of(new SegmentContext(_timestampIndexedSegment), new
SegmentContext(second));
+ try (QueryThreadContext ignored = QueryThreadContext.openForSseTest()) {
+ InstanceResponseBlock result =
+ new InstancePlanMakerImplV2().makeInstancePlan(segments, context,
EXECUTOR_SERVICE).execute();
+ assertEquals(((Number) result.getRows().get(0)[0]).longValue(), 1L);
+ assertEquals(((Number) result.getRows().get(0)[1]).longValue(), 2L *
NUM_RECORDS);
+ }
+ assertEquals(context.getColumns(), allIndexed ? Set.of(derived) :
Set.of(TIMESTAMP_COLUMN));
}
@Test
@@ -238,6 +307,7 @@ public class TimestampQueriesTest extends BaseQueriesTest {
public void tearDown()
throws IOException {
_indexSegment.destroy();
+ _timestampIndexedSegment.destroy();
FileUtils.deleteDirectory(INDEX_DIR);
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
