xiangfu0 commented on code in PR #20035:
URL: https://github.com/apache/pinot/pull/20035#discussion_r4371298131
##########
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:
This segment-level overload shares the same EXPRESSION_HINTS_REWRITTEN_KEY
guard as the instance-level one, so once makeInstancePlan has run the
cross-segment rewrite, these per-segment calls in makeSegmentPlanNode become
no-ops. That is the intended behavior and it is correct, but the two methods
are now quietly coupled through that shared key. A one-line note here that the
segment overload is a no-op after the instance-level rewrite would save a
future reader from assuming this path re-narrows columns per segment.
##########
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:
refreshAfterExpressionRewrite rebuilds every aggregation function whenever
any override column is present, even aggregations whose arguments were not
actually rewritten. This path is only reached when columns is non-empty (the
no-op case returns earlier, which testMissingIndexRetainsCompiledAggregation
confirms with assertSame), so it is just a little redundant recompilation
during planning, not a correctness issue. Fine to leave as is.
--
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]
