From b9456d5f1a7db76dd182ccc4de1974fa245e2ef6 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 29 Jun 2026 11:27:35 -0700 Subject: [PATCH 1/2] [multistage] Support GROUP BY GROUPING SETS / ROLLUP / CUBE in the multi-stage query engine Adds native execution of GROUP BY GROUPING SETS (...) / ROLLUP(...) / CUBE(...) and the GROUPING() / GROUPING_ID() functions to the multi-stage query engine (MSE), in both the default planner and the v2 physical planner (usePhysicalOptimizer). A grouping-set query is represented as the union of all grouping columns plus, per set, a synthetic $groupingId bitmask (bit i set iff union column i is rolled up). The bitmask is carried as an internal key column that keeps the sets distinct and powers GROUPING() / GROUPING_ID(); the shared conventions live in GroupingSets (pinot-common). Execution paths (identical results on both): - Aggregate directly over a table scan: the whole aggregate is pushed to the single-stage (leaf) engine, which expands each row across the sets and appends $groupingId. - Aggregate over any other input (e.g. after a JOIN): a new RepeatOperator in the MSE runtime performs the same per-set row expansion, then an ordinary GROUP BY over [union keys..., $groupingId] runs -- no grouping-set-specific aggregation logic. Planner split (both planners): LEAF -> EXCHANGE(hash union + $groupingId) -> FINAL(group by union + $groupingId) -> PROJECT. The PROJECT computes GROUPING() / GROUPING_ID() from $groupingId (shared GroupingSetsRexUtils) and drops $groupingId. Guards (explicit failure, never wrong results): at most 31 distinct grouping columns (the $groupingId bitmask is a 32-bit int); WITHIN GROUP ordered aggregates under grouping sets are rejected in the v2 planner (run on single-stage). Tests: GroupingSetsQueriesTest exercises grouping-set aggregation, GROUPING() / GROUPING_ID() (SELECT and HAVING), filtered aggregation, composite/nested sets, explicit GROUPING SETS syntax, grouping sets over a JOIN, and the v2 physical planner -- against both engines via data providers. PlanNodeSerDeTest and RepeatOperatorTest cover the wire format and the runtime row expansion. Rolling upgrade: the feature ships within a single release, so brokers and servers are upgraded together; AggregateNode.groupingSets is additive (proto field 8). --- .../common/request/context/GroupingSets.java | 28 ++ pinot-common/src/main/proto/plan.proto | 6 + .../tests/custom/GroupingSetsQueriesTest.java | 336 +++++++++++++----- .../rel/logical/PinotLogicalAggregate.java | 18 + .../rel/rules/GroupingSetsRexUtils.java | 84 +++++ .../PinotAggregateExchangeNodeInsertRule.java | 157 +++++++- .../calcite/sql/fun/PinotOperatorTable.java | 4 + .../logical/RelToPlanNodeConverter.java | 29 +- .../physical/v2/PRelToPlanNodeConverter.java | 6 +- .../physical/v2/nodes/PhysicalAggregate.java | 19 + .../v2/opt/rules/AggregatePushdownRule.java | 161 ++++++++- .../query/planner/plannode/AggregateNode.java | 32 +- .../planner/serde/PlanNodeDeserializer.java | 3 +- .../planner/serde/PlanNodeSerializer.java | 1 + .../RowExpressionValidationVisitor.java | 7 +- .../pinot/query/QueryCompilationTest.java | 17 + .../planner/serde/PlanNodeSerDeTest.java | 15 + .../runtime/operator/AggregateOperator.java | 33 +- .../runtime/operator/MultiStageOperator.java | 10 +- .../runtime/operator/RepeatOperator.java | 163 +++++++++ .../plan/server/ServerPlanRequestVisitor.java | 7 + .../runtime/operator/RepeatOperatorTest.java | 120 +++++++ 22 files changed, 1139 insertions(+), 117 deletions(-) create mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsRexUtils.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java create mode 100644 pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/RepeatOperatorTest.java diff --git a/pinot-common/src/main/java/org/apache/pinot/common/request/context/GroupingSets.java b/pinot-common/src/main/java/org/apache/pinot/common/request/context/GroupingSets.java index ab071b1ecc6b..04452b371f8f 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/request/context/GroupingSets.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/request/context/GroupingSets.java @@ -18,6 +18,12 @@ */ package org.apache.pinot.common.request.context; +import java.util.List; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.sql.type.SqlTypeName; + /// Engine-agnostic conventions for GROUP BY GROUPING SETS / ROLLUP / CUBE, shared so both query engines /// agree on the discriminator column and the GROUPING()/GROUPING_ID() semantics. /// @@ -36,6 +42,28 @@ private GroupingSets() { /// result (never projected unless referenced via GROUPING()/GROUPING_ID()). public static final String GROUPING_ID_COLUMN = "$groupingId"; + /// The grouping-id discriminator is a 32-bit INT bitmask over the union group-by columns, so a grouping-set query + /// may reference at most 31 distinct grouping columns (bit 31 is the sign bit). + public static final int MAX_GROUPING_SET_COLUMNS = 31; + + /// Returns the LEAF row type for a grouping-set aggregate: the synthetic {@link #GROUPING_ID_COLUMN} INT column + /// inserted right after the {@code groupCount} union group-by columns, i.e. {@code [group keys..., $groupingId, + /// aggregates...]}. Both query engines share this single layout so the multi-stage final stage can group on + /// {@code $groupingId}; keeping it here prevents the two planner {@code deriveRowType()} overrides from drifting. + public static RelDataType appendGroupingIdColumn(RelDataTypeFactory typeFactory, RelDataType rowType, + int groupCount) { + RelDataTypeFactory.Builder builder = typeFactory.builder(); + List fields = rowType.getFieldList(); + for (int i = 0; i < groupCount; i++) { + builder.add(fields.get(i)); + } + builder.add(GROUPING_ID_COLUMN, typeFactory.createSqlType(SqlTypeName.INTEGER)); + for (int i = groupCount; i < fields.size(); i++) { + builder.add(fields.get(i)); + } + return builder.build(); + } + /// Computes the {@code GROUPING(args...)} / {@code GROUPING_ID(args...)} value from a row's grouping-id /// bitmask. For each argument column (identified by its index in the union of grouping columns) the /// corresponding bit is read from {@code groupingId} (1 = rolled up), and the bits are packed with the diff --git a/pinot-common/src/main/proto/plan.proto b/pinot-common/src/main/proto/plan.proto index 8c2917d09686..a912fe989a31 100644 --- a/pinot-common/src/main/proto/plan.proto +++ b/pinot-common/src/main/proto/plan.proto @@ -73,6 +73,12 @@ message AggregateNode { bool leafReturnFinalResult = 5; repeated Collation collations = 6; int32 limit = 7; + // GROUP BY GROUPING SETS / ROLLUP / CUBE: one membership bitmask per grouping set, over groupKeys. Empty for a + // plain GROUP BY. See AggregateNode#getGroupingSets. + // Rolling upgrade: like every other multi-stage plan field, this is generated by the broker and executed by the + // servers, which must be on a version that understands it. A pre-grouping-sets server silently ignores this field + // and runs a plain GROUP BY (dropping the subtotal/grand-total rows), so upgrade servers before brokers. + repeated int32 groupingSets = 8; } message FilterNode { diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GroupingSetsQueriesTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GroupingSetsQueriesTest.java index dd774545eaed..86f34f579d58 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GroupingSetsQueriesTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GroupingSetsQueriesTest.java @@ -21,6 +21,7 @@ import com.fasterxml.jackson.databind.JsonNode; import java.io.File; import java.io.IOException; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -38,7 +39,7 @@ /// End-to-end coverage for GROUP BY GROUPING SETS / ROLLUP / CUBE and the GROUPING() / GROUPING_ID() helper -/// functions in the single-stage (v1) query engine. +/// functions, run against both the single-stage and multi-stage query engines. /// /// The dataset has two dimensions {@code d1} (values a/b, never null) and {@code d2} (values x or genuine /// NULL), with a metric {@code met = 1} per row. The genuine NULLs in {@code d2} are the crux of the @@ -47,7 +48,14 @@ /// GROUPING(d2)=1). These must remain distinct rows with independent counts; without the synthetic /// $groupingId discriminator they would incorrectly merge. /// -/// This feature is single-stage only, so every test pins {@code setUseMultiStageQueryEngine(false)}. +/// Every test selects its engine(s) through a data provider. Most run on BOTH via {@code useBothQueryEngines}: the +/// multi-stage engine returns identical results by pushing the per-set expansion down to the single-stage (leaf) +/// engine. This includes null-handling-disabled queries (a rolled-up key still comes back NULL regardless of the +/// null-handling option), the ORDER BY / HAVING aggregation and GROUPING() cases, and the compile-time rejection +/// cases (rejected on both engines). The only test pinned to {@code useV1QueryEngine} is the multi-value grouping +/// column, which the multi-stage engine rejects as an intermediate-stage group key. Tests inherently multi-stage +/// specific use {@code useV2QueryEngine}: grouping sets over a JOIN, and the v2 physical planner +/// ({@code usePhysicalOptimizer}) cases. @Test(suiteName = "CustomClusterIntegrationTest") public class GroupingSetsQueriesTest extends CustomDataQueryClusterIntegrationTest { private static final String DEFAULT_TABLE_NAME = "GroupingSetsQueriesTest"; @@ -138,24 +146,136 @@ private static String cell(JsonNode row, int index) { return value.isNull() ? "NULL" : value.asText(); } - /// Runs the query (single-stage) and returns the result rows, surfacing any broker/server exception. - private JsonNode runRows(String query, boolean nullHandling) + /// --- Grouping sets over a JOIN (multi-stage only: the single-stage engine has no joins). The grouping-set + /// aggregate runs in the multi-stage runtime over the join output, where a RepeatOperator expands each row across the + /// sets. The self-join on d1 multiplies each d1 group 4x4=16, so d1=a -> 16, d1=b -> 16, grand total -> 32. + + @Test(dataProvider = "useBothQueryEngines") + public void testV2PhysicalPlannerRollup(boolean useMultiStageQueryEngine) + throws Exception { + /// With usePhysicalOptimizer the multi-stage v2 physical planner splits the grouping-set aggregate itself; verify + /// it returns the same rows, including the genuine (a, NULL) detail vs the rolled-up (a, NULL) subtotal kept + /// distinct via $groupingId. The single-stage engine ignores usePhysicalOptimizer and runs the same ROLLUP, so + /// this also cross-checks both engines agree. (GROUPING() / GROUPING_ID() are not yet supported on the v2 planner.) + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + String query = "SET usePhysicalOptimizer=true; SET enableNullHandling=true; SELECT " + D1 + ", " + D2 + + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ", " + D2 + ")"; + JsonNode response = postQuery(query); + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions == null || exceptions.isEmpty(), "v2 query failed: " + response.toPrettyString()); + List actual = new ArrayList<>(); + for (JsonNode row : response.get("resultTable").get("rows")) { + actual.add(cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(2).asLong()); + } + actual.sort(null); + List expected = new ArrayList<>(List.of( + "NULL|NULL|8", "a|NULL|2", "a|NULL|4", "a|x|2", "b|NULL|2", "b|NULL|4", "b|x|2")); + expected.sort(null); + assertEquals(actual, expected, "v2 grouping-set result mismatch: " + response.toPrettyString()); + } + + @Test(dataProvider = "useV2QueryEngine") + public void testV2PhysicalPlannerGrouping(boolean useMultiStageQueryEngine) + throws Exception { + /// The v2 physical planner (usePhysicalOptimizer) computes GROUPING() / GROUPING_ID() from the synthetic + /// $groupingId in its final projection, just like the default planner. CUBE(d1, d2) gives every combination of + /// GROUPING(d1)/GROUPING(d2), and GROUPING_ID(d1, d2) packs them as (GROUPING(d1) << 1) | GROUPING(d2). + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + String query = "SET usePhysicalOptimizer=true; SET enableNullHandling=true; SELECT " + D1 + ", " + D2 + + ", GROUPING(" + D1 + "), GROUPING(" + D2 + "), GROUPING_ID(" + D1 + ", " + D2 + "), COUNT(*) FROM " + + getTableName() + " GROUP BY CUBE(" + D1 + ", " + D2 + ")"; + JsonNode response = postQuery(query); + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions == null || exceptions.isEmpty(), "v2 GROUPING query failed: " + response.toPrettyString()); + /// Key: d1|d2|grouping(d1)|grouping(d2)|grouping_id(d1,d2) -> count + Map actual = new HashMap<>(); + for (JsonNode row : response.get("resultTable").get("rows")) { + String key = cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(2).asInt() + "|" + row.get(3).asInt() + "|" + + row.get(4).asInt(); + actual.put(key, row.get(5).asLong()); + } + Map expected = new HashMap<>(); + /// {d1, d2} detail sets (GROUPING=0,0 -> id 0): includes the genuine (a, NULL) / (b, NULL) groups. + expected.put("a|x|0|0|0", 2L); + expected.put("a|NULL|0|0|0", 2L); + expected.put("b|x|0|0|0", 2L); + expected.put("b|NULL|0|0|0", 2L); + /// {d1} subtotals (d2 rolled up -> GROUPING=0,1 -> id 1): the rolled-up (a, NULL) / (b, NULL) stay distinct. + expected.put("a|NULL|0|1|1", 4L); + expected.put("b|NULL|0|1|1", 4L); + /// {d2} subtotals (d1 rolled up -> GROUPING=1,0 -> id 2). + expected.put("NULL|x|1|0|2", 4L); + expected.put("NULL|NULL|1|0|2", 4L); + /// {} grand total (both rolled up -> GROUPING=1,1 -> id 3). + expected.put("NULL|NULL|1|1|3", 8L); + assertEquals(actual, expected, "v2 GROUPING result mismatch: " + response.toPrettyString()); + } + + @Test(dataProvider = "useV2QueryEngine") + public void testV2PhysicalPlannerRejectsWithinGroupRollup(boolean useMultiStageQueryEngine) + throws Exception { + /// A WITHIN GROUP ordered aggregate (LISTAGG ... WITHIN GROUP (ORDER BY ...)) under ROLLUP cannot be leaf/final + /// split without losing the ORDER BY. The v2 physical planner rejects it explicitly (so it can run on the default + /// planner / single-stage) rather than silently producing a wrong list. This also proves the v2 path is engaged: + /// the default planner accepts the query, so a silent fallback would not reject it. + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + String query = "SET usePhysicalOptimizer=true; SELECT " + D1 + ", LISTAGG(" + D2 + ", '|') WITHIN GROUP (ORDER BY " + + D2 + ") FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ")"; + JsonNode response = postQuery(query); + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions != null && !exceptions.isEmpty(), + "expected WITHIN GROUP under ROLLUP to be rejected on the v2 planner: " + response.toPrettyString()); + } + + @Test(dataProvider = "useV2QueryEngine") + public void testMultiStageRollupOverJoin(boolean useMultiStageQueryEngine) throws Exception { - setUseMultiStageQueryEngine(false); - String prefixed = nullHandling ? "SET enableNullHandling=true; " + query : query; - JsonNode response = postQuery(prefixed); + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + String query = "SELECT t1." + D1 + ", COUNT(*) FROM " + getTableName() + " t1 JOIN " + getTableName() + " t2 ON t1." + + D1 + " = t2." + D1 + " GROUP BY ROLLUP(t1." + D1 + ")"; + JsonNode response = postQuery("SET enableNullHandling=true; " + query); JsonNode exceptions = response.get("exceptions"); assertTrue(exceptions == null || exceptions.isEmpty(), "query failed: " + response.toPrettyString()); - return response.get("resultTable").get("rows"); + Map actual = new HashMap<>(); + for (JsonNode row : response.get("resultTable").get("rows")) { + actual.put(cell(row, 0), row.get(1).asLong()); + } + Map expected = new HashMap<>(); + expected.put("a", 16L); + expected.put("b", 16L); + expected.put("NULL", 32L); + assertEquals(actual, expected); } - @Test - public void testRollupWithGenuineAndRolledUpNulls() + @Test(dataProvider = "useV2QueryEngine") + public void testMultiStageGroupingOverJoin(boolean useMultiStageQueryEngine) throws Exception { + /// GROUPING() must also work over a join: 0 for the per-d1 subtotals, 1 for the rolled-up grand total. + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + String query = "SELECT t1." + D1 + ", GROUPING(t1." + D1 + "), COUNT(*) FROM " + getTableName() + " t1 JOIN " + + getTableName() + " t2 ON t1." + D1 + " = t2." + D1 + " GROUP BY ROLLUP(t1." + D1 + ")"; + JsonNode response = postQuery("SET enableNullHandling=true; " + query); + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions == null || exceptions.isEmpty(), "query failed: " + response.toPrettyString()); + Map actual = new HashMap<>(); + for (JsonNode row : response.get("resultTable").get("rows")) { + actual.put(cell(row, 0) + "|" + row.get(1).asInt(), row.get(2).asLong()); + } + Map expected = new HashMap<>(); + expected.put("a|0", 16L); + expected.put("b|0", 16L); + expected.put("NULL|1", 32L); + assertEquals(actual, expected); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testRollupWithGenuineAndRolledUpNulls(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// SELECT layout: d1, d2, COUNT(*), GROUPING(d1), GROUPING(d2) String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*), GROUPING(" + D1 + "), GROUPING(" + D2 + ") FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ", " + D2 + ")"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); /// Key: d1|d2|grouping(d1)|grouping(d2) -> count Map actual = new HashMap<>(); @@ -184,13 +304,14 @@ public void testRollupWithGenuineAndRolledUpNulls() "genuine-NULL detail and rolled-up-NULL subtotal must both be present"); } - @Test - public void testRolledUpNullWithoutNullHandling() + @Test(dataProvider = "useBothQueryEngines") + public void testRolledUpNullWithoutNullHandling(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// ROLLUP(d1) over a column with no genuine nulls. Even with null handling DISABLED, the rolled-up /// grand-total row must come back with d1 = NULL (grouping sets force null-aware key round-trip). String query = "SELECT " + D1 + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ")"; - JsonNode rows = runRows(query, false); + JsonNode rows = postQuery(query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0), row.get(1).asLong()); @@ -203,12 +324,13 @@ public void testRolledUpNullWithoutNullHandling() assertTrue(actual.containsKey("NULL"), "rolled-up grand-total row must have NULL key without null handling"); } - @Test - public void testCube() + @Test(dataProvider = "useBothQueryEngines") + public void testCube(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*), GROUPING(" + D1 + "), GROUPING(" + D2 + ") FROM " + getTableName() + " GROUP BY CUBE(" + D1 + ", " + D2 + ")"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { String key = cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt() + "|" + row.get(4).asInt(); @@ -232,12 +354,13 @@ public void testCube() assertEquals(rows.size(), expected.size()); } - @Test - public void testGroupingSets() + @Test(dataProvider = "useBothQueryEngines") + public void testGroupingSets(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*) FROM " + getTableName() + " GROUP BY GROUPING SETS ((" + D1 + "), (" + D2 + "))"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0) + "|" + cell(row, 1), row.get(2).asLong()); @@ -253,13 +376,14 @@ public void testGroupingSets() assertEquals(rows.size(), 4); } - @Test - public void testGroupingIdMultiArg() + @Test(dataProvider = "useBothQueryEngines") + public void testGroupingIdMultiArg(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// GROUPING_ID(d1, d2) returns a 2-bit mask: bit for d1 is MSB, bit for d2 is LSB. String query = "SELECT " + D1 + ", " + D2 + ", GROUPING_ID(" + D1 + ", " + D2 + "), COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ", " + D2 + ")"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); /// GROUPING_ID values: {d1,d2}->0, {d1}->1, {}->3. Exact shape over the 8-doc dataset: 4 detail groups, /// 2 per-d1 subtotals and 1 grand total (7 rows), each grouping set covering all 8 docs. Map rowsPerGroupingId = new HashMap<>(); @@ -274,9 +398,10 @@ public void testGroupingIdMultiArg() assertEquals(rows.size(), 7); } - @Test - public void testPlainGroupByRegression() + @Test(dataProvider = "useBothQueryEngines") + public void testPlainGroupByRegression(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// A plain GROUP BY (no grouping sets) must be unaffected: 4 detail groups, no synthetic column. String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*) FROM " + getTableName() + " GROUP BY " + D1 + ", " + D2; @@ -297,14 +422,15 @@ public void testPlainGroupByRegression() assertEquals(actual, expected); } - @Test - public void testHavingOnGrouping() + @Test(dataProvider = "useBothQueryEngines") + public void testHavingOnGrouping(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// HAVING GROUPING(d2) = 1 keeps only rows where d2 is rolled up (the {d1} subtotals and the {} grand total /// from a ROLLUP(d1, d2)). String query = "SELECT " + D1 + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ", " + D2 + ") HAVING GROUPING(" + D2 + ") = 1"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0), row.get(1).asLong()); @@ -316,12 +442,13 @@ public void testHavingOnGrouping() assertEquals(actual, expected); } - @Test - public void testLongGroupingColumn() + @Test(dataProvider = "useBothQueryEngines") + public void testLongGroupingColumn(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// Exercises the LONG key-resolution branch of the generator (and rolled-up NULL without null handling). String query = "SELECT " + LNG + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + LNG + ")"; - JsonNode rows = runRows(query, false); + JsonNode rows = postQuery(query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0), row.get(1).asLong()); @@ -333,12 +460,13 @@ public void testLongGroupingColumn() assertEquals(actual, expected); } - @Test - public void testDoubleGroupingColumn() + @Test(dataProvider = "useBothQueryEngines") + public void testDoubleGroupingColumn(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// Exercises the DOUBLE key-resolution branch of the generator. String query = "SELECT " + DBL + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + DBL + ")"; - JsonNode rows = runRows(query, false); + JsonNode rows = postQuery(query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0), row.get(1).asLong()); @@ -350,15 +478,15 @@ public void testDoubleGroupingColumn() assertEquals(actual, expected); } - @Test - public void testOrderByGroupingColumnWithoutNullHandling() + @Test(dataProvider = "useBothQueryEngines") + public void testOrderByGroupingColumnWithoutNullHandling(boolean useMultiStageQueryEngine) throws Exception { /// Regression: ORDER BY a grouping column whose rolled-up grand-total value is NULL must not NPE even /// when null handling is disabled (grouping sets produce NULL keys regardless of the null-handling /// option, so the order-by comparator must be null-safe). String query = "SELECT " + D1 + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ") ORDER BY " + D1; - setUseMultiStageQueryEngine(false); + setUseMultiStageQueryEngine(useMultiStageQueryEngine); JsonNode response = postQuery(query); assertEquals(response.get("exceptions").size(), 0, response.toPrettyString()); JsonNode rows = response.get("resultTable").get("rows"); @@ -373,33 +501,34 @@ public void testOrderByGroupingColumnWithoutNullHandling() assertEquals(actual, expected); } - @Test - public void testGroupingOnNonGroupingColumnRejected() + @Test(dataProvider = "useBothQueryEngines") + public void testGroupingOnNonGroupingColumnRejected(boolean useMultiStageQueryEngine) throws Exception { /// GROUPING(arg) requires arg to be a grouping column; met is a metric, so the query must fail clearly /// rather than produce a wrong value. - setUseMultiStageQueryEngine(false); + setUseMultiStageQueryEngine(useMultiStageQueryEngine); String query = "SELECT " + D1 + ", GROUPING(" + MET + ") FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ")"; JsonNode response = postQuery(query); assertTrue(response.get("exceptions").size() > 0, "GROUPING() over a non-grouping column must be rejected: " + response.toPrettyString()); } - @Test - public void testAggregationOnlyInHavingOrOrderBy() + @Test(dataProvider = "useBothQueryEngines") + public void testAggregationOnlyInHavingOrOrderBy(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// The only aggregation lives in HAVING / ORDER-BY (none in SELECT): the query must execute as an /// aggregation group-by — not be rewritten to DISTINCT nor rejected. ROLLUP(d1) groups over the 8-doc /// dataset: (a) = 4 docs, (b) = 4 docs, grand total = 8 docs. String havingQuery = "SELECT " + D1 + " FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ") HAVING COUNT(*) > 4"; - JsonNode rows = runRows(havingQuery, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + havingQuery).get("resultTable").get("rows"); assertEquals(rows.size(), 1, "only the grand-total group has more than 4 docs"); assertEquals(cell(rows.get(0), 0), "NULL"); String orderByQuery = "SELECT " + D1 + " FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ") ORDER BY COUNT(*) DESC, " + D1; - rows = runRows(orderByQuery, true); + rows = postQuery("SET enableNullHandling=true; " + orderByQuery).get("resultTable").get("rows"); assertEquals(rows.size(), 3); /// Grand total (8 docs) first, then (a) and (b) (4 docs each) tie-broken by d1. assertEquals(cell(rows.get(0), 0), "NULL"); @@ -407,13 +536,13 @@ public void testAggregationOnlyInHavingOrOrderBy() assertEquals(cell(rows.get(2), 0), "b"); } - @Test - public void testAggregationFreeGroupingSetsRejected() + @Test(dataProvider = "useBothQueryEngines") + public void testAggregationFreeGroupingSetsRejected(boolean useMultiStageQueryEngine) throws Exception { /// A grouping-set query without any aggregation must fail with a clear compilation error. It must NOT be /// rewritten to SELECT DISTINCT (which would drop the rolled-up subtotal rows) nor silently executed as a /// selection query (which would ignore the GROUP BY entirely). - setUseMultiStageQueryEngine(false); + setUseMultiStageQueryEngine(useMultiStageQueryEngine); for (String query : new String[]{ "SELECT " + D1 + " FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ")", "SELECT " + D1 + ", " + D2 + " FROM " + getTableName() + " GROUP BY GROUPING SETS ((" + D1 + "), (" + D2 @@ -429,14 +558,17 @@ public void testAggregationFreeGroupingSetsRejected() } } - @Test - public void testMultiValueColumnInGroupingSet() + @Test(dataProvider = "useV1QueryEngine") + public void testMultiValueColumnInGroupingSet(boolean useMultiStageQueryEngine) throws Exception { - /// Phase 2: a multi-value column may participate in a grouping set. Every row has mv = [t1, t2]. - /// In the {mv} set the row expands over its values (contributes to both t1 and t2); in the {} set mv is - /// rolled up so the row counts once toward the grand total. + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + /// Single-stage only: the multi-stage engine rejects a multi-value column as a group key (it cannot serve as an + /// intermediate-stage hash/exchange key), so a multi-value grouping set runs on the single-stage engine. + /// A multi-value column may participate in a grouping set there: every row has mv = [t1, t2]. In the {mv} set + /// the row expands over its values (contributes to both t1 and t2); in the {} set mv is rolled up so the row + /// counts once toward the grand total. String query = "SELECT " + MV + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + MV + ")"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0), row.get(1).asLong()); @@ -449,14 +581,15 @@ public void testMultiValueColumnInGroupingSet() assertEquals(rows.size(), 3); } - @Test - public void testFilteredAggregation() + @Test(dataProvider = "useBothQueryEngines") + public void testFilteredAggregation(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// Phase 2: filtered aggregations combine with grouping sets. COUNT(*) FILTER (WHERE d2 = 'x') alongside /// an unfiltered COUNT(*), under ROLLUP(d1). Layout: d1, cntX, cntAll. String query = "SELECT " + D1 + ", COUNT(*) FILTER (WHERE " + D2 + " = 'x'), COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ")"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0), new long[]{row.get(1).asLong(), row.get(2).asLong()}); @@ -468,15 +601,16 @@ public void testFilteredAggregation() assertEquals(actual.get("NULL"), new long[]{4L, 8L}); } - @Test - public void testOrderByAggregation() + @Test(dataProvider = "useBothQueryEngines") + public void testOrderByAggregation(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// ORDER BY an aggregation in a grouping-set query exercises the order-by aggregation extractor at the /// discriminator-shifted aggregation offset (the $groupingId column sits between the group keys and the /// aggregations). Without that offset fix this throws a ClassCastException at the IndexedTable resize. String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ", " + D2 + ") ORDER BY COUNT(*) DESC, " + D1 + " LIMIT 100"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); /// ROLLUP counts: grand total = 8; the two {d1} subtotals = 4; the four detail groups = 2. assertEquals(rows.get(0).get(2).asLong(), 8L); long previousCount = Long.MAX_VALUE; @@ -487,16 +621,17 @@ public void testOrderByAggregation() } } - @Test - public void testOrderByGroupingFunction() + @Test(dataProvider = "useBothQueryEngines") + public void testOrderByGroupingFunction(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// Regression: ORDER BY GROUPING(...) / GROUPING_ID(...) must not fail when the IndexedTable builds its /// order-by extractors. GROUPING/GROUPING_ID are context-dependent (computed from the $groupingId /// discriminator), so TableResizer needs a dedicated extractor; the generic post-aggregation path cannot /// resolve them and would throw while constructing the ORDER BY comparator. String groupingQuery = "SELECT " + D1 + ", " + D2 + ", COUNT(*), GROUPING(" + D2 + ") FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ", " + D2 + ") ORDER BY GROUPING(" + D2 + "), " + D1 + " LIMIT 100"; - JsonNode rows = runRows(groupingQuery, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + groupingQuery).get("resultTable").get("rows"); assertEquals(rows.size(), 7); /// ORDER BY GROUPING(d2) ASC: every detail row (GROUPING(d2)=0) sorts before any rolled-up row (=1). int previous = Integer.MIN_VALUE; @@ -512,7 +647,7 @@ public void testOrderByGroupingFunction() String groupingIdQuery = "SELECT " + D1 + ", " + D2 + ", COUNT(*), GROUPING_ID(" + D1 + ", " + D2 + ") FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ", " + D2 + ") ORDER BY GROUPING_ID(" + D1 + ", " + D2 + ") DESC LIMIT 100"; - JsonNode idRows = runRows(groupingIdQuery, true); + JsonNode idRows = postQuery("SET enableNullHandling=true; " + groupingIdQuery).get("resultTable").get("rows"); assertEquals(idRows.size(), 7); assertEquals(idRows.get(0).get(3).asInt(), 3, "grand total (GROUPING_ID=3) sorts first under DESC"); int previousId = Integer.MAX_VALUE; @@ -523,28 +658,33 @@ public void testOrderByGroupingFunction() } } - @Test - public void testEmptyMatchRollup() + @Test(dataProvider = "useBothQueryEngines") + public void testEmptyMatchRollup(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// Regression: an empty match (filter matches no rows, so segments return empty results) for a /// grouping-set query must return 0 rows without error. The empty group-by result block must carry the /// same $groupingId column as a non-empty one; otherwise the reducer rejects the narrower schema with a - /// spurious "upgrade servers" error on a fully-upgraded cluster. runRows() asserts no broker exception. + /// spurious "upgrade servers" error on a fully-upgraded cluster, so assert there is no broker exception. String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*) FROM " + getTableName() + " WHERE " + D1 + " = 'no_such_value' GROUP BY ROLLUP(" + D1 + ", " + D2 + ")"; - JsonNode rows = runRows(query, true); + JsonNode response = postQuery("SET enableNullHandling=true; " + query); + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions == null || exceptions.isEmpty(), "query failed: " + response.toPrettyString()); + JsonNode rows = response.get("resultTable").get("rows"); assertEquals(rows.size(), 0); } - @Test - public void testMixedPlainAndRollup() + @Test(dataProvider = "useBothQueryEngines") + public void testMixedPlainAndRollup(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// GROUP BY d1, ROLLUP(d2) == GROUPING SETS ((d1, d2), (d1)): plain keys cross-multiply with grouping /// constructs. Layout: d1, d2, COUNT(*), GROUPING(d2). The (a, NULL) detail group (genuine NULL, /// GROUPING(d2)=0) must stay distinct from the (a) subtotal (GROUPING(d2)=1). String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*), GROUPING(" + D2 + ") FROM " + getTableName() + " GROUP BY " + D1 + ", ROLLUP(" + D2 + ")"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt(), row.get(2).asLong()); @@ -560,14 +700,15 @@ public void testMixedPlainAndRollup() assertEquals(rows.size(), 6); } - @Test - public void testCompositeRollupLevel() + @Test(dataProvider = "useBothQueryEngines") + public void testCompositeRollupLevel(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// ROLLUP((d1, d2)): a parenthesized level rolls up both columns together, so the expansion is /// GROUPING SETS ((d1, d2), ()) — detail rows plus the grand total, with no per-d1 subtotals. String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*), GROUPING_ID(" + D1 + ", " + D2 + ") FROM " + getTableName() + " GROUP BY ROLLUP((" + D1 + ", " + D2 + "))"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt(), row.get(2).asLong()); @@ -582,15 +723,16 @@ public void testCompositeRollupLevel() assertEquals(rows.size(), 5); } - @Test - public void testNestedRollupInsideGroupingSets() + @Test(dataProvider = "useBothQueryEngines") + public void testNestedRollupInsideGroupingSets(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// GROUPING SETS may nest other grouping constructs: GROUPING SETS ((d1), ROLLUP(d2)) expands to /// {d1}, {d2}, {}. The {d2} set includes a genuine-NULL d2 group (GROUPING_ID=2) that must stay distinct /// from the grand total (GROUPING_ID=3) although both render as (NULL, NULL). String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*), GROUPING_ID(" + D1 + ", " + D2 + ") FROM " + getTableName() + " GROUP BY GROUPING SETS ((" + D1 + "), ROLLUP(" + D2 + "))"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt(), row.get(2).asLong()); @@ -605,13 +747,14 @@ public void testNestedRollupInsideGroupingSets() assertEquals(rows.size(), 5); } - @Test - public void testExpressionGroupingColumn() + @Test(dataProvider = "useBothQueryEngines") + public void testExpressionGroupingColumn(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// Grouping columns may be transform expressions, not just identifiers: ROLLUP(UPPER(d1)). String query = "SELECT UPPER(" + D1 + "), COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(UPPER(" + D1 + "))"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0), row.get(1).asLong()); @@ -624,14 +767,15 @@ public void testExpressionGroupingColumn() assertEquals(rows.size(), 3); } - @Test - public void testWhereFilterWithCube() + @Test(dataProvider = "useBothQueryEngines") + public void testWhereFilterWithCube(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// The WHERE filter applies before grouping: with d2 = 'x' only 4 docs remain, and every CUBE(d1, d2) /// subtotal reflects the filtered counts. String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*), GROUPING_ID(" + D1 + ", " + D2 + ") FROM " + getTableName() + " WHERE " + D2 + " = 'x' GROUP BY CUBE(" + D1 + ", " + D2 + ")"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt(), row.get(2).asLong()); @@ -647,14 +791,15 @@ public void testWhereFilterWithCube() assertEquals(rows.size(), 6); } - @Test - public void testCaseWhenGroupingRelabelsSubtotals() + @Test(dataProvider = "useBothQueryEngines") + public void testCaseWhenGroupingRelabelsSubtotals(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// GROUPING() nested inside a post-aggregation transform: the canonical pattern for relabeling subtotal /// rows. CASE WHEN GROUPING(d1) = 1 THEN 'ALL' ELSE d1 END turns the rolled-up NULL into 'ALL'. String query = "SELECT CASE WHEN GROUPING(" + D1 + ") = 1 THEN 'ALL' ELSE " + D1 + " END, COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ")"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0), row.get(1).asLong()); @@ -667,14 +812,15 @@ public void testCaseWhenGroupingRelabelsSubtotals() assertEquals(rows.size(), 3); } - @Test - public void testMultipleAggregationsWithOrderByAndLimit() + @Test(dataProvider = "useBothQueryEngines") + public void testMultipleAggregationsWithOrderByAndLimit(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// Several aggregation types under one ROLLUP, with ORDER BY + LIMIT applied after the grouping-set /// expansion. Full result: (NULL, 8, 200, 8), (a, 4, 100, 4), (b, 4, 200, 4); LIMIT 2 keeps the first two. String query = "SELECT " + D1 + ", SUM(" + MET + "), MAX(" + LNG + "), COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ") ORDER BY COUNT(*) DESC, " + D1 + " LIMIT 2"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); assertEquals(rows.size(), 2); assertEquals(cell(rows.get(0), 0), "NULL"); assertEquals(rows.get(0).get(1).asDouble(), 8.0); @@ -686,14 +832,15 @@ public void testMultipleAggregationsWithOrderByAndLimit() assertEquals(rows.get(1).get(3).asLong(), 4L); } - @Test - public void testDistinctCountUnderRollup() + @Test(dataProvider = "useBothQueryEngines") + public void testDistinctCountUnderRollup(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// DISTINCTCOUNT(d1) recomputes per grouping set: 2 distinct d1 values in every d2 group and in the grand /// total. The genuine-NULL d2 group (GROUPING(d2)=0) stays distinct from the grand total (GROUPING(d2)=1). String query = "SELECT " + D2 + ", DISTINCTCOUNT(" + D1 + "), COUNT(*), GROUPING(" + D2 + ") FROM " + getTableName() + " GROUP BY ROLLUP(" + D2 + ")"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); Map actual = new HashMap<>(); for (JsonNode row : rows) { actual.put(cell(row, 0) + "|" + row.get(3).asInt(), new long[]{row.get(1).asLong(), row.get(2).asLong()}); @@ -704,14 +851,15 @@ public void testDistinctCountUnderRollup() assertEquals(actual.get("NULL|1"), new long[]{2L, 8L}); } - @Test - public void testHavingOnAggregationAndGrouping() + @Test(dataProvider = "useBothQueryEngines") + public void testHavingOnAggregationAndGrouping(boolean useMultiStageQueryEngine) throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); /// HAVING may combine a regular aggregation with GROUPING(): keep only rolled-up-d2 rows with more than /// 4 docs — the per-d1 subtotals have exactly 4, so only the grand total survives. String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ", " + D2 + ") HAVING GROUPING(" + D2 + ") = 1 AND COUNT(*) > 4"; - JsonNode rows = runRows(query, true); + JsonNode rows = postQuery("SET enableNullHandling=true; " + query).get("resultTable").get("rows"); assertEquals(rows.size(), 1); assertEquals(cell(rows.get(0), 0), "NULL"); assertEquals(cell(rows.get(0), 1), "NULL"); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalAggregate.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalAggregate.java index f9edb412c883..df3486e3fcf9 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalAggregate.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalAggregate.java @@ -28,7 +28,9 @@ import org.apache.calcite.rel.core.Aggregate; import org.apache.calcite.rel.core.AggregateCall; import org.apache.calcite.rel.hint.RelHint; +import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.util.ImmutableBitSet; +import org.apache.pinot.common.request.context.GroupingSets; import org.apache.pinot.query.planner.plannode.AggregateNode.AggType; @@ -92,6 +94,22 @@ public PinotLogicalAggregate copy(RelTraitSet traitSet, RelNode input, Immutable _leafReturnFinalResult, _collations, _limit); } + /// Whether this LEAF aggregate must synthesize the {@code $groupingId} discriminator column. Only a LEAF + /// grouping-set aggregate does: it is converted to a single-stage query that expands each row across the + /// grouping sets and appends {@code $groupingId}; the multi-stage final stage then groups on it. + public boolean emitsGroupingId() { + return _aggType == AggType.LEAF && getGroupType() != Group.SIMPLE; + } + + @Override + protected RelDataType deriveRowType() { + RelDataType rowType = super.deriveRowType(); + if (!emitsGroupingId()) { + return rowType; + } + return GroupingSets.appendGroupingIdColumn(getCluster().getTypeFactory(), rowType, getGroupCount()); + } + @Override public RelWriter explainTerms(RelWriter pw) { RelWriter relWriter = super.explainTerms(pw); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsRexUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsRexUtils.java new file mode 100644 index 000000000000..e7e7e0eebb4f --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsRexUtils.java @@ -0,0 +1,84 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.calcite.rel.rules; + +import com.google.common.base.Preconditions; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.List; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlIdentifier; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.SqlSyntax; +import org.apache.calcite.sql.parser.SqlParserPos; +import org.apache.calcite.sql.validate.SqlNameMatchers; +import org.apache.pinot.calcite.sql.fun.PinotOperatorTable; + + +/// Builds the {@code GROUPING(args...)} / {@code GROUPING_ID(args...)} value as a bit expression over the synthetic +/// {@code $groupingId} discriminator column. Shared by the default planner +/// ({@link PinotAggregateExchangeNodeInsertRule}) and the v2 physical planner ({@code AggregatePushdownRule}) so both +/// compute the value identically, mirroring {@link org.apache.pinot.common.request.context.GroupingSets#groupingValue}. +public class GroupingSetsRexUtils { + private GroupingSetsRexUtils() { + } + + /// Builds the GROUPING / GROUPING_ID value: for each argument (identified by its union-column index) extract its bit + /// from {@code groupingIdRef}, and pack the bits with the first argument as the most significant bit. Uses Pinot's + /// bit scalar functions, evaluated in the projection operator. + /// + /// @param groupingIdRef reference to the {@code $groupingId} INT column (bit i set iff union column i is rolled up) + /// @param unionIndexes the union-column index of each GROUPING argument, in argument order + public static RexNode buildGroupingValue(RexBuilder rexBuilder, RelDataType intType, RexNode groupingIdRef, + List unionIndexes) { + SqlOperator bitAnd = pinotScalarOperator("bitAnd"); + SqlOperator bitShiftRightUnsigned = pinotScalarOperator("bitShiftRightUnsigned"); + SqlOperator bitShiftLeft = pinotScalarOperator("bitShiftLeft"); + SqlOperator bitOr = pinotScalarOperator("bitOr"); + int numArgs = unionIndexes.size(); + RexNode result = null; + for (int j = 0; j < numArgs; j++) { + RexNode k = rexBuilder.makeExactLiteral(BigDecimal.valueOf(unionIndexes.get(j)), intType); + /// bit = (groupingId >>> k) & 1 + RexNode shifted = rexBuilder.makeCall(intType, bitShiftRightUnsigned, List.of(groupingIdRef, k)); + RexNode bit = rexBuilder.makeCall(intType, bitAnd, + List.of(shifted, rexBuilder.makeExactLiteral(BigDecimal.ONE, intType))); + /// Pack with the first argument as the MSB: shift left by (numArgs - 1 - j). + int shiftLeftBy = numArgs - 1 - j; + RexNode placed = shiftLeftBy == 0 ? bit : rexBuilder.makeCall(intType, bitShiftLeft, + List.of(bit, rexBuilder.makeExactLiteral(BigDecimal.valueOf(shiftLeftBy), intType))); + result = result == null ? placed : rexBuilder.makeCall(intType, bitOr, List.of(result, placed)); + } + return result; + } + + /// Looks up a Pinot scalar function as a Calcite {@link SqlOperator} by name, for building bit expressions in the + /// grouping-set final projection. + private static SqlOperator pinotScalarOperator(String name) { + List operators = new ArrayList<>(1); + PinotOperatorTable.instance(false).lookupOperatorOverloads(new SqlIdentifier(name, SqlParserPos.ZERO), + SqlFunctionCategory.USER_DEFINED_FUNCTION, SqlSyntax.FUNCTION, operators, + SqlNameMatchers.withCaseSensitive(false)); + Preconditions.checkState(!operators.isEmpty(), "Pinot scalar function not found: %s", name); + return operators.get(0); + } +} 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 fb3d0f437c48..de48c427ac6a 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 @@ -18,6 +18,7 @@ */ package org.apache.pinot.calcite.rel.rules; +import com.google.common.base.Preconditions; import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -37,9 +38,11 @@ import org.apache.calcite.rel.core.Sort; import org.apache.calcite.rel.core.Union; import org.apache.calcite.rel.logical.LogicalAggregate; +import org.apache.calcite.rel.logical.LogicalProject; import org.apache.calcite.rel.rules.AggregateExtractProjectRule; import org.apache.calcite.rel.rules.AggregateReduceFunctionsRule; import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; @@ -50,6 +53,7 @@ import org.apache.calcite.sql.type.ReturnTypes; import org.apache.calcite.sql.type.SqlOperandTypeChecker; import org.apache.calcite.sql.type.SqlReturnTypeInference; +import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.tools.RelBuilder; import org.apache.calcite.tools.RelBuilderFactory; import org.apache.calcite.util.ImmutableBitSet; @@ -64,6 +68,7 @@ import org.apache.pinot.calcite.rel.logical.PinotLogicalExchange; import org.apache.pinot.calcite.rel.logical.PinotLogicalSortExchange; import org.apache.pinot.common.function.sql.PinotSqlAggFunction; +import org.apache.pinot.common.request.context.GroupingSets; import org.apache.pinot.query.QueryEnvironment; import org.apache.pinot.query.planner.plannode.AggregateNode.AggType; import org.apache.pinot.segment.spi.AggregationFunctionType; @@ -160,7 +165,7 @@ public void onMatch(RelOptRuleCall call) { return; } - PinotLogicalAggregate newAggRel = createPlan(call, aggRel, true, hintOptions, newCollations, limit); + RelNode newAggRel = createPlan(call, aggRel, true, hintOptions, newCollations, limit); RelNode newProjectRel = projectRel.copy(projectRel.getTraitSet(), List.of(newAggRel)); call.transformTo(sortRel.copy(sortRel.getTraitSet(), List.of(newProjectRel))); } @@ -202,7 +207,7 @@ public void onMatch(RelOptRuleCall call) { return; } - PinotLogicalAggregate newAggRel = createPlan(call, aggRel, true, hintOptions, collations, limit); + RelNode newAggRel = createPlan(call, aggRel, true, hintOptions, collations, limit); call.transformTo(sortRel.copy(sortRel.getTraitSet(), List.of(newAggRel))); } } @@ -227,7 +232,7 @@ public void onMatch(RelOptRuleCall call) { } } - private static PinotLogicalAggregate createPlan(RelOptRuleCall call, Aggregate aggRel, boolean hasGroupBy, + private static RelNode createPlan(RelOptRuleCall call, Aggregate aggRel, boolean hasGroupBy, Map hintOptions, @Nullable List collations, int limit) { // WITHIN GROUP collation is not supported in leaf stage aggregation. RelCollation withinGroupCollation = extractWithinGroupCollation(aggRel); @@ -288,8 +293,11 @@ private static PinotLogicalAggregate createPlanWithExchangeDirectAggregation(Rel * Aggregate node will be split into LEAF + EXCHANGE + FINAL. * TODO: Add optional INTERMEDIATE stage to reduce hotspot. */ - private static PinotLogicalAggregate createPlanWithLeafExchangeFinalAggregate(Aggregate aggRel, - boolean leafReturnFinalResult, @Nullable List collations, int limit) { + private static RelNode createPlanWithLeafExchangeFinalAggregate(Aggregate aggRel, boolean leafReturnFinalResult, + @Nullable List collations, int limit) { + if (aggRel.getGroupType() != Aggregate.Group.SIMPLE) { + return createGroupingSetsPlanWithLeafExchangeFinalAggregate(aggRel); + } // Create a LEAF aggregate. PinotLogicalAggregate leafAggRel = new PinotLogicalAggregate(aggRel, aggRel.getInput(), buildAggCalls(aggRel, AggType.LEAF, leafReturnFinalResult), @@ -301,6 +309,101 @@ private static PinotLogicalAggregate createPlanWithLeafExchangeFinalAggregate(Ag return convertAggFromIntermediateInput(aggRel, exchange, AggType.FINAL, leafReturnFinalResult, collations, limit); } + /// Splits a GROUP BY GROUPING SETS / ROLLUP / CUBE aggregate into LEAF + EXCHANGE + FINAL, pushing the per-set + /// row expansion down to the single-stage (leaf) engine, which expands each row across the sets at the backend so + /// the rest of the plan is an ordinary aggregate. Layout: + /// - LEAF (carries the grouping sets): the single-stage engine expands each row across the sets and appends the + /// synthetic {@code $groupingId} discriminator, so the leaf output is {@code [union keys..., $groupingId, + /// leaf aggregates...]} (see {@link PinotLogicalAggregate#deriveRowType}). + /// - EXCHANGE: hash-partition by the union keys AND {@code $groupingId} so all rows of a (set, key) co-locate. + /// - FINAL: a plain (SIMPLE) aggregate grouping by {@code [union keys..., $groupingId]}, so rows from different + /// grouping sets stay distinct with no grouping-set-specific merge logic. + /// - PROJECT: drop {@code $groupingId} so the output row type matches the original aggregate. + /// TODO: group trim (collations/limit) is not pushed for grouping sets yet — the single-stage engine disables + /// combine/reduce trim for them anyway; add it when leaf per-set trim is wired through. + private static RelNode createGroupingSetsPlanWithLeafExchangeFinalAggregate(Aggregate aggRel) { + int groupCount = aggRel.getGroupCount(); + /// The single-stage leaf encodes the grouping set as a 32-bit $groupingId bitmask over the union columns. + if (groupCount > GroupingSets.MAX_GROUPING_SET_COLUMNS) { + throw new UnsupportedOperationException( + "GROUP BY GROUPING SETS / ROLLUP / CUBE supports at most " + GroupingSets.MAX_GROUPING_SET_COLUMNS + + " distinct grouping columns in the multi-stage query engine, got " + groupCount); + } + /// GROUPING() / GROUPING_ID() are not real aggregations: they are functions of which grouping set a row belongs + /// to. Like the single-stage post-aggregation handler, they are computed from $groupingId in the final projection, + /// so they are split out of the LEAF/FINAL aggregate calls here. + List orgAggCalls = aggRel.getAggCallList(); + List realAggCalls = new ArrayList<>(orgAggCalls.size()); + /// For each original aggregate call: its index among the real (non-GROUPING) aggregates, or -1 if it is a + /// GROUPING / GROUPING_ID call (computed in the projection instead). + int[] realAggIndex = new int[orgAggCalls.size()]; + for (int i = 0; i < orgAggCalls.size(); i++) { + SqlKind kind = orgAggCalls.get(i).getAggregation().getKind(); + if (kind == SqlKind.GROUPING || kind == SqlKind.GROUPING_ID) { + realAggIndex[i] = -1; + } else { + realAggIndex[i] = realAggCalls.size(); + realAggCalls.add(orgAggCalls.get(i)); + } + } + /// LEAF: carries the original groupSets; deriveRowType() appends $groupingId after the union group keys. Only the + /// real aggregations are pushed down, so the leaf output is [union keys..., $groupingId, real aggregates...]. + PinotLogicalAggregate leafAggRel = new PinotLogicalAggregate(aggRel, aggRel.getInput(), + buildAggCalls(aggRel, realAggCalls, AggType.LEAF, false), AggType.LEAF, false, null, 0); + /// The leaf output adds $groupingId after the union keys, so the final group key count is groupCount + 1. + int finalGroupCount = groupCount + 1; + /// EXCHANGE: hash by union keys + $groupingId. + PinotLogicalExchange exchange = + PinotLogicalExchange.create(leafAggRel, RelDistributions.hash(ImmutableIntList.range(0, finalGroupCount))); + /// FINAL: SIMPLE aggregate over [union keys..., $groupingId]; the real aggregate input refs are shifted by + /// finalGroupCount and typed against the exchange (intermediate) row type, which carries the extra $groupingId. + PinotLogicalAggregate finalAggRel = + convertAggFromIntermediateInput(aggRel, exchange, realAggCalls, ImmutableBitSet.range(finalGroupCount), + finalGroupCount, exchange.getRowType(), AggType.FINAL, false, null, 0); + /// PROJECT: emit the original row type — group keys, then per original aggregate call either a real aggregate + /// result reference or a GROUPING/GROUPING_ID value computed from $groupingId — dropping $groupingId itself. + return buildGroupingSetsProject(aggRel, finalAggRel, groupCount, realAggIndex); + } + + /// Builds the Project on top of the FINAL grouping-set aggregate. The FINAL output is + /// {@code [union keys..., $groupingId, real aggregate results...]}; the project restores the original aggregate row + /// type by emitting, in order: the group keys, then for each original aggregate call either the real aggregate + /// result (referenced from the FINAL) or, for GROUPING() / GROUPING_ID(), the value computed from the {@code + /// $groupingId} discriminator column (mirroring the single-stage post-aggregation handler). {@code $groupingId} + /// itself is dropped. + private static RelNode buildGroupingSetsProject(Aggregate aggRel, PinotLogicalAggregate finalAggRel, int groupCount, + int[] realAggIndex) { + RexBuilder rexBuilder = finalAggRel.getCluster().getRexBuilder(); + RelDataType intType = finalAggRel.getCluster().getTypeFactory().createSqlType(SqlTypeName.INTEGER); + int finalGroupCount = groupCount + 1; + /// $groupingId is the INT discriminator column immediately after the union group keys. + RexNode groupingIdRef = rexBuilder.makeInputRef(finalAggRel, groupCount); + List union = aggRel.getGroupSet().asList(); + List orgAggCalls = aggRel.getAggCallList(); + List projects = new ArrayList<>(groupCount + orgAggCalls.size()); + for (int i = 0; i < groupCount; i++) { + projects.add(rexBuilder.makeInputRef(finalAggRel, i)); + } + for (int i = 0; i < orgAggCalls.size(); i++) { + if (realAggIndex[i] >= 0) { + projects.add(rexBuilder.makeInputRef(finalAggRel, finalGroupCount + realAggIndex[i])); + } else { + AggregateCall groupingCall = orgAggCalls.get(i); + /// Map each GROUPING argument (an input column index) to its position in the union group key list, which is + /// the bit position in $groupingId (bit set iff the column is rolled up in the row's grouping set). + List unionIndexes = new ArrayList<>(groupingCall.getArgList().size()); + for (int arg : groupingCall.getArgList()) { + int unionIndex = union.indexOf(arg); + Preconditions.checkState(unionIndex >= 0, "GROUPING argument must be a grouping column"); + unionIndexes.add(unionIndex); + } + RexNode value = GroupingSetsRexUtils.buildGroupingValue(rexBuilder, intType, groupingIdRef, unionIndexes); + projects.add(rexBuilder.makeCast(groupingCall.getType(), value)); + } + } + return LogicalProject.create(finalAggRel, List.of(), projects, aggRel.getRowType().getFieldNames()); + } + /** * The following is copied from {@link AggregateExtractProjectRule#onMatch(RelOptRuleCall)} with modification to take * aggregate input as input. @@ -354,20 +457,42 @@ private static RelNode generateProjectUnderAggregate(RelOptRuleCall call, Aggreg private static PinotLogicalAggregate convertAggFromIntermediateInput(Aggregate aggRel, PinotLogicalExchange exchange, AggType aggType, boolean leafReturnFinalResult, @Nullable List collations, int limit) { + /// Plain GROUP BY: the FINAL groups by the union keys (range(groupCount)) and the intermediate aggregate columns + /// start at groupCount; the input ref types are taken from the original aggregate row type (same field count). + int groupCount = aggRel.getGroupCount(); + return convertAggFromIntermediateInput(aggRel, exchange, ImmutableBitSet.range(groupCount), groupCount, + aggRel.getRowType(), aggType, leafReturnFinalResult, collations, limit); + } + + /// As above, but with an explicit FINAL group set, the column offset where intermediate aggregate results begin, + /// and the row type to type the aggregate input refs against. The grouping-set split passes + /// {@code finalGroupSet = range(groupCount + 1)} and {@code aggColumnOffset = groupCount + 1} (the union keys plus + /// the synthetic {@code $groupingId}), and types the refs against the exchange (intermediate) row type. + private static PinotLogicalAggregate convertAggFromIntermediateInput(Aggregate aggRel, PinotLogicalExchange exchange, + ImmutableBitSet finalGroupSet, int aggColumnOffset, RelDataType refRowType, AggType aggType, + boolean leafReturnFinalResult, @Nullable List collations, int limit) { + return convertAggFromIntermediateInput(aggRel, exchange, aggRel.getAggCallList(), finalGroupSet, aggColumnOffset, + refRowType, aggType, leafReturnFinalResult, collations, limit); + } + + /// As above, but over an explicit list of aggregate calls (rather than {@code aggRel.getAggCallList()}). The + /// grouping-set split passes only the real (non-GROUPING) aggregates, which the leaf emits in this order starting + /// at {@code aggColumnOffset}. + private static PinotLogicalAggregate convertAggFromIntermediateInput(Aggregate aggRel, PinotLogicalExchange exchange, + List orgAggCalls, ImmutableBitSet finalGroupSet, int aggColumnOffset, RelDataType refRowType, + AggType aggType, boolean leafReturnFinalResult, @Nullable List collations, int limit) { RelNode input = aggRel.getInput(); List projects = findImmediateProjects(input); // Create new AggregateCalls from exchange input. Exchange produces results with group keys followed by intermediate // aggregate results. - int groupCount = aggRel.getGroupCount(); - List orgAggCalls = aggRel.getAggCallList(); int numAggCalls = orgAggCalls.size(); List aggCalls = new ArrayList<>(numAggCalls); for (int i = 0; i < numAggCalls; i++) { AggregateCall orgAggCall = orgAggCalls.get(i); List argList = orgAggCall.getArgList(); - int index = groupCount + i; - RexInputRef inputRef = RexInputRef.of(index, aggRel.getRowType()); + int index = aggColumnOffset + i; + RexInputRef inputRef = RexInputRef.of(index, refRowType); // Generate rexList from argList and replace literal reference with literal. Keep the first argument as is. int numArguments = argList.size(); List rexList; @@ -386,17 +511,23 @@ private static PinotLogicalAggregate convertAggFromIntermediateInput(Aggregate a } } } - aggCalls.add(buildAggCall(exchange, orgAggCall, rexList, groupCount, aggType, leafReturnFinalResult)); + aggCalls.add(buildAggCall(exchange, orgAggCall, rexList, aggColumnOffset, aggType, leafReturnFinalResult)); } - return new PinotLogicalAggregate(aggRel, exchange, ImmutableBitSet.range(groupCount), aggCalls, aggType, - leafReturnFinalResult, collations, limit); + return new PinotLogicalAggregate(aggRel, exchange, finalGroupSet, aggCalls, aggType, leafReturnFinalResult, + collations, limit); } private static List buildAggCalls(Aggregate aggRel, AggType aggType, boolean leafReturnFinalResult) { + return buildAggCalls(aggRel, aggRel.getAggCallList(), aggType, leafReturnFinalResult); + } + + /// As above, but over an explicit list of aggregate calls (rather than {@code aggRel.getAggCallList()}). The + /// grouping-set split passes only the real (non-GROUPING) aggregates to push down to the leaf. + private static List buildAggCalls(Aggregate aggRel, List orgAggCalls, AggType aggType, + boolean leafReturnFinalResult) { RelNode input = aggRel.getInput(); List projects = findImmediateProjects(input); - List orgAggCalls = aggRel.getAggCallList(); List aggCalls = new ArrayList<>(orgAggCalls.size()); for (AggregateCall orgAggCall : orgAggCalls) { // Generate rexList from argList and replace literal reference with literal. Keep the first argument as is. 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 cea0cd934220..65df604e1234 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 @@ -204,6 +204,10 @@ public static PinotOperatorTable instance(boolean nullHandlingEnabled) { SqlStdOperatorTable.VAR_POP, SqlStdOperatorTable.VAR_SAMP, SqlStdOperatorTable.SUM0, + /// GROUPING() / GROUPING_ID() for GROUP BY GROUPING SETS / ROLLUP / CUBE. These are computed from the synthetic + /// $groupingId discriminator in the grouping-set final projection (see PinotAggregateExchangeNodeInsertRule). + SqlStdOperatorTable.GROUPING, + SqlStdOperatorTable.GROUPING_ID, // WINDOW Rank Functions SqlStdOperatorTable.DENSE_RANK, diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java index 66271d5c0067..4c5ae0aee128 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java @@ -22,7 +22,9 @@ import com.google.common.collect.Sets; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Set; import javax.annotation.Nullable; import org.apache.calcite.plan.RelOptTable; @@ -30,6 +32,7 @@ import org.apache.calcite.rel.RelDistribution; import org.apache.calcite.rel.RelFieldCollation; import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.Aggregate; import org.apache.calcite.rel.core.AggregateCall; import org.apache.calcite.rel.core.Exchange; import org.apache.calcite.rel.core.JoinInfo; @@ -58,6 +61,7 @@ import org.apache.calcite.rex.RexWindowExclusion; import org.apache.calcite.sql.SqlLiteral; import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.ImmutableBitSet; import org.apache.pinot.calcite.rel.hint.PinotHintOptions; import org.apache.pinot.calcite.rel.logical.PinotLogicalAggregate; import org.apache.pinot.calcite.rel.logical.PinotLogicalEnrichedJoin; @@ -706,6 +710,29 @@ private SortNode convertLogicalSort(LogicalSort node) { convertInputs(node.getInputs()), node.getCollation().getFieldCollations(), fetch, offset); } + /// Encodes each grouping set as a membership bitmask over the union group-by columns ({@code groupSet.asList()}), + /// mirroring the single-stage engine's {@code PinotQuery.groupingSetMasks} so the per-set expansion can be pushed + /// down to the single-stage (leaf) engine. Returns an empty list for a plain GROUP BY (SIMPLE). + public static List computeGroupingSetMasks(Aggregate node) { + if (node.getGroupType() == Aggregate.Group.SIMPLE) { + return List.of(); + } + List union = node.getGroupSet().asList(); + Map unionIndex = new HashMap<>(); + for (int i = 0; i < union.size(); i++) { + unionIndex.put(union.get(i), i); + } + List masks = new ArrayList<>(node.getGroupSets().size()); + for (ImmutableBitSet set : node.getGroupSets()) { + int mask = 0; + for (int bit : set) { + mask |= 1 << unionIndex.get(bit); + } + masks.add(mask); + } + return masks; + } + private AggregateNode convertLogicalAggregate(PinotLogicalAggregate node) { List aggregateCalls = node.getAggCallList(); int numAggregates = aggregateCalls.size(); @@ -717,7 +744,7 @@ private AggregateNode convertLogicalAggregate(PinotLogicalAggregate node) { } return new AggregateNode(DEFAULT_STAGE_ID, toDataSchema(node.getRowType()), NodeHint.fromRelHints(node.getHints()), convertInputs(node.getInputs()), functionCalls, filterArgs, node.getGroupSet().asList(), node.getAggType(), - node.isLeafReturnFinalResult(), node.getCollations(), node.getLimit()); + node.isLeafReturnFinalResult(), node.getCollations(), node.getLimit(), computeGroupingSetMasks(node)); } private ProjectNode convertLogicalProject(LogicalProject node) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java index 8704529bad97..91ef973476a0 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java @@ -47,6 +47,7 @@ import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.common.utils.DatabaseUtils; import org.apache.pinot.common.utils.request.RequestUtils; +import org.apache.pinot.query.planner.logical.RelToPlanNodeConverter; import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.logical.RexExpressionUtils; import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalAggregate; @@ -199,6 +200,8 @@ public static SortNode convertSort(Sort node) { } public static AggregateNode convertAggregate(PhysicalAggregate node) { + /// GROUP BY GROUPING SETS / ROLLUP / CUBE is split by AggregatePushdownRule (LEAF carrying the grouping sets + + /// $groupingId, FINAL grouping on $groupingId); the masks below carry to the runtime RepeatOperator. List aggregateCalls = node.getAggCallList(); int numAggregates = aggregateCalls.size(); List functionCalls = new ArrayList<>(numAggregates); @@ -209,7 +212,8 @@ public static AggregateNode convertAggregate(PhysicalAggregate node) { } return new AggregateNode(DEFAULT_STAGE_ID, toDataSchema(node.getRowType()), NodeHint.fromRelHints(node.getHints()), new ArrayList<>(), functionCalls, filterArgs, node.getGroupSet().asList(), node.getAggType(), - node.isLeafReturnFinalResult(), node.getCollations(), node.getLimit()); + node.isLeafReturnFinalResult(), node.getCollations(), node.getLimit(), + RelToPlanNodeConverter.computeGroupingSetMasks(node)); } public static ProjectNode convertProject(Project node) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalAggregate.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalAggregate.java index 149795eb465d..596e7658af71 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalAggregate.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalAggregate.java @@ -28,7 +28,9 @@ import org.apache.calcite.rel.core.Aggregate; import org.apache.calcite.rel.core.AggregateCall; import org.apache.calcite.rel.hint.RelHint; +import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.util.ImmutableBitSet; +import org.apache.pinot.common.request.context.GroupingSets; import org.apache.pinot.query.planner.physical.v2.PRelNode; import org.apache.pinot.query.planner.physical.v2.PinotDataDistribution; import org.apache.pinot.query.planner.plannode.AggregateNode; @@ -72,6 +74,23 @@ public Aggregate copy(RelTraitSet traitSet, RelNode input, ImmutableBitSet group (PRelNode) input, _pinotDataDistribution, _leafStage, _aggType, _leafReturnFinalResult, _collations, _limit); } + /// Whether this LEAF aggregate must synthesize the {@code $groupingId} discriminator column — the v2 physical + /// counterpart of {@link org.apache.pinot.calcite.rel.logical.PinotLogicalAggregate#emitsGroupingId()}. A LEAF + /// grouping-set aggregate appends {@code $groupingId} after the union group-by columns so the final stage can group + /// on it (the per-set row expansion itself happens at runtime in the RepeatOperator). + public boolean emitsGroupingId() { + return _aggType == AggregateNode.AggType.LEAF && getGroupType() != Group.SIMPLE; + } + + @Override + protected RelDataType deriveRowType() { + RelDataType rowType = super.deriveRowType(); + if (!emitsGroupingId()) { + return rowType; + } + return GroupingSets.appendGroupingIdColumn(getCluster().getTypeFactory(), rowType, getGroupCount()); + } + @Override public int getNodeId() { return _nodeId; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java index fdb7e7b482c1..5aefc84a7e0b 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java @@ -22,6 +22,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.Supplier; import javax.annotation.Nullable; import org.apache.calcite.plan.RelTraitSet; @@ -35,6 +36,7 @@ import org.apache.calcite.rel.core.Project; import org.apache.calcite.rel.core.Union; import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; @@ -45,12 +47,15 @@ import org.apache.calcite.sql.type.ReturnTypes; import org.apache.calcite.sql.type.SqlOperandTypeChecker; import org.apache.calcite.sql.type.SqlReturnTypeInference; +import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.util.ImmutableBitSet; import org.apache.pinot.calcite.rel.hint.PinotHintOptions; import org.apache.pinot.calcite.rel.hint.PinotHintStrategyTable; +import org.apache.pinot.calcite.rel.rules.GroupingSetsRexUtils; import org.apache.pinot.calcite.rel.rules.PinotRuleUtils; import org.apache.pinot.calcite.rel.traits.PinotExecStrategyTrait; import org.apache.pinot.common.function.sql.PinotSqlAggFunction; +import org.apache.pinot.common.request.context.GroupingSets; import org.apache.pinot.query.context.PhysicalPlannerContext; import org.apache.pinot.query.planner.physical.v2.PRelNode; import org.apache.pinot.query.planner.physical.v2.PinotDataDistribution; @@ -58,6 +63,7 @@ import org.apache.pinot.query.planner.physical.v2.mapping.PinotDistMapping; import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalAggregate; import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalExchange; +import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalProject; import org.apache.pinot.query.planner.physical.v2.opt.PRelOptRule; import org.apache.pinot.query.planner.physical.v2.opt.PRelOptRuleCall; import org.apache.pinot.query.planner.plannode.AggregateNode.AggType; @@ -94,6 +100,34 @@ public PRelNode onMatch(PRelOptRuleCall call) { hintOptions = Map.of(); } boolean isInputExchange = call._currentNode.unwrap().getInput(0) instanceof Exchange; + if (aggRel.getGroupType() != Aggregate.Group.SIMPLE) { + /// GROUP BY GROUPING SETS / ROLLUP / CUBE: split with the synthetic $groupingId discriminator carried through the + /// exchange (the runtime RepeatOperator does the per-set row expansion). Requires an input exchange to + /// repartition by union keys + $groupingId. + /// $groupingId is a 32-bit bitmask over the union columns, so cap the distinct grouping columns (same guard as + /// the single-stage rule PinotAggregateExchangeNodeInsertRule). + if (aggRel.getGroupCount() > GroupingSets.MAX_GROUPING_SET_COLUMNS) { + throw new UnsupportedOperationException( + "GROUP BY GROUPING SETS / ROLLUP / CUBE supports at most " + GroupingSets.MAX_GROUPING_SET_COLUMNS + + " distinct grouping columns in the multi-stage query engine, got " + aggRel.getGroupCount()); + } + /// A WITHIN GROUP ordered aggregate (e.g. LISTAGG ... WITHIN GROUP) under a grouping set cannot be leaf/final + /// split without losing the ORDER BY, and the DIRECT fallback does not preserve the ordering across the per-set + /// expansion either. Reject explicitly (run on the default planner / single-stage) rather than silently + /// returning a wrong result. The default planner checks WITHIN GROUP before any split for the same reason. + if (withinGroupCollation != null) { + throw new UnsupportedOperationException( + "WITHIN GROUP ordered aggregates with GROUP BY GROUPING SETS / ROLLUP / CUBE are not supported in the " + + "multi-stage v2 physical planner. Run the query on the single-stage query engine instead."); + } + if (!isInputExchange) { + throw new UnsupportedOperationException( + "GROUP BY GROUPING SETS / ROLLUP / CUBE without a repartitioning exchange is not yet supported in the " + + "multi-stage v2 physical planner. Run the query on the single-stage query engine instead."); + } + return addPartialAggregateForGroupingSets((PhysicalAggregate) call._currentNode, hintOptions, + _context.getNodeIdGenerator()); + } if (!isInputExchange || withinGroupCollation != null || (hasGroupBy && Boolean.parseBoolean( hintOptions.get(PinotHintOptions.AggregateOptions.IS_SKIP_LEAF_STAGE_GROUP_BY)))) { return skipPartialAggregate(call._currentNode); @@ -110,6 +144,125 @@ private static PRelNode skipPartialAggregate(PRelNode aggPRelNode) { false /* leaf return final agg */, aggRel.getCollations(), aggRel.getLimit()); } + /// Splits a GROUP BY GROUPING SETS / ROLLUP / CUBE aggregate (v2 physical planner). Mirrors the SIMPLE split + /// {@link #addPartialAggregate} but threads the synthetic {@code $groupingId} discriminator column, exactly like the + /// single-stage rule PinotAggregateExchangeNodeInsertRule: + /// - LEAF carries the grouping sets and only the real aggregations; {@link PhysicalAggregate#deriveRowType} appends + /// {@code $groupingId} after the union group keys. The per-set row expansion happens at runtime in the + /// RepeatOperator. + /// - EXCHANGE repartitions by the union keys AND {@code $groupingId} (so all rows of a (set, key) co-locate). + /// - FINAL is a SIMPLE aggregate grouping by {@code [union keys..., $groupingId]}. + /// - PROJECT restores the original aggregate row type: the group keys, then for each original aggregate call either + /// the real aggregate result or, for GROUPING() / GROUPING_ID(), the value computed from {@code $groupingId}; + /// {@code $groupingId} itself is dropped. + private static PRelNode addPartialAggregateForGroupingSets(PhysicalAggregate aggPRelNode, + Map hintOptions, Supplier idGenerator) { + PhysicalAggregate o0 = aggPRelNode; + PhysicalExchange o1 = (PhysicalExchange) o0.getPRelInput(0); + PRelNode o2 = o1.getPRelInput(0); + int groupCount = o0.getGroupCount(); + /// GROUPING() / GROUPING_ID() are not real aggregations: they are functions of which grouping set a row belongs to, + /// computed from $groupingId in the projection. Split them out of the LEAF/FINAL aggregate calls. realAggIndex[i] + /// is the position of original aggregate call i among the real aggregates, or -1 if it is a GROUPING / GROUPING_ID. + List orgAggCalls = o0.getAggCallList(); + List realAggCalls = new ArrayList<>(orgAggCalls.size()); + int[] realAggIndex = new int[orgAggCalls.size()]; + for (int i = 0; i < orgAggCalls.size(); i++) { + SqlKind kind = orgAggCalls.get(i).getAggregation().getKind(); + if (kind == SqlKind.GROUPING || kind == SqlKind.GROUPING_ID) { + realAggIndex[i] = -1; + } else { + realAggIndex[i] = realAggCalls.size(); + realAggCalls.add(orgAggCalls.get(i)); + } + } + /// LEAF: carries the grouping sets and only the real aggregations; deriveRowType() appends $groupingId at position + /// groupCount, so the leaf output is [union keys..., $groupingId, real aggregates...]. + PhysicalAggregate n2 = new PhysicalAggregate(o0.getCluster(), RelTraitSet.createEmpty(), List.of() /* hints */, + o0.getGroupSet(), o0.groupSets, buildAggCalls(o0, realAggCalls, AggType.LEAF, false), idGenerator.get(), o2, + null /* data dist */, o2.isLeafStage(), AggType.LEAF, false, o0.getCollations(), o0.getLimit()); + PinotDistMapping mapFromInputToPartialAgg = DistMappingGenerator.compute(o2.unwrap(), n2, null); + PinotDataDistribution leafAggDataDistribution = + o2.getPinotDataDistributionOrThrow().apply(mapFromInputToPartialAgg); + n2 = (PhysicalAggregate) n2.with(n2.getPRelInputs(), leafAggDataDistribution); + /// EXCHANGE: hash by the union keys plus $groupingId (at position groupCount in the leaf output). + List mappedUnionKeys = mapFromInputToPartialAgg.getMappedKeys(o1.getDistributionKeys()).get(0); + List newDistKeys = new ArrayList<>(mappedUnionKeys); + newDistKeys.add(groupCount); + PhysicalExchange n1 = new PhysicalExchange(o1.getNodeId(), n2, + o1.getPinotDataDistributionOrThrow().apply(mapFromInputToPartialAgg), newDistKeys, o1.getExchangeStrategy(), + null /* collation */, PinotExecStrategyTrait.getDefaultExecStrategy(), o1.getHashFunction()); + /// FINAL: SIMPLE aggregate grouping by [union keys..., $groupingId]; the real aggregate input refs start at + /// groupCount + 1 in the exchange (intermediate) output. + int finalGroupCount = groupCount + 1; + PhysicalAggregate n0 = convertAggForGroupingSets(o0, n1, realAggCalls, finalGroupCount, idGenerator); + /// PROJECT: restore the original aggregate row type. For each original aggregate call, reference the real result or + /// compute the GROUPING() / GROUPING_ID() value from $groupingId (at position groupCount in the final output). + RexBuilder rexBuilder = o0.getCluster().getRexBuilder(); + RelDataType intType = o0.getCluster().getTypeFactory().createSqlType(SqlTypeName.INTEGER); + RexNode groupingIdRef = rexBuilder.makeInputRef(n0, groupCount); + List union = o0.getGroupSet().asList(); + List projects = new ArrayList<>(groupCount + orgAggCalls.size()); + for (int i = 0; i < groupCount; i++) { + projects.add(rexBuilder.makeInputRef(n0, i)); + } + for (int i = 0; i < orgAggCalls.size(); i++) { + if (realAggIndex[i] >= 0) { + projects.add(rexBuilder.makeInputRef(n0, finalGroupCount + realAggIndex[i])); + } else { + AggregateCall groupingCall = orgAggCalls.get(i); + /// Map each GROUPING argument (an input column index) to its bit position in $groupingId (its position in the + /// union group key list). + List unionIndexes = new ArrayList<>(groupingCall.getArgList().size()); + for (int arg : groupingCall.getArgList()) { + int unionIndex = union.indexOf(arg); + Preconditions.checkState(unionIndex >= 0, "GROUPING argument must be a grouping column"); + unionIndexes.add(unionIndex); + } + RexNode value = GroupingSetsRexUtils.buildGroupingValue(rexBuilder, intType, groupingIdRef, unionIndexes); + projects.add(rexBuilder.makeCast(groupingCall.getType(), value)); + } + } + return new PhysicalProject(o0.getCluster(), o0.getTraitSet(), List.of() /* hints */, projects, o0.getRowType(), + Set.of(), idGenerator.get(), n0, n0.getPinotDataDistributionOrThrow(), false); + } + + /// FINAL aggregate for a grouping-set split: like {@link #convertAggFromIntermediateInput} but groups by + /// {@code [union keys..., $groupingId]} ({@code finalGroupCount = groupCount + 1}), aggregates only the real + /// (non-GROUPING) calls, and reads the intermediate aggregate columns starting at {@code finalGroupCount}. + private static PhysicalAggregate convertAggForGroupingSets(PhysicalAggregate physicalAggregate, + PhysicalExchange exchange, List realAggCalls, int finalGroupCount, Supplier nodeId) { + Aggregate aggRel = (Aggregate) physicalAggregate.unwrap(); + List projects = findImmediateProjects(aggRel.getInput()); + List aggCalls = new ArrayList<>(realAggCalls.size()); + for (int i = 0; i < realAggCalls.size(); i++) { + AggregateCall orgAggCall = realAggCalls.get(i); + List argList = orgAggCall.getArgList(); + RexInputRef inputRef = RexInputRef.of(finalGroupCount + i, exchange.getRowType()); + int numArguments = argList.size(); + List rexList; + if (numArguments <= 1) { + rexList = List.of(inputRef); + } else { + rexList = new ArrayList<>(numArguments); + rexList.add(inputRef); + for (int j = 1; j < numArguments; j++) { + int argument = argList.get(j); + if (projects != null && projects.get(argument) instanceof RexLiteral) { + rexList.add(projects.get(argument)); + } else { + rexList.add(inputRef); + } + } + } + aggCalls.add(buildAggCall(exchange, orgAggCall, rexList, finalGroupCount, AggType.FINAL, false)); + } + ImmutableBitSet groupSet = ImmutableBitSet.range(finalGroupCount); + return new PhysicalAggregate(aggRel.getCluster(), aggRel.getTraitSet(), aggRel.getHints(), groupSet, + List.of(groupSet), aggCalls, nodeId.get(), exchange, physicalAggregate.getPinotDataDistributionOrThrow(), + false, AggType.FINAL, false, List.of(), 0); + } + private static PRelNode addPartialAggregate(PhysicalAggregate aggPRelNode, Map hintOptions, Supplier idGenerator) { // Old: Aggregate (o0) > Exchange (o1) > Input (o2) @@ -205,9 +358,15 @@ private static PhysicalAggregate convertAggFromIntermediateInput(PhysicalAggrega } public static List buildAggCalls(Aggregate aggRel, AggType aggType, boolean leafReturnFinalResult) { + return buildAggCalls(aggRel, aggRel.getAggCallList(), aggType, leafReturnFinalResult); + } + + /// As {@link #buildAggCalls(Aggregate, AggType, boolean)} but over an explicit list of aggregate calls (rather than + /// {@code aggRel.getAggCallList()}). The grouping-set split passes only the real (non-GROUPING) aggregates. + public static List buildAggCalls(Aggregate aggRel, List orgAggCalls, AggType aggType, + boolean leafReturnFinalResult) { RelNode input = aggRel.getInput(); List projects = findImmediateProjects(input); - List orgAggCalls = aggRel.getAggCallList(); List aggCalls = new ArrayList<>(orgAggCalls.size()); for (AggregateCall orgAggCall : orgAggCalls) { // Generate rexList from argList and replace literal reference with literal. Keep the first argument as is. diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/AggregateNode.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/AggregateNode.java index 5e6fda1e1b6e..e2d6fcf68c3b 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/AggregateNode.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/AggregateNode.java @@ -33,6 +33,12 @@ public class AggregateNode extends BasePlanNode { private final AggType _aggType; private final boolean _leafReturnFinalResult; + /// GROUP BY GROUPING SETS / ROLLUP / CUBE: one membership bitmask per grouping set, over _groupKeys (the union of + /// all grouping columns). Bit i is set iff _groupKeys[i] participates in (is grouped by) that set; a mask of 0 is + /// the grand-total set (). Empty for a plain GROUP BY. Mirrors the single-stage engine's PinotQuery.groupingSetMasks + /// so the per-set row expansion can be pushed down to the single-stage (leaf) engine. + private final List _groupingSets; + // The following fields are set when group trim is enabled, and are extracted from the Sort on top of this Aggregate. // The group trim behavior at leaf stage is shared with single-stage engine. private final List _collations; @@ -41,6 +47,14 @@ public class AggregateNode extends BasePlanNode { public AggregateNode(int stageId, DataSchema dataSchema, NodeHint nodeHint, List inputs, List aggCalls, List filterArgs, List groupKeys, AggType aggType, boolean leafReturnFinalResult, @Nullable List collations, int limit) { + this(stageId, dataSchema, nodeHint, inputs, aggCalls, filterArgs, groupKeys, aggType, leafReturnFinalResult, + collations, limit, List.of()); + } + + public AggregateNode(int stageId, DataSchema dataSchema, NodeHint nodeHint, List inputs, + List aggCalls, List filterArgs, List groupKeys, AggType aggType, + boolean leafReturnFinalResult, @Nullable List collations, int limit, + List groupingSets) { super(stageId, dataSchema, nodeHint, inputs); _aggCalls = aggCalls; _filterArgs = filterArgs; @@ -49,6 +63,7 @@ public AggregateNode(int stageId, DataSchema dataSchema, NodeHint nodeHint, List _leafReturnFinalResult = leafReturnFinalResult; _collations = collations != null ? collations : List.of(); _limit = limit; + _groupingSets = groupingSets; } public List getAggCalls() { @@ -71,6 +86,16 @@ public boolean isLeafReturnFinalResult() { return _leafReturnFinalResult; } + /// One membership bitmask per grouping set over {@link #getGroupKeys()}, or empty for a plain GROUP BY. + public List getGroupingSets() { + return _groupingSets; + } + + /// Whether this is a GROUP BY GROUPING SETS / ROLLUP / CUBE aggregate. + public boolean isGroupingSets() { + return !_groupingSets.isEmpty(); + } + public List getCollations() { return _collations; } @@ -92,7 +117,7 @@ public T visit(PlanNodeVisitor visitor, C context) { @Override public PlanNode withInputs(List inputs) { return new AggregateNode(_stageId, _dataSchema, _nodeHint, inputs, _aggCalls, _filterArgs, _groupKeys, _aggType, - _leafReturnFinalResult, _collations, _limit); + _leafReturnFinalResult, _collations, _limit, _groupingSets); } @Override @@ -109,13 +134,14 @@ public boolean equals(Object o) { AggregateNode that = (AggregateNode) o; return _leafReturnFinalResult == that._leafReturnFinalResult && _limit == that._limit && Objects.equals(_aggCalls, that._aggCalls) && Objects.equals(_filterArgs, that._filterArgs) && Objects.equals(_groupKeys, that._groupKeys) - && _aggType == that._aggType && Objects.equals(_collations, that._collations); + && _aggType == that._aggType && Objects.equals(_collations, that._collations) && Objects.equals(_groupingSets, + that._groupingSets); } @Override public int hashCode() { return Objects.hash(super.hashCode(), _aggCalls, _filterArgs, _groupKeys, _aggType, _leafReturnFinalResult, - _collations, _limit); + _collations, _limit, _groupingSets); } /** diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java index c91e46d8f5fb..540fd6060822 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java @@ -95,7 +95,8 @@ private static AggregateNode deserializeAggregateNode(Plan.PlanNode protoNode) { extractInputs(protoNode), convertFunctionCalls(protoAggregateNode.getAggCallsList()), protoAggregateNode.getFilterArgsList(), protoAggregateNode.getGroupKeysList(), convertAggType(protoAggregateNode.getAggType()), protoAggregateNode.getLeafReturnFinalResult(), - convertCollations(protoAggregateNode.getCollationsList()), protoAggregateNode.getLimit()); + convertCollations(protoAggregateNode.getCollationsList()), protoAggregateNode.getLimit(), + protoAggregateNode.getGroupingSetsList()); } private static FilterNode deserializeFilterNode(Plan.PlanNode protoNode) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java index f2283ed09f51..c329152fc142 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java @@ -102,6 +102,7 @@ public Void visitAggregate(AggregateNode node, Plan.PlanNode.Builder builder) { .setLeafReturnFinalResult(node.isLeafReturnFinalResult()) .addAllCollations(convertCollations(node.getCollations())) .setLimit(node.getLimit()) + .addAllGroupingSets(node.getGroupingSets()) .build(); builder.setAggregateNode(aggregateNode); return null; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/validate/RowExpressionValidationVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/validate/RowExpressionValidationVisitor.java index 6d23d86bf11c..3840988b38e2 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/validate/RowExpressionValidationVisitor.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/validate/RowExpressionValidationVisitor.java @@ -157,7 +157,12 @@ private boolean isRowExpression(SqlNode node) { private boolean isAllowedRowContext(SqlKind kind) { return kind == SqlKind.VALUES || kind == SqlKind.INSERT - || kind == SqlKind.ARRAY_VALUE_CONSTRUCTOR; + || kind == SqlKind.ARRAY_VALUE_CONSTRUCTOR + /// GROUP BY GROUPING SETS ((a, b), (a), ()) parses the parenthesized grouping sets as ROW expressions; they + /// are grouping-set tuples, not row constructors, so allow ROW operands directly under these constructs. + || kind == SqlKind.GROUPING_SETS + || kind == SqlKind.ROLLUP + || kind == SqlKind.CUBE; } private boolean hasRowOperands(SqlCall call) { diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java index 1bf31abd4314..237a9c8ff95d 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java @@ -142,6 +142,23 @@ public void testUnsignedBigintCastIsRejected() { } } + @Test + public void testGroupingSetsSupportedInMultiStage() { + /// GROUP BY GROUPING SETS / ROLLUP / CUBE plan successfully in the multi-stage engine: the per-set expansion is + /// pushed down to the single-stage leaf. They must not throw, nor silently collapse to a plain GROUP BY (which + /// would drop the subtotal/grand-total rows). GROUPING() / GROUPING_ID() and the explicit GROUPING SETS tuple + /// syntax are included. + List queries = List.of( + "SELECT col1, SUM(col3) FROM a GROUP BY ROLLUP(col1)", + "SELECT col1, col2, SUM(col3) FROM a GROUP BY CUBE(col1, col2)", + "SELECT col1, col2, SUM(col3) FROM a GROUP BY GROUPING SETS ((col1), (col2))", + "SELECT col1, col2, SUM(col3) FROM a GROUP BY GROUPING SETS ((col1, col2), (col1), ())", + "SELECT col1, GROUPING(col1), GROUPING_ID(col1, col2), SUM(col3) FROM a GROUP BY ROLLUP(col1, col2)"); + for (String sql : queries) { + assertNotNull(_queryEnvironment.planQuery(sql), "expected a multi-stage plan for: " + sql); + } + } + @Test public void testUnsignedLiteralCastIsFolded() { // Companion to testUnsignedTypeCastIsAccepted using literal (constant-foldable) casts, so the unsigned diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java index 98745869d709..a92ddd575387 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java @@ -26,6 +26,8 @@ import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; import org.apache.pinot.query.planner.physical.DispatchableSubPlan; +import org.apache.pinot.query.planner.plannode.AggregateNode; +import org.apache.pinot.query.planner.plannode.AggregateNode.AggType; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.planner.plannode.UnnestNode; import org.testng.annotations.Test; @@ -78,4 +80,17 @@ public void testLegacyUnnestNodeSerDe() { assertEquals(deserialized.isPrunedPassthrough(), false); assertEquals(deserialized.getPassthroughInputIndexes(), List.of()); } + + @Test + public void testAggregateGroupingSetsSerDe() { + /// The grouping-set masks (over the union group keys) must survive serialization to the worker. ROLLUP(g0, g1) + /// over the union {g0, g1} expands to masks {0b11, 0b01, 0b00}. + DataSchema schema = new DataSchema(new String[]{"g0", "g1", "sum"}, + new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.INT, ColumnDataType.DOUBLE}); + AggregateNode node = new AggregateNode(0, schema, PlanNode.NodeHint.EMPTY, List.of(), List.of(), List.of(), + List.of(0, 1), AggType.DIRECT, false, List.of(), 0, List.of(0b11, 0b01, 0b00)); + AggregateNode deserialized = (AggregateNode) PlanNodeDeserializer.process(PlanNodeSerializer.process(node)); + assertEquals(deserialized.getGroupingSets(), List.of(0b11, 0b01, 0b00)); + assertEquals(deserialized, node); + } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java index d47a6f60548e..1296cf71ea06 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java @@ -32,6 +32,7 @@ import org.apache.pinot.common.datatable.StatMap; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.request.context.FunctionContext; +import org.apache.pinot.common.request.context.GroupingSets; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.core.common.BlockValSet; @@ -95,7 +96,6 @@ public class AggregateOperator extends MultiStageOperator { public AggregateOperator(OpChainExecutionContext context, MultiStageOperator input, AggregateNode node) { super(context); - _input = input; _resultSchema = node.getDataSchema(); _aggFunctions = getAggFunctions(node.getAggCalls()); int numFunctions = _aggFunctions.length; @@ -111,6 +111,21 @@ public AggregateOperator(OpChainExecutionContext context, MultiStageOperator inp List groupKeys = node.getGroupKeys(); + /// GROUP BY GROUPING SETS / ROLLUP / CUBE in the multi-stage runtime: expand each input row across the grouping + /// sets (NULLing the rolled-up columns and appending the $groupingId discriminator) via a RepeatOperator, then run + /// an ordinary GROUP BY over the union columns plus $groupingId. This handles grouping sets over any input (e.g. + /// above a JOIN); when the aggregate sits directly on a table scan it is instead pushed down to the single-stage + /// leaf and this operator is not used. + List groupingSets = node.getGroupingSets(); + if (!groupingSets.isEmpty()) { + int groupingIdColumnIndex = node.getInputs().get(0).getDataSchema().size(); + input = new RepeatOperator(context, input, getGroupKeyIds(groupKeys), groupingSets, repeatResultSchema(node)); + /// The group keys are the union columns (unchanged input positions) plus $groupingId, appended by the Repeat. + groupKeys = new ArrayList<>(groupKeys); + groupKeys.add(groupingIdColumnIndex); + } + _input = input; + int groupTrimSize = Integer.MAX_VALUE; Comparator comparator = null; int limit = node.getLimit(); @@ -346,6 +361,22 @@ private int[] getGroupKeyIds(List groupKeys) { return groupKeyIds; } + /// Output schema of the {@link RepeatOperator} that feeds a grouping-set aggregate: the aggregate input schema with + /// the synthetic {@code $groupingId} INT discriminator column appended. + private static DataSchema repeatResultSchema(AggregateNode node) { + DataSchema inputSchema = node.getInputs().get(0).getDataSchema(); + int numInputColumns = inputSchema.size(); + String[] columnNames = new String[numInputColumns + 1]; + DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[numInputColumns + 1]; + for (int i = 0; i < numInputColumns; i++) { + columnNames[i] = inputSchema.getColumnName(i); + columnDataTypes[i] = inputSchema.getColumnDataType(i); + } + columnNames[numInputColumns] = GroupingSets.GROUPING_ID_COLUMN; + columnDataTypes[numInputColumns] = DataSchema.ColumnDataType.INT; + return new DataSchema(columnNames, columnDataTypes); + } + static RoaringBitmap getMatchedBitmap(MseBlock.Data block, int filterArgId) { Preconditions.checkArgument(filterArgId >= 0, "Got negative filter argument id: %s", filterArgId); RoaringBitmap matchedBitmap = new RoaringBitmap(); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java index 0a8b6bc94d61..dcde2f2dff8c 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java @@ -439,10 +439,18 @@ public void mergeInto(BrokerResponseNativeV2 response, StatMap map) { StatMap stats = (StatMap) map; response.mergeMaxRowsInOperator(stats.getLong(UnnestOperator.StatKey.EMITTED_ROWS)); } + }, + REPEAT(16, RepeatOperator.StatKey.class) { + @Override + public void mergeInto(BrokerResponseNativeV2 response, StatMap map) { + @SuppressWarnings("unchecked") + StatMap stats = (StatMap) map; + response.mergeMaxRowsInOperator(stats.getLong(RepeatOperator.StatKey.EMITTED_ROWS)); + } }; // When adding new operator types, update MAX_ID if the new ID exceeds the current max - private static final int MAX_ID = 15; + private static final int MAX_ID = 16; private static final Type[] ID_TO_TYPE = new Type[MAX_ID + 1]; static { diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java new file mode 100644 index 000000000000..460c19ede4fa --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java @@ -0,0 +1,163 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.operator; + +import java.util.ArrayList; +import java.util.List; +import org.apache.pinot.common.datatable.StatMap; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Expands each input row across the grouping sets of a GROUP BY GROUPING SETS / ROLLUP / CUBE query — the +/// multi-stage equivalent of the single-stage per-set row expansion. For every input row and every grouping set it +/// emits one output row in which the columns NOT grouped by that set (the +/// "rolled up" columns) are set to NULL, and appends the synthetic INT discriminator column +/// {@link org.apache.pinot.common.request.context.GroupingSets#GROUPING_ID_COLUMN} whose bit i is set iff union +/// group-by column i is rolled up in that set (matching the single-stage convention so GROUPING() / GROUPING_ID() +/// agree across engines). +/// +/// With the rows expanded and tagged, the downstream aggregate is an ordinary GROUP BY over the union group-by +/// columns plus {@code $groupingId} — no grouping-set-specific aggregation logic is needed. This lets grouping sets +/// run over any input (e.g. above a JOIN), not just a leaf table scan. +public class RepeatOperator extends MultiStageOperator { + private static final Logger LOGGER = LoggerFactory.getLogger(RepeatOperator.class); + private static final String EXPLAIN_NAME = "REPEAT"; + + private final MultiStageOperator _input; + private final DataSchema _resultSchema; + /// Input column index of each union group-by column, in union order; bit i of a grouping-set mask refers to + /// _unionGroupKeyIds[i]. + private final int[] _unionGroupKeyIds; + /// Per grouping set: the rolled-up bitmask over the union columns (bit i set iff column i is NOT in the set), in the + /// order the output rows are emitted. Precomputed from the participation masks. + private final int[] _groupingIds; + private final int _inputColumnCount; + private final StatMap _statMap = new StatMap<>(StatKey.class); + + /// @param unionGroupKeyIds input column index of each union group-by column (in union order) + /// @param groupingSetMasks one participation bitmask per grouping set over the union columns (bit i set iff + /// {@code unionGroupKeyIds[i]} is grouped by that set) + /// @param resultSchema the input schema with the {@code $groupingId} INT column appended + public RepeatOperator(OpChainExecutionContext context, MultiStageOperator input, int[] unionGroupKeyIds, + List groupingSetMasks, DataSchema resultSchema) { + super(context); + _input = input; + _resultSchema = resultSchema; + _unionGroupKeyIds = unionGroupKeyIds; + _inputColumnCount = resultSchema.size() - 1; + int fullMask = (1 << unionGroupKeyIds.length) - 1; + _groupingIds = new int[groupingSetMasks.size()]; + for (int s = 0; s < groupingSetMasks.size(); s++) { + /// Rolled-up bits are the complement of the participation mask over the union columns. + _groupingIds[s] = ~groupingSetMasks.get(s) & fullMask; + } + } + + @Override + public void registerExecution(long time, int numRows, long memoryUsedBytes, long gcTimeMs) { + _statMap.merge(StatKey.EXECUTION_TIME_MS, time); + _statMap.merge(StatKey.EMITTED_ROWS, numRows); + _statMap.merge(StatKey.ALLOCATED_MEMORY_BYTES, memoryUsedBytes); + _statMap.merge(StatKey.GC_TIME_MS, gcTimeMs); + } + + @Override + public Type getOperatorType() { + return Type.REPEAT; + } + + @Override + protected Logger logger() { + return LOGGER; + } + + @Override + public List getChildOperators() { + return List.of(_input); + } + + @Override + public String toExplainString() { + return EXPLAIN_NAME; + } + + @Override + protected MseBlock getNextBlock() { + MseBlock block = _input.nextBlock(); + if (block.isEos()) { + return block; + } + List inputRows = ((MseBlock.Data) block).asRowHeap().getRows(); + int groupingIdIndex = _inputColumnCount; + List expanded = new ArrayList<>(inputRows.size() * _groupingIds.length); + for (Object[] inputRow : inputRows) { + for (int groupingId : _groupingIds) { + Object[] row = new Object[_inputColumnCount + 1]; + System.arraycopy(inputRow, 0, row, 0, _inputColumnCount); + /// NULL out the union columns that are rolled up in this set (their bit is set in groupingId). + for (int i = 0; i < _unionGroupKeyIds.length; i++) { + if ((groupingId & (1 << i)) != 0) { + row[_unionGroupKeyIds[i]] = null; + } + } + row[groupingIdIndex] = groupingId; + expanded.add(row); + } + } + return new RowHeapDataBlock(expanded, _resultSchema); + } + + @Override + public StatMap copyStatMaps() { + return new StatMap<>(_statMap); + } + + public enum StatKey implements StatMap.Key { + EXECUTION_TIME_MS(StatMap.Type.LONG) { + @Override + public boolean includeDefaultInJson() { + return true; + } + }, + EMITTED_ROWS(StatMap.Type.LONG) { + @Override + public boolean includeDefaultInJson() { + return true; + } + }, + ALLOCATED_MEMORY_BYTES(StatMap.Type.LONG), + GC_TIME_MS(StatMap.Type.LONG); + + private final StatMap.Type _type; + + StatKey(StatMap.Type type) { + _type = type; + } + + @Override + public StatMap.Type getType() { + return _type; + } + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java index 29b369c3f2fd..49e2a622f112 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java @@ -78,6 +78,13 @@ public Void visitAggregate(AggregateNode node, ServerPlanRequestContext context) if (!groupByList.isEmpty()) { pinotQuery.setGroupByList(groupByList); } + /// GROUP BY GROUPING SETS / ROLLUP / CUBE: push the per-set row expansion down to the single-stage engine. The + /// masks are over groupByList (== node.getGroupKeys()), so they transfer directly. The single-stage leaf expands + /// each row across the sets and appends the synthetic $groupingId discriminator column; the multi-stage final + /// stage then groups on it (it is one of the group keys), keeping the grouping sets distinct. + if (node.isGroupingSets()) { + pinotQuery.setGroupingSetMasks(node.getGroupingSets()); + } List selectList = CalciteRexExpressionParser.convertAggregateList(groupByList, node.getAggCalls(), node.getFilterArgs(), pinotQuery.getSelectList()); for (Expression expression : selectList) { diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/RepeatOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/RepeatOperatorTest.java new file mode 100644 index 000000000000..55103bd3dbba --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/RepeatOperatorTest.java @@ -0,0 +1,120 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.operator; + +import java.util.List; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.mockito.Mock; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.when; +import static org.mockito.MockitoAnnotations.openMocks; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; + + +/// Unit coverage for [RepeatOperator]'s per-set row expansion: it must NULL the rolled-up union columns and append the +/// `$groupingId` discriminator (the rolled-up complement of each grouping set's participation mask). +public class RepeatOperatorTest { + private AutoCloseable _mocks; + @Mock + private MultiStageOperator _input; + + @BeforeMethod + public void setUp() { + _mocks = openMocks(this); + } + + @AfterMethod + public void tearDown() + throws Exception { + _mocks.close(); + } + + /// Input [d1, d2, cnt]; union group keys are d1 (idx 0) and d2 (idx 1); result appends $groupingId. + private static final DataSchema INPUT_SCHEMA = new DataSchema(new String[]{"d1", "d2", "cnt"}, + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT}); + private static final DataSchema RESULT_SCHEMA = new DataSchema(new String[]{"d1", "d2", "cnt", "$groupingId"}, + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT, ColumnDataType.INT}); + + @Test + public void shouldExpandRollupAndNullRolledUpColumns() { + /// ROLLUP(d1, d2) => sets {d1,d2}, {d1}, {} with participation masks 0b11, 0b01, 0b00. + when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(INPUT_SCHEMA, new Object[]{"a", "x", 5})); + RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0, 1}, + List.of(0b11, 0b01, 0b00), RESULT_SCHEMA); + + List rows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + assertEquals(rows.size(), 3); /// one input row x three sets + + /// {d1,d2}: nothing rolled up -> groupingId 0, both keys kept. + assertEquals(rows.get(0), new Object[]{"a", "x", 5, 0}); + /// {d1}: d2 rolled up -> groupingId 0b10 = 2, d2 NULLed. + assertEquals(rows.get(1)[0], "a"); + assertNull(rows.get(1)[1]); + assertEquals(rows.get(1)[3], 0b10); + /// {}: both rolled up -> groupingId 0b11 = 3, both keys NULLed. + assertNull(rows.get(2)[0]); + assertNull(rows.get(2)[1]); + assertEquals(rows.get(2)[3], 0b11); + /// cnt (non-key column) is never touched. + assertEquals(rows.get(2)[2], 5); + } + + @Test + public void shouldNullTheCorrectKeyForCube() { + /// CUBE(d1, d2) => the {d2}-only set rolls up d1 (groupingId 0b01), which a wrong bit order would get backwards. + when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(INPUT_SCHEMA, new Object[]{"a", "x", 5})); + RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0, 1}, + List.of(0b11, 0b01, 0b10, 0b00), RESULT_SCHEMA); + + List rows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + assertEquals(rows.size(), 4); + /// {d2}: participation 0b10 -> groupingId 0b01 = 1, d1 NULLed, d2 kept. + Object[] d2Only = rows.get(2); + assertNull(d2Only[0]); + assertEquals(d2Only[1], "x"); + assertEquals(d2Only[3], 0b01); + } + + @Test + public void shouldHandleNonContiguousMaskAndMultipleRows() { + /// Three union columns, a single explicit set {c0, c2} (participation 0b101); only c1 must be NULLed. + DataSchema inputSchema = new DataSchema(new String[]{"c0", "c1", "c2", "cnt"}, + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT}); + DataSchema resultSchema = new DataSchema(new String[]{"c0", "c1", "c2", "cnt", "$groupingId"}, + new ColumnDataType[]{ + ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT, ColumnDataType.INT + }); + when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inputSchema, + new Object[]{"a", "m", "x", 1}, new Object[]{"b", "n", "y", 2})); + RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0, 1, 2}, + List.of(0b101), resultSchema); + + List rows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + assertEquals(rows.size(), 2); /// two input rows x one set + /// groupingId = ~0b101 & 0b111 = 0b010; only c1 (idx 1) NULLed, c0/c2 kept. + assertEquals(rows.get(0), new Object[]{"a", null, "x", 1, 0b010}); + assertEquals(rows.get(1), new Object[]{"b", null, "y", 2, 0b010}); + } +} From ffaebbfe126474fd4a1c0502ba406dfeb1c67ee3 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 1 Jul 2026 16:51:17 -0700 Subject: [PATCH 2/2] [multistage] Encode grouping sets as per-set column-index lists, removing the 31-column limit Encode grouping sets the way Calcite does: each set is the list of participating grouping-column indexes (an unbounded per-set membership list) rather than a 32-bit membership bitmask, so the number of grouping columns is no longer capped at 31. A set's position in the plan's set list is its ordinal, carried as the synthetic $groupingId discriminator through the leaf/exchange/final split and the wire. Also addresses review feedback: - RepeatOperator now appends NULLable per-set group-key copies and leaves the original input columns untouched, so aggregation arguments that reference a grouping column (e.g. SUM(b) under ROLLUP(a, b)) aggregate the real values; expansion is streamed one set per block with a termination/usage check. - The per-set expansion is wired in PlanNodeToOpChain (not inside a specific AggregateOperator), so every AggregateOperatorFactory receives already-expanded input and needs no grouping-set awareness. - The GROUPING()/GROUPING_ID() split, projection, LEAF row type, and set encoding are shared between the default and v2 physical planners via GroupingSetsPlanUtils (replaces GroupingSetsRexUtils). - Both planners reject combinations that would otherwise produce broken plans or silently wrong results: WITHIN GROUP / skip-leaf / partitioned-by hints with grouping sets, aggregation-free grouping sets, GROUPING() over a plain GROUP BY, and (v2) grouping sets over an input with no repartitioning exchange. The v2 grouping-set LEAF no longer forwards group-trim collations/limit indexed against the pre-$groupingId layout. - The single-stage engine treats GROUP BY GROUPING SETS (()) (grand-total-only, no grouping columns) as a plain aggregation instead of crashing the reducer. - Wire: plan.proto grouping sets field renumbered to 8; query.thrift/plan.proto rolling-upgrade notes reworded. --- .../pinot/common/request/PinotQuery.java | 317 ++++++++++-------- .../common/request/context/GroupingSets.java | 123 ++++--- .../pinot/sql/parsers/CalciteSqlParser.java | 83 +++-- ...egationGroupByToDistinctQueryRewriter.java | 2 +- pinot-common/src/main/proto/plan.proto | 20 +- .../request/context/GroupingSetsTest.java | 66 ++++ .../sql/parsers/GroupingSetsParserTest.java | 59 +++- pinot-common/src/thrift/query.thrift | 17 +- .../pinot/core/data/table/TableResizer.java | 30 +- .../GroupingSetsGroupKeyGenerator.java | 36 +- .../query/reduce/PostAggregationHandler.java | 36 +- .../utils/QueryContextConverterUtils.java | 53 ++- ...kerRequestToQueryContextConverterTest.java | 44 ++- .../tests/custom/GroupingSetsQueriesTest.java | 110 +++++- .../rel/logical/PinotLogicalAggregate.java | 4 +- .../rel/rules/GroupingSetsPlanUtils.java | 198 +++++++++++ .../rel/rules/GroupingSetsRexUtils.java | 84 ----- .../PinotAggregateExchangeNodeInsertRule.java | 104 +++--- .../logical/RelToPlanNodeConverter.java | 31 +- .../physical/v2/PRelToPlanNodeConverter.java | 6 +- .../physical/v2/nodes/PhysicalAggregate.java | 4 +- .../v2/opt/rules/AggregatePushdownRule.java | 95 +++--- .../query/planner/plannode/AggregateNode.java | 19 +- .../planner/serde/PlanNodeDeserializer.java | 7 +- .../planner/serde/PlanNodeSerializer.java | 11 +- .../pinot/query/QueryCompilationTest.java | 42 +++ .../planner/serde/PlanNodeSerDeTest.java | 10 +- .../runtime/operator/AggregateOperator.java | 33 +- .../runtime/operator/RepeatOperator.java | 109 +++--- .../factory/AggregateOperatorFactory.java | 5 + .../query/runtime/plan/PlanNodeToOpChain.java | 61 ++++ .../plan/server/ServerPlanRequestVisitor.java | 8 +- .../runtime/operator/RepeatOperatorTest.java | 209 +++++++++--- 33 files changed, 1340 insertions(+), 696 deletions(-) create mode 100644 pinot-common/src/test/java/org/apache/pinot/common/request/context/GroupingSetsTest.java create mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsPlanUtils.java delete mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsRexUtils.java diff --git a/pinot-common/src/main/java/org/apache/pinot/common/request/PinotQuery.java b/pinot-common/src/main/java/org/apache/pinot/common/request/PinotQuery.java index 45f0de1bbf85..5dff7ee89574 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/request/PinotQuery.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/request/PinotQuery.java @@ -41,7 +41,7 @@ public class PinotQuery implements org.apache.thrift.TBase queryOptions; // optional private boolean explain; // optional private @org.apache.thrift.annotation.Nullable java.util.Map expressionOverrideHints; // optional - private @org.apache.thrift.annotation.Nullable java.util.List groupingSetMasks; // optional + private @org.apache.thrift.annotation.Nullable java.util.List> groupingSets; // optional /** The set of fields this struct contains, along with convenience methods for finding and manipulating them. */ public enum _Fields implements org.apache.thrift.TFieldIdEnum { @@ -74,7 +74,7 @@ public enum _Fields implements org.apache.thrift.TFieldIdEnum { QUERY_OPTIONS((short)11, "queryOptions"), EXPLAIN((short)12, "explain"), EXPRESSION_OVERRIDE_HINTS((short)13, "expressionOverrideHints"), - GROUPING_SET_MASKS((short)14, "groupingSetMasks"); + GROUPING_SETS((short)15, "groupingSets"); private static final java.util.Map byName = new java.util.HashMap(); @@ -114,8 +114,8 @@ public static _Fields findByThriftId(int fieldId) { return EXPLAIN; case 13: // EXPRESSION_OVERRIDE_HINTS return EXPRESSION_OVERRIDE_HINTS; - case 14: // GROUPING_SET_MASKS - return GROUPING_SET_MASKS; + case 15: // GROUPING_SETS + return GROUPING_SETS; default: return null; } @@ -164,7 +164,7 @@ public java.lang.String getFieldName() { private static final int __OFFSET_ISSET_ID = 2; private static final int __EXPLAIN_ISSET_ID = 3; private byte __isset_bitfield = 0; - private static final _Fields optionals[] = {_Fields.VERSION,_Fields.DATA_SOURCE,_Fields.SELECT_LIST,_Fields.FILTER_EXPRESSION,_Fields.GROUP_BY_LIST,_Fields.ORDER_BY_LIST,_Fields.HAVING_EXPRESSION,_Fields.LIMIT,_Fields.OFFSET,_Fields.QUERY_OPTIONS,_Fields.EXPLAIN,_Fields.EXPRESSION_OVERRIDE_HINTS,_Fields.GROUPING_SET_MASKS}; + private static final _Fields optionals[] = {_Fields.VERSION,_Fields.DATA_SOURCE,_Fields.SELECT_LIST,_Fields.FILTER_EXPRESSION,_Fields.GROUP_BY_LIST,_Fields.ORDER_BY_LIST,_Fields.HAVING_EXPRESSION,_Fields.LIMIT,_Fields.OFFSET,_Fields.QUERY_OPTIONS,_Fields.EXPLAIN,_Fields.EXPRESSION_OVERRIDE_HINTS,_Fields.GROUPING_SETS}; public static final java.util.Map<_Fields, org.apache.thrift.meta_data.FieldMetaData> metaDataMap; static { java.util.Map<_Fields, org.apache.thrift.meta_data.FieldMetaData> tmpMap = new java.util.EnumMap<_Fields, org.apache.thrift.meta_data.FieldMetaData>(_Fields.class); @@ -199,9 +199,10 @@ public java.lang.String getFieldName() { new org.apache.thrift.meta_data.MapMetaData(org.apache.thrift.protocol.TType.MAP, new org.apache.thrift.meta_data.StructMetaData(org.apache.thrift.protocol.TType.STRUCT, Expression.class), new org.apache.thrift.meta_data.StructMetaData(org.apache.thrift.protocol.TType.STRUCT, Expression.class)))); - tmpMap.put(_Fields.GROUPING_SET_MASKS, new org.apache.thrift.meta_data.FieldMetaData("groupingSetMasks", org.apache.thrift.TFieldRequirementType.OPTIONAL, + tmpMap.put(_Fields.GROUPING_SETS, new org.apache.thrift.meta_data.FieldMetaData("groupingSets", org.apache.thrift.TFieldRequirementType.OPTIONAL, new org.apache.thrift.meta_data.ListMetaData(org.apache.thrift.protocol.TType.LIST, - new org.apache.thrift.meta_data.FieldValueMetaData(org.apache.thrift.protocol.TType.I32)))); + new org.apache.thrift.meta_data.ListMetaData(org.apache.thrift.protocol.TType.LIST, + new org.apache.thrift.meta_data.FieldValueMetaData(org.apache.thrift.protocol.TType.I32))))); metaDataMap = java.util.Collections.unmodifiableMap(tmpMap); org.apache.thrift.meta_data.FieldMetaData.addStructMetaDataMap(PinotQuery.class, metaDataMap); } @@ -271,9 +272,13 @@ public PinotQuery(PinotQuery other) { } this.expressionOverrideHints = __this__expressionOverrideHints; } - if (other.isSetGroupingSetMasks()) { - java.util.List __this__groupingSetMasks = new java.util.ArrayList(other.groupingSetMasks); - this.groupingSetMasks = __this__groupingSetMasks; + if (other.isSetGroupingSets()) { + java.util.List> __this__groupingSets = new java.util.ArrayList>(other.groupingSets.size()); + for (java.util.List other_element : other.groupingSets) { + java.util.List __this__groupingSets_copy = new java.util.ArrayList(other_element); + __this__groupingSets.add(__this__groupingSets_copy); + } + this.groupingSets = __this__groupingSets; } } @@ -300,7 +305,7 @@ public void clear() { setExplainIsSet(false); this.explain = false; this.expressionOverrideHints = null; - this.groupingSetMasks = null; + this.groupingSets = null; } public int getVersion() { @@ -653,43 +658,43 @@ public void setExpressionOverrideHintsIsSet(boolean value) { } } - public int getGroupingSetMasksSize() { - return (this.groupingSetMasks == null) ? 0 : this.groupingSetMasks.size(); + public int getGroupingSetsSize() { + return (this.groupingSets == null) ? 0 : this.groupingSets.size(); } @org.apache.thrift.annotation.Nullable - public java.util.Iterator getGroupingSetMasksIterator() { - return (this.groupingSetMasks == null) ? null : this.groupingSetMasks.iterator(); + public java.util.Iterator> getGroupingSetsIterator() { + return (this.groupingSets == null) ? null : this.groupingSets.iterator(); } - public void addToGroupingSetMasks(int elem) { - if (this.groupingSetMasks == null) { - this.groupingSetMasks = new java.util.ArrayList(); + public void addToGroupingSets(java.util.List elem) { + if (this.groupingSets == null) { + this.groupingSets = new java.util.ArrayList>(); } - this.groupingSetMasks.add(elem); + this.groupingSets.add(elem); } @org.apache.thrift.annotation.Nullable - public java.util.List getGroupingSetMasks() { - return this.groupingSetMasks; + public java.util.List> getGroupingSets() { + return this.groupingSets; } - public void setGroupingSetMasks(@org.apache.thrift.annotation.Nullable java.util.List groupingSetMasks) { - this.groupingSetMasks = groupingSetMasks; + public void setGroupingSets(@org.apache.thrift.annotation.Nullable java.util.List> groupingSets) { + this.groupingSets = groupingSets; } - public void unsetGroupingSetMasks() { - this.groupingSetMasks = null; + public void unsetGroupingSets() { + this.groupingSets = null; } - /** Returns true if field groupingSetMasks is set (has been assigned a value) and false otherwise */ - public boolean isSetGroupingSetMasks() { - return this.groupingSetMasks != null; + /** Returns true if field groupingSets is set (has been assigned a value) and false otherwise */ + public boolean isSetGroupingSets() { + return this.groupingSets != null; } - public void setGroupingSetMasksIsSet(boolean value) { + public void setGroupingSetsIsSet(boolean value) { if (!value) { - this.groupingSetMasks = null; + this.groupingSets = null; } } @@ -792,11 +797,11 @@ public void setFieldValue(_Fields field, @org.apache.thrift.annotation.Nullable } break; - case GROUPING_SET_MASKS: + case GROUPING_SETS: if (value == null) { - unsetGroupingSetMasks(); + unsetGroupingSets(); } else { - setGroupingSetMasks((java.util.List)value); + setGroupingSets((java.util.List>)value); } break; @@ -843,8 +848,8 @@ public java.lang.Object getFieldValue(_Fields field) { case EXPRESSION_OVERRIDE_HINTS: return getExpressionOverrideHints(); - case GROUPING_SET_MASKS: - return getGroupingSetMasks(); + case GROUPING_SETS: + return getGroupingSets(); } throw new java.lang.IllegalStateException(); @@ -882,8 +887,8 @@ public boolean isSet(_Fields field) { return isSetExplain(); case EXPRESSION_OVERRIDE_HINTS: return isSetExpressionOverrideHints(); - case GROUPING_SET_MASKS: - return isSetGroupingSetMasks(); + case GROUPING_SETS: + return isSetGroupingSets(); } throw new java.lang.IllegalStateException(); } @@ -1009,12 +1014,12 @@ public boolean equals(PinotQuery that) { return false; } - boolean this_present_groupingSetMasks = true && this.isSetGroupingSetMasks(); - boolean that_present_groupingSetMasks = true && that.isSetGroupingSetMasks(); - if (this_present_groupingSetMasks || that_present_groupingSetMasks) { - if (!(this_present_groupingSetMasks && that_present_groupingSetMasks)) + boolean this_present_groupingSets = true && this.isSetGroupingSets(); + boolean that_present_groupingSets = true && that.isSetGroupingSets(); + if (this_present_groupingSets || that_present_groupingSets) { + if (!(this_present_groupingSets && that_present_groupingSets)) return false; - if (!this.groupingSetMasks.equals(that.groupingSetMasks)) + if (!this.groupingSets.equals(that.groupingSets)) return false; } @@ -1073,9 +1078,9 @@ public int hashCode() { if (isSetExpressionOverrideHints()) hashCode = hashCode * 8191 + expressionOverrideHints.hashCode(); - hashCode = hashCode * 8191 + ((isSetGroupingSetMasks()) ? 131071 : 524287); - if (isSetGroupingSetMasks()) - hashCode = hashCode * 8191 + groupingSetMasks.hashCode(); + hashCode = hashCode * 8191 + ((isSetGroupingSets()) ? 131071 : 524287); + if (isSetGroupingSets()) + hashCode = hashCode * 8191 + groupingSets.hashCode(); return hashCode; } @@ -1208,12 +1213,12 @@ public int compareTo(PinotQuery other) { return lastComparison; } } - lastComparison = java.lang.Boolean.compare(isSetGroupingSetMasks(), other.isSetGroupingSetMasks()); + lastComparison = java.lang.Boolean.compare(isSetGroupingSets(), other.isSetGroupingSets()); if (lastComparison != 0) { return lastComparison; } - if (isSetGroupingSetMasks()) { - lastComparison = org.apache.thrift.TBaseHelper.compareTo(this.groupingSetMasks, other.groupingSetMasks); + if (isSetGroupingSets()) { + lastComparison = org.apache.thrift.TBaseHelper.compareTo(this.groupingSets, other.groupingSets); if (lastComparison != 0) { return lastComparison; } @@ -1345,13 +1350,13 @@ public java.lang.String toString() { } first = false; } - if (isSetGroupingSetMasks()) { + if (isSetGroupingSets()) { if (!first) sb.append(", "); - sb.append("groupingSetMasks:"); - if (this.groupingSetMasks == null) { + sb.append("groupingSets:"); + if (this.groupingSets == null) { sb.append("null"); } else { - sb.append(this.groupingSetMasks); + sb.append(this.groupingSets); } first = false; } @@ -1569,20 +1574,30 @@ public void read(org.apache.thrift.protocol.TProtocol iprot, PinotQuery struct) org.apache.thrift.protocol.TProtocolUtil.skip(iprot, schemeField.type); } break; - case 14: // GROUPING_SET_MASKS + case 15: // GROUPING_SETS if (schemeField.type == org.apache.thrift.protocol.TType.LIST) { { org.apache.thrift.protocol.TList _list17 = iprot.readListBegin(); - struct.groupingSetMasks = new java.util.ArrayList(_list17.size); - int _elem18; + struct.groupingSets = new java.util.ArrayList>(_list17.size); + @org.apache.thrift.annotation.Nullable java.util.List _elem18; for (int _i19 = 0; _i19 < _list17.size; ++_i19) { - _elem18 = iprot.readI32(); - struct.groupingSetMasks.add(_elem18); + { + org.apache.thrift.protocol.TList _list20 = iprot.readListBegin(); + _elem18 = new java.util.ArrayList(_list20.size); + int _elem21; + for (int _i22 = 0; _i22 < _list20.size; ++_i22) + { + _elem21 = iprot.readI32(); + _elem18.add(_elem21); + } + iprot.readListEnd(); + } + struct.groupingSets.add(_elem18); } iprot.readListEnd(); } - struct.setGroupingSetMasksIsSet(true); + struct.setGroupingSetsIsSet(true); } else { org.apache.thrift.protocol.TProtocolUtil.skip(iprot, schemeField.type); } @@ -1618,9 +1633,9 @@ public void write(org.apache.thrift.protocol.TProtocol oprot, PinotQuery struct) oprot.writeFieldBegin(SELECT_LIST_FIELD_DESC); { oprot.writeListBegin(new org.apache.thrift.protocol.TList(org.apache.thrift.protocol.TType.STRUCT, struct.selectList.size())); - for (Expression _iter20 : struct.selectList) + for (Expression _iter23 : struct.selectList) { - _iter20.write(oprot); + _iter23.write(oprot); } oprot.writeListEnd(); } @@ -1639,9 +1654,9 @@ public void write(org.apache.thrift.protocol.TProtocol oprot, PinotQuery struct) oprot.writeFieldBegin(GROUP_BY_LIST_FIELD_DESC); { oprot.writeListBegin(new org.apache.thrift.protocol.TList(org.apache.thrift.protocol.TType.STRUCT, struct.groupByList.size())); - for (Expression _iter21 : struct.groupByList) + for (Expression _iter24 : struct.groupByList) { - _iter21.write(oprot); + _iter24.write(oprot); } oprot.writeListEnd(); } @@ -1653,9 +1668,9 @@ public void write(org.apache.thrift.protocol.TProtocol oprot, PinotQuery struct) oprot.writeFieldBegin(ORDER_BY_LIST_FIELD_DESC); { oprot.writeListBegin(new org.apache.thrift.protocol.TList(org.apache.thrift.protocol.TType.STRUCT, struct.orderByList.size())); - for (Expression _iter22 : struct.orderByList) + for (Expression _iter25 : struct.orderByList) { - _iter22.write(oprot); + _iter25.write(oprot); } oprot.writeListEnd(); } @@ -1684,10 +1699,10 @@ public void write(org.apache.thrift.protocol.TProtocol oprot, PinotQuery struct) oprot.writeFieldBegin(QUERY_OPTIONS_FIELD_DESC); { oprot.writeMapBegin(new org.apache.thrift.protocol.TMap(org.apache.thrift.protocol.TType.STRING, org.apache.thrift.protocol.TType.STRING, struct.queryOptions.size())); - for (java.util.Map.Entry _iter23 : struct.queryOptions.entrySet()) + for (java.util.Map.Entry _iter26 : struct.queryOptions.entrySet()) { - oprot.writeString(_iter23.getKey()); - oprot.writeString(_iter23.getValue()); + oprot.writeString(_iter26.getKey()); + oprot.writeString(_iter26.getValue()); } oprot.writeMapEnd(); } @@ -1704,24 +1719,31 @@ public void write(org.apache.thrift.protocol.TProtocol oprot, PinotQuery struct) oprot.writeFieldBegin(EXPRESSION_OVERRIDE_HINTS_FIELD_DESC); { oprot.writeMapBegin(new org.apache.thrift.protocol.TMap(org.apache.thrift.protocol.TType.STRUCT, org.apache.thrift.protocol.TType.STRUCT, struct.expressionOverrideHints.size())); - for (java.util.Map.Entry _iter24 : struct.expressionOverrideHints.entrySet()) + for (java.util.Map.Entry _iter27 : struct.expressionOverrideHints.entrySet()) { - _iter24.getKey().write(oprot); - _iter24.getValue().write(oprot); + _iter27.getKey().write(oprot); + _iter27.getValue().write(oprot); } oprot.writeMapEnd(); } oprot.writeFieldEnd(); } } - if (struct.groupingSetMasks != null) { - if (struct.isSetGroupingSetMasks()) { - oprot.writeFieldBegin(GROUPING_SET_MASKS_FIELD_DESC); + if (struct.groupingSets != null) { + if (struct.isSetGroupingSets()) { + oprot.writeFieldBegin(GROUPING_SETS_FIELD_DESC); { - oprot.writeListBegin(new org.apache.thrift.protocol.TList(org.apache.thrift.protocol.TType.I32, struct.groupingSetMasks.size())); - for (int _iter25 : struct.groupingSetMasks) + oprot.writeListBegin(new org.apache.thrift.protocol.TList(org.apache.thrift.protocol.TType.LIST, struct.groupingSets.size())); + for (java.util.List _iter28 : struct.groupingSets) { - oprot.writeI32(_iter25); + { + oprot.writeListBegin(new org.apache.thrift.protocol.TList(org.apache.thrift.protocol.TType.I32, _iter28.size())); + for (int _iter29 : _iter28) + { + oprot.writeI32(_iter29); + } + oprot.writeListEnd(); + } } oprot.writeListEnd(); } @@ -1783,7 +1805,7 @@ public void write(org.apache.thrift.protocol.TProtocol prot, PinotQuery struct) if (struct.isSetExpressionOverrideHints()) { optionals.set(11); } - if (struct.isSetGroupingSetMasks()) { + if (struct.isSetGroupingSets()) { optionals.set(12); } oprot.writeBitSet(optionals, 13); @@ -1796,9 +1818,9 @@ public void write(org.apache.thrift.protocol.TProtocol prot, PinotQuery struct) if (struct.isSetSelectList()) { { oprot.writeI32(struct.selectList.size()); - for (Expression _iter26 : struct.selectList) + for (Expression _iter30 : struct.selectList) { - _iter26.write(oprot); + _iter30.write(oprot); } } } @@ -1808,18 +1830,18 @@ public void write(org.apache.thrift.protocol.TProtocol prot, PinotQuery struct) if (struct.isSetGroupByList()) { { oprot.writeI32(struct.groupByList.size()); - for (Expression _iter27 : struct.groupByList) + for (Expression _iter31 : struct.groupByList) { - _iter27.write(oprot); + _iter31.write(oprot); } } } if (struct.isSetOrderByList()) { { oprot.writeI32(struct.orderByList.size()); - for (Expression _iter28 : struct.orderByList) + for (Expression _iter32 : struct.orderByList) { - _iter28.write(oprot); + _iter32.write(oprot); } } } @@ -1835,10 +1857,10 @@ public void write(org.apache.thrift.protocol.TProtocol prot, PinotQuery struct) if (struct.isSetQueryOptions()) { { oprot.writeI32(struct.queryOptions.size()); - for (java.util.Map.Entry _iter29 : struct.queryOptions.entrySet()) + for (java.util.Map.Entry _iter33 : struct.queryOptions.entrySet()) { - oprot.writeString(_iter29.getKey()); - oprot.writeString(_iter29.getValue()); + oprot.writeString(_iter33.getKey()); + oprot.writeString(_iter33.getValue()); } } } @@ -1848,19 +1870,25 @@ public void write(org.apache.thrift.protocol.TProtocol prot, PinotQuery struct) if (struct.isSetExpressionOverrideHints()) { { oprot.writeI32(struct.expressionOverrideHints.size()); - for (java.util.Map.Entry _iter30 : struct.expressionOverrideHints.entrySet()) + for (java.util.Map.Entry _iter34 : struct.expressionOverrideHints.entrySet()) { - _iter30.getKey().write(oprot); - _iter30.getValue().write(oprot); + _iter34.getKey().write(oprot); + _iter34.getValue().write(oprot); } } } - if (struct.isSetGroupingSetMasks()) { + if (struct.isSetGroupingSets()) { { - oprot.writeI32(struct.groupingSetMasks.size()); - for (int _iter31 : struct.groupingSetMasks) + oprot.writeI32(struct.groupingSets.size()); + for (java.util.List _iter35 : struct.groupingSets) { - oprot.writeI32(_iter31); + { + oprot.writeI32(_iter35.size()); + for (int _iter36 : _iter35) + { + oprot.writeI32(_iter36); + } + } } } } @@ -1881,14 +1909,14 @@ public void read(org.apache.thrift.protocol.TProtocol prot, PinotQuery struct) t } if (incoming.get(2)) { { - org.apache.thrift.protocol.TList _list32 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT); - struct.selectList = new java.util.ArrayList(_list32.size); - @org.apache.thrift.annotation.Nullable Expression _elem33; - for (int _i34 = 0; _i34 < _list32.size; ++_i34) + org.apache.thrift.protocol.TList _list37 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT); + struct.selectList = new java.util.ArrayList(_list37.size); + @org.apache.thrift.annotation.Nullable Expression _elem38; + for (int _i39 = 0; _i39 < _list37.size; ++_i39) { - _elem33 = new Expression(); - _elem33.read(iprot); - struct.selectList.add(_elem33); + _elem38 = new Expression(); + _elem38.read(iprot); + struct.selectList.add(_elem38); } } struct.setSelectListIsSet(true); @@ -1900,28 +1928,28 @@ public void read(org.apache.thrift.protocol.TProtocol prot, PinotQuery struct) t } if (incoming.get(4)) { { - org.apache.thrift.protocol.TList _list35 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT); - struct.groupByList = new java.util.ArrayList(_list35.size); - @org.apache.thrift.annotation.Nullable Expression _elem36; - for (int _i37 = 0; _i37 < _list35.size; ++_i37) + org.apache.thrift.protocol.TList _list40 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT); + struct.groupByList = new java.util.ArrayList(_list40.size); + @org.apache.thrift.annotation.Nullable Expression _elem41; + for (int _i42 = 0; _i42 < _list40.size; ++_i42) { - _elem36 = new Expression(); - _elem36.read(iprot); - struct.groupByList.add(_elem36); + _elem41 = new Expression(); + _elem41.read(iprot); + struct.groupByList.add(_elem41); } } struct.setGroupByListIsSet(true); } if (incoming.get(5)) { { - org.apache.thrift.protocol.TList _list38 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT); - struct.orderByList = new java.util.ArrayList(_list38.size); - @org.apache.thrift.annotation.Nullable Expression _elem39; - for (int _i40 = 0; _i40 < _list38.size; ++_i40) + org.apache.thrift.protocol.TList _list43 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT); + struct.orderByList = new java.util.ArrayList(_list43.size); + @org.apache.thrift.annotation.Nullable Expression _elem44; + for (int _i45 = 0; _i45 < _list43.size; ++_i45) { - _elem39 = new Expression(); - _elem39.read(iprot); - struct.orderByList.add(_elem39); + _elem44 = new Expression(); + _elem44.read(iprot); + struct.orderByList.add(_elem44); } } struct.setOrderByListIsSet(true); @@ -1941,15 +1969,15 @@ public void read(org.apache.thrift.protocol.TProtocol prot, PinotQuery struct) t } if (incoming.get(9)) { { - org.apache.thrift.protocol.TMap _map41 = iprot.readMapBegin(org.apache.thrift.protocol.TType.STRING, org.apache.thrift.protocol.TType.STRING); - struct.queryOptions = new java.util.HashMap(2*_map41.size); - @org.apache.thrift.annotation.Nullable java.lang.String _key42; - @org.apache.thrift.annotation.Nullable java.lang.String _val43; - for (int _i44 = 0; _i44 < _map41.size; ++_i44) + org.apache.thrift.protocol.TMap _map46 = iprot.readMapBegin(org.apache.thrift.protocol.TType.STRING, org.apache.thrift.protocol.TType.STRING); + struct.queryOptions = new java.util.HashMap(2*_map46.size); + @org.apache.thrift.annotation.Nullable java.lang.String _key47; + @org.apache.thrift.annotation.Nullable java.lang.String _val48; + for (int _i49 = 0; _i49 < _map46.size; ++_i49) { - _key42 = iprot.readString(); - _val43 = iprot.readString(); - struct.queryOptions.put(_key42, _val43); + _key47 = iprot.readString(); + _val48 = iprot.readString(); + struct.queryOptions.put(_key47, _val48); } } struct.setQueryOptionsIsSet(true); @@ -1960,33 +1988,42 @@ public void read(org.apache.thrift.protocol.TProtocol prot, PinotQuery struct) t } if (incoming.get(11)) { { - org.apache.thrift.protocol.TMap _map45 = iprot.readMapBegin(org.apache.thrift.protocol.TType.STRUCT, org.apache.thrift.protocol.TType.STRUCT); - struct.expressionOverrideHints = new java.util.HashMap(2*_map45.size); - @org.apache.thrift.annotation.Nullable Expression _key46; - @org.apache.thrift.annotation.Nullable Expression _val47; - for (int _i48 = 0; _i48 < _map45.size; ++_i48) + org.apache.thrift.protocol.TMap _map50 = iprot.readMapBegin(org.apache.thrift.protocol.TType.STRUCT, org.apache.thrift.protocol.TType.STRUCT); + struct.expressionOverrideHints = new java.util.HashMap(2*_map50.size); + @org.apache.thrift.annotation.Nullable Expression _key51; + @org.apache.thrift.annotation.Nullable Expression _val52; + for (int _i53 = 0; _i53 < _map50.size; ++_i53) { - _key46 = new Expression(); - _key46.read(iprot); - _val47 = new Expression(); - _val47.read(iprot); - struct.expressionOverrideHints.put(_key46, _val47); + _key51 = new Expression(); + _key51.read(iprot); + _val52 = new Expression(); + _val52.read(iprot); + struct.expressionOverrideHints.put(_key51, _val52); } } struct.setExpressionOverrideHintsIsSet(true); } if (incoming.get(12)) { { - org.apache.thrift.protocol.TList _list49 = iprot.readListBegin(org.apache.thrift.protocol.TType.I32); - struct.groupingSetMasks = new java.util.ArrayList(_list49.size); - int _elem50; - for (int _i51 = 0; _i51 < _list49.size; ++_i51) + org.apache.thrift.protocol.TList _list54 = iprot.readListBegin(org.apache.thrift.protocol.TType.LIST); + struct.groupingSets = new java.util.ArrayList>(_list54.size); + @org.apache.thrift.annotation.Nullable java.util.List _elem55; + for (int _i56 = 0; _i56 < _list54.size; ++_i56) { - _elem50 = iprot.readI32(); - struct.groupingSetMasks.add(_elem50); + { + org.apache.thrift.protocol.TList _list57 = iprot.readListBegin(org.apache.thrift.protocol.TType.I32); + _elem55 = new java.util.ArrayList(_list57.size); + int _elem58; + for (int _i59 = 0; _i59 < _list57.size; ++_i59) + { + _elem58 = iprot.readI32(); + _elem55.add(_elem58); + } + } + struct.groupingSets.add(_elem55); } } - struct.setGroupingSetMasksIsSet(true); + struct.setGroupingSetsIsSet(true); } } } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/request/context/GroupingSets.java b/pinot-common/src/main/java/org/apache/pinot/common/request/context/GroupingSets.java index 04452b371f8f..400fcea20486 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/request/context/GroupingSets.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/request/context/GroupingSets.java @@ -18,68 +18,105 @@ */ package org.apache.pinot.common.request.context; +import java.util.ArrayList; import java.util.List; -import org.apache.calcite.rel.type.RelDataType; -import org.apache.calcite.rel.type.RelDataTypeFactory; -import org.apache.calcite.rel.type.RelDataTypeField; -import org.apache.calcite.sql.type.SqlTypeName; /// Engine-agnostic conventions for GROUP BY GROUPING SETS / ROLLUP / CUBE, shared so both query engines /// agree on the discriminator column and the GROUPING()/GROUPING_ID() semantics. /// -/// A grouping-set query is represented as the union of all grouping columns plus, per set, a "grouping-id" -/// bitmask over that union where bit `i` is set iff union column `i` is **rolled up** (aggregated away) in -/// that set — matching the SQL/PostgreSQL convention that GROUPING(col) returns 1 for an aggregated-away -/// column. The bitmask is carried through the result as the synthetic key column named {@link -/// #GROUPING_ID_COLUMN}, which the discriminator keeps distinct across sets and which powers -/// GROUPING()/GROUPING_ID(). +/// A grouping-set query is represented as the ordered union of all grouping columns plus an ordered list of +/// grouping sets, each set given as the indexes of its participating union columns (mirroring Calcite's +/// per-set column bitset, so the number of grouping columns is unlimited). A set's position in that list is +/// its **ordinal**, carried through the result as the synthetic INT key column named {@link +/// #GROUPING_ID_COLUMN}: the discriminator keeps rows from different sets distinct — e.g. set `{a}` with `b` +/// rolled up to NULL stays apart from set `{a, b}` where `b` is genuinely NULL — and identifies the set that +/// produced a row, which is what powers GROUPING()/GROUPING_ID(). public class GroupingSets { private GroupingSets() { } - /// Name of the synthetic INT key column that carries the per-set grouping-id bitmask. It is internal: it is - /// appended after the union group-by columns in intermediate results and dropped from the user-facing - /// result (never projected unless referenced via GROUPING()/GROUPING_ID()). + /// Name of the synthetic INT key column that carries the grouping-set ordinal (the set's position in the + /// query's grouping-set list). It is internal: it is appended after the union group-by columns in + /// intermediate results and dropped from the user-facing result (never projected unless referenced via + /// GROUPING()/GROUPING_ID()). Note: despite the name, the value identifies the grouping set (its ordinal); + /// it is NOT the SQL {@code GROUPING_ID()} bitmask, which is computed from it per call. public static final String GROUPING_ID_COLUMN = "$groupingId"; - /// The grouping-id discriminator is a 32-bit INT bitmask over the union group-by columns, so a grouping-set query - /// may reference at most 31 distinct grouping columns (bit 31 is the sign bit). - public static final int MAX_GROUPING_SET_COLUMNS = 31; + /// GROUPING()/GROUPING_ID() pack one bit per argument into an INT (first argument = most significant bit), + /// so a single call accepts at most 31 arguments. This is a per-call limit on the function's arguments — + /// the number of grouping columns in the query is unlimited. + public static final int MAX_GROUPING_FUNCTION_ARGS = 31; - /// Returns the LEAF row type for a grouping-set aggregate: the synthetic {@link #GROUPING_ID_COLUMN} INT column - /// inserted right after the {@code groupCount} union group-by columns, i.e. {@code [group keys..., $groupingId, - /// aggregates...]}. Both query engines share this single layout so the multi-stage final stage can group on - /// {@code $groupingId}; keeping it here prevents the two planner {@code deriveRowType()} overrides from drifting. - public static RelDataType appendGroupingIdColumn(RelDataTypeFactory typeFactory, RelDataType rowType, - int groupCount) { - RelDataTypeFactory.Builder builder = typeFactory.builder(); - List fields = rowType.getFieldList(); - for (int i = 0; i < groupCount; i++) { - builder.add(fields.get(i)); - } - builder.add(GROUPING_ID_COLUMN, typeFactory.createSqlType(SqlTypeName.INTEGER)); - for (int i = groupCount; i < fields.size(); i++) { - builder.add(fields.get(i)); - } - return builder.build(); - } + /// Upper bound on the number of grouping sets a single query may expand to. Guards against CUBE blow-up: + /// every set is materialized in the plan, expanded per input row at runtime, and enumerated by the + /// GROUPING()/GROUPING_ID() projection, so the set count bounds plan size and per-row work. Shared by the + /// single-stage parser and the multi-stage plan conversion. + public static final int MAX_GROUPING_SETS = 4096; - /// Computes the {@code GROUPING(args...)} / {@code GROUPING_ID(args...)} value from a row's grouping-id - /// bitmask. For each argument column (identified by its index in the union of grouping columns) the - /// corresponding bit is read from {@code groupingId} (1 = rolled up), and the bits are packed with the - /// first argument as the most significant bit, per the SQL standard. + /// Computes the {@code GROUPING(args...)} / {@code GROUPING_ID(args...)} value for one grouping set. For + /// each argument column (identified by its index in the union of grouping columns) the bit is 1 iff the + /// column is **rolled up** (does not participate) in the set, and the bits are packed with the first + /// argument as the most significant bit, per the SQL/PostgreSQL convention. /// - /// @param groupingId the row's grouping-id bitmask (bit i set iff union column i is rolled up) - /// @param unionColumnIndexes the union-column index of each GROUPING argument, in argument order - /// @return the GROUPING / GROUPING_ID integer value - public static int groupingValue(int groupingId, int[] unionColumnIndexes) { + /// @param setColumnIndexes the union-column indexes participating in the grouping set + /// @param argUnionIndexes the union-column index of each GROUPING argument, in argument order + /// @return the GROUPING / GROUPING_ID integer value for rows of that set + public static int groupingValue(int[] setColumnIndexes, int[] argUnionIndexes) { int result = 0; - for (int unionColumnIndex : unionColumnIndexes) { - result = (result << 1) | ((groupingId >>> unionColumnIndex) & 1); + for (int argUnionIndex : argUnionIndexes) { + int rolledUp = 1; + for (int setColumnIndex : setColumnIndexes) { + if (setColumnIndex == argUnionIndex) { + rolledUp = 0; + break; + } + } + result = (result << 1) | rolledUp; } return result; } + /// Precomputes the {@code GROUPING(args...)} / {@code GROUPING_ID(args...)} value for every grouping set of + /// a query, indexed by the set's ordinal (the value of the {@link #GROUPING_ID_COLUMN} discriminator). The + /// value of a GROUPING call is fully determined by the grouping set a row belongs to, so evaluation per row + /// is a single lookup: {@code values[ordinal]}. + /// + /// @param groupingSets per grouping set (in ordinal order), the union-column indexes participating in it + /// @param argUnionIndexes the union-column index of each GROUPING argument, in argument order + /// @return the GROUPING / GROUPING_ID value per grouping-set ordinal + public static int[] groupingValuesByOrdinal(List groupingSets, int[] argUnionIndexes) { + if (argUnionIndexes.length > MAX_GROUPING_FUNCTION_ARGS) { + throw new IllegalStateException( + "GROUPING / GROUPING_ID supports at most " + MAX_GROUPING_FUNCTION_ARGS + " arguments, got " + + argUnionIndexes.length); + } + int numSets = groupingSets.size(); + int[] values = new int[numSets]; + for (int ordinal = 0; ordinal < numSets; ordinal++) { + values[ordinal] = groupingValue(groupingSets.get(ordinal), argUnionIndexes); + } + return values; + } + + /// Boxed-list overload of {@link #groupingValuesByOrdinal(List, int[])} for callers holding the wire / plan + /// representation ({@code List>} sets and {@code List} argument indexes). + public static int[] groupingValuesByOrdinal(List> groupingSets, + List argUnionIndexes) { + List sets = new ArrayList<>(groupingSets.size()); + for (List groupingSet : groupingSets) { + int[] set = new int[groupingSet.size()]; + for (int i = 0; i < set.length; i++) { + set[i] = groupingSet.get(i); + } + sets.add(set); + } + int[] args = new int[argUnionIndexes.size()]; + for (int i = 0; i < args.length; i++) { + args[i] = argUnionIndexes.get(i); + } + return groupingValuesByOrdinal(sets, args); + } + /// Returns true if the given function name denotes {@code GROUPING} or {@code GROUPING_ID}. The argument is /// the post-canonicalization form (lower-cased, underscores removed), so {@code GROUPING_ID} appears as /// {@code "groupingid"}. Centralizes the name contract shared by the broker post-aggregation handler and the diff --git a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java index 932ff800637b..0d19d6200f6d 100644 --- a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java +++ b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java @@ -23,6 +23,7 @@ import com.google.common.collect.ImmutableSet; import java.io.StringReader; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.LinkedHashMap; @@ -65,6 +66,7 @@ import org.apache.pinot.common.request.Join; import org.apache.pinot.common.request.JoinType; import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.request.context.GroupingSets; import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.segment.spi.AggregationFunctionType; @@ -97,11 +99,9 @@ private CalciteSqlParser() { public static final List QUERY_REWRITERS = new ArrayList<>(QueryRewriterFactory.getQueryRewriters()); // TODO: Add the ability to configure the parser's maximum identifier length via configuration if needed in the future public static final int CALCITE_SQL_PARSER_IDENTIFIER_MAX_LENGTH = 1024; - /// The grouping-set discriminator is encoded as a 32-bit INT bitmask over the union group-by columns, so a - /// GROUPING SETS / ROLLUP / CUBE query may reference at most 31 distinct grouping columns. - public static final int MAX_GROUPING_SETS_COLUMNS = 31; - /// Upper bound on the number of grouping sets a single query may expand to (guards against CUBE blow-up). - public static final int MAX_GROUPING_SETS = 4096; + /// Upper bound on the number of grouping sets a single query may expand to (guards against CUBE blow-up); + /// shared with the multi-stage plan conversion. + public static final int MAX_GROUPING_SETS = GroupingSets.MAX_GROUPING_SETS; private static final Logger LOGGER = LoggerFactory.getLogger(CalciteSqlParser.class); // To Keep the backward compatibility with 'OPTION' Functionality in PQL, which is used to @@ -256,11 +256,47 @@ private static void validateGroupByClause(PinotQuery pinotQuery) // sets. Note that GROUPING() / GROUPING_ID() are not aggregation functions. (Plain non-aggregation GROUP BY // queries are rewritten to DISTINCT by NonAggregationGroupByToDistinctQueryRewriter, but that rewrite cannot // represent multiple grouping sets.) - if (pinotQuery.getGroupingSetMasks() != null && aggregateExprCount == 0 && !hasAggregationOutsideSelect( + if (pinotQuery.getGroupingSets() != null && aggregateExprCount == 0 && !hasAggregationOutsideSelect( pinotQuery)) { throw new SqlCompilationException( "GROUP BY GROUPING SETS / ROLLUP / CUBE requires at least one aggregation function in the query"); } + // GROUPING()/GROUPING_ID() pack one bit per argument into an INT, so a single call accepts at most 31 + // arguments (a per-call limit — the number of grouping columns is unlimited). Reject at compile time so the + // user gets a clear error instead of a mid-execution failure from the shared value computation. + if (pinotQuery.getGroupingSets() != null) { + for (Expression selectExpression : pinotQuery.getSelectList()) { + validateGroupingFunctionArgs(selectExpression); + } + if (pinotQuery.getHavingExpression() != null) { + validateGroupingFunctionArgs(pinotQuery.getHavingExpression()); + } + if (pinotQuery.getOrderByList() != null) { + for (Expression orderBy : pinotQuery.getOrderByList()) { + validateGroupingFunctionArgs(orderBy); + } + } + } + } + + /// Recursively rejects GROUPING() / GROUPING_ID() calls with more than + /// {@link GroupingSets#MAX_GROUPING_FUNCTION_ARGS} arguments (the packed INT bit width). + private static void validateGroupingFunctionArgs(Expression expression) { + Function function = expression.getFunctionCall(); + if (function == null) { + return; + } + if (GroupingSets.isGroupingFunction(function.getOperator()) + && function.getOperandsSize() > GroupingSets.MAX_GROUPING_FUNCTION_ARGS) { + throw new SqlCompilationException( + "GROUPING / GROUPING_ID supports at most " + GroupingSets.MAX_GROUPING_FUNCTION_ARGS + " arguments, got " + + function.getOperandsSize()); + } + if (function.getOperands() != null) { + for (Expression operand : function.getOperands()) { + validateGroupingFunctionArgs(operand); + } + } } /// Returns true if the HAVING clause or any ORDER-BY expression contains an aggregation. @@ -725,16 +761,16 @@ private static List convertSelectList(SqlNodeList selectList) { } /// Converts the GROUP BY clause into {@link PinotQuery#groupByList} and (when grouping constructs are - /// present) {@link PinotQuery#groupingSetMasks}. + /// present) {@link PinotQuery#groupingSets}. /// /// For a plain GROUP BY (no ROLLUP / CUBE / GROUPING SETS) this behaves exactly like - /// {@link #convertSelectList} and leaves {@code groupingSetMasks} unset, so non-grouping-set queries are + /// {@link #convertSelectList} and leaves {@code groupingSets} unset, so non-grouping-set queries are /// unchanged. When grouping constructs are present, the grouping elements are cross-multiplied into the /// canonical, de-duplicated list of grouping sets (standard SQL semantics, e.g. /// {@code GROUP BY a, ROLLUP(b, c)} produces {@code {a,b,c}, {a,b}, {a}}); the ordered, de-duplicated /// union of all participating columns is stored as {@code groupByList}, and each grouping set is stored - /// as a membership bitmask over that union — bit {@code i} set iff union column {@code i} participates, - /// a mask of {@code 0} being the grand-total set {@code ()}. + /// as the sorted list of its participating union-column indexes, an empty list being the grand-total set + /// {@code ()}. private static void setGroupByListAndGroupingSets(PinotQuery pinotQuery, SqlNodeList groupByNodeList) { boolean hasGroupingConstruct = false; for (SqlNode node : groupByNodeList) { @@ -776,28 +812,25 @@ private static void setGroupByListAndGroupingSets(PinotQuery pinotQuery, SqlNode unionIndex.computeIfAbsent(expr, k -> unionIndex.size()); } } - if (unionIndex.size() > MAX_GROUPING_SETS_COLUMNS) { - throw new SqlCompilationException( - "GROUPING SETS / ROLLUP / CUBE support at most " + MAX_GROUPING_SETS_COLUMNS - + " distinct grouping columns, got " + unionIndex.size()); - } pinotQuery.setGroupByList(new ArrayList<>(unionIndex.keySet())); - /// Encode each grouping set as a membership bitmask over the union columns (bit i set iff union column i - /// participates), de-duplicating overlapping sets produced by CUBE/ROLLUP. A mask of 0 is the grand-total - /// set (). The union is capped at MAX_GROUPING_SETS_COLUMNS (31) above, so a 32-bit mask never overflows. - Set seen = new HashSet<>(); - List groupingSetMasks = new ArrayList<>(); + /// Encode each grouping set as the sorted list of its participating union-column indexes (mirroring + /// Calcite's per-set column bitset, so the number of grouping columns is unlimited), de-duplicating + /// overlapping sets produced by CUBE/ROLLUP. An empty list is the grand-total set (). A set's position in + /// the list is its ordinal — the value the engine carries in the synthetic $groupingId discriminator. + Set> seen = new HashSet<>(); + List> groupingSets = new ArrayList<>(); for (LinkedHashSet set : combinedSets) { - int mask = 0; + List columnIndexes = new ArrayList<>(set.size()); for (Expression expr : set) { - mask |= 1 << unionIndex.get(expr); + columnIndexes.add(unionIndex.get(expr)); } - if (seen.add(mask)) { - groupingSetMasks.add(mask); + Collections.sort(columnIndexes); + if (seen.add(columnIndexes)) { + groupingSets.add(columnIndexes); } } - pinotQuery.setGroupingSetMasks(groupingSetMasks); + pinotQuery.setGroupingSets(groupingSets); } private static boolean isGroupingConstruct(SqlKind kind) { diff --git a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/NonAggregationGroupByToDistinctQueryRewriter.java b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/NonAggregationGroupByToDistinctQueryRewriter.java index 35f0e68ab223..504dcaa3e80b 100644 --- a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/NonAggregationGroupByToDistinctQueryRewriter.java +++ b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/NonAggregationGroupByToDistinctQueryRewriter.java @@ -54,7 +54,7 @@ public PinotQuery rewrite(PinotQuery pinotQuery) { // The DISTINCT rewrite is only valid for a plain GROUP BY: a GROUPING SETS / ROLLUP / CUBE query produces // one row per group per grouping set, which DISTINCT over the union of grouping columns cannot represent. // Such queries are rejected by CalciteSqlParser.validateGroupByClause() when they have no aggregation. - if (pinotQuery.getGroupingSetMasks() != null) { + if (pinotQuery.getGroupingSets() != null) { return pinotQuery; } for (Expression select : pinotQuery.getSelectList()) { diff --git a/pinot-common/src/main/proto/plan.proto b/pinot-common/src/main/proto/plan.proto index a912fe989a31..490a4dd7c605 100644 --- a/pinot-common/src/main/proto/plan.proto +++ b/pinot-common/src/main/proto/plan.proto @@ -65,6 +65,13 @@ enum AggType { FINAL = 3; } +// One grouping set of a GROUP BY GROUPING SETS / ROLLUP / CUBE aggregate: the indexes into +// AggregateNode.groupKeys (the union of all grouping columns) that participate in (are grouped by) the set. +// An empty list is the grand-total set (). +message GroupingSet { + repeated int32 groupKeyIndexes = 1; +} + message AggregateNode { repeated FunctionCall aggCalls = 1; repeated int32 filterArgs = 2; @@ -73,12 +80,13 @@ message AggregateNode { bool leafReturnFinalResult = 5; repeated Collation collations = 6; int32 limit = 7; - // GROUP BY GROUPING SETS / ROLLUP / CUBE: one membership bitmask per grouping set, over groupKeys. Empty for a - // plain GROUP BY. See AggregateNode#getGroupingSets. - // Rolling upgrade: like every other multi-stage plan field, this is generated by the broker and executed by the - // servers, which must be on a version that understands it. A pre-grouping-sets server silently ignores this field - // and runs a plain GROUP BY (dropping the subtotal/grand-total rows), so upgrade servers before brokers. - repeated int32 groupingSets = 8; + // GROUP BY GROUPING SETS / ROLLUP / CUBE: one entry per grouping set, in plan order; a set's position in + // this list is its ordinal, carried as the synthetic $groupingId discriminator, so the number of grouping + // columns is unlimited (mirroring Calcite's per-set column bitset). Empty for a plain GROUP BY. + // Rolling upgrade: a pre-grouping-sets server ignores this field and runs a plain GROUP BY; its response is + // then rejected by the broker-side schema validation (missing $groupingId column) with an actionable error + // rather than returning collapsed results. Avoid grouping-set queries until all servers are upgraded. + repeated GroupingSet groupingSets = 8; } message FilterNode { diff --git a/pinot-common/src/test/java/org/apache/pinot/common/request/context/GroupingSetsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/request/context/GroupingSetsTest.java new file mode 100644 index 000000000000..4a01f4ee0ce0 --- /dev/null +++ b/pinot-common/src/test/java/org/apache/pinot/common/request/context/GroupingSetsTest.java @@ -0,0 +1,66 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.common.request.context; + +import java.util.List; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertThrows; + + +/// Pins the GROUPING()/GROUPING_ID() value contract shared by the single-stage reduce extractors and the +/// multi-stage planner projection: bit = 1 iff the argument column is rolled up (absent from the set), packed +/// with the first argument as the most significant bit, indexed by the grouping set's ordinal. +public class GroupingSetsTest { + + @Test + public void testGroupingValuePacksRolledUpBitsMsbFirst() { + /// ROLLUP(c0, c1) sets: (c0, c1), (c0), () — args (c0, c1) pack as GROUPING(c0) << 1 | GROUPING(c1). + List rollupSets = List.of(new int[]{0, 1}, new int[]{0}, new int[]{}); + assertEquals(GroupingSets.groupingValuesByOrdinal(rollupSets, new int[]{0, 1}), new int[]{0b00, 0b01, 0b11}); + /// Argument order defines bit order: reversed args produce the mirrored packing. + assertEquals(GroupingSets.groupingValuesByOrdinal(rollupSets, new int[]{1, 0}), new int[]{0b00, 0b10, 0b11}); + /// Single-argument GROUPING(c1): 0 where c1 participates, 1 where rolled up. + assertEquals(GroupingSets.groupingValuesByOrdinal(rollupSets, new int[]{1}), new int[]{0, 1, 1}); + /// The boxed (wire / plan representation) overload computes identical values. + assertEquals( + GroupingSets.groupingValuesByOrdinal(List.of(List.of(0, 1), List.of(0), List.of()), List.of(0, 1)), + new int[]{0b00, 0b01, 0b11}); + } + + @Test + public void testGroupingValueUnboundedColumns() { + /// Grouping columns are unlimited; only the per-call argument count is capped. A set over high column + /// indexes (way past the retired 31-column bitmask range) computes fine. + List sets = List.of(new int[]{0, 40}, new int[]{40}, new int[]{}); + assertEquals(GroupingSets.groupingValuesByOrdinal(sets, new int[]{40}), new int[]{0, 0, 1}); + assertEquals(GroupingSets.groupingValuesByOrdinal(sets, new int[]{0, 40}), new int[]{0b00, 0b10, 0b11}); + } + + @Test + public void testGroupingValueRejectsTooManyArguments() { + /// The packed value is an INT, so a single GROUPING/GROUPING_ID call accepts at most 31 arguments. + List sets = List.of(new int[]{0}); + assertEquals(GroupingSets.groupingValuesByOrdinal(sets, new int[GroupingSets.MAX_GROUPING_FUNCTION_ARGS]).length, + 1); + assertThrows(IllegalStateException.class, + () -> GroupingSets.groupingValuesByOrdinal(sets, new int[GroupingSets.MAX_GROUPING_FUNCTION_ARGS + 1])); + } +} diff --git a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/GroupingSetsParserTest.java b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/GroupingSetsParserTest.java index ebfaaa3adabf..5ad99c44bb8e 100644 --- a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/GroupingSetsParserTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/GroupingSetsParserTest.java @@ -32,20 +32,18 @@ /// Tests that GROUP BY GROUPING SETS / ROLLUP / CUBE are expanded by {@link CalciteSqlParser} into the union -/// of grouping columns ({@link PinotQuery#getGroupByList()}) plus a per-set membership bitmask -/// ({@link PinotQuery#getGroupingSetMasks()}). +/// of grouping columns ({@link PinotQuery#getGroupByList()}) plus, per grouping set, the list of participating +/// union-column indexes ({@link PinotQuery#getGroupingSets()}). public class GroupingSetsParserTest { /// Returns the grouping sets as a set of column-name sets, decoupled from union/set ordering. private static Set> groupingSetsByName(PinotQuery query) { List union = query.getGroupByList(); Set> result = new HashSet<>(); - for (int mask : query.getGroupingSetMasks()) { + for (List set : query.getGroupingSets()) { Set names = new HashSet<>(); - for (int bit = 0; bit < Integer.SIZE; bit++) { - if ((mask & (1 << bit)) != 0) { - names.add(union.get(bit).getIdentifier().getName()); - } + for (int columnIndex : set) { + names.add(union.get(columnIndex).getIdentifier().getName()); } result.add(names); } @@ -60,8 +58,8 @@ private static Set names(String... names) { public void testPlainGroupByLeavesGroupingSetsUnset() { PinotQuery query = CalciteSqlParser.compileToPinotQuery("SELECT a, SUM(c) FROM t GROUP BY a, b"); assertEquals(query.getGroupByList().size(), 2); - /// Plain GROUP BY must be byte-for-byte unchanged: no grouping-set masks emitted. - assertNull(query.getGroupingSetMasks()); + /// Plain GROUP BY must be byte-for-byte unchanged: no grouping sets emitted. + assertNull(query.getGroupingSets()); assertTrue(query.getGroupByListSize() == 2); } @@ -117,7 +115,7 @@ public void testDuplicateGroupingSetsDeduplicated() { "SELECT a, SUM(c) FROM t GROUP BY GROUPING SETS ((a), (a), ())"); assertEquals(groupingSetsByName(query), Set.of(names("a"), names())); /// {a} and {} only -> 2 distinct grouping sets. - assertEquals(query.getGroupingSetMasks().size(), 2); + assertEquals(query.getGroupingSets().size(), 2); } @Test @@ -155,15 +153,46 @@ public void testGroupingSetsWithAggregationOnlyInHavingOrOrderBy() { } } - @Test(expectedExceptions = SqlCompilationException.class) - public void testTooManyGroupingColumnsRejected() { - /// 32 distinct grouping columns exceed the 31-column limit imposed by the INT grouping-id bitmask. + @Test + public void testManyGroupingColumnsSupported() { + /// The number of grouping columns is unlimited (each set is a column-index list and the discriminator is + /// the set ordinal, mirroring Calcite's per-set column bitset): 40 distinct grouping columns parse fine, + /// where the retired 32-bit bitmask encoding capped the union at 31. StringBuilder sets = new StringBuilder(); - for (int i = 0; i < 32; i++) { + for (int i = 0; i < 40; i++) { sets.append(i == 0 ? "(c" : ", (c").append(i).append(')'); } - CalciteSqlParser.compileToPinotQuery( + PinotQuery query = CalciteSqlParser.compileToPinotQuery( "SELECT SUM(m) FROM t GROUP BY GROUPING SETS (" + sets + ")"); + assertEquals(query.getGroupByList().size(), 40); + assertEquals(query.getGroupingSets().size(), 40); + /// Each set holds exactly one distinct column index within range. + Set seenColumns = new HashSet<>(); + for (List set : query.getGroupingSets()) { + assertEquals(set.size(), 1); + int columnIndex = set.get(0); + assertTrue(columnIndex >= 0 && columnIndex < 40, "column index out of range: " + columnIndex); + assertTrue(seenColumns.add(columnIndex), "duplicate column index: " + columnIndex); + } + } + + @Test + public void testTooManyGroupingFunctionArgsRejected() { + /// GROUPING()/GROUPING_ID() pack one bit per argument into an INT, so a single call accepts at most 31 + /// arguments — rejected at compile time (not mid-execution), in SELECT and in HAVING/ORDER BY alike. + StringBuilder sets = new StringBuilder(); + StringBuilder args = new StringBuilder(); + for (int i = 0; i < 32; i++) { + sets.append(i == 0 ? "(c" : ", (c").append(i).append(')'); + args.append(i == 0 ? "c" : ", c").append(i); + } + for (String sql : List.of( + "SELECT GROUPING_ID(" + args + "), SUM(m) FROM t GROUP BY GROUPING SETS (" + sets + ")", + "SELECT SUM(m) FROM t GROUP BY GROUPING SETS (" + sets + ") ORDER BY GROUPING_ID(" + args + ")")) { + SqlCompilationException e = + expectThrows(SqlCompilationException.class, () -> CalciteSqlParser.compileToPinotQuery(sql)); + assertTrue(e.getMessage().contains("at most 31 arguments"), e.getMessage()); + } } @Test(expectedExceptions = SqlCompilationException.class) diff --git a/pinot-common/src/thrift/query.thrift b/pinot-common/src/thrift/query.thrift index 0b49b4a86af2..d3aa18928adb 100644 --- a/pinot-common/src/thrift/query.thrift +++ b/pinot-common/src/thrift/query.thrift @@ -32,12 +32,17 @@ struct PinotQuery { 11: optional map queryOptions; 12: optional bool explain; 13: optional map expressionOverrideHints; - // GROUP BY GROUPING SETS / ROLLUP / CUBE: one membership bitmask per grouping set, over groupByList. Bit i - // is set iff groupByList[i] participates in (is grouped by) that set; a mask of 0 is the grand-total set (). - // Unset for a plain GROUP BY query. Only set when grouping constructs are used, in which case groupByList (the - // union of all grouping columns) is capped at 31 entries so a mask always fits in an i32 (enforced at compile - // time by CalciteSqlParser.MAX_GROUPING_SETS_COLUMNS); plain GROUP BY queries have no such limit. - 14: optional list groupingSetMasks; +//14: retired before release (was list groupingSetMasks, a bitmask encoding capped at 31 grouping columns); +// do not reuse the id. + // GROUP BY GROUPING SETS / ROLLUP / CUBE: one entry per grouping set, in query order; each entry lists the + // indexes into groupByList (the ordered union of all grouping columns) that participate in (are grouped by) + // that set. An empty inner list is the grand-total set (). A set's position in this list is its ordinal, which + // the engine carries as the synthetic $groupingId discriminator key, so the number of grouping columns is + // unlimited (mirroring Calcite's per-set column bitset). Unset for a plain GROUP BY query. + // Rolling upgrade: a pre-grouping-sets server ignores this field and runs a plain GROUP BY; its response is + // then rejected by the broker-side schema validation (missing $groupingId column) with an actionable error + // rather than returning collapsed results. Avoid grouping-set queries until all servers are upgraded. + 15: optional list> groupingSets; } struct DataSource { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java b/pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java index 985bf261dc41..9d7c0341ba41 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java @@ -29,6 +29,7 @@ import java.util.List; import java.util.Map; import java.util.PriorityQueue; +import javax.annotation.Nullable; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.request.context.FilterContext; @@ -57,6 +58,10 @@ public class TableResizer { /// Number of key columns preceding the aggregation columns in a record: the group-by (union) columns plus /// the synthetic $groupingId column for grouping-set queries. Equals _numGroupByExpressions otherwise. private final int _numKeyColumns; + /// Per grouping set (in ordinal order), the participating union-column indexes; null when the query has no + /// grouping sets. Read by the GROUPING()/GROUPING_ID() ORDER BY extractor. + @Nullable + private final List _groupingSets; private final Map _groupByExpressionIndexMap; private final AggregationFunction[] _aggregationFunctions; private final Map, Integer> _filteredAggregationIndexMap; @@ -79,6 +84,7 @@ public TableResizer(DataSchema dataSchema, boolean hasFinalInput, QueryContext q assert groupByExpressions != null; _numGroupByExpressions = groupByExpressions.size(); _numKeyColumns = queryContext.getNumGroupByKeyColumns(); + _groupingSets = queryContext.getGroupingSets(); _groupByExpressionIndexMap = new HashMap<>(); for (int i = 0; i < _numGroupByExpressions; i++) { _groupByExpressionIndexMap.put(groupByExpressions.get(i), i); @@ -537,26 +543,30 @@ public Comparable extract(Record record) { } /// Extractor for {@code GROUPING(col, ...)} / {@code GROUPING_ID(col, ...)} in ORDER BY. Mirrors - /// {@code PostAggregationHandler.GroupingValueExtractor}: reads the synthetic {@code $groupingId} - /// discriminator column (immediately after the union group-by columns) and computes the bitmask over the - /// argument columns, with the first argument as the most significant bit (PostgreSQL semantics). Returns 0 - /// when the query has no grouping sets (no column is rolled up). + /// {@code PostAggregationHandler.GroupingValueExtractor}: the value is fully determined by the grouping set + /// a row belongs to, so it is precomputed per grouping-set ordinal at construction (first argument = most + /// significant bit, PostgreSQL semantics) and extraction is a single lookup on the synthetic + /// {@code $groupingId} ordinal column (immediately after the union group-by columns). Returns 0 when the + /// query has no grouping sets (no column is rolled up). private class GroupingExtractor implements OrderByValueExtractor { - final int[] _unionColumnIndexes; + @Nullable + final int[] _valuesByOrdinal; final int _discriminatorIndex; GroupingExtractor(FunctionContext function) { List arguments = function.getArguments(); Preconditions.checkState(!arguments.isEmpty(), "GROUPING requires at least one argument"); - _unionColumnIndexes = new int[arguments.size()]; + int[] argUnionIndexes = new int[arguments.size()]; for (int i = 0; i < arguments.size(); i++) { Integer index = _groupByExpressionIndexMap.get(arguments.get(i)); Preconditions.checkState(index != null, "GROUPING argument must be a grouping column, got: %s", arguments.get(i)); - _unionColumnIndexes[i] = index; + argUnionIndexes[i] = index; } /// The $groupingId column sits right after the union columns; it is absent for non-grouping-set queries. _discriminatorIndex = _numKeyColumns > _numGroupByExpressions ? _numGroupByExpressions : -1; + _valuesByOrdinal = + _groupingSets != null ? GroupingSets.groupingValuesByOrdinal(_groupingSets, argUnionIndexes) : null; } @Override @@ -566,11 +576,11 @@ public ColumnDataType getValueType() { @Override public Comparable extract(Record record) { - if (_discriminatorIndex < 0) { + if (_discriminatorIndex < 0 || _valuesByOrdinal == null) { return 0; } - int groupingId = ((Number) record.getValues()[_discriminatorIndex]).intValue(); - return GroupingSets.groupingValue(groupingId, _unionColumnIndexes); + int ordinal = ((Number) record.getValues()[_discriminatorIndex]).intValue(); + return _valuesByOrdinal[ordinal]; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupingSetsGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupingSetsGroupKeyGenerator.java index a9d93563aefd..ecfc812ad28c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupingSetsGroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupingSetsGroupKeyGenerator.java @@ -48,10 +48,11 @@ /// ``` /// /// where a column `c_i` that does NOT participate in `S` (rolled up) is pinned to the NULL sentinel, and the -/// trailing `groupingId` is a bitmask over the union with bit `i` set iff column `i` is rolled up in `S` -/// (PostgreSQL `GROUPING` semantics: 1 = aggregated away). The `groupingId` is appended as a synthetic key -/// column so that rows from different grouping sets never collide -- e.g. set `{a}` with `b` rolled up to -/// NULL stays distinct from set `{a, b}` where `b` is genuinely NULL, even though both render `(a, NULL)`. +/// trailing `groupingId` is the ordinal of `S` (its index in the query's grouping-set list). The ordinal is +/// appended as a synthetic key column so that rows from different grouping sets never collide -- e.g. set +/// `{a}` with `b` rolled up to NULL stays distinct from set `{a, b}` where `b` is genuinely NULL, even though +/// both render `(a, NULL)` -- and it identifies the producing set for GROUPING() / GROUPING_ID(). Because the +/// ordinal is not a per-column bitmask, the number of grouping columns is unlimited. /// /// This generator always emits multiple keys per row, so it is only used via the multi-value /// ({@code int[][]}) executor path. The union columns are resolved through per-column on-the-fly @@ -75,11 +76,10 @@ public class GroupingSetsGroupKeyGenerator implements GroupKeyGenerator { private final boolean _nullHandlingEnabled; private final int _numGroupsLimit; - /// Per grouping set: membership over the union columns, and the grouping-id bitmask (bit i set iff column i - /// is rolled up / excluded from the set). The bitmask is also the value stored in the synthetic - /// $groupingId key column and consumed by GROUPING() / GROUPING_ID(). + /// Per grouping set: membership over the union columns. The set's ordinal (its index in the query's + /// grouping-set list) is the value stored in the synthetic $groupingId key column, keeping rows from + /// different sets distinct and identifying the set for GROUPING() / GROUPING_ID(). private final boolean[][] _setContains; - private final int[] _setBitmasks; private final int _numSets; private final Object2IntOpenHashMap _groupKeyMap; @@ -110,18 +110,10 @@ public GroupingSetsGroupKeyGenerator(BaseProjectOperator projectOperator, _numSets = groupingSets.size(); _setContains = new boolean[_numSets][_numGroupByExpressions]; - _setBitmasks = new int[_numSets]; for (int s = 0; s < _numSets; s++) { for (int columnIndex : groupingSets.get(s)) { _setContains[s][columnIndex] = true; } - int bitmask = 0; - for (int i = 0; i < _numGroupByExpressions; i++) { - if (!_setContains[s][i]) { - bitmask |= 1 << i; - } - } - _setBitmasks[s] = bitmask; } _groupKeyMap = new Object2IntOpenHashMap<>(); @@ -164,7 +156,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { for (int col = 0; col < _numGroupByExpressions; col++) { keyValues[col] = _setContains[s][col] ? columnIds[col][row] : ID_FOR_NULL; } - keyValues[_numGroupByExpressions] = _setBitmasks[s]; + keyValues[_numGroupByExpressions] = s; int groupId = getGroupIdForKey(flyweightKey); if (groupId == numGroups) { /// A new group was inserted, so the map now retains this buffer; allocate a fresh one to reuse. @@ -188,15 +180,15 @@ private void generateKeysForBlockWithMv(ValueBlock valueBlock, int[][] groupKeys columnValueIds[col] = resolveColumnValueIds(valueBlock, col, numDocs); } int[][] keyComponents = new int[_numGroupByExpressions + 1][]; - int[] bitmaskComponent = new int[1]; - keyComponents[_numGroupByExpressions] = bitmaskComponent; + int[] ordinalComponent = new int[1]; + keyComponents[_numGroupByExpressions] = ordinalComponent; for (int row = 0; row < numDocs; row++) { IntArrayList rowGroupIds = new IntArrayList(_numSets); for (int s = 0; s < _numSets; s++) { for (int col = 0; col < _numGroupByExpressions; col++) { keyComponents[col] = _setContains[s][col] ? columnValueIds[col][row] : _nullComponent; } - bitmaskComponent[0] = _setBitmasks[s]; + ordinalComponent[0] = s; expandGroupIds(keyComponents, new int[_numGroupByExpressions + 1], 0, rowGroupIds); } groupKeys[row] = rowGroupIds.toIntArray(); @@ -381,14 +373,14 @@ public Iterator getGroupKeys() { } /// Reconstructs the output row for a composite key: the union column values (NULL where rolled up) followed - /// by the integer grouping-id discriminator. + /// by the integer grouping-set-ordinal discriminator. private Object[] buildKeysFromIds(FixedIntArray keyList) { int[] ids = keyList.elements(); Object[] keys = new Object[_numGroupByExpressions + 1]; for (int i = 0; i < _numGroupByExpressions; i++) { keys[i] = ids[i] == ID_FOR_NULL ? null : _onTheFlyDictionaries[i].get(ids[i]); } - /// The trailing slot stores the grouping-id bitmask directly (not a dictionary id). + /// The trailing slot stores the grouping-set ordinal directly (not a dictionary id). keys[_numGroupByExpressions] = ids[_numGroupByExpressions]; return keys; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/PostAggregationHandler.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/PostAggregationHandler.java index 0f6d9b16f42e..d57833e1cef2 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/PostAggregationHandler.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/PostAggregationHandler.java @@ -22,6 +22,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import javax.annotation.Nullable; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.request.context.FilterContext; @@ -48,9 +49,13 @@ public class PostAggregationHandler implements ValueExtractorFactory { /// Number of key columns (union group-by columns + the synthetic $groupingId column for grouping sets); /// aggregation columns start at this offset in the row. private final int _numKeyColumns; - /// Row index of the synthetic $groupingId discriminator column (read by GROUPING()/GROUPING_ID()), or -1 - /// when the query has no grouping sets. + /// Row index of the synthetic $groupingId discriminator column (the grouping-set ordinal, read by + /// GROUPING()/GROUPING_ID()), or -1 when the query has no grouping sets. private final int _groupingIdColumnIndex; + /// Per grouping set (in ordinal order), the participating union-column indexes; null when the query has no + /// grouping sets. + @Nullable + private final List _groupingSets; private final Map _groupByExpressionIndexMap; private final DataSchema _dataSchema; private final ValueExtractor[] _valueExtractors; @@ -72,6 +77,7 @@ public PostAggregationHandler(QueryContext queryContext, DataSchema dataSchema) } _numKeyColumns = _numGroupByExpressions + queryContext.getNumExtraGroupByKeyColumns(); _groupingIdColumnIndex = queryContext.isGroupingSets() ? _numGroupByExpressions : -1; + _groupingSets = queryContext.getGroupingSets(); // NOTE: The data schema will always have group-by expressions in the front, followed by aggregation functions of // the same order as in the query context. This is handled in AggregationGroupByOrderByOperator. @@ -201,26 +207,30 @@ public Object extract(Object[] row) { } } - /// Value extractor for {@code GROUPING(col, ...)} / {@code GROUPING_ID(col, ...)}. Computes the bitmask - /// from the synthetic {@code $groupingId} discriminator column: a bit is set (1) when the corresponding - /// argument column is rolled up (aggregated away) in the row's grouping set, following PostgreSQL - /// semantics with the first argument as the most significant bit. Returns 0 for every argument when the - /// query has no grouping sets (no column is rolled up). + /// Value extractor for {@code GROUPING(col, ...)} / {@code GROUPING_ID(col, ...)}. The value is fully + /// determined by the grouping set a row belongs to, so it is precomputed per grouping-set ordinal at + /// construction (a bit is 1 when the corresponding argument column is rolled up / aggregated away in that + /// set, PostgreSQL semantics, first argument = most significant bit); extraction is a single lookup on the + /// synthetic {@code $groupingId} ordinal column. Returns 0 for every argument when the query has no + /// grouping sets (no column is rolled up). private class GroupingValueExtractor implements ValueExtractor { private final FunctionContext _function; - private final int[] _unionColumnIndexes; + @Nullable + private final int[] _valuesByOrdinal; GroupingValueExtractor(FunctionContext function) { _function = function; List arguments = function.getArguments(); Preconditions.checkState(!arguments.isEmpty(), "GROUPING requires at least one argument"); - _unionColumnIndexes = new int[arguments.size()]; + int[] argUnionIndexes = new int[arguments.size()]; for (int i = 0; i < arguments.size(); i++) { ExpressionContext argument = arguments.get(i); Integer index = _groupByExpressionIndexMap != null ? _groupByExpressionIndexMap.get(argument) : null; Preconditions.checkState(index != null, "GROUPING argument must be a grouping column, got: %s", argument); - _unionColumnIndexes[i] = index; + argUnionIndexes[i] = index; } + _valuesByOrdinal = + _groupingSets != null ? GroupingSets.groupingValuesByOrdinal(_groupingSets, argUnionIndexes) : null; } @Override @@ -236,11 +246,11 @@ public ColumnDataType getColumnDataType() { @Override public Object extract(Object[] row) { /// Without grouping sets no column is rolled up, so GROUPING is always 0. - if (_groupingIdColumnIndex < 0) { + if (_groupingIdColumnIndex < 0 || _valuesByOrdinal == null) { return 0; } - int groupingId = ((Number) row[_groupingIdColumnIndex]).intValue(); - return GroupingSets.groupingValue(groupingId, _unionColumnIndexes); + int ordinal = ((Number) row[_groupingIdColumnIndex]).intValue(); + return _valuesByOrdinal[ordinal]; } } } 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 ec165cf8f7ea..07b2b899709e 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 @@ -114,40 +114,33 @@ public static QueryContext getQueryContext(PinotQuery pinotQuery) { } } - /// GROUP BY GROUPING SETS / ROLLUP / CUBE: the wire carries one membership bitmask per set over - /// groupByExpressions (the union of all grouping columns). Decode each mask into the sorted list of - /// participating column indexes; a mask of 0 yields the empty (grand-total) set (). Null when this is a - /// plain GROUP BY query. + /// GROUP BY GROUPING SETS / ROLLUP / CUBE: the wire carries, per grouping set (in ordinal order), the + /// list of participating column indexes into groupByExpressions (the union of all grouping columns). An + /// empty inner list is the grand-total set (). Null when this is a plain GROUP BY query. List groupingSets = null; - List groupingSetMasks = pinotQuery.getGroupingSetMasks(); - if (groupingSetMasks != null) { - // Guard against malformed requests (the parser always emits at least one set over a union of at most 31 - // columns): an empty set list would silently produce empty results, and a union of more than 31 columns - // cannot be represented in the synthetic INT grouping-id bitmask (bit 31 is the sign bit). - if (groupingSetMasks.isEmpty()) { - throw new IllegalStateException("Grouping set masks must not be empty"); + List> thriftGroupingSets = pinotQuery.getGroupingSets(); + // A grouping-sets clause with no grouping columns is the grand-total-only set () — semantically identical to a + // plain aggregation (a single group over all rows, with GROUPING() resolving to 0). Drop the marker so the + // no-group-by aggregation path runs; otherwise the reducer would expect a $groupingId key column the server + // never emits, and index past the result schema. + if (thriftGroupingSets != null && groupByExpressions != null) { + // Guard against malformed requests: an empty set list would silently produce empty results, and an + // out-of-range column index would cause out-of-bounds access deep in the group key generator. + if (thriftGroupingSets.isEmpty()) { + throw new IllegalStateException("Grouping sets must not be empty"); } - int numGroupByExpressions = groupByExpressions != null ? groupByExpressions.size() : 0; - if (numGroupByExpressions > CalciteSqlParser.MAX_GROUPING_SETS_COLUMNS) { - throw new IllegalStateException( - "Cannot use grouping sets over " + numGroupByExpressions + " group-by expressions (max: " - + CalciteSqlParser.MAX_GROUPING_SETS_COLUMNS + ")"); - } - groupingSets = new ArrayList<>(groupingSetMasks.size()); - for (int mask : groupingSetMasks) { - // Every bit must reference an existing union column, otherwise the decoded indexes would cause - // out-of-bounds access deep in the group key generator. The union size is capped at 31 above, so the - // unsigned shift also rejects negative masks (bit 31 set). - if ((mask >>> numGroupByExpressions) != 0) { - throw new IllegalStateException( - "Invalid grouping set mask: " + mask + " for " + numGroupByExpressions + " group-by expressions"); - } - int[] indexes = new int[Integer.bitCount(mask)]; + int numGroupByExpressions = groupByExpressions.size(); + groupingSets = new ArrayList<>(thriftGroupingSets.size()); + for (List set : thriftGroupingSets) { + int[] indexes = new int[set.size()]; int idx = 0; - for (int bit = 0; bit < Integer.SIZE && idx < indexes.length; bit++) { - if ((mask & (1 << bit)) != 0) { - indexes[idx++] = bit; + for (int columnIndex : set) { + if (columnIndex < 0 || columnIndex >= numGroupByExpressions) { + throw new IllegalStateException( + "Invalid grouping set column index: " + columnIndex + " for " + numGroupByExpressions + + " group-by expressions"); } + indexes[idx++] = columnIndex; } groupingSets.add(indexes); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/request/context/utils/BrokerRequestToQueryContextConverterTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/request/context/utils/BrokerRequestToQueryContextConverterTest.java index 3639d470c347..620919359667 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/request/context/utils/BrokerRequestToQueryContextConverterTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/request/context/utils/BrokerRequestToQueryContextConverterTest.java @@ -22,6 +22,7 @@ import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.common.request.PinotQuery; import org.apache.pinot.common.request.context.ExpressionContext; @@ -811,36 +812,49 @@ public void testConstantFilter() { } @Test - public void testInvalidGroupingSetMasksRejected() { + public void testInvalidGroupingSetsRejected() { /// A well-formed grouping-set query decodes into per-set column-index arrays. PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery("SELECT COUNT(*) FROM testTable GROUP BY ROLLUP(a, b)"); assertNotNull(QueryContextConverterUtils.getQueryContext(pinotQuery).getGroupingSets()); - /// A mask referencing a non-existent union column (bit >= groupByList size) or with the sign bit set must be - /// rejected at conversion instead of causing out-of-bounds access in the group key generator. - for (int badMask : new int[]{1 << 2, -1}) { + /// A set referencing a non-existent union column (index >= groupByList size, or negative) must be rejected + /// at conversion instead of causing out-of-bounds access in the group key generator. + for (int badIndex : new int[]{2, -1}) { PinotQuery corrupted = CalciteSqlParser.compileToPinotQuery("SELECT COUNT(*) FROM testTable GROUP BY ROLLUP(a, b)"); - corrupted.setGroupingSetMasks(List.of(3, badMask)); + corrupted.setGroupingSets(List.of(List.of(0, 1), List.of(badIndex))); assertThrows(IllegalStateException.class, () -> QueryContextConverterUtils.getQueryContext(corrupted)); } /// An empty set list must be rejected rather than silently producing zero groups. - PinotQuery emptyMasks = + PinotQuery emptySets = CalciteSqlParser.compileToPinotQuery("SELECT COUNT(*) FROM testTable GROUP BY ROLLUP(a, b)"); - emptyMasks.setGroupingSetMasks(List.of()); - assertThrows(IllegalStateException.class, () -> QueryContextConverterUtils.getQueryContext(emptyMasks)); + emptySets.setGroupingSets(List.of()); + assertThrows(IllegalStateException.class, () -> QueryContextConverterUtils.getQueryContext(emptySets)); + } - /// A group-by union of more than 31 columns cannot be represented in the INT grouping-id bitmask; the - /// parser never emits this (it caps the union at 31), so it must be rejected as malformed. + @Test + public void testManyGroupingColumnsSupported() { + /// The number of grouping columns is unlimited (the discriminator is the set ordinal, not a per-column + /// bitmask): a 40-column ROLLUP compiles and converts, where the retired 32-bit mask encoding capped the + /// union at 31 columns. StringBuilder columns = new StringBuilder(); - for (int i = 0; i < 32; i++) { + for (int i = 0; i < 40; i++) { columns.append(i == 0 ? "c0" : ", c" + i); } - PinotQuery wideUnion = - CalciteSqlParser.compileToPinotQuery("SELECT COUNT(*) FROM testTable GROUP BY " + columns); - wideUnion.setGroupingSetMasks(List.of(0)); - assertThrows(IllegalStateException.class, () -> QueryContextConverterUtils.getQueryContext(wideUnion)); + PinotQuery pinotQuery = + CalciteSqlParser.compileToPinotQuery("SELECT COUNT(*) FROM testTable GROUP BY ROLLUP(" + columns + ")"); + List groupingSets = QueryContextConverterUtils.getQueryContext(pinotQuery).getGroupingSets(); + assertNotNull(groupingSets); + /// ROLLUP(c0..c39) -> the 41 prefixes (c0..c39), (c0..c38), ..., (c0), (). + assertEquals(groupingSets.size(), 41); + Set seenSizes = new HashSet<>(); + for (int[] set : groupingSets) { + assertTrue(seenSizes.add(set.length), "duplicate prefix length: " + set.length); + for (int columnIndex : set) { + assertTrue(columnIndex >= 0 && columnIndex < 40, "column index out of range: " + columnIndex); + } + } } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GroupingSetsQueriesTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GroupingSetsQueriesTest.java index 86f34f579d58..5ae01cf7faa9 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GroupingSetsQueriesTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GroupingSetsQueriesTest.java @@ -156,7 +156,8 @@ public void testV2PhysicalPlannerRollup(boolean useMultiStageQueryEngine) /// With usePhysicalOptimizer the multi-stage v2 physical planner splits the grouping-set aggregate itself; verify /// it returns the same rows, including the genuine (a, NULL) detail vs the rolled-up (a, NULL) subtotal kept /// distinct via $groupingId. The single-stage engine ignores usePhysicalOptimizer and runs the same ROLLUP, so - /// this also cross-checks both engines agree. (GROUPING() / GROUPING_ID() are not yet supported on the v2 planner.) + /// this also cross-checks both engines agree. (GROUPING() / GROUPING_ID() on the v2 planner are covered by + /// testV2PhysicalPlannerGrouping.) setUseMultiStageQueryEngine(useMultiStageQueryEngine); String query = "SET usePhysicalOptimizer=true; SET enableNullHandling=true; SELECT " + D1 + ", " + D2 + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ", " + D2 + ")"; @@ -268,6 +269,113 @@ public void testMultiStageGroupingOverJoin(boolean useMultiStageQueryEngine) assertEquals(actual, expected); } + @Test(dataProvider = "useV2QueryEngine") + public void testMultiStageAggregationOfGroupingColumnOverJoin(boolean useMultiStageQueryEngine) + throws Exception { + /// Regression: an aggregation argument may BE a grouping column. Under ROLLUP(d1, met) the (d1) subtotals + /// and the grand total roll up met, but SUM(met) must still aggregate the real met values — the runtime + /// per-set expansion must not NULL the aggregation input in place (it groups by separate key copies). + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + String query = "SELECT t1." + D1 + ", t1." + MET + ", SUM(t1." + MET + ") FROM " + getTableName() + " t1 JOIN " + + getTableName() + " t2 ON t1." + D1 + " = t2." + D1 + " GROUP BY ROLLUP(t1." + D1 + ", t1." + MET + ")"; + JsonNode response = postQuery("SET enableNullHandling=true; " + query); + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions == null || exceptions.isEmpty(), "query failed: " + response.toPrettyString()); + Map actual = new HashMap<>(); + for (JsonNode row : response.get("resultTable").get("rows")) { + actual.put(cell(row, 0) + "|" + cell(row, 1), row.get(2).asLong()); + } + /// The self-join on d1 multiplies each d1 group 4x4=16; met = 1 on every row. + Map expected = new HashMap<>(); + expected.put("a|1", 16L); + expected.put("b|1", 16L); + /// The (d1) subtotals: met rolled up to NULL in the key, but SUM(met) still 16 — NOT null/0. + expected.put("a|NULL", 16L); + expected.put("b|NULL", 16L); + expected.put("NULL|NULL", 32L); + assertEquals(actual, expected, response.toPrettyString()); + } + + @Test(dataProvider = "useV2QueryEngine") + public void testV2PhysicalPlannerRejectsGroupingSetsOverJoin(boolean useMultiStageQueryEngine) + throws Exception { + /// The v2 physical planner only splits a grouping-set aggregate when its input is a repartitioning exchange; + /// over a JOIN there is none, so it rejects explicitly (a clean plan-time error) rather than producing a broken + /// plan. The default planner handles the same query (see testMultiStageRollupOverJoin). + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + String query = "SELECT t1." + D1 + ", COUNT(*) FROM " + getTableName() + " t1 JOIN " + getTableName() + " t2 ON t1." + + D1 + " = t2." + D1 + " GROUP BY ROLLUP(t1." + D1 + ")"; + JsonNode response = postQuery("SET usePhysicalOptimizer=true; SET enableNullHandling=true; " + query); + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions != null && !exceptions.isEmpty(), "expected a rejection: " + response.toPrettyString()); + assertTrue(exceptions.toString().contains("GROUP BY GROUPING SETS"), + "expected a grouping-sets rejection: " + response.toPrettyString()); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testManyGroupingColumnsEndToEnd(boolean useMultiStageQueryEngine) + throws Exception { + /// The 31-column limit is gone end-to-end: ROLLUP over 33 grouping expressions (d1, d2 and 31 met-derived + /// expressions) executes on both engines without hitting a column-count cap. Assert only what is robust + /// regardless of how the engine folds the constant expressions: it does not error, and the grand-total set () + /// produces the max-count row (all 8 docs). ORDER BY count DESC pins that row first (it is the unique maximum), + /// so the default result limit cannot truncate it away. + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + StringBuilder rollup = new StringBuilder(D1 + ", " + D2); + for (int i = 0; i < 31; i++) { + rollup.append(", ").append(MET).append(" + ").append(i); + } + String query = + "SELECT COUNT(*) AS c FROM " + getTableName() + " GROUP BY ROLLUP(" + rollup + ") ORDER BY c DESC"; + JsonNode response = postQuery("SET enableNullHandling=true; " + query); + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions == null || exceptions.isEmpty(), "query failed: " + response.toPrettyString()); + JsonNode rows = response.get("resultTable").get("rows"); + assertTrue(rows.size() > 0, response.toPrettyString()); + assertEquals(rows.get(0).get(0).asLong(), 8L, + "grand-total row counting all 8 docs must sort first: " + response.toPrettyString()); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testDuplicateGroupingSetsValues(boolean useMultiStageQueryEngine) + throws Exception { + /// Explicit duplicate grouping sets do not corrupt the aggregation values: {d1} yields a=4/b=4 and () yields + /// NULL=8 on both engines. NOTE: whether the duplicate (d1) set is de-duplicated (affecting the row COUNT) is + /// not pinned here — the single-stage parser collapses duplicates while the multi-stage path may keep them + /// (standard SQL keeps them); reconciling that is tracked separately. Asserting the value map (keyed by d1) + /// is robust to either choice. + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + String query = "SELECT " + D1 + ", COUNT(*) FROM " + getTableName() + " GROUP BY GROUPING SETS ((" + D1 + "), (" + + D1 + "), ())"; + JsonNode response = postQuery("SET enableNullHandling=true; " + query); + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions == null || exceptions.isEmpty(), "query failed: " + response.toPrettyString()); + Map actual = new HashMap<>(); + for (JsonNode row : response.get("resultTable").get("rows")) { + actual.put(cell(row, 0), row.get(1).asLong()); + } + Map expected = new HashMap<>(); + expected.put("a", 4L); + expected.put("b", 4L); + expected.put("NULL", 8L); + assertEquals(actual, expected, response.toPrettyString()); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testGrandTotalOnlyGroupingSet(boolean useMultiStageQueryEngine) + throws Exception { + /// GROUP BY GROUPING SETS (()) — the grand total as the only set, with an empty grouping-column union — is + /// a QueryContext shape no other query produces (grouping sets present, no group-by expressions). + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + String query = "SELECT COUNT(*) FROM " + getTableName() + " GROUP BY GROUPING SETS (())"; + JsonNode response = postQuery("SET enableNullHandling=true; " + query); + JsonNode exceptions = response.get("exceptions"); + assertTrue(exceptions == null || exceptions.isEmpty(), "query failed: " + response.toPrettyString()); + JsonNode rows = response.get("resultTable").get("rows"); + assertEquals(rows.size(), 1, response.toPrettyString()); + assertEquals(rows.get(0).get(0).asLong(), 8L); + } + @Test(dataProvider = "useBothQueryEngines") public void testRollupWithGenuineAndRolledUpNulls(boolean useMultiStageQueryEngine) throws Exception { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalAggregate.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalAggregate.java index df3486e3fcf9..c7782a102497 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalAggregate.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalAggregate.java @@ -30,7 +30,7 @@ import org.apache.calcite.rel.hint.RelHint; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.util.ImmutableBitSet; -import org.apache.pinot.common.request.context.GroupingSets; +import org.apache.pinot.calcite.rel.rules.GroupingSetsPlanUtils; import org.apache.pinot.query.planner.plannode.AggregateNode.AggType; @@ -107,7 +107,7 @@ protected RelDataType deriveRowType() { if (!emitsGroupingId()) { return rowType; } - return GroupingSets.appendGroupingIdColumn(getCluster().getTypeFactory(), rowType, getGroupCount()); + return GroupingSetsPlanUtils.appendGroupingIdColumn(getCluster().getTypeFactory(), rowType, getGroupCount()); } @Override diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsPlanUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsPlanUtils.java new file mode 100644 index 000000000000..60900e54bd7a --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsPlanUtils.java @@ -0,0 +1,198 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.calcite.rel.rules; + +import com.google.common.base.Preconditions; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.Aggregate; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.ImmutableBitSet; +import org.apache.pinot.common.request.context.GroupingSets; + + +/// The shared planner-side logic for GROUP BY GROUPING SETS / ROLLUP / CUBE, used by both the default planner +/// ({@link PinotAggregateExchangeNodeInsertRule}) and the v2 physical planner ({@code AggregatePushdownRule}) +/// so the two cannot drift on the encoding: the grouping sets of a query are the ordered, de-duplicated list +/// produced by {@link #computeGroupingSets}, a set's position in that list is its ordinal (the value of the +/// synthetic {@link GroupingSets#GROUPING_ID_COLUMN} discriminator), and GROUPING()/GROUPING_ID() values are +/// derived from set membership via {@link GroupingSets#groupingValuesByOrdinal}. +public class GroupingSetsPlanUtils { + private GroupingSetsPlanUtils() { + } + + /// Encodes each grouping set as the sorted list of its participating union group-by column indexes + /// (positions in {@code groupSet.asList()}), in {@code getGroupSets()} order, de-duplicated (first + /// occurrence wins, matching the single-stage parser: duplicate sets — reachable via explicit + /// {@code GROUPING SETS ((a), (a))} — collapse to one). A set's position in the returned list is its + /// ordinal, the value carried as the synthetic $groupingId discriminator. Mirrors the single-stage engine's + /// {@code PinotQuery.groupingSets} so the per-set expansion can be pushed down to the single-stage (leaf) + /// engine, and mirrors Calcite's per-set column bitset so the number of grouping columns is unlimited. + /// Returns an empty list for a plain GROUP BY (SIMPLE). This is the single conversion point for both + /// planners: every consumer of the ordinal (planner GROUPING projections, RepeatOperator, leaf pushdown) + /// derives the set order from the list produced here. + public static List> computeGroupingSets(Aggregate node) { + if (node.getGroupType() == Aggregate.Group.SIMPLE) { + return List.of(); + } + /// Same cap as the single-stage parser: every set is materialized in the plan, expanded per input row at + /// runtime, and enumerated by the GROUPING()/GROUPING_ID() projection. + if (node.getGroupSets().size() > GroupingSets.MAX_GROUPING_SETS) { + throw new UnsupportedOperationException( + "GROUPING SETS / ROLLUP / CUBE expands to more than " + GroupingSets.MAX_GROUPING_SETS + + " grouping sets"); + } + List union = node.getGroupSet().asList(); + Map unionIndex = new HashMap<>(); + for (int i = 0; i < union.size(); i++) { + unionIndex.put(union.get(i), i); + } + Set> seen = new LinkedHashSet<>(); + for (ImmutableBitSet set : node.getGroupSets()) { + List columnIndexes = new ArrayList<>(set.cardinality()); + for (int bit : set) { + columnIndexes.add(unionIndex.get(bit)); + } + Collections.sort(columnIndexes); + seen.add(columnIndexes); + } + return new ArrayList<>(seen); + } + + /// Returns the LEAF row type for a grouping-set aggregate: the synthetic + /// {@link GroupingSets#GROUPING_ID_COLUMN} INT column inserted right after the {@code groupCount} union + /// group-by columns, i.e. {@code [group keys..., $groupingId, aggregates...]}. Shared by + /// {@code PinotLogicalAggregate} and {@code PhysicalAggregate} so the two planners' {@code deriveRowType()} + /// overrides cannot drift. + public static RelDataType appendGroupingIdColumn(RelDataTypeFactory typeFactory, RelDataType rowType, + int groupCount) { + RelDataTypeFactory.Builder builder = typeFactory.builder(); + List fields = rowType.getFieldList(); + for (int i = 0; i < groupCount; i++) { + builder.add(fields.get(i)); + } + builder.add(GroupingSets.GROUPING_ID_COLUMN, typeFactory.createSqlType(SqlTypeName.INTEGER)); + for (int i = groupCount; i < fields.size(); i++) { + builder.add(fields.get(i)); + } + return builder.build(); + } + + /// Splits GROUPING() / GROUPING_ID() out of the aggregate calls: they are not real aggregations (they are + /// functions of which grouping set a row belongs to, computed from $groupingId in the final projection). + /// Adds the real (non-GROUPING) calls to {@code outRealAggCalls} and returns, per original call, its index + /// among the real calls or -1 for a GROUPING / GROUPING_ID call. + public static int[] splitOutGroupingCalls(List orgAggCalls, List outRealAggCalls) { + int[] realAggIndex = new int[orgAggCalls.size()]; + for (int i = 0; i < orgAggCalls.size(); i++) { + SqlKind kind = orgAggCalls.get(i).getAggregation().getKind(); + if (kind == SqlKind.GROUPING || kind == SqlKind.GROUPING_ID) { + realAggIndex[i] = -1; + } else { + realAggIndex[i] = outRealAggCalls.size(); + outRealAggCalls.add(orgAggCalls.get(i)); + } + } + return realAggIndex; + } + + /// Builds the projection expressions restoring the original grouping-set aggregate row type on top of the + /// FINAL aggregate, whose output is {@code [union keys..., $groupingId, real aggregate results...]}: the + /// group keys, then per original aggregate call either the real aggregate result reference or the + /// GROUPING() / GROUPING_ID() value computed from $groupingId. {@code $groupingId} itself is dropped. + /// + /// @param finalAggRel the FINAL aggregate node the projection reads from + /// @param aggRel the original grouping-set aggregate (provides union order, aggregate calls, and sets) + /// @param realAggIndex per original aggregate call, its index among the real calls or -1 for GROUPING calls + /// (from {@link #splitOutGroupingCalls}) + public static List buildGroupingSetsProjects(RexBuilder rexBuilder, RelNode finalAggRel, Aggregate aggRel, + int[] realAggIndex) { + int groupCount = aggRel.getGroupCount(); + int finalGroupCount = groupCount + 1; + RelDataType intType = finalAggRel.getCluster().getTypeFactory().createSqlType(SqlTypeName.INTEGER); + /// $groupingId is the INT discriminator column immediately after the union group keys. + RexNode groupingIdRef = rexBuilder.makeInputRef(finalAggRel, groupCount); + List union = aggRel.getGroupSet().asList(); + /// The ordinal order of the sets — the same list the runtime receives on the wire. + List> groupingSets = computeGroupingSets(aggRel); + List orgAggCalls = aggRel.getAggCallList(); + List projects = new ArrayList<>(groupCount + orgAggCalls.size()); + for (int i = 0; i < groupCount; i++) { + projects.add(rexBuilder.makeInputRef(finalAggRel, i)); + } + for (int i = 0; i < orgAggCalls.size(); i++) { + if (realAggIndex[i] >= 0) { + projects.add(rexBuilder.makeInputRef(finalAggRel, finalGroupCount + realAggIndex[i])); + } else { + AggregateCall groupingCall = orgAggCalls.get(i); + /// Map each GROUPING argument (an input column index) to its position in the union group key list (the + /// index space the grouping sets are expressed in). + List unionIndexes = new ArrayList<>(groupingCall.getArgList().size()); + for (int arg : groupingCall.getArgList()) { + int unionIndex = union.indexOf(arg); + Preconditions.checkState(unionIndex >= 0, "GROUPING argument must be a grouping column"); + unionIndexes.add(unionIndex); + } + RexNode value = buildGroupingValue(rexBuilder, intType, groupingIdRef, groupingSets, unionIndexes); + projects.add(rexBuilder.makeCast(groupingCall.getType(), value)); + } + } + return projects; + } + + /// Builds the GROUPING / GROUPING_ID value expression. The value is fully determined by the grouping set a + /// row belongs to, so it is a plan-time constant per grouping-set ordinal: the expression is + /// {@code CASE WHEN $groupingId = 0 THEN v_0 WHEN $groupingId = 1 THEN v_1 ... ELSE null END}, with each + /// {@code v_k} packed from the argument columns' rolled-up bits in set {@code k} (first argument = most + /// significant bit, PostgreSQL semantics, see {@link GroupingSets#groupingValuesByOrdinal}). + /// + /// The projection evaluates the CASE per output group row, so the per-row cost is O(number of grouping + /// sets), bounded by {@link GroupingSets#MAX_GROUPING_SETS} (enforced in {@link #computeGroupingSets}). + /// TODO: replace the CASE with an O(1) per-ordinal lookup (an INT-array literal indexed by $groupingId, + /// mirroring the single-stage reduce-side precompute) once array literals are supported in the plan + /// expression conversion. + static RexNode buildGroupingValue(RexBuilder rexBuilder, RelDataType intType, RexNode groupingIdRef, + List> groupingSets, List argUnionIndexes) { + int[] valuesByOrdinal = GroupingSets.groupingValuesByOrdinal(groupingSets, argUnionIndexes); + int numSets = valuesByOrdinal.length; + List caseOperands = new ArrayList<>(2 * numSets + 1); + for (int ordinal = 0; ordinal < numSets; ordinal++) { + caseOperands.add(rexBuilder.makeCall(SqlStdOperatorTable.EQUALS, groupingIdRef, + rexBuilder.makeExactLiteral(BigDecimal.valueOf(ordinal), intType))); + caseOperands.add(rexBuilder.makeExactLiteral(BigDecimal.valueOf(valuesByOrdinal[ordinal]), intType)); + } + caseOperands.add(rexBuilder.makeNullLiteral(intType)); + return rexBuilder.makeCall(intType, SqlStdOperatorTable.CASE, caseOperands); + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsRexUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsRexUtils.java deleted file mode 100644 index e7e7e0eebb4f..000000000000 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/GroupingSetsRexUtils.java +++ /dev/null @@ -1,84 +0,0 @@ -/** - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -package org.apache.pinot.calcite.rel.rules; - -import com.google.common.base.Preconditions; -import java.math.BigDecimal; -import java.util.ArrayList; -import java.util.List; -import org.apache.calcite.rel.type.RelDataType; -import org.apache.calcite.rex.RexBuilder; -import org.apache.calcite.rex.RexNode; -import org.apache.calcite.sql.SqlFunctionCategory; -import org.apache.calcite.sql.SqlIdentifier; -import org.apache.calcite.sql.SqlOperator; -import org.apache.calcite.sql.SqlSyntax; -import org.apache.calcite.sql.parser.SqlParserPos; -import org.apache.calcite.sql.validate.SqlNameMatchers; -import org.apache.pinot.calcite.sql.fun.PinotOperatorTable; - - -/// Builds the {@code GROUPING(args...)} / {@code GROUPING_ID(args...)} value as a bit expression over the synthetic -/// {@code $groupingId} discriminator column. Shared by the default planner -/// ({@link PinotAggregateExchangeNodeInsertRule}) and the v2 physical planner ({@code AggregatePushdownRule}) so both -/// compute the value identically, mirroring {@link org.apache.pinot.common.request.context.GroupingSets#groupingValue}. -public class GroupingSetsRexUtils { - private GroupingSetsRexUtils() { - } - - /// Builds the GROUPING / GROUPING_ID value: for each argument (identified by its union-column index) extract its bit - /// from {@code groupingIdRef}, and pack the bits with the first argument as the most significant bit. Uses Pinot's - /// bit scalar functions, evaluated in the projection operator. - /// - /// @param groupingIdRef reference to the {@code $groupingId} INT column (bit i set iff union column i is rolled up) - /// @param unionIndexes the union-column index of each GROUPING argument, in argument order - public static RexNode buildGroupingValue(RexBuilder rexBuilder, RelDataType intType, RexNode groupingIdRef, - List unionIndexes) { - SqlOperator bitAnd = pinotScalarOperator("bitAnd"); - SqlOperator bitShiftRightUnsigned = pinotScalarOperator("bitShiftRightUnsigned"); - SqlOperator bitShiftLeft = pinotScalarOperator("bitShiftLeft"); - SqlOperator bitOr = pinotScalarOperator("bitOr"); - int numArgs = unionIndexes.size(); - RexNode result = null; - for (int j = 0; j < numArgs; j++) { - RexNode k = rexBuilder.makeExactLiteral(BigDecimal.valueOf(unionIndexes.get(j)), intType); - /// bit = (groupingId >>> k) & 1 - RexNode shifted = rexBuilder.makeCall(intType, bitShiftRightUnsigned, List.of(groupingIdRef, k)); - RexNode bit = rexBuilder.makeCall(intType, bitAnd, - List.of(shifted, rexBuilder.makeExactLiteral(BigDecimal.ONE, intType))); - /// Pack with the first argument as the MSB: shift left by (numArgs - 1 - j). - int shiftLeftBy = numArgs - 1 - j; - RexNode placed = shiftLeftBy == 0 ? bit : rexBuilder.makeCall(intType, bitShiftLeft, - List.of(bit, rexBuilder.makeExactLiteral(BigDecimal.valueOf(shiftLeftBy), intType))); - result = result == null ? placed : rexBuilder.makeCall(intType, bitOr, List.of(result, placed)); - } - return result; - } - - /// Looks up a Pinot scalar function as a Calcite {@link SqlOperator} by name, for building bit expressions in the - /// grouping-set final projection. - private static SqlOperator pinotScalarOperator(String name) { - List operators = new ArrayList<>(1); - PinotOperatorTable.instance(false).lookupOperatorOverloads(new SqlIdentifier(name, SqlParserPos.ZERO), - SqlFunctionCategory.USER_DEFINED_FUNCTION, SqlSyntax.FUNCTION, operators, - SqlNameMatchers.withCaseSensitive(false)); - Preconditions.checkState(!operators.isEmpty(), "Pinot scalar function not found: %s", name); - return operators.get(0); - } -} 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 de48c427ac6a..95efc01d1b7a 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 @@ -18,7 +18,6 @@ */ package org.apache.pinot.calcite.rel.rules; -import com.google.common.base.Preconditions; import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -53,7 +52,6 @@ import org.apache.calcite.sql.type.ReturnTypes; import org.apache.calcite.sql.type.SqlOperandTypeChecker; import org.apache.calcite.sql.type.SqlReturnTypeInference; -import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.tools.RelBuilder; import org.apache.calcite.tools.RelBuilderFactory; import org.apache.calcite.util.ImmutableBitSet; @@ -68,7 +66,6 @@ import org.apache.pinot.calcite.rel.logical.PinotLogicalExchange; import org.apache.pinot.calcite.rel.logical.PinotLogicalSortExchange; import org.apache.pinot.common.function.sql.PinotSqlAggFunction; -import org.apache.pinot.common.request.context.GroupingSets; import org.apache.pinot.query.QueryEnvironment; import org.apache.pinot.query.planner.plannode.AggregateNode.AggType; import org.apache.pinot.segment.spi.AggregationFunctionType; @@ -236,6 +233,36 @@ private static RelNode createPlan(RelOptRuleCall call, Aggregate aggRel, boolean Map hintOptions, @Nullable List collations, int limit) { // WITHIN GROUP collation is not supported in leaf stage aggregation. RelCollation withinGroupCollation = extractWithinGroupCollation(aggRel); + if (aggRel.getGroupType() != Aggregate.Group.SIMPLE) { + /// GROUP BY GROUPING SETS / ROLLUP / CUBE only supports the LEAF+EXCHANGE+FINAL split: the DIRECT paths + /// (WITHIN GROUP ordered aggregates and the skip-leaf-stage / partitioned-by hints) would carry the grouping + /// sets into a plan shape whose schema and exchange keys do not account for the per-set expansion, producing + /// broken plans at runtime. Reject explicitly; the v2 physical planner rejects the same combinations. + if (withinGroupCollation != null) { + throw new UnsupportedOperationException( + "WITHIN GROUP ordered aggregates with GROUP BY GROUPING SETS / ROLLUP / CUBE are not supported in the " + + "multi-stage query engine. Run the query on the single-stage query engine instead."); + } + if (hasGroupBy && (Boolean.parseBoolean( + hintOptions.get(PinotHintOptions.AggregateOptions.IS_SKIP_LEAF_STAGE_GROUP_BY)) || Boolean.parseBoolean( + hintOptions.get(PinotHintOptions.AggregateOptions.IS_PARTITIONED_BY_GROUP_BY_KEYS)))) { + throw new UnsupportedOperationException( + "Aggregate hints are not supported with GROUP BY GROUPING SETS / ROLLUP / CUBE in the multi-stage " + + "query engine"); + } + return createPlanWithLeafExchangeFinalAggregate(aggRel, false, collations, limit); + } + /// GROUPING() / GROUPING_ID() with a plain GROUP BY is constant 0 per the SQL standard; it is not wired into + /// the multi-stage runtime (it is not a real aggregation function), so reject it explicitly rather than failing + /// with an obscure error at execution. TODO: fold it to the literal 0 for single-stage parity. + for (AggregateCall aggCall : aggRel.getAggCallList()) { + SqlKind kind = aggCall.getAggregation().getKind(); + if (kind == SqlKind.GROUPING || kind == SqlKind.GROUPING_ID) { + throw new UnsupportedOperationException( + "GROUPING() / GROUPING_ID() requires GROUP BY GROUPING SETS / ROLLUP / CUBE in the multi-stage query " + + "engine. Run the query on the single-stage query engine instead."); + } + } if (withinGroupCollation != null || (hasGroupBy && Boolean.parseBoolean( hintOptions.get(PinotHintOptions.AggregateOptions.IS_SKIP_LEAF_STAGE_GROUP_BY)))) { return createPlanWithExchangeDirectAggregation(call, aggRel, withinGroupCollation, collations, limit); @@ -323,28 +350,17 @@ private static RelNode createPlanWithLeafExchangeFinalAggregate(Aggregate aggRel /// combine/reduce trim for them anyway; add it when leaf per-set trim is wired through. private static RelNode createGroupingSetsPlanWithLeafExchangeFinalAggregate(Aggregate aggRel) { int groupCount = aggRel.getGroupCount(); - /// The single-stage leaf encodes the grouping set as a 32-bit $groupingId bitmask over the union columns. - if (groupCount > GroupingSets.MAX_GROUPING_SET_COLUMNS) { - throw new UnsupportedOperationException( - "GROUP BY GROUPING SETS / ROLLUP / CUBE supports at most " + GroupingSets.MAX_GROUPING_SET_COLUMNS - + " distinct grouping columns in the multi-stage query engine, got " + groupCount); - } - /// GROUPING() / GROUPING_ID() are not real aggregations: they are functions of which grouping set a row belongs - /// to. Like the single-stage post-aggregation handler, they are computed from $groupingId in the final projection, - /// so they are split out of the LEAF/FINAL aggregate calls here. + /// GROUPING() / GROUPING_ID() are not real aggregations: they are computed from $groupingId in the final + /// projection, so they are split out of the LEAF/FINAL aggregate calls here (shared with the v2 planner). List orgAggCalls = aggRel.getAggCallList(); List realAggCalls = new ArrayList<>(orgAggCalls.size()); - /// For each original aggregate call: its index among the real (non-GROUPING) aggregates, or -1 if it is a - /// GROUPING / GROUPING_ID call (computed in the projection instead). - int[] realAggIndex = new int[orgAggCalls.size()]; - for (int i = 0; i < orgAggCalls.size(); i++) { - SqlKind kind = orgAggCalls.get(i).getAggregation().getKind(); - if (kind == SqlKind.GROUPING || kind == SqlKind.GROUPING_ID) { - realAggIndex[i] = -1; - } else { - realAggIndex[i] = realAggCalls.size(); - realAggCalls.add(orgAggCalls.get(i)); - } + int[] realAggIndex = GroupingSetsPlanUtils.splitOutGroupingCalls(orgAggCalls, realAggCalls); + /// Without a real aggregation the LEAF would push an aggregation-free query down to the single-stage engine, + /// which would execute it as a plain selection and ignore the grouping sets. Mirrors the single-stage parser + /// rejection (GROUPING()/GROUPING_ID() are not aggregations). + if (realAggCalls.isEmpty()) { + throw new UnsupportedOperationException( + "GROUP BY GROUPING SETS / ROLLUP / CUBE requires at least one aggregation function in the query"); } /// LEAF: carries the original groupSets; deriveRowType() appends $groupingId after the union group keys. Only the /// real aggregations are pushed down, so the leaf output is [union keys..., $groupingId, real aggregates...]. @@ -362,45 +378,17 @@ private static RelNode createGroupingSetsPlanWithLeafExchangeFinalAggregate(Aggr finalGroupCount, exchange.getRowType(), AggType.FINAL, false, null, 0); /// PROJECT: emit the original row type — group keys, then per original aggregate call either a real aggregate /// result reference or a GROUPING/GROUPING_ID value computed from $groupingId — dropping $groupingId itself. - return buildGroupingSetsProject(aggRel, finalAggRel, groupCount, realAggIndex); + return buildGroupingSetsProject(aggRel, finalAggRel, realAggIndex); } - /// Builds the Project on top of the FINAL grouping-set aggregate. The FINAL output is - /// {@code [union keys..., $groupingId, real aggregate results...]}; the project restores the original aggregate row - /// type by emitting, in order: the group keys, then for each original aggregate call either the real aggregate - /// result (referenced from the FINAL) or, for GROUPING() / GROUPING_ID(), the value computed from the {@code - /// $groupingId} discriminator column (mirroring the single-stage post-aggregation handler). {@code $groupingId} - /// itself is dropped. - private static RelNode buildGroupingSetsProject(Aggregate aggRel, PinotLogicalAggregate finalAggRel, int groupCount, + /// Builds the Project on top of the FINAL grouping-set aggregate via the shared + /// {@link GroupingSetsPlanUtils#buildGroupingSetsProjects} (also used by the v2 physical planner), restoring + /// the original aggregate row type and dropping {@code $groupingId}. + private static RelNode buildGroupingSetsProject(Aggregate aggRel, PinotLogicalAggregate finalAggRel, int[] realAggIndex) { RexBuilder rexBuilder = finalAggRel.getCluster().getRexBuilder(); - RelDataType intType = finalAggRel.getCluster().getTypeFactory().createSqlType(SqlTypeName.INTEGER); - int finalGroupCount = groupCount + 1; - /// $groupingId is the INT discriminator column immediately after the union group keys. - RexNode groupingIdRef = rexBuilder.makeInputRef(finalAggRel, groupCount); - List union = aggRel.getGroupSet().asList(); - List orgAggCalls = aggRel.getAggCallList(); - List projects = new ArrayList<>(groupCount + orgAggCalls.size()); - for (int i = 0; i < groupCount; i++) { - projects.add(rexBuilder.makeInputRef(finalAggRel, i)); - } - for (int i = 0; i < orgAggCalls.size(); i++) { - if (realAggIndex[i] >= 0) { - projects.add(rexBuilder.makeInputRef(finalAggRel, finalGroupCount + realAggIndex[i])); - } else { - AggregateCall groupingCall = orgAggCalls.get(i); - /// Map each GROUPING argument (an input column index) to its position in the union group key list, which is - /// the bit position in $groupingId (bit set iff the column is rolled up in the row's grouping set). - List unionIndexes = new ArrayList<>(groupingCall.getArgList().size()); - for (int arg : groupingCall.getArgList()) { - int unionIndex = union.indexOf(arg); - Preconditions.checkState(unionIndex >= 0, "GROUPING argument must be a grouping column"); - unionIndexes.add(unionIndex); - } - RexNode value = GroupingSetsRexUtils.buildGroupingValue(rexBuilder, intType, groupingIdRef, unionIndexes); - projects.add(rexBuilder.makeCast(groupingCall.getType(), value)); - } - } + List projects = GroupingSetsPlanUtils.buildGroupingSetsProjects(rexBuilder, finalAggRel, aggRel, + realAggIndex); return LogicalProject.create(finalAggRel, List.of(), projects, aggRel.getRowType().getFieldNames()); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java index 4c5ae0aee128..cb166dedd9d3 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java @@ -22,9 +22,7 @@ import com.google.common.collect.Sets; import java.util.ArrayList; import java.util.Arrays; -import java.util.HashMap; import java.util.List; -import java.util.Map; import java.util.Set; import javax.annotation.Nullable; import org.apache.calcite.plan.RelOptTable; @@ -32,7 +30,6 @@ import org.apache.calcite.rel.RelDistribution; import org.apache.calcite.rel.RelFieldCollation; import org.apache.calcite.rel.RelNode; -import org.apache.calcite.rel.core.Aggregate; import org.apache.calcite.rel.core.AggregateCall; import org.apache.calcite.rel.core.Exchange; import org.apache.calcite.rel.core.JoinInfo; @@ -61,7 +58,6 @@ import org.apache.calcite.rex.RexWindowExclusion; import org.apache.calcite.sql.SqlLiteral; import org.apache.calcite.sql.type.SqlTypeName; -import org.apache.calcite.util.ImmutableBitSet; import org.apache.pinot.calcite.rel.hint.PinotHintOptions; import org.apache.pinot.calcite.rel.logical.PinotLogicalAggregate; import org.apache.pinot.calcite.rel.logical.PinotLogicalEnrichedJoin; @@ -69,6 +65,7 @@ import org.apache.pinot.calcite.rel.logical.PinotLogicalSortExchange; import org.apache.pinot.calcite.rel.logical.PinotLogicalTableScan; import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType; +import org.apache.pinot.calcite.rel.rules.GroupingSetsPlanUtils; import org.apache.pinot.calcite.rel.rules.PinotRuleUtils; import org.apache.pinot.common.metrics.BrokerMeter; import org.apache.pinot.common.metrics.BrokerMetrics; @@ -710,29 +707,6 @@ private SortNode convertLogicalSort(LogicalSort node) { convertInputs(node.getInputs()), node.getCollation().getFieldCollations(), fetch, offset); } - /// Encodes each grouping set as a membership bitmask over the union group-by columns ({@code groupSet.asList()}), - /// mirroring the single-stage engine's {@code PinotQuery.groupingSetMasks} so the per-set expansion can be pushed - /// down to the single-stage (leaf) engine. Returns an empty list for a plain GROUP BY (SIMPLE). - public static List computeGroupingSetMasks(Aggregate node) { - if (node.getGroupType() == Aggregate.Group.SIMPLE) { - return List.of(); - } - List union = node.getGroupSet().asList(); - Map unionIndex = new HashMap<>(); - for (int i = 0; i < union.size(); i++) { - unionIndex.put(union.get(i), i); - } - List masks = new ArrayList<>(node.getGroupSets().size()); - for (ImmutableBitSet set : node.getGroupSets()) { - int mask = 0; - for (int bit : set) { - mask |= 1 << unionIndex.get(bit); - } - masks.add(mask); - } - return masks; - } - private AggregateNode convertLogicalAggregate(PinotLogicalAggregate node) { List aggregateCalls = node.getAggCallList(); int numAggregates = aggregateCalls.size(); @@ -744,7 +718,8 @@ private AggregateNode convertLogicalAggregate(PinotLogicalAggregate node) { } return new AggregateNode(DEFAULT_STAGE_ID, toDataSchema(node.getRowType()), NodeHint.fromRelHints(node.getHints()), convertInputs(node.getInputs()), functionCalls, filterArgs, node.getGroupSet().asList(), node.getAggType(), - node.isLeafReturnFinalResult(), node.getCollations(), node.getLimit(), computeGroupingSetMasks(node)); + node.isLeafReturnFinalResult(), node.getCollations(), node.getLimit(), + GroupingSetsPlanUtils.computeGroupingSets(node)); } private ProjectNode convertLogicalProject(LogicalProject node) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java index 91ef973476a0..b0a2f0032b89 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java @@ -43,11 +43,11 @@ import org.apache.calcite.sql.SqlLiteral; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.pinot.calcite.rel.hint.PinotHintOptions; +import org.apache.pinot.calcite.rel.rules.GroupingSetsPlanUtils; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.common.utils.DatabaseUtils; import org.apache.pinot.common.utils.request.RequestUtils; -import org.apache.pinot.query.planner.logical.RelToPlanNodeConverter; import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.logical.RexExpressionUtils; import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalAggregate; @@ -201,7 +201,7 @@ public static SortNode convertSort(Sort node) { public static AggregateNode convertAggregate(PhysicalAggregate node) { /// GROUP BY GROUPING SETS / ROLLUP / CUBE is split by AggregatePushdownRule (LEAF carrying the grouping sets + - /// $groupingId, FINAL grouping on $groupingId); the masks below carry to the runtime RepeatOperator. + /// $groupingId, FINAL grouping on $groupingId); the grouping sets below carry to the runtime RepeatOperator. List aggregateCalls = node.getAggCallList(); int numAggregates = aggregateCalls.size(); List functionCalls = new ArrayList<>(numAggregates); @@ -213,7 +213,7 @@ public static AggregateNode convertAggregate(PhysicalAggregate node) { return new AggregateNode(DEFAULT_STAGE_ID, toDataSchema(node.getRowType()), NodeHint.fromRelHints(node.getHints()), new ArrayList<>(), functionCalls, filterArgs, node.getGroupSet().asList(), node.getAggType(), node.isLeafReturnFinalResult(), node.getCollations(), node.getLimit(), - RelToPlanNodeConverter.computeGroupingSetMasks(node)); + GroupingSetsPlanUtils.computeGroupingSets(node)); } public static ProjectNode convertProject(Project node) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalAggregate.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalAggregate.java index 596e7658af71..28ebea03b8ad 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalAggregate.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalAggregate.java @@ -30,7 +30,7 @@ import org.apache.calcite.rel.hint.RelHint; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.util.ImmutableBitSet; -import org.apache.pinot.common.request.context.GroupingSets; +import org.apache.pinot.calcite.rel.rules.GroupingSetsPlanUtils; import org.apache.pinot.query.planner.physical.v2.PRelNode; import org.apache.pinot.query.planner.physical.v2.PinotDataDistribution; import org.apache.pinot.query.planner.plannode.AggregateNode; @@ -88,7 +88,7 @@ protected RelDataType deriveRowType() { if (!emitsGroupingId()) { return rowType; } - return GroupingSets.appendGroupingIdColumn(getCluster().getTypeFactory(), rowType, getGroupCount()); + return GroupingSetsPlanUtils.appendGroupingIdColumn(getCluster().getTypeFactory(), rowType, getGroupCount()); } @Override diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java index 5aefc84a7e0b..6b79d934deee 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/AggregatePushdownRule.java @@ -36,7 +36,6 @@ import org.apache.calcite.rel.core.Project; import org.apache.calcite.rel.core.Union; import org.apache.calcite.rel.type.RelDataType; -import org.apache.calcite.rex.RexBuilder; import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; @@ -47,15 +46,13 @@ import org.apache.calcite.sql.type.ReturnTypes; import org.apache.calcite.sql.type.SqlOperandTypeChecker; import org.apache.calcite.sql.type.SqlReturnTypeInference; -import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.util.ImmutableBitSet; import org.apache.pinot.calcite.rel.hint.PinotHintOptions; import org.apache.pinot.calcite.rel.hint.PinotHintStrategyTable; -import org.apache.pinot.calcite.rel.rules.GroupingSetsRexUtils; +import org.apache.pinot.calcite.rel.rules.GroupingSetsPlanUtils; import org.apache.pinot.calcite.rel.rules.PinotRuleUtils; import org.apache.pinot.calcite.rel.traits.PinotExecStrategyTrait; import org.apache.pinot.common.function.sql.PinotSqlAggFunction; -import org.apache.pinot.common.request.context.GroupingSets; import org.apache.pinot.query.context.PhysicalPlannerContext; import org.apache.pinot.query.planner.physical.v2.PRelNode; import org.apache.pinot.query.planner.physical.v2.PinotDataDistribution; @@ -104,17 +101,10 @@ public PRelNode onMatch(PRelOptRuleCall call) { /// GROUP BY GROUPING SETS / ROLLUP / CUBE: split with the synthetic $groupingId discriminator carried through the /// exchange (the runtime RepeatOperator does the per-set row expansion). Requires an input exchange to /// repartition by union keys + $groupingId. - /// $groupingId is a 32-bit bitmask over the union columns, so cap the distinct grouping columns (same guard as - /// the single-stage rule PinotAggregateExchangeNodeInsertRule). - if (aggRel.getGroupCount() > GroupingSets.MAX_GROUPING_SET_COLUMNS) { - throw new UnsupportedOperationException( - "GROUP BY GROUPING SETS / ROLLUP / CUBE supports at most " + GroupingSets.MAX_GROUPING_SET_COLUMNS - + " distinct grouping columns in the multi-stage query engine, got " + aggRel.getGroupCount()); - } /// A WITHIN GROUP ordered aggregate (e.g. LISTAGG ... WITHIN GROUP) under a grouping set cannot be leaf/final /// split without losing the ORDER BY, and the DIRECT fallback does not preserve the ordering across the per-set - /// expansion either. Reject explicitly (run on the default planner / single-stage) rather than silently - /// returning a wrong result. The default planner checks WITHIN GROUP before any split for the same reason. + /// expansion either. Reject explicitly (run on the single-stage engine) rather than silently returning a wrong + /// result; the default planner rejects the same combination. if (withinGroupCollation != null) { throw new UnsupportedOperationException( "WITHIN GROUP ordered aggregates with GROUP BY GROUPING SETS / ROLLUP / CUBE are not supported in the " @@ -125,9 +115,13 @@ public PRelNode onMatch(PRelOptRuleCall call) { "GROUP BY GROUPING SETS / ROLLUP / CUBE without a repartitioning exchange is not yet supported in the " + "multi-stage v2 physical planner. Run the query on the single-stage query engine instead."); } - return addPartialAggregateForGroupingSets((PhysicalAggregate) call._currentNode, hintOptions, - _context.getNodeIdGenerator()); + /// Aggregate hints (e.g. is_leaf_return_final_result) are not applied to grouping-set splits. + return addPartialAggregateForGroupingSets((PhysicalAggregate) call._currentNode, _context.getNodeIdGenerator()); } + /// GROUPING() / GROUPING_ID() with a plain GROUP BY is constant 0 per the SQL standard; it is not wired into + /// the multi-stage runtime (it is not a real aggregation function), so reject it explicitly rather than failing + /// with an obscure error at execution. TODO: fold it to the literal 0 for single-stage parity. + rejectGroupingFunctionWithPlainGroupBy(aggRel); if (!isInputExchange || withinGroupCollation != null || (hasGroupBy && Boolean.parseBoolean( hintOptions.get(PinotHintOptions.AggregateOptions.IS_SKIP_LEAF_STAGE_GROUP_BY)))) { return skipPartialAggregate(call._currentNode); @@ -135,6 +129,18 @@ public PRelNode onMatch(PRelOptRuleCall call) { return addPartialAggregate((PhysicalAggregate) call._currentNode, hintOptions, _context.getNodeIdGenerator()); } + /// Rejects GROUPING() / GROUPING_ID() calls on a plain (single-grouping-set) aggregate. + private static void rejectGroupingFunctionWithPlainGroupBy(Aggregate aggRel) { + for (AggregateCall aggCall : aggRel.getAggCallList()) { + SqlKind kind = aggCall.getAggregation().getKind(); + if (kind == SqlKind.GROUPING || kind == SqlKind.GROUPING_ID) { + throw new UnsupportedOperationException( + "GROUPING() / GROUPING_ID() requires GROUP BY GROUPING SETS / ROLLUP / CUBE in the multi-stage query " + + "engine. Run the query on the single-stage query engine instead."); + } + } + } + private static PRelNode skipPartialAggregate(PRelNode aggPRelNode) { PhysicalAggregate aggRel = (PhysicalAggregate) aggPRelNode.unwrap(); List newAggCalls = buildAggCalls(aggRel, AggType.DIRECT, false); @@ -156,31 +162,31 @@ private static PRelNode skipPartialAggregate(PRelNode aggPRelNode) { /// the real aggregate result or, for GROUPING() / GROUPING_ID(), the value computed from {@code $groupingId}; /// {@code $groupingId} itself is dropped. private static PRelNode addPartialAggregateForGroupingSets(PhysicalAggregate aggPRelNode, - Map hintOptions, Supplier idGenerator) { + Supplier idGenerator) { PhysicalAggregate o0 = aggPRelNode; PhysicalExchange o1 = (PhysicalExchange) o0.getPRelInput(0); PRelNode o2 = o1.getPRelInput(0); int groupCount = o0.getGroupCount(); - /// GROUPING() / GROUPING_ID() are not real aggregations: they are functions of which grouping set a row belongs to, - /// computed from $groupingId in the projection. Split them out of the LEAF/FINAL aggregate calls. realAggIndex[i] - /// is the position of original aggregate call i among the real aggregates, or -1 if it is a GROUPING / GROUPING_ID. + /// GROUPING() / GROUPING_ID() are not real aggregations: they are computed from $groupingId in the final + /// projection, so they are split out of the LEAF/FINAL aggregate calls here (shared with the default planner). List orgAggCalls = o0.getAggCallList(); List realAggCalls = new ArrayList<>(orgAggCalls.size()); - int[] realAggIndex = new int[orgAggCalls.size()]; - for (int i = 0; i < orgAggCalls.size(); i++) { - SqlKind kind = orgAggCalls.get(i).getAggregation().getKind(); - if (kind == SqlKind.GROUPING || kind == SqlKind.GROUPING_ID) { - realAggIndex[i] = -1; - } else { - realAggIndex[i] = realAggCalls.size(); - realAggCalls.add(orgAggCalls.get(i)); - } + int[] realAggIndex = GroupingSetsPlanUtils.splitOutGroupingCalls(orgAggCalls, realAggCalls); + /// Without a real aggregation the LEAF would push an aggregation-free query down to the single-stage engine, + /// which would execute it as a plain selection and ignore the grouping sets. Mirrors the single-stage parser + /// rejection (GROUPING()/GROUPING_ID() are not aggregations). + if (realAggCalls.isEmpty()) { + throw new UnsupportedOperationException( + "GROUP BY GROUPING SETS / ROLLUP / CUBE requires at least one aggregation function in the query"); } /// LEAF: carries the grouping sets and only the real aggregations; deriveRowType() appends $groupingId at position /// groupCount, so the leaf output is [union keys..., $groupingId, real aggregates...]. + /// Group trim (collations/limit) is not applied to the grouping-set LEAF: the collation indexes were + /// extracted against the original [union keys..., aggs...] layout and would be off by one past the inserted + /// $groupingId column (the default planner disables trim for grouping sets for the same reason). PhysicalAggregate n2 = new PhysicalAggregate(o0.getCluster(), RelTraitSet.createEmpty(), List.of() /* hints */, o0.getGroupSet(), o0.groupSets, buildAggCalls(o0, realAggCalls, AggType.LEAF, false), idGenerator.get(), o2, - null /* data dist */, o2.isLeafStage(), AggType.LEAF, false, o0.getCollations(), o0.getLimit()); + null /* data dist */, o2.isLeafStage(), AggType.LEAF, false, List.of() /* collations */, 0 /* limit */); PinotDistMapping mapFromInputToPartialAgg = DistMappingGenerator.compute(o2.unwrap(), n2, null); PinotDataDistribution leafAggDataDistribution = o2.getPinotDataDistributionOrThrow().apply(mapFromInputToPartialAgg); @@ -196,33 +202,10 @@ private static PRelNode addPartialAggregateForGroupingSets(PhysicalAggregate agg /// groupCount + 1 in the exchange (intermediate) output. int finalGroupCount = groupCount + 1; PhysicalAggregate n0 = convertAggForGroupingSets(o0, n1, realAggCalls, finalGroupCount, idGenerator); - /// PROJECT: restore the original aggregate row type. For each original aggregate call, reference the real result or - /// compute the GROUPING() / GROUPING_ID() value from $groupingId (at position groupCount in the final output). - RexBuilder rexBuilder = o0.getCluster().getRexBuilder(); - RelDataType intType = o0.getCluster().getTypeFactory().createSqlType(SqlTypeName.INTEGER); - RexNode groupingIdRef = rexBuilder.makeInputRef(n0, groupCount); - List union = o0.getGroupSet().asList(); - List projects = new ArrayList<>(groupCount + orgAggCalls.size()); - for (int i = 0; i < groupCount; i++) { - projects.add(rexBuilder.makeInputRef(n0, i)); - } - for (int i = 0; i < orgAggCalls.size(); i++) { - if (realAggIndex[i] >= 0) { - projects.add(rexBuilder.makeInputRef(n0, finalGroupCount + realAggIndex[i])); - } else { - AggregateCall groupingCall = orgAggCalls.get(i); - /// Map each GROUPING argument (an input column index) to its bit position in $groupingId (its position in the - /// union group key list). - List unionIndexes = new ArrayList<>(groupingCall.getArgList().size()); - for (int arg : groupingCall.getArgList()) { - int unionIndex = union.indexOf(arg); - Preconditions.checkState(unionIndex >= 0, "GROUPING argument must be a grouping column"); - unionIndexes.add(unionIndex); - } - RexNode value = GroupingSetsRexUtils.buildGroupingValue(rexBuilder, intType, groupingIdRef, unionIndexes); - projects.add(rexBuilder.makeCast(groupingCall.getType(), value)); - } - } + /// PROJECT: restore the original aggregate row type via the shared projection builder (also used by the + /// default planner), dropping $groupingId. + List projects = + GroupingSetsPlanUtils.buildGroupingSetsProjects(o0.getCluster().getRexBuilder(), n0, o0, realAggIndex); return new PhysicalProject(o0.getCluster(), o0.getTraitSet(), List.of() /* hints */, projects, o0.getRowType(), Set.of(), idGenerator.get(), n0, n0.getPinotDataDistributionOrThrow(), false); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/AggregateNode.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/AggregateNode.java index e2d6fcf68c3b..b606bd59f87a 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/AggregateNode.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/AggregateNode.java @@ -33,11 +33,13 @@ public class AggregateNode extends BasePlanNode { private final AggType _aggType; private final boolean _leafReturnFinalResult; - /// GROUP BY GROUPING SETS / ROLLUP / CUBE: one membership bitmask per grouping set, over _groupKeys (the union of - /// all grouping columns). Bit i is set iff _groupKeys[i] participates in (is grouped by) that set; a mask of 0 is - /// the grand-total set (). Empty for a plain GROUP BY. Mirrors the single-stage engine's PinotQuery.groupingSetMasks - /// so the per-set row expansion can be pushed down to the single-stage (leaf) engine. - private final List _groupingSets; + /// GROUP BY GROUPING SETS / ROLLUP / CUBE: one entry per grouping set, each the list of indexes into _groupKeys + /// (the union of all grouping columns) participating in (grouped by) that set; an empty inner list is the + /// grand-total set (). A set's position in this list is its ordinal, carried as the synthetic $groupingId + /// discriminator, so the number of grouping columns is unlimited. Empty for a plain GROUP BY. Mirrors the + /// single-stage engine's PinotQuery.groupingSets so the per-set row expansion can be pushed down to the + /// single-stage (leaf) engine. + private final List> _groupingSets; // The following fields are set when group trim is enabled, and are extracted from the Sort on top of this Aggregate. // The group trim behavior at leaf stage is shared with single-stage engine. @@ -54,7 +56,7 @@ public AggregateNode(int stageId, DataSchema dataSchema, NodeHint nodeHint, List public AggregateNode(int stageId, DataSchema dataSchema, NodeHint nodeHint, List inputs, List aggCalls, List filterArgs, List groupKeys, AggType aggType, boolean leafReturnFinalResult, @Nullable List collations, int limit, - List groupingSets) { + List> groupingSets) { super(stageId, dataSchema, nodeHint, inputs); _aggCalls = aggCalls; _filterArgs = filterArgs; @@ -86,8 +88,9 @@ public boolean isLeafReturnFinalResult() { return _leafReturnFinalResult; } - /// One membership bitmask per grouping set over {@link #getGroupKeys()}, or empty for a plain GROUP BY. - public List getGroupingSets() { + /// Per grouping set (in ordinal order), the indexes into {@link #getGroupKeys()} participating in it, or + /// empty for a plain GROUP BY. + public List> getGroupingSets() { return _groupingSets; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java index 540fd6060822..f0ebfd3c8742 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java @@ -91,12 +91,15 @@ public static PlanNode process(Plan.PlanNode protoNode) { private static AggregateNode deserializeAggregateNode(Plan.PlanNode protoNode) { Plan.AggregateNode protoAggregateNode = protoNode.getAggregateNode(); + List> groupingSets = new ArrayList<>(protoAggregateNode.getGroupingSetsCount()); + for (Plan.GroupingSet protoGroupingSet : protoAggregateNode.getGroupingSetsList()) { + groupingSets.add(protoGroupingSet.getGroupKeyIndexesList()); + } return new AggregateNode(protoNode.getStageId(), extractDataSchema(protoNode), extractNodeHint(protoNode), extractInputs(protoNode), convertFunctionCalls(protoAggregateNode.getAggCallsList()), protoAggregateNode.getFilterArgsList(), protoAggregateNode.getGroupKeysList(), convertAggType(protoAggregateNode.getAggType()), protoAggregateNode.getLeafReturnFinalResult(), - convertCollations(protoAggregateNode.getCollationsList()), protoAggregateNode.getLimit(), - protoAggregateNode.getGroupingSetsList()); + convertCollations(protoAggregateNode.getCollationsList()), protoAggregateNode.getLimit(), groupingSets); } private static FilterNode deserializeFilterNode(Plan.PlanNode protoNode) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java index c329152fc142..bec743cbcd45 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java @@ -94,17 +94,18 @@ private SerializationVisitor() { @Override public Void visitAggregate(AggregateNode node, Plan.PlanNode.Builder builder) { - Plan.AggregateNode aggregateNode = Plan.AggregateNode.newBuilder() + Plan.AggregateNode.Builder aggregateNodeBuilder = Plan.AggregateNode.newBuilder() .addAllAggCalls(convertFunctionCalls(node.getAggCalls())) .addAllFilterArgs(node.getFilterArgs()) .addAllGroupKeys(node.getGroupKeys()) .setAggType(convertAggType(node.getAggType())) .setLeafReturnFinalResult(node.isLeafReturnFinalResult()) .addAllCollations(convertCollations(node.getCollations())) - .setLimit(node.getLimit()) - .addAllGroupingSets(node.getGroupingSets()) - .build(); - builder.setAggregateNode(aggregateNode); + .setLimit(node.getLimit()); + for (List groupingSet : node.getGroupingSets()) { + aggregateNodeBuilder.addGroupingSets(Plan.GroupingSet.newBuilder().addAllGroupKeyIndexes(groupingSet).build()); + } + builder.setAggregateNode(aggregateNodeBuilder.build()); return null; } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java index 237a9c8ff95d..2e7eb73ca9a0 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java @@ -159,6 +159,48 @@ public void testGroupingSetsSupportedInMultiStage() { } } + @Test + public void testGroupingSetsRejectionsInMultiStage() { + /// Combinations the multi-stage planners reject explicitly (instead of producing broken plans or silently + /// wrong results): too many expanded sets (CUBE blow-up past the 4096 cap), WITHIN GROUP ordered aggregates + /// under a grouping set, aggregate hints with grouping sets, aggregation-free grouping sets (GROUPING() is + /// not an aggregation), and GROUPING() with a plain GROUP BY. + StringBuilder cubeColumns = new StringBuilder("col1, col2"); + for (int i = 0; i < 11; i++) { + cubeColumns.append(", col3 + ").append(i); + } + List queries = List.of( + /// CUBE over 13 grouping expressions expands to 2^13 = 8192 grouping sets, exceeding the 4096 cap. + "SELECT COUNT(*) FROM a GROUP BY CUBE(" + cubeColumns + ")", + "SELECT col1, LISTAGG(col2, ',') WITHIN GROUP (ORDER BY col2) FROM a GROUP BY ROLLUP(col1)", + "SELECT /*+ aggOptions(is_skip_leaf_stage_group_by='true') */ col1, COUNT(*) FROM a GROUP BY ROLLUP(col1)", + "SELECT col1 FROM a GROUP BY ROLLUP(col1)", + "SELECT col1, GROUPING(col1) FROM a GROUP BY ROLLUP(col1)", + "SELECT col1, GROUPING(col1), COUNT(*) FROM a GROUP BY col1"); + for (String sql : queries) { + try { + _queryEnvironment.planQuery(sql); + fail("expected rejection for: " + sql); + } catch (RuntimeException e) { + // expected + } + } + } + + @Test + public void testGroupingSetsUnlimitedColumnsInMultiStage() { + /// The number of distinct grouping columns is unlimited (each grouping set is carried as a member-index + /// list and the discriminator is the set ordinal, mirroring Calcite's per-set column bitset). 40 distinct + /// grouping expressions — past the retired 31-column bitmask cap — must plan successfully, including with a + /// GROUPING() call. + StringBuilder columns = new StringBuilder("col1, col2"); + for (int i = 0; i < 38; i++) { + columns.append(", col3 + ").append(i); + } + String sql = "SELECT col1, GROUPING(col1), SUM(col3) FROM a GROUP BY ROLLUP(" + columns + ")"; + assertNotNull(_queryEnvironment.planQuery(sql), "expected a multi-stage plan for a 40-column ROLLUP"); + } + @Test public void testUnsignedLiteralCastIsFolded() { // Companion to testUnsignedTypeCastIsAccepted using literal (constant-foldable) casts, so the unsigned diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java index a92ddd575387..cdafce4fee7d 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java @@ -83,14 +83,16 @@ public void testLegacyUnnestNodeSerDe() { @Test public void testAggregateGroupingSetsSerDe() { - /// The grouping-set masks (over the union group keys) must survive serialization to the worker. ROLLUP(g0, g1) - /// over the union {g0, g1} expands to masks {0b11, 0b01, 0b00}. + /// The grouping sets (member indexes over the union group keys, in ordinal order) must survive serialization + /// to the worker. ROLLUP(g0, g1) over the union {g0, g1} expands to the sets (g0, g1), (g0), (). The empty + /// grand-total set in particular must round-trip as an entry (not vanish as a proto default). DataSchema schema = new DataSchema(new String[]{"g0", "g1", "sum"}, new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.INT, ColumnDataType.DOUBLE}); + List> groupingSets = List.of(List.of(0, 1), List.of(0), List.of()); AggregateNode node = new AggregateNode(0, schema, PlanNode.NodeHint.EMPTY, List.of(), List.of(), List.of(), - List.of(0, 1), AggType.DIRECT, false, List.of(), 0, List.of(0b11, 0b01, 0b00)); + List.of(0, 1), AggType.DIRECT, false, List.of(), 0, groupingSets); AggregateNode deserialized = (AggregateNode) PlanNodeDeserializer.process(PlanNodeSerializer.process(node)); - assertEquals(deserialized.getGroupingSets(), List.of(0b11, 0b01, 0b00)); + assertEquals(deserialized.getGroupingSets(), groupingSets); assertEquals(deserialized, node); } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java index 1296cf71ea06..c40022ed4fee 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java @@ -32,7 +32,6 @@ import org.apache.pinot.common.datatable.StatMap; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.request.context.FunctionContext; -import org.apache.pinot.common.request.context.GroupingSets; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.core.common.BlockValSet; @@ -109,21 +108,9 @@ public AggregateOperator(OpChainExecutionContext context, MultiStageOperator inp maxFilterArgId = Math.max(maxFilterArgId, filterArgIds[i]); } + /// Grouping-set aggregates never reach this operator directly: PlanNodeToOpChain pre-wraps the input in a + /// RepeatOperator and rewrites the node into the equivalent plain GROUP BY over the expanded input. List groupKeys = node.getGroupKeys(); - - /// GROUP BY GROUPING SETS / ROLLUP / CUBE in the multi-stage runtime: expand each input row across the grouping - /// sets (NULLing the rolled-up columns and appending the $groupingId discriminator) via a RepeatOperator, then run - /// an ordinary GROUP BY over the union columns plus $groupingId. This handles grouping sets over any input (e.g. - /// above a JOIN); when the aggregate sits directly on a table scan it is instead pushed down to the single-stage - /// leaf and this operator is not used. - List groupingSets = node.getGroupingSets(); - if (!groupingSets.isEmpty()) { - int groupingIdColumnIndex = node.getInputs().get(0).getDataSchema().size(); - input = new RepeatOperator(context, input, getGroupKeyIds(groupKeys), groupingSets, repeatResultSchema(node)); - /// The group keys are the union columns (unchanged input positions) plus $groupingId, appended by the Repeat. - groupKeys = new ArrayList<>(groupKeys); - groupKeys.add(groupingIdColumnIndex); - } _input = input; int groupTrimSize = Integer.MAX_VALUE; @@ -361,22 +348,6 @@ private int[] getGroupKeyIds(List groupKeys) { return groupKeyIds; } - /// Output schema of the {@link RepeatOperator} that feeds a grouping-set aggregate: the aggregate input schema with - /// the synthetic {@code $groupingId} INT discriminator column appended. - private static DataSchema repeatResultSchema(AggregateNode node) { - DataSchema inputSchema = node.getInputs().get(0).getDataSchema(); - int numInputColumns = inputSchema.size(); - String[] columnNames = new String[numInputColumns + 1]; - DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[numInputColumns + 1]; - for (int i = 0; i < numInputColumns; i++) { - columnNames[i] = inputSchema.getColumnName(i); - columnDataTypes[i] = inputSchema.getColumnDataType(i); - } - columnNames[numInputColumns] = GroupingSets.GROUPING_ID_COLUMN; - columnDataTypes[numInputColumns] = DataSchema.ColumnDataType.INT; - return new DataSchema(columnNames, columnDataTypes); - } - static RoaringBitmap getMatchedBitmap(MseBlock.Data block, int filterArgId) { Preconditions.checkArgument(filterArgId >= 0, "Got negative filter argument id: %s", filterArgId); RoaringBitmap matchedBitmap = new RoaringBitmap(); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java index 460c19ede4fa..8f87a6a462ef 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java @@ -20,6 +20,7 @@ import java.util.ArrayList; import java.util.List; +import javax.annotation.Nullable; import org.apache.pinot.common.datatable.StatMap; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.query.runtime.blocks.MseBlock; @@ -31,46 +32,65 @@ /// Expands each input row across the grouping sets of a GROUP BY GROUPING SETS / ROLLUP / CUBE query — the /// multi-stage equivalent of the single-stage per-set row expansion. For every input row and every grouping set it -/// emits one output row in which the columns NOT grouped by that set (the -/// "rolled up" columns) are set to NULL, and appends the synthetic INT discriminator column -/// {@link org.apache.pinot.common.request.context.GroupingSets#GROUPING_ID_COLUMN} whose bit i is set iff union -/// group-by column i is rolled up in that set (matching the single-stage convention so GROUPING() / GROUPING_ID() -/// agree across engines). +/// emits one output row of the shape {@code [input columns..., group-key copies..., $groupingId]}: the original +/// input columns are carried through UNTOUCHED (aggregation arguments may reference a grouping column, so nulling +/// in place would corrupt them — e.g. SUM(b) under ROLLUP(a, b) must still aggregate the real b values for the +/// (a) subtotal), one appended copy per union group-by column holds the group-key value — NULL when the column is +/// rolled up (not grouped by) in the set — and the trailing synthetic INT discriminator column +/// {@link org.apache.pinot.common.request.context.GroupingSets#GROUPING_ID_COLUMN} carries the grouping set's +/// ordinal (its index in the plan's grouping-set list, matching the single-stage convention so GROUPING() / +/// GROUPING_ID() agree across engines). Because the ordinal is not a per-column bitmask, the number of grouping +/// columns is unlimited. The downstream aggregate groups by the appended copies plus the ordinal. /// -/// With the rows expanded and tagged, the downstream aggregate is an ordinary GROUP BY over the union group-by -/// columns plus {@code $groupingId} — no grouping-set-specific aggregation logic is needed. This lets grouping sets +/// With the rows expanded and tagged, the downstream aggregate is an ordinary GROUP BY over the appended key +/// copies plus {@code $groupingId} — no grouping-set-specific aggregation logic is needed. This lets grouping sets /// run over any input (e.g. above a JOIN), not just a leaf table scan. +/// +/// The expansion is streamed one grouping set at a time: each {@link #getNextBlock()} call emits the current input +/// block expanded for a single set, so the transient materialization is bounded by the input block size rather than +/// multiplied by the set count. Not thread-safe: like all multi-stage operators, it executes on the single OpChain +/// thread. public class RepeatOperator extends MultiStageOperator { private static final Logger LOGGER = LoggerFactory.getLogger(RepeatOperator.class); private static final String EXPLAIN_NAME = "REPEAT"; private final MultiStageOperator _input; private final DataSchema _resultSchema; - /// Input column index of each union group-by column, in union order; bit i of a grouping-set mask refers to - /// _unionGroupKeyIds[i]. + /// Input column index of each union group-by column, in union order; output column _inputColumnCount + i is the + /// group-key copy of _unionGroupKeyIds[i]. private final int[] _unionGroupKeyIds; - /// Per grouping set: the rolled-up bitmask over the union columns (bit i set iff column i is NOT in the set), in the - /// order the output rows are emitted. Precomputed from the participation masks. - private final int[] _groupingIds; + /// Per grouping set (in ordinal order): membership over the union columns (_setContains[s][i] iff union column i + /// participates in set s; its key copy is NULLed otherwise). Output rows of set s carry the ordinal s in the + /// discriminator column. + private final boolean[][] _setContains; + private final int _numSets; private final int _inputColumnCount; private final StatMap _statMap = new StatMap<>(StatKey.class); + /// Rows of the input block currently being expanded, or null when the next input block must be pulled. + @Nullable + private List _currentRows; + /// The grouping set (ordinal) the next getNextBlock() call will expand the current input block for. + private int _currentSet; + /// @param unionGroupKeyIds input column index of each union group-by column (in union order) - /// @param groupingSetMasks one participation bitmask per grouping set over the union columns (bit i set iff - /// {@code unionGroupKeyIds[i]} is grouped by that set) - /// @param resultSchema the input schema with the {@code $groupingId} INT column appended + /// @param groupingSets per grouping set (in ordinal order), the union-column indexes participating in + /// (grouped by) that set — {@code unionGroupKeyIds[i]} for member index i + /// @param resultSchema the input schema with one group-key copy column per union group-by column and the + /// {@code $groupingId} INT column appended public RepeatOperator(OpChainExecutionContext context, MultiStageOperator input, int[] unionGroupKeyIds, - List groupingSetMasks, DataSchema resultSchema) { + List> groupingSets, DataSchema resultSchema) { super(context); _input = input; _resultSchema = resultSchema; _unionGroupKeyIds = unionGroupKeyIds; - _inputColumnCount = resultSchema.size() - 1; - int fullMask = (1 << unionGroupKeyIds.length) - 1; - _groupingIds = new int[groupingSetMasks.size()]; - for (int s = 0; s < groupingSetMasks.size(); s++) { - /// Rolled-up bits are the complement of the participation mask over the union columns. - _groupingIds[s] = ~groupingSetMasks.get(s) & fullMask; + _inputColumnCount = resultSchema.size() - unionGroupKeyIds.length - 1; + _numSets = groupingSets.size(); + _setContains = new boolean[_numSets][unionGroupKeyIds.length]; + for (int s = 0; s < _numSets; s++) { + for (int memberIndex : groupingSets.get(s)) { + _setContains[s][memberIndex] = true; + } } } @@ -104,26 +124,35 @@ public String toExplainString() { @Override protected MseBlock getNextBlock() { - MseBlock block = _input.nextBlock(); - if (block.isEos()) { - return block; + if (_currentRows == null) { + MseBlock block = _input.nextBlock(); + if (block.isEos()) { + return block; + } + _currentRows = ((MseBlock.Data) block).asRowHeap().getRows(); + _currentSet = 0; } - List inputRows = ((MseBlock.Data) block).asRowHeap().getRows(); - int groupingIdIndex = _inputColumnCount; - List expanded = new ArrayList<>(inputRows.size() * _groupingIds.length); - for (Object[] inputRow : inputRows) { - for (int groupingId : _groupingIds) { - Object[] row = new Object[_inputColumnCount + 1]; - System.arraycopy(inputRow, 0, row, 0, _inputColumnCount); - /// NULL out the union columns that are rolled up in this set (their bit is set in groupingId). - for (int i = 0; i < _unionGroupKeyIds.length; i++) { - if ((groupingId & (1 << i)) != 0) { - row[_unionGroupKeyIds[i]] = null; - } - } - row[groupingIdIndex] = groupingId; - expanded.add(row); + /// Expand the current input block for one grouping set per call. + int set = _currentSet; + boolean[] contains = _setContains[set]; + int numUnionKeys = _unionGroupKeyIds.length; + List expanded = new ArrayList<>(_currentRows.size()); + for (Object[] inputRow : _currentRows) { + Object[] row = new Object[_inputColumnCount + numUnionKeys + 1]; + System.arraycopy(inputRow, 0, row, 0, _inputColumnCount); + /// Append the group-key copies: the union column's value where it participates in this set, NULL where it + /// is rolled up. The original input columns stay untouched for aggregation arguments. + for (int i = 0; i < numUnionKeys; i++) { + row[_inputColumnCount + i] = contains[i] ? inputRow[_unionGroupKeyIds[i]] : null; } + row[_inputColumnCount + numUnionKeys] = set; + expanded.add(row); + /// Honor query cancellation/deadline and sample resource usage during the row amplification, like the + /// other row-amplifying operators (joins, window). + checkTerminationAndSampleUsagePeriodically(expanded.size(), EXPLAIN_NAME); + } + if (++_currentSet == _numSets) { + _currentRows = null; } return new RowHeapDataBlock(expanded, _resultSchema); } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/AggregateOperatorFactory.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/AggregateOperatorFactory.java index 19dd9969e7d7..70dcf9e56501 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/AggregateOperatorFactory.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/AggregateOperatorFactory.java @@ -30,6 +30,11 @@ */ public interface AggregateOperatorFactory { + /// Creates the aggregate operator. Grouping-set aggregates are normalized before this is called: the plan + /// visitor wraps {@code inputOperator} in a RepeatOperator (whose output schema is {@code inputPlanNode}'s + /// schema plus a trailing $groupingId INT column) and passes the equivalent plain GROUP BY + /// {@code aggregateNode} (group keys extended with the $groupingId column, grouping sets cleared), so + /// implementations need no grouping-set awareness. MultiStageOperator createAggregateOperator(OpChainExecutionContext context, MultiStageOperator inputOperator, PlanNode inputPlanNode, AggregateNode aggregateNode); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanNodeToOpChain.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanNodeToOpChain.java index f8ee41aae95e..ad262a03bf95 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanNodeToOpChain.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanNodeToOpChain.java @@ -25,6 +25,8 @@ import java.util.stream.Collectors; import java.util.stream.Stream; import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.GroupingSets; +import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.query.planner.plannode.AggregateNode; import org.apache.pinot.query.planner.plannode.EnrichedJoinNode; import org.apache.pinot.query.planner.plannode.ExchangeNode; @@ -50,6 +52,7 @@ import org.apache.pinot.query.runtime.operator.MailboxSendOperator; import org.apache.pinot.query.runtime.operator.MultiStageOperator; import org.apache.pinot.query.runtime.operator.OpChain; +import org.apache.pinot.query.runtime.operator.RepeatOperator; import org.apache.pinot.query.runtime.operator.SortOperator; import org.apache.pinot.query.runtime.operator.SortedMailboxReceiveOperator; import org.apache.pinot.query.runtime.operator.TransformOperator; @@ -220,6 +223,20 @@ public MultiStageOperator visitAggregate(AggregateNode node, OpChainExecutionCon try { PlanNode input = node.getInputs().get(0); child = visit(input, context); + /// GROUP BY GROUPING SETS / ROLLUP / CUBE in the multi-stage runtime: expand each input row across the + /// grouping sets via a RepeatOperator — appending per-set group-key copies (NULL where rolled up) and the + /// $groupingId ordinal while leaving the original input columns untouched (aggregation arguments may + /// reference a grouping column) — then hand the factory an ordinary GROUP BY over the appended key copies + /// plus $groupingId. Wrapping here — rather than inside a specific AggregateOperator implementation — means + /// every AggregateOperatorFactory receives already-expanded input and needs no grouping-set awareness. This + /// path handles grouping sets over any input (e.g. above a JOIN); when the aggregate sits directly on a + /// table scan the whole expansion is pushed down to the single-stage leaf instead and the plan reaching here + /// has no grouping sets. + if (node.isGroupingSets()) { + child = new RepeatOperator(context, child, groupKeyIds(node.getGroupKeys()), node.getGroupingSets(), + repeatResultSchema(input.getDataSchema(), node.getGroupKeys())); + node = asPlainGroupByOverExpandedInput(node, input.getDataSchema().size()); + } AggregateOperatorFactory aggregateOperatorFactory = context.getQueryOperatorFactoryProvider().getAggregateOperatorFactory(); return aggregateOperatorFactory.createAggregateOperator(context, child, input, node); @@ -228,6 +245,50 @@ public MultiStageOperator visitAggregate(AggregateNode node, OpChainExecutionCon } } + /// Output schema of the {@link RepeatOperator} that feeds a grouping-set aggregate: the aggregate input schema + /// with one group-key copy column per union group-by column and the synthetic {@code $groupingId} INT + /// discriminator column appended. + private static DataSchema repeatResultSchema(DataSchema inputSchema, List unionGroupKeyIds) { + int numInputColumns = inputSchema.size(); + int numUnionKeys = unionGroupKeyIds.size(); + String[] columnNames = new String[numInputColumns + numUnionKeys + 1]; + DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[numInputColumns + numUnionKeys + 1]; + for (int i = 0; i < numInputColumns; i++) { + columnNames[i] = inputSchema.getColumnName(i); + columnDataTypes[i] = inputSchema.getColumnDataType(i); + } + for (int i = 0; i < numUnionKeys; i++) { + columnNames[numInputColumns + i] = "$groupingSetKey" + i; + columnDataTypes[numInputColumns + i] = inputSchema.getColumnDataType(unionGroupKeyIds.get(i)); + } + columnNames[numInputColumns + numUnionKeys] = GroupingSets.GROUPING_ID_COLUMN; + columnDataTypes[numInputColumns + numUnionKeys] = DataSchema.ColumnDataType.INT; + return new DataSchema(columnNames, columnDataTypes); + } + + /// The equivalent plain GROUP BY over the RepeatOperator-expanded input: the group keys are the appended + /// group-key copies plus the $groupingId column (the original input columns are left untouched for aggregation + /// arguments), and the grouping sets are cleared. + private static AggregateNode asPlainGroupByOverExpandedInput(AggregateNode node, int numInputColumns) { + int numUnionKeys = node.getGroupKeys().size(); + List groupKeys = new ArrayList<>(numUnionKeys + 1); + for (int i = 0; i <= numUnionKeys; i++) { + groupKeys.add(numInputColumns + i); + } + return new AggregateNode(node.getStageId(), node.getDataSchema(), node.getNodeHint(), node.getInputs(), + node.getAggCalls(), node.getFilterArgs(), groupKeys, node.getAggType(), node.isLeafReturnFinalResult(), + node.getCollations(), node.getLimit()); + } + + private static int[] groupKeyIds(List groupKeys) { + int numKeys = groupKeys.size(); + int[] groupKeyIds = new int[numKeys]; + for (int i = 0; i < numKeys; i++) { + groupKeyIds[i] = groupKeys.get(i); + } + return groupKeyIds; + } + @Override public MultiStageOperator visitWindow(WindowNode node, OpChainExecutionContext context) { MultiStageOperator child = null; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java index 49e2a622f112..02930e1cdd0e 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java @@ -79,11 +79,11 @@ public Void visitAggregate(AggregateNode node, ServerPlanRequestContext context) pinotQuery.setGroupByList(groupByList); } /// GROUP BY GROUPING SETS / ROLLUP / CUBE: push the per-set row expansion down to the single-stage engine. The - /// masks are over groupByList (== node.getGroupKeys()), so they transfer directly. The single-stage leaf expands - /// each row across the sets and appends the synthetic $groupingId discriminator column; the multi-stage final - /// stage then groups on it (it is one of the group keys), keeping the grouping sets distinct. + /// per-set column-index lists are over groupByList (== node.getGroupKeys()), so they transfer directly. The + /// single-stage leaf expands each row across the sets and appends the synthetic $groupingId ordinal column; the + /// multi-stage final stage then groups on it (it is one of the group keys), keeping the grouping sets distinct. if (node.isGroupingSets()) { - pinotQuery.setGroupingSetMasks(node.getGroupingSets()); + pinotQuery.setGroupingSets(node.getGroupingSets()); } List selectList = CalciteRexExpressionParser.convertAggregateList(groupByList, node.getAggCalls(), node.getFilterArgs(), pinotQuery.getSelectList()); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/RepeatOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/RepeatOperatorTest.java index 55103bd3dbba..2170bc993873 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/RepeatOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/RepeatOperatorTest.java @@ -18,10 +18,13 @@ */ package org.apache.pinot.query.runtime.operator; +import java.util.ArrayList; import java.util.List; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.query.runtime.blocks.ErrorMseBlock; import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; import org.mockito.Mock; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; @@ -31,10 +34,14 @@ import static org.mockito.MockitoAnnotations.openMocks; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; -/// Unit coverage for [RepeatOperator]'s per-set row expansion: it must NULL the rolled-up union columns and append the -/// `$groupingId` discriminator (the rolled-up complement of each grouping set's participation mask). +/// Unit coverage for [RepeatOperator]'s per-set row expansion. Output rows have the shape +/// `[input columns..., group-key copies..., $groupingId]`: the original input columns must stay UNTOUCHED +/// (aggregation arguments may reference a grouping column), the appended key copies carry the group-key values +/// with NULL where the column is rolled up in the set, and `$groupingId` is the set's ordinal in the plan's set +/// list. The expansion is streamed — one grouping set per output block — so tests collect blocks until EOS. public class RepeatOperatorTest { private AutoCloseable _mocks; @Mock @@ -51,70 +58,188 @@ public void tearDown() _mocks.close(); } - /// Input [d1, d2, cnt]; union group keys are d1 (idx 0) and d2 (idx 1); result appends $groupingId. + /// Input [d1, d2, cnt]; union group keys are d1 (idx 0) and d2 (idx 1); result appends the two key copies and + /// $groupingId: [d1, d2, cnt, key0, key1, $groupingId]. private static final DataSchema INPUT_SCHEMA = new DataSchema(new String[]{"d1", "d2", "cnt"}, new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT}); - private static final DataSchema RESULT_SCHEMA = new DataSchema(new String[]{"d1", "d2", "cnt", "$groupingId"}, - new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT, ColumnDataType.INT}); + private static final DataSchema RESULT_SCHEMA = new DataSchema( + new String[]{"d1", "d2", "cnt", "$groupingSetKey0", "$groupingSetKey1", "$groupingId"}, + new ColumnDataType[]{ + ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT, ColumnDataType.STRING, + ColumnDataType.STRING, ColumnDataType.INT + }); + + /// Drains the operator until EOS, returning all emitted rows in block order (one block per grouping set). + private static List collectRows(RepeatOperator operator) { + List rows = new ArrayList<>(); + MseBlock block = operator.nextBlock(); + while (!block.isEos()) { + rows.addAll(((MseBlock.Data) block).asRowHeap().getRows()); + block = operator.nextBlock(); + } + return rows; + } @Test - public void shouldExpandRollupAndNullRolledUpColumns() { - /// ROLLUP(d1, d2) => sets {d1,d2}, {d1}, {} with participation masks 0b11, 0b01, 0b00. - when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(INPUT_SCHEMA, new Object[]{"a", "x", 5})); + public void shouldExpandRollupAndNullRolledUpKeyCopies() { + /// ROLLUP(d1, d2) => sets {d1,d2}, {d1}, {} — member union-column indexes (0,1), (0), (); ordinal = position. + when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(INPUT_SCHEMA, new Object[]{"a", "x", 5}), + SuccessMseBlock.INSTANCE); RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0, 1}, - List.of(0b11, 0b01, 0b00), RESULT_SCHEMA); + List.of(List.of(0, 1), List.of(0), List.of()), RESULT_SCHEMA); - List rows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + List rows = collectRows(operator); assertEquals(rows.size(), 3); /// one input row x three sets - /// {d1,d2}: nothing rolled up -> groupingId 0, both keys kept. - assertEquals(rows.get(0), new Object[]{"a", "x", 5, 0}); - /// {d1}: d2 rolled up -> groupingId 0b10 = 2, d2 NULLed. - assertEquals(rows.get(1)[0], "a"); - assertNull(rows.get(1)[1]); - assertEquals(rows.get(1)[3], 0b10); - /// {}: both rolled up -> groupingId 0b11 = 3, both keys NULLed. - assertNull(rows.get(2)[0]); - assertNull(rows.get(2)[1]); - assertEquals(rows.get(2)[3], 0b11); - /// cnt (non-key column) is never touched. - assertEquals(rows.get(2)[2], 5); + /// {d1,d2} (ordinal 0): both key copies carry the values. + assertEquals(rows.get(0), new Object[]{"a", "x", 5, "a", "x", 0}); + /// {d1} (ordinal 1): the d2 key copy is NULLed — but the ORIGINAL d2 column must stay untouched, because + /// aggregation arguments may reference it (e.g. SUM(d2) under ROLLUP(d1, d2)). + assertEquals(rows.get(1), new Object[]{"a", "x", 5, "a", null, 1}); + /// {} (ordinal 2): both key copies NULLed, originals untouched. + assertEquals(rows.get(2), new Object[]{"a", "x", 5, null, null, 2}); } @Test - public void shouldNullTheCorrectKeyForCube() { - /// CUBE(d1, d2) => the {d2}-only set rolls up d1 (groupingId 0b01), which a wrong bit order would get backwards. - when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(INPUT_SCHEMA, new Object[]{"a", "x", 5})); + public void shouldNullTheCorrectKeyCopyForCube() { + /// CUBE(d1, d2) => the {d2}-only set (ordinal 2) rolls up d1, which a wrong member/column mapping would get + /// backwards. + when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(INPUT_SCHEMA, new Object[]{"a", "x", 5}), + SuccessMseBlock.INSTANCE); RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0, 1}, - List.of(0b11, 0b01, 0b10, 0b00), RESULT_SCHEMA); + List.of(List.of(0, 1), List.of(0), List.of(1), List.of()), RESULT_SCHEMA); - List rows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + List rows = collectRows(operator); assertEquals(rows.size(), 4); - /// {d2}: participation 0b10 -> groupingId 0b01 = 1, d1 NULLed, d2 kept. - Object[] d2Only = rows.get(2); - assertNull(d2Only[0]); - assertEquals(d2Only[1], "x"); - assertEquals(d2Only[3], 0b01); + /// {d2} (ordinal 2): the d1 key copy NULLed, the d2 key copy kept. + assertEquals(rows.get(2), new Object[]{"a", "x", 5, null, "x", 2}); } @Test - public void shouldHandleNonContiguousMaskAndMultipleRows() { - /// Three union columns, a single explicit set {c0, c2} (participation 0b101); only c1 must be NULLed. + public void shouldHandleNonContiguousSetAndMultipleRows() { + /// Three union columns, a single explicit set {c0, c2} (member indexes 0 and 2); only the c1 key copy is NULLed. DataSchema inputSchema = new DataSchema(new String[]{"c0", "c1", "c2", "cnt"}, new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT}); - DataSchema resultSchema = new DataSchema(new String[]{"c0", "c1", "c2", "cnt", "$groupingId"}, + DataSchema resultSchema = new DataSchema( + new String[]{"c0", "c1", "c2", "cnt", "$groupingSetKey0", "$groupingSetKey1", "$groupingSetKey2", + "$groupingId"}, new ColumnDataType[]{ - ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT, ColumnDataType.INT + ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT, + ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.INT }); when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inputSchema, - new Object[]{"a", "m", "x", 1}, new Object[]{"b", "n", "y", 2})); + new Object[]{"a", "m", "x", 1}, new Object[]{"b", "n", "y", 2}), SuccessMseBlock.INSTANCE); RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0, 1, 2}, - List.of(0b101), resultSchema); + List.of(List.of(0, 2)), resultSchema); - List rows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + List rows = collectRows(operator); assertEquals(rows.size(), 2); /// two input rows x one set - /// groupingId = ~0b101 & 0b111 = 0b010; only c1 (idx 1) NULLed, c0/c2 kept. - assertEquals(rows.get(0), new Object[]{"a", null, "x", 1, 0b010}); - assertEquals(rows.get(1), new Object[]{"b", null, "y", 2, 0b010}); + assertEquals(rows.get(0), new Object[]{"a", "m", "x", 1, "a", null, "x", 0}); + assertEquals(rows.get(1), new Object[]{"b", "n", "y", 2, "b", null, "y", 0}); + } + + @Test + public void shouldEmitOneBlockPerGroupingSet() { + /// The expansion is streamed one set per output block, bounding the transient materialization by the input + /// block size: 2 input rows x 2 sets arrive as two 2-row blocks (set-major), not one 4-row block. + when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(INPUT_SCHEMA, + new Object[]{"a", "x", 1}, new Object[]{"b", "y", 2}), SuccessMseBlock.INSTANCE); + RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0, 1}, + List.of(List.of(0, 1), List.of()), RESULT_SCHEMA); + + List block0 = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + assertEquals(block0.size(), 2); + assertEquals(block0.get(0), new Object[]{"a", "x", 1, "a", "x", 0}); + assertEquals(block0.get(1), new Object[]{"b", "y", 2, "b", "y", 0}); + List block1 = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + assertEquals(block1.size(), 2); + assertEquals(block1.get(0), new Object[]{"a", "x", 1, null, null, 1}); + assertEquals(block1.get(1), new Object[]{"b", "y", 2, null, null, 1}); + assertTrue(operator.nextBlock().isEos()); + } + + @Test + public void shouldPropagateUpstreamErrorBlock() { + when(_input.nextBlock()).thenReturn(ErrorMseBlock.fromException(new Exception("boom!"))); + RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0, 1}, + List.of(List.of(0, 1), List.of()), RESULT_SCHEMA); + + MseBlock block = operator.nextBlock(); + assertTrue(block.isError(), "upstream error must pass through unexpanded"); + assertTrue(((ErrorMseBlock) block).getErrorMessages().values().toString().contains("boom!")); + } + + @Test + public void shouldPropagateUpstreamEos() { + when(_input.nextBlock()).thenReturn(SuccessMseBlock.INSTANCE); + RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0, 1}, + List.of(List.of(0, 1), List.of()), RESULT_SCHEMA); + + assertTrue(operator.nextBlock().isEos(), "EOS must pass through with no expansion"); + } + + @Test + public void shouldExpandEmptyDataBlock() { + /// An empty data block (e.g. a filtered-out segment) expands to empty per-set blocks, not an error. + when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(INPUT_SCHEMA), SuccessMseBlock.INSTANCE); + RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0, 1}, + List.of(List.of(0, 1), List.of()), RESULT_SCHEMA); + + assertEquals(collectRows(operator).size(), 0); + } + + @Test + public void shouldHandleMoreThanThirtyTwoUnionColumns() { + /// The retired bitmask encoding computed (1 << numUnionColumns) - 1, which overflows past 31 columns; the + /// ordinal encoding has no such limit. 34 union columns, ROLLUP-style sets (all columns) and (): the + /// grand-total set must NULL all 34 key copies and both ordinals must be emitted. + int numColumns = 34; + int numResultColumns = numColumns + 1 + numColumns + 1; + String[] columnNames = new String[numResultColumns]; + ColumnDataType[] resultTypes = new ColumnDataType[numResultColumns]; + ColumnDataType[] inputTypes = new ColumnDataType[numColumns + 1]; + Object[] inputRow = new Object[numColumns + 1]; + int[] unionGroupKeyIds = new int[numColumns]; + List allColumns = new ArrayList<>(numColumns); + for (int i = 0; i < numColumns; i++) { + columnNames[i] = "c" + i; + inputTypes[i] = ColumnDataType.STRING; + resultTypes[i] = ColumnDataType.STRING; + inputRow[i] = "v" + i; + unionGroupKeyIds[i] = i; + allColumns.add(i); + columnNames[numColumns + 1 + i] = "$groupingSetKey" + i; + resultTypes[numColumns + 1 + i] = ColumnDataType.STRING; + } + columnNames[numColumns] = "cnt"; + inputTypes[numColumns] = ColumnDataType.INT; + resultTypes[numColumns] = ColumnDataType.INT; + inputRow[numColumns] = 7; + columnNames[numResultColumns - 1] = "$groupingId"; + resultTypes[numResultColumns - 1] = ColumnDataType.INT; + String[] inputNames = new String[numColumns + 1]; + System.arraycopy(columnNames, 0, inputNames, 0, numColumns + 1); + DataSchema inputSchema = new DataSchema(inputNames, inputTypes); + DataSchema resultSchema = new DataSchema(columnNames, resultTypes); + + when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inputSchema, inputRow), SuccessMseBlock.INSTANCE); + RepeatOperator operator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, unionGroupKeyIds, + List.of(allColumns, List.of()), resultSchema); + + List rows = collectRows(operator); + assertEquals(rows.size(), 2); + /// Set (all columns), ordinal 0: every key copy carries the value. + for (int i = 0; i < numColumns; i++) { + assertEquals(rows.get(0)[numColumns + 1 + i], "v" + i); + } + assertEquals(rows.get(0)[numResultColumns - 1], 0); + /// Grand-total set (), ordinal 1: every key copy NULLed — including the ones past the 32-bit boundary — while + /// the original columns stay untouched. + for (int i = 0; i < numColumns; i++) { + assertNull(rows.get(1)[numColumns + 1 + i]); + assertEquals(rows.get(1)[i], "v" + i); + } + assertEquals(rows.get(1)[numColumns], 7); + assertEquals(rows.get(1)[numResultColumns - 1], 1); } }