Support GROUP BY GROUPING SETS / ROLLUP / CUBE in the multi-stage query engine

pinot-common/src/main/java/org/apache/pinot/common/request/context/GroupingSets.java

@@ -18,40 +18,105 @@
  */
 package org.apache.pinot.common.request.context;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /// 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";
 
-  /// 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.
+  /// 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;
+
+  /// 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 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<int[]> 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<List<Integer>>} sets and {@code List<Integer>} argument indexes).
+  public static int[] groupingValuesByOrdinal(List<? extends List<Integer>> groupingSets,
+      List<Integer> argUnionIndexes) {
+    List<int[]> sets = new ArrayList<>(groupingSets.size());
+    for (List<Integer> 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

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<QueryRewriter> 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<Expression> 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<Integer> seen = new HashSet<>();
-    List<Integer> 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<List<Integer>> seen = new HashSet<>();
+    List<List<Integer>> groupingSets = new ArrayList<>();
     for (LinkedHashSet<Expression> set : combinedSets) {
-      int mask = 0;
+      List<Integer> 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) {

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()) {

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,6 +80,13 @@ message AggregateNode {
   bool leafReturnFinalResult = 5;
   repeated Collation collations = 6;
   int32 limit = 7;
+  // 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 {