gortiz opened a new pull request, #20061: URL: https://github.com/apache/pinot/pull/20061
## Summary **What.** The multi-stage engine (MSE) can now run `GROUP BY ARRAY_TO_MV(col)` in every stage of a query. For example: - directly on a table scan; - above a join; - in a subquery; - with group trimming. Before this change, the query worked only when the whole aggregate ran directly on a table scan. **Why.** In the single-stage engine (SSE), `GROUP BY mvCol` groups each row by every element of the multi-value (MV) column. The MSE has no equivalent above the leaf stage. `ARRAY_TO_MV` is only a type trick that the leaf stage removes. In an intermediate stage, the planner rejects it. **How.** This PR adds a new plan node and runtime operator, the *MV aggregate*: - `PinotAggregateExchangeNodeInsertRule` finds the group keys that are exactly `ARRAY_TO_MV(expr)`. It rewrites each key to the array `expr` and splits the aggregate into a LEAF `PinotLogicalMvAggregate`, a hash exchange, and a plain FINAL aggregate. - In the leaf stage, `ServerPlanRequestVisitor` sends the MV aggregate to the SSE as `GROUP BY expr`. The SSE then groups by each element, as it does today. - In an intermediate stage, the new `MvAggregateOperator` expands the MV keys of each row. It calls `AggregationFunction.aggregateGroupByMV`, so it never creates one row for each element. - Above a lookup join, the node runs as `MvAggregateOperator` inside the leaf stage, above the leaf boundary. A plain aggregate above a lookup join runs the same way. - `ArrayToMvValidationVisitor` rejects every other use of `ARRAY_TO_MV` with one error message. The broker uses the instance capabilities of [apache/pinot#20058](https://github.com/apache/pinot/pull/20058) to know which servers can run the new plan node. This PR includes the commit of that PR as its first commit. Review that commit in [apache/pinot#20058](https://github.com/apache/pinot/pull/20058). The broker config `pinot.broker.multistage.use.mv.aggregate` has three values: `true`, `false` and `auto`. The default is `auto`. In `auto` mode, a rolling upgrade breaks no query that works today, and no manual step is necessary. ## Background ### SSE semantics The MV aggregate uses the SSE rules for `GROUP BY` on an MV column: - A row goes into the group of each of its elements. - If an element occurs two times in one row, the row counts two times in that group. - An empty array goes into no group. A stored MV value is never empty, but a function such as `filterMv` can return an empty array. - When null handling is on, a null MV value goes into one null group. - In an intermediate stage, a null MV value from an outer join goes into the null group, also when null handling is off. `MvKeyExpander` maps a null value to one null element. - Two or more MV keys give, for each row, the cartesian product of their elements. `CROSS JOIN UNNEST(col)` gives the same groups, but it drops a null MV value. The order in which the operator creates the combinations is not a contract. Without `ORDER BY`, no rule tells which groups stay when the query reaches the groups limit. ### Why the exchange must carry scalars The MV aggregate always runs as LEAF + exchange + FINAL, and the LEAF side expands the MV keys before the exchange. The exchange must partition by element: with whole arrays, rows that share an element but have different arrays go to different workers, and the FINAL side returns the same group more than once. Also, the MSE has no array equality: `Key` compares arrays by reference, and the exchange hashes `toString()`. For the same reason, a plain `GROUP BY arrayCol` stays rejected. That syntax is reserved for SQL array semantics (one group for each whole array). ### Why a new node `MvAggregateNode` does not extend `AggregateNode`. A `PlanNodeVisitor` that handles only `visitAggregate` must fail at compile time. It must not treat the MV aggregate as a plain aggregate. The proto gets a new `MvAggregateNode` message (field 19 of `PlanNode.node`). The runtime gets the new operator type `MV_AGGREGATE` (id 17). ## Plan shapes `SELECT ARRAY_TO_MV(mcol), COUNT(*) FROM t GROUP BY ARRAY_TO_MV(mcol)` Before (works only on a scan): ``` Aggregate FINAL Exchange hash[0] Aggregate LEAF Project ARRAY_TO_MV(mcol) TableScan t ``` After: ``` Aggregate FINAL <- plain aggregate over scalar keys Exchange hash[0] <- scalar elements, never arrays MvAggregate LEAF <- mvKeys=[0], runs inside the SSE leaf query Project mcol <- ARRAY_TO_MV removed, the key is the array TableScan t ``` Above a join, the same shape runs in an intermediate stage. There, `MvAggregateOperator` runs the LEAF side: ``` Aggregate FINAL Exchange hash[0] MvAggregate LEAF <- MvAggregateOperator Project t.mcol Join t.k = d.k Exchange hash[0] TableScan t Exchange hash[0] TableScan d ``` Notes on the rule: - The hints `is_skip_leaf_stage_group_by` and `is_partitioned_by_group_by_keys` are ignored for an MV aggregate. A skip-leaf plan hashes whole arrays. A partitioned plan assumes a partitioning by element that no partition function gives. - The hint `is_leaf_return_final_result` is honored. Group trimming applies. - If an MV key column is also an aggregate argument, as in `SUM(ARRAY_TO_MV(c)) ... GROUP BY ARRAY_TO_MV(c)`, the rule keeps the plain aggregate. The same occurs when the MV key column is in the `FILTER (WHERE ...)` clause of an aggregate call. - An aggregate with `GROUPING SETS`, `ROLLUP` or `CUBE` is never an MV aggregate. It keeps the plan that it has before this PR. - `PinotProjectJoinTransposeRule` no longer pushes an expression with `ARRAY_TO_MV` below a join. The group expression stays in the `Project` under the aggregate. - The physical optimizer (`usePhysicalOptimizer=true`) does not use the MV aggregate. It keeps the old behavior. ## Supported and unsupported uses `ARRAY_TO_MV` is supported in two places only: 1. As a complete GROUP BY key: `GROUP BY ARRAY_TO_MV(expr)`, where `expr` is an MV column or an array expression such as `VALUE_IN(col, 'a', 'b')`. 2. In a WHERE predicate that runs in the leaf stage, directly on a table scan. You can repeat the group key in the SELECT list and in ORDER BY. There, it refers to the group key. These uses are not supported: | Use | Leaf stage | Intermediate stage | |---|---|---| | Nested in a group key, for example `LOWER(ARRAY_TO_MV(c))` | The SSE decides. `CAST` applies to each element. Most other functions fail when the query runs. | Error | | In the SELECT list, not as a group key | Error | Error | | In an aggregate argument, for example `SUM(ARRAY_TO_MV(c))` | Works, as today | Error | | Plain `GROUP BY mvCol` | Error (new) | Error | | In a JOIN condition | Error | Error | | In a filter above a join or an aggregate, in a window function | n/a | Error | | With `GROUPING SETS`, `ROLLUP` or `CUBE` | Works, as today | Error | | With `WITHIN GROUP` | Error | Error | In a leaf fragment with a lookup join, the join and the nodes above it count as intermediate stage. They run as multi-stage operators. Only the right input of the join counts as leaf stage. Most of these queries fail at planning time with a `QUERY_PLANNING` error. The message ends with `ArrayToMvValidationVisitor.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 predicates that run in the leaf stage (directly on a table scan). It is not supported nested inside another expression (for example CASE WHEN ARRAY_TO_MV(col) IS NOT NULL ...), in the SELECT list, in aggregate arguments, in filters above a join or an aggregate, in JOIN conditions, or in window functions. ``` There are two exceptions: - `WITHIN GROUP` fails at planning time with a `QUERY_PLANNING` error, but only when the rule plans an MV aggregate. The message is `GROUP BY ARRAY_TO_MV() is not supported with WITHIN GROUP ordered aggregates in the multi-stage query engine`. - A nested group key directly on a table scan fails in the SSE when the query runs, with the message of the SSE. ## Gate and rolling upgrade ### The problem An old server does not know the new proto field. It fails with `Unsupported PlanNode type`. Pinot upgrades brokers before servers. So, for a time, a new broker sends queries to old servers. A leaf-only `GROUP BY ARRAY_TO_MV(x)` works today on old servers. If the new broker plans it as an MV aggregate, the query fails until the servers are upgraded. ### Instance capabilities [apache/pinot#20058](https://github.com/apache/pinot/pull/20058) adds the generic mechanism: each instance writes a map of capabilities with integer revisions to its instance config, and `InstanceCapabilityTracker` gives the minimum revision over all the instances of one type. That PR also explains the rollback guard and the ZooKeeper cost. This PR adds: - the first capability, `MSE_MV_AGGREGATE` (server, revision 1); - the broker wiring. In `auto` mode, `BaseBrokerStarter` creates a `MultiClusterInstanceCapabilityTracker` for the servers of the local cluster. `MultiClusterHelixBrokerStarter` adds the servers of each remote cluster. A remote cluster that does not connect counts as revision 0 until the broker restarts. **ZooKeeper cost.** The servers add no write: the capabilities are part of the instance config update that each server already does at startup, in its own znode. Each broker in `auto` mode adds one instance config listener for each cluster. When an instance is added or removed, the broker reads the config of each server, about 2 ZooKeeper operations each. The broker routing manager already reads all instance configs on every change, so the order of magnitude does not change. A possible follow-up removes this cost: the routing manager can feed the tracker from the records that it already reads. ### The gate | Setting | Name | Values | Default | |---|---|---|---| | Broker config | `pinot.broker.multistage.use.mv.aggregate` | `true`, `false`, `auto` | `auto` | | Query option | `useMvAggregate` | `true`, `false`, `auto` | the broker config | - `true`: the rule plans the MV aggregate everywhere. - `false`: the rule does not fire. `GROUP BY ARRAY_TO_MV(col)` plans as before this PR. - `auto`: the rule plans the MV aggregate everywhere, with one exception. If the aggregate runs whole in the leaf stage, the rule keeps the plain plan of today until every server advertises `MSE_MV_AGGREGATE` with revision 1 or higher. The query option `true` or `false` overrides the broker config for one query. The value `auto` uses the broker config. Any other value fails the query with a `QUERY_VALIDATION` error. A broker config value that is not valid stops the broker at startup. An aggregate runs whole in the leaf stage when its input is only `Project` and `Filter` nodes over a table scan. A semi-join that becomes a dynamic broadcast also stays in the leaf stage. `PinotJoinToDynamicBroadcastRule` runs after the aggregate rule, so the aggregate rule predicts its result: - Both rules use `PinotJoinToDynamicBroadcastRule.hasConvertibleShape` for the conditions on the join itself. - The aggregate rule then follows the left input of the join to the table scan, as the broadcast rule does. Nested semi-joins are converted from the bottom up, so this walk also accepts inner joins that the broadcast rule converts later. - If the broker config disables dynamic filtering for semi-joins, no join counts as leaf stage. With `auto`, a rolling upgrade needs no manual step: - Before all servers are upgraded, a leaf-only query keeps the plan of today, which old servers run. - After all servers are upgraded, every query uses the MV aggregate. - Future upgrades do not go back to the old plan, because the upgraded servers already advertise the capability. - A query that runs in an intermediate stage always uses the MV aggregate. Such a query fails before this PR on every version, so nothing regresses. **Stage stats.** Old servers do not know the operator type id 17. If an MV aggregate runs on a new server and the next stage runs on an old server, the query can lose its stage stats. The query results stay correct. ### Changes that do not depend on the gate - `PinotProjectJoinTransposeRule` keeps an expression with `ARRAY_TO_MV` above a join. Before, the rule pushed it into a leaf-stage `Project`, where the validator rejected it. These queries failed before and still fail with the gate off. - `ArrayToMvValidationVisitor` has new error texts. Each one ends with `SUPPORTED_USES`. - `ArrayToMvValidationVisitor` treats a lookup join and the nodes above it as intermediate stage. Before, it treated them as leaf stage. These nodes run as multi-stage operators, which have no implementation of `ARRAY_TO_MV`. Such queries now fail at planning time, not when they run. - `ArrayToMvValidationVisitor` rejects a plain `GROUP BY mvCol` in the leaf stage. Before this PR, only one plan skipped the check: the hint `is_partitioned_by_group_by_keys='true'`. That hint gives one DIRECT aggregate in a leaf fragment, with no FINAL aggregate. In every other plan, the existing check rejected the array key in the FINAL aggregate. - `ARRAY_TO_MV(col)` is nullable when the array or its elements are nullable. Before, it was never nullable. As a result, with column-based null handling, `WHERE ARRAY_TO_MV(col) IS NOT NULL` now removes the rows with a null MV value, as the SSE does for `WHERE col IS NOT NULL`. Before, the planner replaced the predicate with `TRUE`. ## Testing Planner (`pinot-query-planner`): - `MvAggregatePlans.json` in `ResourceBasedQueryPlansTest` has plans for: - an MV key on a scan, and two MV keys with a single-value key; - an MV key above a join, above a lookup join, and in a subquery; - group trimming, `DISTINCT`, and the three hints; - an array-function key and a nullable MV column; - grouping sets on a scan and above a join; - the gate off and the aggregate-argument fallback; - the error cases. - `MvAggregateGateTest`: the broker default, the query option in both directions, and the `auto` mode with and without the server capability. It covers the leaf-only cases (scan, filter, dynamic-broadcast semi-joins, hints) and the other cases (hash join, lookup join, sorted subquery, hash semi-join, `CROSS JOIN UNNEST`). - `ValidationErrorPlan.json`: updated error messages, the gate off, and `usePhysicalOptimizer=true`. - `ArrayToMvGroupKeysTest`, `ArrayToMvValidationVisitorTest`, `ArrayToMvFunctionNameTest`, `PinotLogicalMvAggregateTest`, `MvAggregateNodeTest`. - `PlanNodeSerDeTest`, `PlanNodeToRelConverterTest`, and `NodeEquivalenceFieldCoverageTest` for the new node. Runtime (`pinot-query-runtime`): - `MvAggregateOperatorTest` covers: - the SSE rules above: one group for each element, repeated elements, empty arrays, the null group, and the cartesian product; - FILTER clauses, group trimming, and the groups limit; - one MV key of each stored type, `BYTES` keys, and list and primitive-array values; - a row with more key combinations than an `int` can hold. Most cases run with row-heap and with serialized input blocks. - `MvKeyExpanderTest`, `OperatorTypeRegistryTest`, `MultiStageQueryStatsTest`. - `ServerPlanRequestVisitorMvAggregateTest`: the leaf query and the leaf boundary for a scan, an array-function key, a lookup join, and the group-trim and FILTER options. It also runs the leaf query on a real segment and compares the stored types with the node schema. - `MvAggregate.json` in `ResourceBasedQueriesTest`: query results for the SSE rules on a scan and above a join, two MV keys, MV and single-value keys, subqueries, `DISTINCT`, `HAVING`, `ORDER BY ... LIMIT`, FILTER clauses, lookup joins, and the `CROSS JOIN UNNEST` equivalent. A second table with column-based null handling covers the null group and `IS NOT NULL`. Broker (`pinot-broker`): - `MultiClusterServerCapabilityTest`: the broker tracks the servers of the local and remote clusters, and a remote cluster that does not connect. - `MultiStageBrokerRequestHandlerTest`: the broker config values and the `auto` mode with a tracker. Integration (`pinot-integration-tests`): - `MultiStageEngineIntegrationTest`: - runs the existing `ARRAY_TO_MV` group-by tests with the gate on and off; - compares a leaf-stage `GROUP BY ARRAY_TO_MV(col)` with the SSE `GROUP BY col`; - compares the MV aggregate above a join with the `CROSS JOIN UNNEST` plan; - covers hash and lookup joins to a dimension table, two MV keys, empty arrays, the skip-leaf hint, nested keys on a scan, and the unsupported uses; - makes sure that the servers advertise `MSE_MV_AGGREGATE`; - removes the capability from one server and makes sure that the plan goes back to the plain aggregate. Docs: a follow-up PR in pinot-contrib/pinot-docs. -- 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]
