yashmayya commented on PR #20061: URL: https://github.com/apache/pinot/pull/20061#issuecomment-6488344728
**Suggestion: turn the MV aggregate on by default, with no rolling-upgrade risk** The gate is off by default because the gate-on rule also rewrites leaf-only `GROUP BY ARRAY_TO_MV`, which old servers run today. The rule cannot skip only those aggregates. It does not know the final stage, because `PinotJoinToDynamicBroadcastRule` runs after it. Idea: decide after the plan is final, in an `auto` mode that is the default: 1. Plan the query without the MV aggregate, as master does. 2. If `ArrayToMvValidationVisitor` rejects that plan, optimize the converted plan again with the MV aggregate. Every query that plans on master keeps its plan. Only queries that fail on master get `MvAggregateNode`. This needs no stage prediction and no server capabilities (#20058). **Test.** I built a prototype on `f77feaca85` (diff below) and compared it with master (`eeebb52011`). The harness sets no `useMvAggregate` option. It ran 328 queries in v1, v2 and lite: our review probes, `MvAggregate.json`, and a copy of `MvAggregate.json` without `SET useMvAggregate=true`. | Check | Result | |---|---| | 714 runs that work on master (v1 245, v2 240, lite 229) | Same results. Byte-identical stage plans (SHA-256 of each serialized stage `PlanNode`). | | The same 714 plans, read with the `PlanNodeDeserializer` of master (an old server) | All deserialize and serialize back to the same bytes. | | Queries that fail on master | 53 now work through the MV aggregate (v1). The master deserializer rejects each of these plans with `Unsupported PlanNode type: NODE_NOT_SET`. | | The other 19 v1 failures | They also fail with the gate on: unsupported uses, explicit `useMvAggregate=false`, or runtime errors. | | Expected outputs of the `MvAggregate.json` copy | All pass, except `ARRAY_TO_MV(col) IS NOT NULL`, which now gives the master answer ([blocker 1](https://github.com/apache/pinot/pull/20061#discussion_r4406726392)). | | `BasicQueryPlans.json` of master | Passes. | In v1, the only master-working queries with a new plan are the 11 `MvAggregate.json` queries that set `useMvAggregate=true` themselves. With `SET useMvAggregate=false`, every query gives the master result. **How the prototype works** - `CompiledQuery` keeps the unoptimized `RelRoot`. If `planQuery` or `explain` gets an `ArrayToMvValidationException` (a new `QUERY_PLANNING` subclass) and the query does not set `useMvAggregate`, it optimizes that `RelRoot` again with the MV aggregate on. `PlannerContext.enableMvAggregateFallback()` supplies new planners for that pass. - The second pass does not parse the SQL again, so the broker IdSet rewrite and the RLS options stay. - The rules read the gate from `PlannerContext.isUseMvAggregate()`. `PinotProjectJoinTransposeRule` keeps `ARRAY_TO_MV` above a join only when the gate is on ([blocker 2](https://github.com/apache/pinot/pull/20061#discussion_r4406726399)). - `ARRAY_TO_MV` has the master type in both passes, because the second pass reuses the converted plan ([blocker 1](https://github.com/apache/pinot/pull/20061#discussion_r4406726392)). A null MV value still goes into the null group at runtime. - The validator treats `UNNEST` like a lookup join: the `UNNEST` input is the leaf boundary. Without this, `CROSS JOIN UNNEST` with `GROUP BY ARRAY_TO_MV` passes the validation, fails at runtime, and never falls back. - Cost: one more optimizer pass, only for queries that fail on master. **Not in the prototype:** the broker config values (`auto` as the default, plus `true` and `false`), `EXPLAIN PLAN WITHOUT IMPLEMENTATION` (it skips the validation, so it shows the first plan), and the [`HAVING` blocker](https://github.com/apache/pinot/pull/20061#discussion_r4406726405). <details><summary>Prototype diff (planner only)</summary> ```diff diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java index 380bd0163c..356ec6744a 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java @@ -65,9 +65,7 @@ import org.apache.pinot.calcite.rel.logical.PinotLogicalExchange; import org.apache.pinot.calcite.rel.logical.PinotLogicalMvAggregate; import org.apache.pinot.calcite.rel.logical.PinotLogicalSortExchange; import org.apache.pinot.common.function.sql.PinotSqlAggFunction; -import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.query.QueryEnvironment; -import org.apache.pinot.query.context.PlannerContext; import org.apache.pinot.query.planner.plannode.AggregateNode.AggType; import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.spi.utils.CommonConstants; @@ -356,20 +354,9 @@ public class PinotAggregateExchangeNodeInsertRule { return convertAggFromIntermediateInput(aggRel, exchange, AggType.FINAL, leafReturnFinalResult, collations, limit); } - /// Whether `GROUP BY ARRAY_TO_MV(col)` is planned as an MV aggregate. The query option `useMvAggregate` decides when - /// the query sets it. Otherwise [QueryEnvironment.Config#defaultUseMvAggregate()] decides. + /// Whether `GROUP BY ARRAY_TO_MV(col)` is planned as an MV aggregate, see [PinotRuleUtils#isMvAggregateEnabled]. private static boolean isMvAggregateEnabled(RelOptRuleCall call) { - Context genericContext = call.getPlanner().getContext(); - if (genericContext == null) { - return true; - } - PlannerContext plannerContext = genericContext.unwrap(PlannerContext.class); - if (plannerContext != null) { - return QueryOptionsUtils.isUseMvAggregate(plannerContext.getOptions(), - plannerContext.getEnvConfig().defaultUseMvAggregate()); - } - QueryEnvironment.Config envConfig = genericContext.unwrap(QueryEnvironment.Config.class); - return envConfig == null || envConfig.defaultUseMvAggregate(); + return PinotRuleUtils.isMvAggregateEnabled(call); } // TODO: Currently it only handles one WITHIN GROUP collation across all AggregateCalls. diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotProjectJoinTransposeRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotProjectJoinTransposeRule.java index 3b4e310bcd..a49caace8d 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotProjectJoinTransposeRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotProjectJoinTransposeRule.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.calcite.rel.rules; +import javax.annotation.Nullable; import org.apache.calcite.plan.RelOptRuleCall; import org.apache.calcite.rel.core.Join; import org.apache.calcite.rel.rules.ProjectJoinTransposeRule; @@ -29,24 +30,33 @@ import org.apache.pinot.calcite.sql.fun.ArrayToMvFunctionName; /// Similar to [ProjectJoinTransposeRule] but do not transpose project into right side of lookup join. /// -/// An expression that contains `ARRAY_TO_MV` stays above the join; only the input references it uses are pushed down. -/// `GROUP BY ARRAY_TO_MV(col)` above a join is planned as an MV aggregate only when the group expression is in the -/// Project under the aggregate (see [ArrayToMvGroupKeys]). Below the join, the expression would be in a leaf-stage -/// Project, where `ARRAY_TO_MV` is not supported. +/// When the MV aggregate is on (see [PinotRuleUtils#isMvAggregateEnabled]), an expression that contains `ARRAY_TO_MV` +/// stays above the join; only the input references it uses are pushed down. `GROUP BY ARRAY_TO_MV(col)` above a join +/// is planned as an MV aggregate only when the group expression is in the Project under the aggregate (see +/// [ArrayToMvGroupKeys]). Below the join, the expression would be in a leaf-stage Project, where `ARRAY_TO_MV` is not +/// supported. When the MV aggregate is off, the rule pushes such an expression down as before the MV aggregate, so +/// that the plan does not change. /// /// TODO: Allow transposing project into left side of lookup join. public class PinotProjectJoinTransposeRule extends ProjectJoinTransposeRule { - private static final Config DEFAULT_CONFIG = Config.DEFAULT.withPreserveExprCondition( + private static final Config MV_AGGREGATE_CONFIG = Config.DEFAULT.withPreserveExprCondition( expr -> Config.DEFAULT.preserveExprCondition().test(expr) && !containsArrayToMv(expr)); - public static final PinotProjectJoinTransposeRule INSTANCE = new PinotProjectJoinTransposeRule(DEFAULT_CONFIG); + public static final PinotProjectJoinTransposeRule INSTANCE = + new PinotProjectJoinTransposeRule(MV_AGGREGATE_CONFIG, new PinotProjectJoinTransposeRule(Config.DEFAULT, null)); public static PinotProjectJoinTransposeRule instanceWithDescription(String description) { - return new PinotProjectJoinTransposeRule((Config) DEFAULT_CONFIG.withDescription(description)); + return new PinotProjectJoinTransposeRule((Config) MV_AGGREGATE_CONFIG.withDescription(description), + new PinotProjectJoinTransposeRule((Config) Config.DEFAULT.withDescription(description), null)); } - private PinotProjectJoinTransposeRule(Config config) { + // The rule as it was before the MV aggregate. It runs while the MV aggregate is off. + @Nullable + private final PinotProjectJoinTransposeRule _plainRule; + + private PinotProjectJoinTransposeRule(Config config, @Nullable PinotProjectJoinTransposeRule plainRule) { super(config); + _plainRule = plainRule; } @Override @@ -55,6 +65,15 @@ public class PinotProjectJoinTransposeRule extends ProjectJoinTransposeRule { return !PinotHintOptions.JoinHintOptions.useLookupJoinStrategy(join); } + @Override + public void onMatch(RelOptRuleCall call) { + if (_plainRule != null && !PinotRuleUtils.isMvAggregateEnabled(call)) { + _plainRule.onMatch(call); + } else { + super.onMatch(call); + } + } + private static boolean containsArrayToMv(RexNode expr) { if (!(expr instanceof RexCall)) { return false; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java index 6bbb888eba..cfd59c5a1f 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRuleUtils.java @@ -24,7 +24,9 @@ import java.util.ArrayList; import java.util.EnumSet; import java.util.List; import javax.annotation.Nullable; +import org.apache.calcite.plan.Context; import org.apache.calcite.plan.Contexts; +import org.apache.calcite.plan.RelOptRuleCall; import org.apache.calcite.plan.RelTraitSet; import org.apache.calcite.plan.hep.HepRelVertex; import org.apache.calcite.rel.RelNode; @@ -55,6 +57,8 @@ import org.apache.calcite.util.Util; import org.apache.pinot.calcite.rel.hint.PinotHintStrategyTable; import org.apache.pinot.calcite.sql.fun.PinotOperatorTable; import org.apache.pinot.common.function.sql.PinotSqlFunction; +import org.apache.pinot.query.QueryEnvironment; +import org.apache.pinot.query.context.PlannerContext; public class PinotRuleUtils { @@ -186,6 +190,22 @@ public class PinotRuleUtils { } } + /// Whether `GROUP BY ARRAY_TO_MV(col)` is planned as an MV aggregate in this planner run, see + /// [PlannerContext#isUseMvAggregate()]. Without a planner context, [QueryEnvironment.Config#defaultUseMvAggregate()] + /// decides. + public static boolean isMvAggregateEnabled(RelOptRuleCall call) { + Context genericContext = call.getPlanner().getContext(); + if (genericContext == null) { + return true; + } + PlannerContext plannerContext = genericContext.unwrap(PlannerContext.class); + if (plannerContext != null) { + return plannerContext.isUseMvAggregate(); + } + QueryEnvironment.Config envConfig = genericContext.unwrap(QueryEnvironment.Config.class); + return envConfig == null || envConfig.defaultUseMvAggregate(); + } + public static class WindowUtils { // Supported window functions // OTHER_FUNCTION supported are: BOOL_AND, BOOL_OR diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java index f2e44fe345..f77dbe8e60 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java @@ -146,15 +146,17 @@ public class PinotOperatorTable implements SqlOperatorTable { opBinding -> arrayToMvReturnType(opBinding.getTypeFactory(), opBinding.getOperandType(0)), OperandTypes.ARRAY, false); - /// The type of `ARRAY_TO_MV(array)`: the component type of `arrayType`, nullable when the array or its component - /// is nullable. The MV aggregate gives its MV keys the same type, so that it can replace a plain aggregate over - /// `ARRAY_TO_MV(array)` without changing the row type. + /// The type of `ARRAY_TO_MV(array)`: the component type of `arrayType`, as before the MV aggregate. The MV aggregate + /// gives its MV keys the same type, so that it can replace a plain aggregate over `ARRAY_TO_MV(array)` without + /// changing the row type. + /// + /// PROTOTYPE (auto mode): the type is the same with the MV aggregate on and off, because the fallback optimizes the + /// already converted plan again. A null MV value still goes into the null group at runtime. public static RelDataType arrayToMvReturnType(RelDataTypeFactory typeFactory, RelDataType arrayType) { RelDataType componentType = arrayType.getComponentType(); Preconditions.checkArgument(componentType != null, "ARRAY_TO_MV requires an array, got: %s", arrayType.getFullTypeString()); - return typeFactory.createTypeWithNullability(componentType, - arrayType.isNullable() || componentType.isNullable()); + return componentType; } /// This list includes the supported standard [SqlOperator]s defined in [SqlStdOperatorTable]. diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java index 053b206893..426b1492dd 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java @@ -94,6 +94,7 @@ import org.apache.pinot.query.planner.physical.v2.RelToPRelConverter; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.planner.rules.PinotRuleSet; import org.apache.pinot.query.planner.spi.Phase; +import org.apache.pinot.query.planner.validation.ArrayToMvValidationVisitor; import org.apache.pinot.query.routing.WorkerManager; import org.apache.pinot.query.type.TypeFactory; import org.apache.pinot.query.validate.BytesCastVisitor; @@ -327,8 +328,10 @@ public class QueryEnvironment { } else { queryNode = sqlNode; } - RelRoot relRoot = compileQuery(queryNode, plannerContext); - return new CompiledQuery(_envConfig.getDatabase(), sqlQuery, relRoot, plannerContext, sqlNodeAndOptions); + RelRoot unoptimizedRelRoot = toRelation(validate(queryNode, plannerContext), plannerContext); + RelRoot relRoot = unoptimizedRelRoot.withRel(optimizeAndConvert(unoptimizedRelRoot, plannerContext)); + return new CompiledQuery(_envConfig.getDatabase(), sqlQuery, relRoot, unoptimizedRelRoot, plannerContext, + sqlNodeAndOptions); } catch (QueryException e) { throw e; } catch (Throwable t) { @@ -418,9 +421,8 @@ public class QueryEnvironment { // steps // -------------------------------------------------------------------------- - private RelRoot compileQuery(SqlNode sqlNode, PlannerContext plannerContext) { - SqlNode validated = validate(sqlNode, plannerContext); - RelRoot relation = toRelation(validated, plannerContext); + /// Optimizes the converted query and, with the physical optimizer, converts it to the physical plan. + private RelNode optimizeAndConvert(RelRoot relation, PlannerContext plannerContext) { RelNode optimized = optimize(relation, plannerContext); if (plannerContext.isUsePhysicalOptimizer()) { Preconditions.checkNotNull(plannerContext.getPhysicalPlannerContext(), "Physical planner context is null"); @@ -428,7 +430,7 @@ public class QueryEnvironment { _envConfig.getTableCache()).unwrap(); PRelNodeTreeValidator.validate((PRelNode) optimized, plannerContext.getPhysicalPlannerContext()); } - return relation.withRel(optimized); + return optimized; } /// Query validation is a transformation from SqlNode to SqlNode where each node is validated. @@ -886,6 +888,16 @@ public class QueryEnvironment { return CommonConstants.Broker.DEFAULT_USE_MV_AGGREGATE; } + /// PROTOTYPE: the `auto` mode of the MV aggregate. When the query does not set `useMvAggregate` and + /// [#defaultUseMvAggregate()] is false, the planner first plans the query without the MV aggregate. Only when that + /// plan fails the `ARRAY_TO_MV` validation does it optimize the query again with the MV aggregate. So a query that + /// plans without the MV aggregate keeps the plan of a broker that predates it, and servers that predate it can run + /// it. + @Value.Default + default boolean defaultMvAggregateFallback() { + return true; + } + /// Whether to use lite mode by default. /// /// This is treated as the default value for the broker and it is expected to be obtained from a Pinot @@ -1007,16 +1019,20 @@ public class QueryEnvironment { public class CompiledQuery implements Closeable { private final String _database; private final String _textQuery; - private final RelRoot _relRoot; + // The optimized plan. The MV aggregate fallback replaces it, see #replanWithMvAggregate + private RelRoot _relRoot; + // The plan before optimization, which the MV aggregate fallback optimizes again + private final RelRoot _unoptimizedRelRoot; private final PlannerContext _plannerContext; private final SqlNodeAndOptions _sqlNodeAndOptions; private final Set<String> _tableNames; - private CompiledQuery(String database, String textQuery, RelRoot relRoot, PlannerContext plannerContext, - SqlNodeAndOptions sqlNodeAndOptions) { + private CompiledQuery(String database, String textQuery, RelRoot relRoot, RelRoot unoptimizedRelRoot, + PlannerContext plannerContext, SqlNodeAndOptions sqlNodeAndOptions) { _database = database; _textQuery = textQuery; _relRoot = relRoot; + _unoptimizedRelRoot = unoptimizedRelRoot; _plannerContext = plannerContext; _sqlNodeAndOptions = sqlNodeAndOptions; // Important & tricky: RelToPlanNodeConverter uses thread local. Therefore we need to get the table names here @@ -1036,6 +1052,34 @@ public class QueryEnvironment { return _plannerContext; } + /// Builds the dispatchable plan. In the auto mode of the MV aggregate, a plan that fails the `ARRAY_TO_MV` + /// validation is optimized again with the MV aggregate (see [#replanWithMvAggregate]). + private DispatchableSubPlan toDispatchableSubPlanWithMvAggregateFallback() { + try { + return toDispatchableSubPlan(_relRoot, _plannerContext); + } catch (QueryException e) { + replanWithMvAggregate(e); + return toDispatchableSubPlan(_relRoot, _plannerContext); + } + } + + /// The auto mode of the MV aggregate ([QueryEnvironment.Config#defaultMvAggregateFallback()]): when the plan + /// without the MV aggregate fails the `ARRAY_TO_MV` validation, optimizes the unoptimized plan again with the MV + /// aggregate on and replaces [#_relRoot]. Rethrows `e` when the fallback does not apply. + /// + /// Only a query that fails without the MV aggregate changes. A query that passes the validation keeps the plan of + /// a broker that predates the MV aggregate, so servers that predate it can run it. + private void replanWithMvAggregate(QueryException e) { + if (!(e instanceof ArrayToMvValidationVisitor.ArrayToMvValidationException) + || !_envConfig.defaultMvAggregateFallback() || _plannerContext.isUseMvAggregate() + || _plannerContext.isUsePhysicalOptimizer() + || _sqlNodeAndOptions.getOptions().containsKey(QueryOptionKey.USE_MV_AGGREGATE)) { + throw e; + } + _plannerContext.enableMvAggregateFallback(); + _relRoot = _unoptimizedRelRoot.withRel(optimizeAndConvert(_unoptimizedRelRoot, _plannerContext)); + } + /// Explain the query plan. /// The original query must be an EXPLAIN query and way it will be explained depends on the options of the EXPLAIN /// query and the [QueryEnvironment.Config] used to create the [QueryEnvironment] that compiled this query. @@ -1050,7 +1094,7 @@ public class QueryEnvironment { SqlExplainFormat format = _plannerContext.getSqlExplainFormat(); if (explain instanceof SqlPhysicalExplain) { // get the physical plan for query. - DispatchableSubPlan dispatchableSubPlan = toDispatchableSubPlan(_relRoot, _plannerContext); + DispatchableSubPlan dispatchableSubPlan = toDispatchableSubPlanWithMvAggregateFallback(); return getQueryPlannerResult(_plannerContext, dispatchableSubPlan, PhysicalExplainPlanVisitor.explain(dispatchableSubPlan), dispatchableSubPlan.getTableNames()); } else { @@ -1066,7 +1110,7 @@ public class QueryEnvironment { // Build the dispatchable subplan even though only the logical plan is rendered: some planning errors are // raised while converting the plan or assigning workers (e.g. a `tableOptions` partition hint that // disagrees with the table's actual partitioning), and EXPLAIN must fail wherever execution would. - DispatchableSubPlan dispatchableSubPlan = toDispatchableSubPlan(_relRoot, _plannerContext); + DispatchableSubPlan dispatchableSubPlan = toDispatchableSubPlanWithMvAggregateFallback(); return getQueryPlannerResult(_plannerContext, dispatchableSubPlan, PlannerUtils.explainPlan(_relRoot.rel, format, level), dispatchableSubPlan.getTableNames()); } else { @@ -1077,8 +1121,14 @@ public class QueryEnvironment { TransformationTracker.ByIdentity.Builder<PlanNode, RelNode> nodeTracker = new TransformationTracker.ByIdentity.Builder<>(); // Transform RelNodes into DispatchableSubPlan - DispatchableSubPlan dispatchableSubPlan = - toDispatchableSubPlan(_relRoot, _plannerContext, nodeTracker); + DispatchableSubPlan dispatchableSubPlan; + try { + dispatchableSubPlan = toDispatchableSubPlan(_relRoot, _plannerContext, nodeTracker); + } catch (QueryException e) { + replanWithMvAggregate(e); + nodeTracker = new TransformationTracker.ByIdentity.Builder<>(); + dispatchableSubPlan = toDispatchableSubPlan(_relRoot, _plannerContext, nodeTracker); + } AskingServerStageExplainer serversExplainer = new AskingServerStageExplainer( onServerExplainer, explainPlanVerbose, RelBuilder.create(_config)); @@ -1103,7 +1153,7 @@ public class QueryEnvironment { // TODO: current code only assume one SubPlan per query, but we should support multiple SubPlans per query. // Each SubPlan should be able to run independently from Broker then set the results into the dependent // SubPlan for further processing. - DispatchableSubPlan dispatchableSubPlan = toDispatchableSubPlan(_relRoot, _plannerContext); + DispatchableSubPlan dispatchableSubPlan = toDispatchableSubPlanWithMvAggregateFallback(); return getQueryPlannerResult(_plannerContext, dispatchableSubPlan, null, dispatchableSubPlan.getTableNames()); } catch (QueryException e) { throw e; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/context/PlannerContext.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/context/PlannerContext.java index 33f95acc16..183e23ed9c 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/context/PlannerContext.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/context/PlannerContext.java @@ -25,6 +25,7 @@ import java.util.Map; import javax.annotation.Nullable; import org.apache.calcite.plan.Context; import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.RelTraitDef; import org.apache.calcite.plan.hep.HepProgram; import org.apache.calcite.plan.hep.HepProgramBuilder; import org.apache.calcite.prepare.PlannerImpl; @@ -54,8 +55,15 @@ public class PlannerContext implements AutoCloseable, Context { private final SqlValidator _validator; - private final RelOptPlanner _relOptPlanner; - private final LogicalPlanner _relTraitPlanner; + // Kept to create new planners in #enableMvAggregateFallback() + private final HepProgram _optProgram; + private final HepProgram _traitProgram; + @Nullable + private final List<RelTraitDef> _optTraitDefs; + private RelOptPlanner _relOptPlanner; + private LogicalPlanner _relTraitPlanner; + // Set when the query is optimized again with the MV aggregate, see #enableMvAggregateFallback() + private boolean _mvAggregateFallback; private final Map<String, String> _options; private final QueryEnvironment.Config _envConfig; @@ -72,7 +80,10 @@ public class PlannerContext implements AutoCloseable, Context { _validator = new Validator(config.getOperatorTable(), catalogReader, typeFactory); _options = options; _envConfig = envConfig; - _relOptPlanner = new LogicalPlanner(optProgram, this, config.getTraitDefs()); + _optProgram = optProgram; + _traitProgram = traitProgram; + _optTraitDefs = config.getTraitDefs(); + _relOptPlanner = new LogicalPlanner(optProgram, this, _optTraitDefs); _relTraitPlanner = new LogicalPlanner(traitProgram, this, List.of(RelDistributionTraitDef.INSTANCE)); _plannerOutput = new HashMap<>(); @@ -96,6 +107,9 @@ public class PlannerContext implements AutoCloseable, Context { _options = options; _envConfig = envConfig; HepProgram emptyProgram = new HepProgramBuilder().build(); + _optProgram = emptyProgram; + _traitProgram = emptyProgram; + _optTraitDefs = null; _relOptPlanner = new LogicalPlanner(emptyProgram, this); _relTraitPlanner = new LogicalPlanner(emptyProgram, this); _plannerOutput = new HashMap<>(); @@ -176,4 +190,20 @@ public class PlannerContext implements AutoCloseable, Context { public boolean isUsePhysicalOptimizer() { return _physicalPlannerContext != null; } + + /// Whether the planner plans `GROUP BY ARRAY_TO_MV(expr)` as an MV aggregate: the query option `useMvAggregate`, + /// else [QueryEnvironment.Config#defaultUseMvAggregate()], unless the MV aggregate fallback turned it on (see + /// [#enableMvAggregateFallback()]). + public boolean isUseMvAggregate() { + return _mvAggregateFallback || QueryOptionsUtils.isUseMvAggregate(_options, _envConfig.defaultUseMvAggregate()); + } + + /// Turns the MV aggregate on and replaces the planners with new ones, so that the query can be optimized again from + /// its unoptimized plan. A new planner starts without the vertices of the plan that the old one optimized. + public void enableMvAggregateFallback() { + _mvAggregateFallback = true; + _relOptPlanner = _optTraitDefs != null ? new LogicalPlanner(_optProgram, this, _optTraitDefs) + : new LogicalPlanner(_optProgram, this); + _relTraitPlanner = new LogicalPlanner(_traitProgram, this, List.of(RelDistributionTraitDef.INSTANCE)); + } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/ArrayToMvValidationVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/ArrayToMvValidationVisitor.java index b4ab63a56e..247621b168 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/ArrayToMvValidationVisitor.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/ArrayToMvValidationVisitor.java @@ -78,7 +78,8 @@ public class ArrayToMvValidationVisitor implements PlanNodeVisitor<Void, Boolean /// @param scansTables whether the fragment scans a table (a leaf fragment) public static void validate(PlanNode fragmentRoot, boolean scansTables) { // In a leaf fragment with a lookup join, the nodes from the root down to the join run above the leaf boundary. - boolean isIntermediateStage = !scansTables || containsLookupJoin(fragmentRoot); + // PROTOTYPE: the same holds for an UNNEST, whose input is the leaf boundary. + boolean isIntermediateStage = !scansTables || containsLookupJoin(fragmentRoot) || containsUnnest(fragmentRoot); fragmentRoot.visit(INSTANCE, isIntermediateStage); } @@ -98,6 +99,36 @@ public class ArrayToMvValidationVisitor implements PlanNodeVisitor<Void, Boolean return false; } + private static boolean containsUnnest(PlanNode node) { + if (node instanceof UnnestNode) { + return true; + } + if (node instanceof MailboxReceiveNode) { + return false; + } + for (PlanNode input : node.getInputs()) { + if (containsUnnest(input)) { + return true; + } + } + return false; + } + + private static boolean scansTable(PlanNode node) { + if (node instanceof TableScanNode) { + return true; + } + if (node instanceof MailboxReceiveNode) { + return false; + } + for (PlanNode input : node.getInputs()) { + if (scansTable(input)) { + return true; + } + } + return false; + } + /// The tail of every error message of this visitor. public static final String SUPPORTED_USES = "ARRAY_TO_MV is supported only: (1) as a complete GROUP BY key, " + "GROUP BY ARRAY_TO_MV(expr), where expr is a multi-value column or an array expression; (2) in WHERE " @@ -253,7 +284,8 @@ public class ArrayToMvValidationVisitor implements PlanNodeVisitor<Void, Boolean @Override public Void visitUnnest(UnnestNode node, Boolean isIntermediateStage) { - node.getInputs().forEach(e -> e.visit(this, isIntermediateStage)); + // PROTOTYPE: in a leaf fragment, the input of the UNNEST is the leaf boundary: it runs in the single-stage engine. + node.getInputs().forEach(e -> e.visit(this, isIntermediateStage && !scansTable(e))); return null; } @@ -275,7 +307,15 @@ public class ArrayToMvValidationVisitor implements PlanNodeVisitor<Void, Boolean } private static QueryException error(String context) { - return new QueryException(QueryErrorCode.QUERY_PLANNING, context + ". " + SUPPORTED_USES); + return new ArrayToMvValidationException(context + ". " + SUPPORTED_USES); + } + + /// The [QueryErrorCode#QUERY_PLANNING] error of a plan that uses `ARRAY_TO_MV` where the plan cannot run it. The + /// auto mode of the MV aggregate plans such a query again with the MV aggregate. + public static class ArrayToMvValidationException extends QueryException { + ArrayToMvValidationException(String message) { + super(QueryErrorCode.QUERY_PLANNING, message); + } } private boolean containsArrayToMv(RexExpression expression) { ``` </details> -- 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]
