xiangfu0 commented on code in PR #20035:
URL: https://github.com/apache/pinot/pull/20035#discussion_r4384033185
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java:
##########
@@ -657,86 +660,130 @@ 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.
+ @VisibleForTesting
+ static void rewriteQueryContextWithHints(QueryContext queryContext,
List<SegmentContext> segmentContexts) {
+ Map<ExpressionContext, ExpressionContext> hints =
queryContext.getExpressionOverrideHints();
+ if (MapUtils.isEmpty(hints) || segmentContexts.isEmpty()) {
+ return;
+ }
+ queryContext.getOrComputeSharedValue(Boolean.class,
EXPRESSION_HINTS_REWRITTEN_KEY, ignored -> {
+ Set<String> columns = getOverrideColumns(hints);
+ for (SegmentContext segmentContext : segmentContexts) {
+ columns.retainAll(segmentContext.getIndexSegment().getColumnNames());
+ }
+ rewriteQueryContextWithHints(queryContext, columns);
+ return true;
+ });
+ }
+
+ /// Rewrites a standalone segment plan once; instance plans first constrain
the rewrite across all their segments.
@VisibleForTesting
public static void rewriteQueryContextWithHints(QueryContext queryContext,
IndexSegment indexSegment) {
- Map<ExpressionContext, ExpressionContext> expressionOverrideHints =
queryContext.getExpressionOverrideHints();
- if (MapUtils.isEmpty(expressionOverrideHints)) {
+ if (MapUtils.isEmpty(queryContext.getExpressionOverrideHints())) {
return;
}
+ queryContext.getOrComputeSharedValue(Boolean.class,
EXPRESSION_HINTS_REWRITTEN_KEY, ignored -> {
Review Comment:
Documented in
[af7d522ed8](https://github.com/apache/pinot/commit/af7d522ed8d0e208911f4d1ed325f2c2ee192c68):
instance planning first rewrites against columns common to every selected
segment, so this overload becomes a no-op afterward. The standalone path
explicitly requires a single-segment context; it must not be reused across
segments with different available columns. Both instance entry points and the
mixed-index regression retain the shared guard.
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java:
##########
@@ -657,86 +660,130 @@ 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.
+ @VisibleForTesting
+ static void rewriteQueryContextWithHints(QueryContext queryContext,
List<SegmentContext> segmentContexts) {
+ Map<ExpressionContext, ExpressionContext> hints =
queryContext.getExpressionOverrideHints();
+ if (MapUtils.isEmpty(hints) || segmentContexts.isEmpty()) {
+ return;
+ }
+ queryContext.getOrComputeSharedValue(Boolean.class,
EXPRESSION_HINTS_REWRITTEN_KEY, ignored -> {
+ Set<String> columns = getOverrideColumns(hints);
+ for (SegmentContext segmentContext : segmentContexts) {
+ columns.retainAll(segmentContext.getIndexSegment().getColumnNames());
+ }
+ rewriteQueryContextWithHints(queryContext, columns);
+ return true;
+ });
+ }
+
+ /// Rewrites a standalone segment plan once; instance plans first constrain
the rewrite across all their segments.
@VisibleForTesting
public static void rewriteQueryContextWithHints(QueryContext queryContext,
IndexSegment indexSegment) {
- Map<ExpressionContext, ExpressionContext> expressionOverrideHints =
queryContext.getExpressionOverrideHints();
- if (MapUtils.isEmpty(expressionOverrideHints)) {
+ if (MapUtils.isEmpty(queryContext.getExpressionOverrideHints())) {
return;
}
+ queryContext.getOrComputeSharedValue(Boolean.class,
EXPRESSION_HINTS_REWRITTEN_KEY, ignored -> {
+ Set<String> columns =
getOverrideColumns(queryContext.getExpressionOverrideHints());
+ columns.retainAll(indexSegment.getColumnNames());
+ rewriteQueryContextWithHints(queryContext, columns);
+ return true;
+ });
+ }
+
+ private static Set<String> getOverrideColumns(Map<ExpressionContext,
ExpressionContext> hints) {
+ Set<String> columns = new HashSet<>();
+ for (ExpressionContext expression : hints.values()) {
+ if (expression.getIdentifier() != null) {
+ columns.add(expression.getIdentifier());
+ }
+ }
+ return columns;
+ }
+ private static void rewriteQueryContextWithHints(QueryContext queryContext,
Set<String> columns) {
+ if (columns.isEmpty()) {
+ return;
+ }
+ Map<ExpressionContext, ExpressionContext> expressionOverrideHints =
queryContext.getExpressionOverrideHints();
List<ExpressionContext> selectExpressions =
queryContext.getSelectExpressions();
selectExpressions.replaceAll(
- expression -> overrideWithExpressionHints(expression, indexSegment,
expressionOverrideHints));
+ expression -> overrideWithExpressionHints(expression, columns,
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);
+ overrideWithExpressionHints(right, columns, expressionOverrideHints);
}
}
}
List<ExpressionContext> groupByExpressions =
queryContext.getGroupByExpressions();
if (CollectionUtils.isNotEmpty(groupByExpressions)) {
groupByExpressions.replaceAll(
- expression -> overrideWithExpressionHints(expression, indexSegment,
expressionOverrideHints));
+ expression -> overrideWithExpressionHints(expression, columns,
expressionOverrideHints));
}
List<OrderByExpressionContext> orderByExpressions =
queryContext.getOrderByExpressions();
if (CollectionUtils.isNotEmpty(orderByExpressions)) {
orderByExpressions.replaceAll(expression -> new OrderByExpressionContext(
- overrideWithExpressionHints(expression.getExpression(),
indexSegment, expressionOverrideHints),
+ overrideWithExpressionHints(expression.getExpression(), columns,
expressionOverrideHints),
expression.isAsc()));
}
- // In-place override
FilterContext filter = queryContext.getFilter();
if (filter != null) {
- overrideWithExpressionHints(filter, indexSegment,
expressionOverrideHints);
+ overrideWithExpressionHints(filter, columns, expressionOverrideHints);
}
-
- // In-place override
FilterContext havingFilter = queryContext.getHavingFilter();
if (havingFilter != null) {
- overrideWithExpressionHints(havingFilter, indexSegment,
expressionOverrideHints);
+ overrideWithExpressionHints(havingFilter, columns,
expressionOverrideHints);
}
+ // Aggregation constructors capture their input expressions; AST
replacement alone leaves those inputs stale.
+ // Rebuilding also repairs hash-map keys whose FunctionContext arguments
were changed in place.
+ queryContext.refreshAfterExpressionRewrite();
Review Comment:
Keeping the once-per-context rebuild here. It retains the original
aggregation slots and rebuilds the filtered-aggregation lookup keys
consistently after the mutable expressions change. The empty-column case still
returns before rebuilding; the existing missing-index regression verifies the
aggregation array is retained in that case. Avoiding recompilation of
individual unchanged aggregates would add change tracking to this correctness
fix for a small planning-only saving.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java:
##########
@@ -353,6 +356,28 @@ public Map<ExpressionContext, ExpressionContext>
getExpressionOverrideHints() {
return _expressionOverrideHints;
}
+ /// Recompiles execution inputs after an in-place expression rewrite, before
any segment operators are built.
+ /// The original aggregation slots are retained even when formerly distinct
HAVING/ORDER-BY inputs become equal.
+ public void refreshAfterExpressionRewrite() {
Review Comment:
Marked this explicitly as an internal planning hook in
[af7d522ed8](https://github.com/apache/pinot/commit/af7d522ed8d0e208911f4d1ed325f2c2ee192c68),
including the requirement to run before fetch planning or segment-operator
construction. It remains public because the caller is in
`org.apache.pinot.core.plan.maker`, while `QueryContext` is in
`org.apache.pinot.core.query.request.context`; package-private would prevent
that call.
--
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]
