You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2026/06/20 00:27:37 UTC
(pinot) branch master updated: Support GROUP BY GROUPING SETS / ROLLUP / CUBE in the single-stage engine (#18662)
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 0bceb35a8de Support GROUP BY GROUPING SETS / ROLLUP / CUBE in the single-stage engine (#18662)
0bceb35a8de is described below
commit 0bceb35a8de6f3dc787ebbea7b3123aad876450d
Author: Xiang Fu <xi...@apache.org>
AuthorDate: Fri Jun 19 17:27:31 2026 -0700
Support GROUP BY GROUPING SETS / ROLLUP / CUBE in the single-stage engine (#18662)
Add native single-scan GROUPING SETS / ROLLUP / CUBE plus the GROUPING() /
GROUPING_ID() functions (PostgreSQL bitmask semantics) to Pinot's
single-stage (v1) query engine.
Approach: a grouping-sets query is represented as the union of all grouping
columns plus, per set, a membership bitmask over that union (ROLLUP/CUBE are
expanded to grouping sets at parse time). Each input row is expanded in a
single scan into one group per set via the multi-value aggregation path. A
synthetic $grouping_id key column (the per-set grouping-id bitmask) is
appended so rows from different sets never merge -- keeping a genuine
(a, NULL) detail row distinct from a rolled-up (a, NULL) subtotal -- and
powers GROUPING()/GROUPING_ID(). Grouping-set key columns are serialized
null-aware regardless of the query's null-handling option, since rolled-up
columns are always NULL.
Highlights:
- Parser/wire: CalciteSqlParser expands and de-duplicates grouping sets;
PinotQuery carries an optional list<i32> groupingSetMasks (one membership
bitmask per set). Plain GROUP BY leaves it unset, so the PinotQuery is
byte-identical for existing queries and a not-yet-upgraded server decodes
it unchanged (broker-first rolling upgrade stays safe).
- Execution: GroupingSetsGroupKeyGenerator expands each row across sets (and
across multi-value column values), reusing on-the-fly dictionaries for
null-representable keys.
- Reduce: the $grouping_id discriminator flows through the IndexedTable merge
key; GROUPING()/GROUPING_ID() are evaluated at the broker as a
post-aggregation extractor over that column.
- Trim safety: per-segment group trim is bucketed by grouping set so a global
top-K cannot starve a low-magnitude set such as the grand total, and the
intermediate combine and broker-reduce trim is disabled for grouping-set
queries so no group is dropped before its partial aggregates merge; the
broker applies the final ORDER BY + LIMIT over the fully-merged table.
- Supports multi-value grouping columns and filtered aggregations. Star-tree
and server-final-result are bypassed for grouping-set queries.
Tests: parser expansion + limit unit tests; TableResizer per-set trim and
GroupByUtils trim-disabled unit tests; a ReducerDataSchemaUtils mixed-version
guard test; a 16-test end-to-end integration suite (real cluster, >=2
segments) covering the NULL discriminator, NULL round-trip with and without
null handling, INT/LONG/DOUBLE/STRING key types, GROUPING in
SELECT/HAVING/ORDER BY, multi-value columns, filtered aggregations, ORDER BY
on an aggregation, and a plain-GROUP-BY regression.
Co-authored-by: Claude Opus 4.8 (1M context) <no...@anthropic.com>
---
.../apache/pinot/common/request/PinotQuery.java | 290 +++++++--
.../pinot/common/request/context/GroupingSets.java | 62 ++
.../apache/pinot/sql/parsers/CalciteSqlParser.java | 203 +++++-
...nAggregationGroupByToDistinctQueryRewriter.java | 6 +
.../pinot/sql/parsers/GroupingSetsParserTest.java | 178 +++++
pinot-common/src/thrift/query.thrift | 6 +
.../apache/pinot/core/data/table/IndexedTable.java | 4 +-
.../pinot/core/data/table/SortedRecordTable.java | 2 +-
.../pinot/core/data/table/SortedRecordsMerger.java | 2 +-
.../apache/pinot/core/data/table/TableResizer.java | 103 ++-
.../blocks/results/GroupByResultsBlock.java | 10 +-
.../operator/blocks/results/ResultsBlockUtils.java | 12 +-
.../operator/combine/GroupByCombineOperator.java | 10 +-
.../operator/query/FilteredGroupByOperator.java | 24 +-
.../pinot/core/operator/query/GroupByOperator.java | 26 +-
.../function/AggregationFunctionUtils.java | 5 +
.../groupby/DefaultGroupByExecutor.java | 8 +-
.../groupby/GroupingSetsGroupKeyGenerator.java | 423 ++++++++++++
.../core/query/reduce/GroupByDataTableReducer.java | 40 +-
.../core/query/reduce/PostAggregationHandler.java | 64 +-
.../core/query/reduce/ReducerDataSchemaUtils.java | 26 +-
.../core/query/request/context/QueryContext.java | 89 ++-
.../context/utils/QueryContextConverterUtils.java | 40 ++
.../org/apache/pinot/core/util/GroupByUtils.java | 61 ++
.../pinot/core/data/table/TableResizerTest.java | 54 ++
.../query/reduce/ReducerDataSchemaUtilsTest.java | 29 +
.../BrokerRequestToQueryContextConverterTest.java | 36 ++
.../apache/pinot/core/util/GroupByUtilsTest.java | 62 ++
.../tests/custom/GroupingSetsQueriesTest.java | 720 +++++++++++++++++++++
29 files changed, 2482 insertions(+), 113 deletions(-)
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 5b3d79f0764..45f0de1bbf8 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,6 +41,7 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
private static final org.apache.thrift.protocol.TField QUERY_OPTIONS_FIELD_DESC = new org.apache.thrift.protocol.TField("queryOptions", org.apache.thrift.protocol.TType.MAP, (short)11);
private static final org.apache.thrift.protocol.TField EXPLAIN_FIELD_DESC = new org.apache.thrift.protocol.TField("explain", org.apache.thrift.protocol.TType.BOOL, (short)12);
private static final org.apache.thrift.protocol.TField EXPRESSION_OVERRIDE_HINTS_FIELD_DESC = new org.apache.thrift.protocol.TField("expressionOverrideHints", org.apache.thrift.protocol.TType.MAP, (short)13);
+ private static final org.apache.thrift.protocol.TField GROUPING_SET_MASKS_FIELD_DESC = new org.apache.thrift.protocol.TField("groupingSetMasks", org.apache.thrift.protocol.TType.LIST, (short)14);
private static final org.apache.thrift.scheme.SchemeFactory STANDARD_SCHEME_FACTORY = new PinotQueryStandardSchemeFactory();
private static final org.apache.thrift.scheme.SchemeFactory TUPLE_SCHEME_FACTORY = new PinotQueryTupleSchemeFactory();
@@ -57,6 +58,7 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
private @org.apache.thrift.annotation.Nullable java.util.Map<java.lang.String,java.lang.String> queryOptions; // optional
private boolean explain; // optional
private @org.apache.thrift.annotation.Nullable java.util.Map<Expression,Expression> expressionOverrideHints; // optional
+ private @org.apache.thrift.annotation.Nullable java.util.List<java.lang.Integer> groupingSetMasks; // 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 {
@@ -71,7 +73,8 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
OFFSET((short)9, "offset"),
QUERY_OPTIONS((short)11, "queryOptions"),
EXPLAIN((short)12, "explain"),
- EXPRESSION_OVERRIDE_HINTS((short)13, "expressionOverrideHints");
+ EXPRESSION_OVERRIDE_HINTS((short)13, "expressionOverrideHints"),
+ GROUPING_SET_MASKS((short)14, "groupingSetMasks");
private static final java.util.Map<java.lang.String, _Fields> byName = new java.util.HashMap<java.lang.String, _Fields>();
@@ -111,6 +114,8 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
return EXPLAIN;
case 13: // EXPRESSION_OVERRIDE_HINTS
return EXPRESSION_OVERRIDE_HINTS;
+ case 14: // GROUPING_SET_MASKS
+ return GROUPING_SET_MASKS;
default:
return null;
}
@@ -159,7 +164,7 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
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};
+ 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};
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);
@@ -194,6 +199,9 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
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,
+ 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);
}
@@ -263,6 +271,10 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
}
this.expressionOverrideHints = __this__expressionOverrideHints;
}
+ if (other.isSetGroupingSetMasks()) {
+ java.util.List<java.lang.Integer> __this__groupingSetMasks = new java.util.ArrayList<java.lang.Integer>(other.groupingSetMasks);
+ this.groupingSetMasks = __this__groupingSetMasks;
+ }
}
@Override
@@ -288,6 +300,7 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
setExplainIsSet(false);
this.explain = false;
this.expressionOverrideHints = null;
+ this.groupingSetMasks = null;
}
public int getVersion() {
@@ -640,6 +653,46 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
}
}
+ public int getGroupingSetMasksSize() {
+ return (this.groupingSetMasks == null) ? 0 : this.groupingSetMasks.size();
+ }
+
+ @org.apache.thrift.annotation.Nullable
+ public java.util.Iterator<java.lang.Integer> getGroupingSetMasksIterator() {
+ return (this.groupingSetMasks == null) ? null : this.groupingSetMasks.iterator();
+ }
+
+ public void addToGroupingSetMasks(int elem) {
+ if (this.groupingSetMasks == null) {
+ this.groupingSetMasks = new java.util.ArrayList<java.lang.Integer>();
+ }
+ this.groupingSetMasks.add(elem);
+ }
+
+ @org.apache.thrift.annotation.Nullable
+ public java.util.List<java.lang.Integer> getGroupingSetMasks() {
+ return this.groupingSetMasks;
+ }
+
+ public void setGroupingSetMasks(@org.apache.thrift.annotation.Nullable java.util.List<java.lang.Integer> groupingSetMasks) {
+ this.groupingSetMasks = groupingSetMasks;
+ }
+
+ public void unsetGroupingSetMasks() {
+ this.groupingSetMasks = null;
+ }
+
+ /** Returns true if field groupingSetMasks is set (has been assigned a value) and false otherwise */
+ public boolean isSetGroupingSetMasks() {
+ return this.groupingSetMasks != null;
+ }
+
+ public void setGroupingSetMasksIsSet(boolean value) {
+ if (!value) {
+ this.groupingSetMasks = null;
+ }
+ }
+
@Override
public void setFieldValue(_Fields field, @org.apache.thrift.annotation.Nullable java.lang.Object value) {
switch (field) {
@@ -739,6 +792,14 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
}
break;
+ case GROUPING_SET_MASKS:
+ if (value == null) {
+ unsetGroupingSetMasks();
+ } else {
+ setGroupingSetMasks((java.util.List<java.lang.Integer>)value);
+ }
+ break;
+
}
}
@@ -782,6 +843,9 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
case EXPRESSION_OVERRIDE_HINTS:
return getExpressionOverrideHints();
+ case GROUPING_SET_MASKS:
+ return getGroupingSetMasks();
+
}
throw new java.lang.IllegalStateException();
}
@@ -818,6 +882,8 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
return isSetExplain();
case EXPRESSION_OVERRIDE_HINTS:
return isSetExpressionOverrideHints();
+ case GROUPING_SET_MASKS:
+ return isSetGroupingSetMasks();
}
throw new java.lang.IllegalStateException();
}
@@ -943,6 +1009,15 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
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))
+ return false;
+ if (!this.groupingSetMasks.equals(that.groupingSetMasks))
+ return false;
+ }
+
return true;
}
@@ -998,6 +1073,10 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
if (isSetExpressionOverrideHints())
hashCode = hashCode * 8191 + expressionOverrideHints.hashCode();
+ hashCode = hashCode * 8191 + ((isSetGroupingSetMasks()) ? 131071 : 524287);
+ if (isSetGroupingSetMasks())
+ hashCode = hashCode * 8191 + groupingSetMasks.hashCode();
+
return hashCode;
}
@@ -1129,6 +1208,16 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
return lastComparison;
}
}
+ lastComparison = java.lang.Boolean.compare(isSetGroupingSetMasks(), other.isSetGroupingSetMasks());
+ if (lastComparison != 0) {
+ return lastComparison;
+ }
+ if (isSetGroupingSetMasks()) {
+ lastComparison = org.apache.thrift.TBaseHelper.compareTo(this.groupingSetMasks, other.groupingSetMasks);
+ if (lastComparison != 0) {
+ return lastComparison;
+ }
+ }
return 0;
}
@@ -1256,6 +1345,16 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
}
first = false;
}
+ if (isSetGroupingSetMasks()) {
+ if (!first) sb.append(", ");
+ sb.append("groupingSetMasks:");
+ if (this.groupingSetMasks == null) {
+ sb.append("null");
+ } else {
+ sb.append(this.groupingSetMasks);
+ }
+ first = false;
+ }
sb.append(")");
return sb.toString();
}
@@ -1470,6 +1569,24 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
org.apache.thrift.protocol.TProtocolUtil.skip(iprot, schemeField.type);
}
break;
+ case 14: // GROUPING_SET_MASKS
+ if (schemeField.type == org.apache.thrift.protocol.TType.LIST) {
+ {
+ org.apache.thrift.protocol.TList _list17 = iprot.readListBegin();
+ struct.groupingSetMasks = new java.util.ArrayList<java.lang.Integer>(_list17.size);
+ int _elem18;
+ for (int _i19 = 0; _i19 < _list17.size; ++_i19)
+ {
+ _elem18 = iprot.readI32();
+ struct.groupingSetMasks.add(_elem18);
+ }
+ iprot.readListEnd();
+ }
+ struct.setGroupingSetMasksIsSet(true);
+ } else {
+ org.apache.thrift.protocol.TProtocolUtil.skip(iprot, schemeField.type);
+ }
+ break;
default:
org.apache.thrift.protocol.TProtocolUtil.skip(iprot, schemeField.type);
}
@@ -1501,9 +1618,9 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
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 _iter17 : struct.selectList)
+ for (Expression _iter20 : struct.selectList)
{
- _iter17.write(oprot);
+ _iter20.write(oprot);
}
oprot.writeListEnd();
}
@@ -1522,9 +1639,9 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
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 _iter18 : struct.groupByList)
+ for (Expression _iter21 : struct.groupByList)
{
- _iter18.write(oprot);
+ _iter21.write(oprot);
}
oprot.writeListEnd();
}
@@ -1536,9 +1653,9 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
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 _iter19 : struct.orderByList)
+ for (Expression _iter22 : struct.orderByList)
{
- _iter19.write(oprot);
+ _iter22.write(oprot);
}
oprot.writeListEnd();
}
@@ -1567,10 +1684,10 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
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<java.lang.String, java.lang.String> _iter20 : struct.queryOptions.entrySet())
+ for (java.util.Map.Entry<java.lang.String, java.lang.String> _iter23 : struct.queryOptions.entrySet())
{
- oprot.writeString(_iter20.getKey());
- oprot.writeString(_iter20.getValue());
+ oprot.writeString(_iter23.getKey());
+ oprot.writeString(_iter23.getValue());
}
oprot.writeMapEnd();
}
@@ -1587,16 +1704,30 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
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<Expression, Expression> _iter21 : struct.expressionOverrideHints.entrySet())
+ for (java.util.Map.Entry<Expression, Expression> _iter24 : struct.expressionOverrideHints.entrySet())
{
- _iter21.getKey().write(oprot);
- _iter21.getValue().write(oprot);
+ _iter24.getKey().write(oprot);
+ _iter24.getValue().write(oprot);
}
oprot.writeMapEnd();
}
oprot.writeFieldEnd();
}
}
+ if (struct.groupingSetMasks != null) {
+ if (struct.isSetGroupingSetMasks()) {
+ oprot.writeFieldBegin(GROUPING_SET_MASKS_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.writeI32(_iter25);
+ }
+ oprot.writeListEnd();
+ }
+ oprot.writeFieldEnd();
+ }
+ }
oprot.writeFieldStop();
oprot.writeStructEnd();
}
@@ -1652,7 +1783,10 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
if (struct.isSetExpressionOverrideHints()) {
optionals.set(11);
}
- oprot.writeBitSet(optionals, 12);
+ if (struct.isSetGroupingSetMasks()) {
+ optionals.set(12);
+ }
+ oprot.writeBitSet(optionals, 13);
if (struct.isSetVersion()) {
oprot.writeI32(struct.version);
}
@@ -1662,9 +1796,9 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
if (struct.isSetSelectList()) {
{
oprot.writeI32(struct.selectList.size());
- for (Expression _iter22 : struct.selectList)
+ for (Expression _iter26 : struct.selectList)
{
- _iter22.write(oprot);
+ _iter26.write(oprot);
}
}
}
@@ -1674,18 +1808,18 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
if (struct.isSetGroupByList()) {
{
oprot.writeI32(struct.groupByList.size());
- for (Expression _iter23 : struct.groupByList)
+ for (Expression _iter27 : struct.groupByList)
{
- _iter23.write(oprot);
+ _iter27.write(oprot);
}
}
}
if (struct.isSetOrderByList()) {
{
oprot.writeI32(struct.orderByList.size());
- for (Expression _iter24 : struct.orderByList)
+ for (Expression _iter28 : struct.orderByList)
{
- _iter24.write(oprot);
+ _iter28.write(oprot);
}
}
}
@@ -1701,10 +1835,10 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
if (struct.isSetQueryOptions()) {
{
oprot.writeI32(struct.queryOptions.size());
- for (java.util.Map.Entry<java.lang.String, java.lang.String> _iter25 : struct.queryOptions.entrySet())
+ for (java.util.Map.Entry<java.lang.String, java.lang.String> _iter29 : struct.queryOptions.entrySet())
{
- oprot.writeString(_iter25.getKey());
- oprot.writeString(_iter25.getValue());
+ oprot.writeString(_iter29.getKey());
+ oprot.writeString(_iter29.getValue());
}
}
}
@@ -1714,10 +1848,19 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
if (struct.isSetExpressionOverrideHints()) {
{
oprot.writeI32(struct.expressionOverrideHints.size());
- for (java.util.Map.Entry<Expression, Expression> _iter26 : struct.expressionOverrideHints.entrySet())
+ for (java.util.Map.Entry<Expression, Expression> _iter30 : struct.expressionOverrideHints.entrySet())
{
- _iter26.getKey().write(oprot);
- _iter26.getValue().write(oprot);
+ _iter30.getKey().write(oprot);
+ _iter30.getValue().write(oprot);
+ }
+ }
+ }
+ if (struct.isSetGroupingSetMasks()) {
+ {
+ oprot.writeI32(struct.groupingSetMasks.size());
+ for (int _iter31 : struct.groupingSetMasks)
+ {
+ oprot.writeI32(_iter31);
}
}
}
@@ -1726,7 +1869,7 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
@Override
public void read(org.apache.thrift.protocol.TProtocol prot, PinotQuery struct) throws org.apache.thrift.TException {
org.apache.thrift.protocol.TTupleProtocol iprot = (org.apache.thrift.protocol.TTupleProtocol) prot;
- java.util.BitSet incoming = iprot.readBitSet(12);
+ java.util.BitSet incoming = iprot.readBitSet(13);
if (incoming.get(0)) {
struct.version = iprot.readI32();
struct.setVersionIsSet(true);
@@ -1738,14 +1881,14 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
}
if (incoming.get(2)) {
{
- org.apache.thrift.protocol.TList _list27 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT);
- struct.selectList = new java.util.ArrayList<Expression>(_list27.size);
- @org.apache.thrift.annotation.Nullable Expression _elem28;
- for (int _i29 = 0; _i29 < _list27.size; ++_i29)
+ org.apache.thrift.protocol.TList _list32 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT);
+ struct.selectList = new java.util.ArrayList<Expression>(_list32.size);
+ @org.apache.thrift.annotation.Nullable Expression _elem33;
+ for (int _i34 = 0; _i34 < _list32.size; ++_i34)
{
- _elem28 = new Expression();
- _elem28.read(iprot);
- struct.selectList.add(_elem28);
+ _elem33 = new Expression();
+ _elem33.read(iprot);
+ struct.selectList.add(_elem33);
}
}
struct.setSelectListIsSet(true);
@@ -1757,28 +1900,28 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
}
if (incoming.get(4)) {
{
- org.apache.thrift.protocol.TList _list30 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT);
- struct.groupByList = new java.util.ArrayList<Expression>(_list30.size);
- @org.apache.thrift.annotation.Nullable Expression _elem31;
- for (int _i32 = 0; _i32 < _list30.size; ++_i32)
+ org.apache.thrift.protocol.TList _list35 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT);
+ struct.groupByList = new java.util.ArrayList<Expression>(_list35.size);
+ @org.apache.thrift.annotation.Nullable Expression _elem36;
+ for (int _i37 = 0; _i37 < _list35.size; ++_i37)
{
- _elem31 = new Expression();
- _elem31.read(iprot);
- struct.groupByList.add(_elem31);
+ _elem36 = new Expression();
+ _elem36.read(iprot);
+ struct.groupByList.add(_elem36);
}
}
struct.setGroupByListIsSet(true);
}
if (incoming.get(5)) {
{
- org.apache.thrift.protocol.TList _list33 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT);
- struct.orderByList = new java.util.ArrayList<Expression>(_list33.size);
- @org.apache.thrift.annotation.Nullable Expression _elem34;
- for (int _i35 = 0; _i35 < _list33.size; ++_i35)
+ org.apache.thrift.protocol.TList _list38 = iprot.readListBegin(org.apache.thrift.protocol.TType.STRUCT);
+ struct.orderByList = new java.util.ArrayList<Expression>(_list38.size);
+ @org.apache.thrift.annotation.Nullable Expression _elem39;
+ for (int _i40 = 0; _i40 < _list38.size; ++_i40)
{
- _elem34 = new Expression();
- _elem34.read(iprot);
- struct.orderByList.add(_elem34);
+ _elem39 = new Expression();
+ _elem39.read(iprot);
+ struct.orderByList.add(_elem39);
}
}
struct.setOrderByListIsSet(true);
@@ -1798,15 +1941,15 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
}
if (incoming.get(9)) {
{
- org.apache.thrift.protocol.TMap _map36 = iprot.readMapBegin(org.apache.thrift.protocol.TType.STRING, org.apache.thrift.protocol.TType.STRING);
- struct.queryOptions = new java.util.HashMap<java.lang.String,java.lang.String>(2*_map36.size);
- @org.apache.thrift.annotation.Nullable java.lang.String _key37;
- @org.apache.thrift.annotation.Nullable java.lang.String _val38;
- for (int _i39 = 0; _i39 < _map36.size; ++_i39)
+ 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<java.lang.String,java.lang.String>(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)
{
- _key37 = iprot.readString();
- _val38 = iprot.readString();
- struct.queryOptions.put(_key37, _val38);
+ _key42 = iprot.readString();
+ _val43 = iprot.readString();
+ struct.queryOptions.put(_key42, _val43);
}
}
struct.setQueryOptionsIsSet(true);
@@ -1817,21 +1960,34 @@ public class PinotQuery implements org.apache.thrift.TBase<PinotQuery, PinotQuer
}
if (incoming.get(11)) {
{
- org.apache.thrift.protocol.TMap _map40 = iprot.readMapBegin(org.apache.thrift.protocol.TType.STRUCT, org.apache.thrift.protocol.TType.STRUCT);
- struct.expressionOverrideHints = new java.util.HashMap<Expression,Expression>(2*_map40.size);
- @org.apache.thrift.annotation.Nullable Expression _key41;
- @org.apache.thrift.annotation.Nullable Expression _val42;
- for (int _i43 = 0; _i43 < _map40.size; ++_i43)
+ 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<Expression,Expression>(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)
{
- _key41 = new Expression();
- _key41.read(iprot);
- _val42 = new Expression();
- _val42.read(iprot);
- struct.expressionOverrideHints.put(_key41, _val42);
+ _key46 = new Expression();
+ _key46.read(iprot);
+ _val47 = new Expression();
+ _val47.read(iprot);
+ struct.expressionOverrideHints.put(_key46, _val47);
}
}
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<java.lang.Integer>(_list49.size);
+ int _elem50;
+ for (int _i51 = 0; _i51 < _list49.size; ++_i51)
+ {
+ _elem50 = iprot.readI32();
+ struct.groupingSetMasks.add(_elem50);
+ }
+ }
+ struct.setGroupingSetMasksIsSet(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
new file mode 100644
index 00000000000..ab071b1ecc6
--- /dev/null
+++ b/pinot-common/src/main/java/org/apache/pinot/common/request/context/GroupingSets.java
@@ -0,0 +1,62 @@
+/**
+ * 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;
+
+/// 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().
+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()).
+ 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.
+ ///
+ /// @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) {
+ int result = 0;
+ for (int unionColumnIndex : unionColumnIndexes) {
+ result = (result << 1) | ((groupingId >>> unionColumnIndex) & 1);
+ }
+ return result;
+ }
+
+ /// 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
+ /// ORDER BY comparator so the two sites cannot drift.
+ public static boolean isGroupingFunction(String functionName) {
+ return functionName.equalsIgnoreCase("grouping") || functionName.equalsIgnoreCase("groupingid");
+ }
+}
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 d3eb407480d..932ff800637 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
@@ -25,6 +25,8 @@ import java.io.StringReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -33,6 +35,7 @@ import java.util.regex.Pattern;
import javax.annotation.Nullable;
import org.apache.calcite.avatica.util.Casing;
import org.apache.calcite.sql.SqlBasicCall;
+import org.apache.calcite.sql.SqlCall;
import org.apache.calcite.sql.SqlDataTypeSpec;
import org.apache.calcite.sql.SqlExplain;
import org.apache.calcite.sql.SqlIdentifier;
@@ -94,6 +97,11 @@ public class 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;
private static final Logger LOGGER = LoggerFactory.getLogger(CalciteSqlParser.class);
// To Keep the backward compatibility with 'OPTION' Functionality in PQL, which is used to
@@ -242,6 +250,32 @@ public class CalciteSqlParser {
}
}
}
+
+ // A GROUPING SETS / ROLLUP / CUBE query must contain at least one aggregation (in SELECT, HAVING or
+ // ORDER-BY): without one the engine would execute it as a selection query and silently ignore the grouping
+ // 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(
+ pinotQuery)) {
+ throw new SqlCompilationException(
+ "GROUP BY GROUPING SETS / ROLLUP / CUBE requires at least one aggregation function in the query");
+ }
+ }
+
+ /// Returns true if the HAVING clause or any ORDER-BY expression contains an aggregation.
+ private static boolean hasAggregationOutsideSelect(PinotQuery pinotQuery) {
+ if (pinotQuery.getHavingExpression() != null && isAggregateExpression(pinotQuery.getHavingExpression())) {
+ return true;
+ }
+ if (pinotQuery.getOrderByList() != null) {
+ for (Expression orderBy : pinotQuery.getOrderByList()) {
+ if (isAggregateExpression(orderBy)) {
+ return true;
+ }
+ }
+ }
+ return false;
}
/*
@@ -537,7 +571,7 @@ public class CalciteSqlParser {
// GROUP-BY
SqlNodeList groupByNodeList = selectNode.getGroup();
if (groupByNodeList != null) {
- pinotQuery.setGroupByList(convertSelectList(groupByNodeList));
+ setGroupByListAndGroupingSets(pinotQuery, groupByNodeList);
}
// HAVING
SqlNode havingNode = selectNode.getHaving();
@@ -690,6 +724,173 @@ public class CalciteSqlParser {
return selectExpr;
}
+ /// Converts the GROUP BY clause into {@link PinotQuery#groupByList} and (when grouping constructs are
+ /// present) {@link PinotQuery#groupingSetMasks}.
+ ///
+ /// 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
+ /// 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 ()}.
+ private static void setGroupByListAndGroupingSets(PinotQuery pinotQuery, SqlNodeList groupByNodeList) {
+ boolean hasGroupingConstruct = false;
+ for (SqlNode node : groupByNodeList) {
+ if (isGroupingConstruct(node.getKind())) {
+ hasGroupingConstruct = true;
+ break;
+ }
+ }
+ if (!hasGroupingConstruct) {
+ pinotQuery.setGroupByList(convertSelectList(groupByNodeList));
+ return;
+ }
+
+ /// Cross-multiply the grouping elements: the overall grouping sets are the union of one chosen set from
+ /// each grouping element. Start with a single empty set (the multiplicative identity).
+ List<LinkedHashSet<Expression>> combinedSets = new ArrayList<>();
+ combinedSets.add(new LinkedHashSet<>());
+ for (SqlNode element : groupByNodeList) {
+ List<List<Expression>> elementSets = parseGroupingElement(element);
+ List<LinkedHashSet<Expression>> next = new ArrayList<>(combinedSets.size() * elementSets.size());
+ for (LinkedHashSet<Expression> prefix : combinedSets) {
+ for (List<Expression> choice : elementSets) {
+ LinkedHashSet<Expression> merged = new LinkedHashSet<>(prefix);
+ merged.addAll(choice);
+ next.add(merged);
+ }
+ }
+ if (next.size() > MAX_GROUPING_SETS) {
+ throw new SqlCompilationException(
+ "GROUPING SETS / ROLLUP / CUBE expands to more than " + MAX_GROUPING_SETS + " grouping sets");
+ }
+ combinedSets = next;
+ }
+
+ /// Build the ordered, de-duplicated union of all participating columns (first-appearance order).
+ LinkedHashMap<Expression, Integer> unionIndex = new LinkedHashMap<>();
+ for (LinkedHashSet<Expression> set : combinedSets) {
+ for (Expression expr : set) {
+ 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<>();
+ for (LinkedHashSet<Expression> set : combinedSets) {
+ int mask = 0;
+ for (Expression expr : set) {
+ mask |= 1 << unionIndex.get(expr);
+ }
+ if (seen.add(mask)) {
+ groupingSetMasks.add(mask);
+ }
+ }
+ pinotQuery.setGroupingSetMasks(groupingSetMasks);
+ }
+
+ private static boolean isGroupingConstruct(SqlKind kind) {
+ return kind == SqlKind.ROLLUP || kind == SqlKind.CUBE || kind == SqlKind.GROUPING_SETS;
+ }
+
+ /// Expands a single grouping element into the list of grouping sets it represents (each set is an ordered
+ /// list of column expressions; the empty list is the grand-total set).
+ /// - {@code ROLLUP(l1, ..., ln)} -> the n+1 prefixes {@code {l1..ln}, {l1..ln-1}, ..., {l1}, {}}
+ /// - {@code CUBE(l1, ..., ln)} -> the power set of the n levels
+ /// - {@code GROUPING SETS(g1, ..., gm)} -> the concatenation of each operand's expansion (operands may be
+ /// nested ROLLUP/CUBE/GROUPING SETS or ordinary sets)
+ /// - an ordinary grouping element (a single column or a parenthesized list) -> a single set
+ private static List<List<Expression>> parseGroupingElement(SqlNode node) {
+ switch (node.getKind()) {
+ case ROLLUP: {
+ List<List<Expression>> levels = parseLevels((SqlCall) node);
+ List<List<Expression>> sets = new ArrayList<>(levels.size() + 1);
+ for (int numLevels = levels.size(); numLevels >= 0; numLevels--) {
+ List<Expression> set = new ArrayList<>();
+ for (int i = 0; i < numLevels; i++) {
+ set.addAll(levels.get(i));
+ }
+ sets.add(set);
+ }
+ return sets;
+ }
+ case CUBE: {
+ List<List<Expression>> levels = parseLevels((SqlCall) node);
+ int numLevels = levels.size();
+ /// Guard the shift against overflow (1L << 64 wraps to 1) before comparing against the set-count cap.
+ if (numLevels >= Integer.SIZE - 1 || (1L << numLevels) > MAX_GROUPING_SETS) {
+ throw new SqlCompilationException(
+ "CUBE expands to more than " + MAX_GROUPING_SETS + " grouping sets");
+ }
+ List<List<Expression>> sets = new ArrayList<>(1 << numLevels);
+ for (int mask = (1 << numLevels) - 1; mask >= 0; mask--) {
+ List<Expression> set = new ArrayList<>();
+ for (int i = 0; i < numLevels; i++) {
+ if ((mask & (1 << i)) != 0) {
+ set.addAll(levels.get(i));
+ }
+ }
+ sets.add(set);
+ }
+ return sets;
+ }
+ case GROUPING_SETS: {
+ List<List<Expression>> sets = new ArrayList<>();
+ for (SqlNode operand : ((SqlCall) node).getOperandList()) {
+ sets.addAll(parseGroupingElement(operand));
+ }
+ return sets;
+ }
+ default:
+ /// Ordinary grouping element: a single column expression or a parenthesized list of columns.
+ return List.of(parseLevel(node));
+ }
+ }
+
+ /// Parses each operand of a ROLLUP/CUBE call into a "level" (a level may be a single column or a
+ /// parenthesized list of columns that roll up together).
+ private static List<List<Expression>> parseLevels(SqlCall call) {
+ List<List<Expression>> levels = new ArrayList<>(call.getOperandList().size());
+ for (SqlNode operand : call.getOperandList()) {
+ levels.add(parseLevel(operand));
+ }
+ return levels;
+ }
+
+ /// Parses a single grouping level/set node into its column expressions. Handles a parenthesized list
+ /// (modeled by Calcite as either a {@link SqlNodeList} or a ROW call), and a bare single column. An empty
+ /// parenthesized list yields an empty column list (the grand-total set).
+ private static List<Expression> parseLevel(SqlNode node) {
+ if (node instanceof SqlNodeList) {
+ SqlNodeList list = (SqlNodeList) node;
+ List<Expression> columns = new ArrayList<>(list.size());
+ for (SqlNode column : list) {
+ columns.add(toExpression(column));
+ }
+ return columns;
+ }
+ if (node.getKind() == SqlKind.ROW) {
+ List<SqlNode> operands = ((SqlCall) node).getOperandList();
+ List<Expression> columns = new ArrayList<>(operands.size());
+ for (SqlNode column : operands) {
+ columns.add(toExpression(column));
+ }
+ return columns;
+ }
+ return List.of(toExpression(node));
+ }
+
private static List<Expression> convertOrderByList(SqlNodeList orderList) {
List<Expression> orderByExpr = new ArrayList<>(orderList.size());
for (SqlNode sqlNode : orderList) {
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 ef64deb5434..35f0e68ab22 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
@@ -51,6 +51,12 @@ public class NonAggregationGroupByToDistinctQueryRewriter implements QueryRewrit
if (pinotQuery.getGroupByListSize() == 0) {
return 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) {
+ return pinotQuery;
+ }
for (Expression select : pinotQuery.getSelectList()) {
if (CalciteSqlParser.isAggregateExpression(select)) {
return pinotQuery;
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
new file mode 100644
index 00000000000..ebfaaa3adab
--- /dev/null
+++ b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/GroupingSetsParserTest.java
@@ -0,0 +1,178 @@
+/**
+ * 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.sql.parsers;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import org.apache.pinot.common.request.Expression;
+import org.apache.pinot.common.request.PinotQuery;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
+
+
+/// 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()}).
+public class GroupingSetsParserTest {
+
+ /// Returns the grouping sets as a set of column-name sets, decoupled from union/set ordering.
+ private static Set<Set<String>> groupingSetsByName(PinotQuery query) {
+ List<Expression> union = query.getGroupByList();
+ Set<Set<String>> result = new HashSet<>();
+ for (int mask : query.getGroupingSetMasks()) {
+ Set<String> names = new HashSet<>();
+ for (int bit = 0; bit < Integer.SIZE; bit++) {
+ if ((mask & (1 << bit)) != 0) {
+ names.add(union.get(bit).getIdentifier().getName());
+ }
+ }
+ result.add(names);
+ }
+ return result;
+ }
+
+ private static Set<String> names(String... names) {
+ return new HashSet<>(List.of(names));
+ }
+
+ @Test
+ 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());
+ assertTrue(query.getGroupByListSize() == 2);
+ }
+
+ @Test
+ public void testRollup() {
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery("SELECT a, b, SUM(c) FROM t GROUP BY ROLLUP(a, b)");
+ assertEquals(query.getGroupByList().size(), 2);
+ assertEquals(groupingSetsByName(query), Set.of(names("a", "b"), names("a"), names()));
+ }
+
+ @Test
+ public void testCube() {
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery("SELECT a, b, SUM(c) FROM t GROUP BY CUBE(a, b)");
+ assertEquals(query.getGroupByList().size(), 2);
+ assertEquals(groupingSetsByName(query), Set.of(names("a", "b"), names("a"), names("b"), names()));
+ }
+
+ @Test
+ public void testGroupingSets() {
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery(
+ "SELECT a, b, SUM(c) FROM t GROUP BY GROUPING SETS ((a, b), (a), ())");
+ assertEquals(query.getGroupByList().size(), 2);
+ assertEquals(groupingSetsByName(query), Set.of(names("a", "b"), names("a"), names()));
+ }
+
+ @Test
+ public void testMixedPlainAndRollup() {
+ /// GROUP BY a, ROLLUP(b, c) == GROUPING SETS ((a, b, c), (a, b), (a))
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery(
+ "SELECT a, b, c, SUM(d) FROM t GROUP BY a, ROLLUP(b, c)");
+ assertEquals(query.getGroupByList().size(), 3);
+ assertEquals(groupingSetsByName(query), Set.of(names("a", "b", "c"), names("a", "b"), names("a")));
+ }
+
+ @Test
+ public void testCompositeRollupLevel() {
+ /// ROLLUP((a, b)) treats the parenthesized pair as one level: GROUPING SETS ((a, b), ()).
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery("SELECT a, b, SUM(c) FROM t GROUP BY ROLLUP((a, b))");
+ assertEquals(groupingSetsByName(query), Set.of(names("a", "b"), names()));
+ }
+
+ @Test
+ public void testNestedRollupInsideGroupingSets() {
+ /// Grouping constructs may nest: GROUPING SETS ((a), ROLLUP(b)) -> {a}, {b}, {}.
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery(
+ "SELECT a, b, SUM(c) FROM t GROUP BY GROUPING SETS ((a), ROLLUP(b))");
+ assertEquals(groupingSetsByName(query), Set.of(names("a"), names("b"), names()));
+ }
+
+ @Test
+ public void testDuplicateGroupingSetsDeduplicated() {
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery(
+ "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);
+ }
+
+ @Test
+ public void testGroupingSetsWithEmptyOnlyIsGrandTotal() {
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery("SELECT SUM(c) FROM t GROUP BY GROUPING SETS (())");
+ assertEquals(groupingSetsByName(query), Set.of(names()));
+ }
+
+ @Test
+ public void testAggregationFreeGroupingSetsRejected() {
+ /// Without an aggregation the query must fail loudly: it must neither be rewritten to DISTINCT (which
+ /// would drop the rollup rows) nor fall through to the selection path (which would ignore GROUP BY).
+ for (String sql : List.of(
+ "SELECT a FROM t GROUP BY ROLLUP(a)",
+ "SELECT a, b FROM t GROUP BY CUBE(a, b)",
+ "SELECT a, b FROM t GROUP BY GROUPING SETS ((a), (b))",
+ "SELECT a, b FROM t GROUP BY GROUPING SETS ((a, b))",
+ "SELECT a, GROUPING(a) FROM t GROUP BY ROLLUP(a)")) {
+ SqlCompilationException e =
+ expectThrows(SqlCompilationException.class, () -> CalciteSqlParser.compileToPinotQuery(sql));
+ assertTrue(e.getMessage().contains("requires at least one aggregation function"), e.getMessage());
+ }
+ }
+
+ @Test
+ public void testGroupingSetsWithAggregationOnlyInHavingOrOrderBy() {
+ /// Aggregations in HAVING / ORDER-BY alone make it an aggregation group-by query; the grouping sets must
+ /// survive compilation (no DISTINCT rewrite, no rejection).
+ for (String sql : List.of(
+ "SELECT a FROM t GROUP BY ROLLUP(a) HAVING COUNT(*) > 1",
+ "SELECT a FROM t GROUP BY ROLLUP(a) ORDER BY COUNT(*)")) {
+ PinotQuery query = CalciteSqlParser.compileToPinotQuery(sql);
+ assertEquals(query.getGroupByList().size(), 1);
+ assertEquals(groupingSetsByName(query), Set.of(names("a"), names()));
+ }
+ }
+
+ @Test(expectedExceptions = SqlCompilationException.class)
+ public void testTooManyGroupingColumnsRejected() {
+ /// 32 distinct grouping columns exceed the 31-column limit imposed by the INT grouping-id bitmask.
+ StringBuilder sets = new StringBuilder();
+ for (int i = 0; i < 32; i++) {
+ sets.append(i == 0 ? "(c" : ", (c").append(i).append(')');
+ }
+ CalciteSqlParser.compileToPinotQuery(
+ "SELECT SUM(m) FROM t GROUP BY GROUPING SETS (" + sets + ")");
+ }
+
+ @Test(expectedExceptions = SqlCompilationException.class)
+ public void testCubeTooManySetsRejected() {
+ /// CUBE over 13 columns expands to 2^13 = 8192 grouping sets, exceeding the 4096 cap.
+ StringBuilder columns = new StringBuilder();
+ for (int i = 0; i < 13; i++) {
+ columns.append(i == 0 ? "c" : ", c").append(i);
+ }
+ CalciteSqlParser.compileToPinotQuery("SELECT SUM(m) FROM t GROUP BY CUBE(" + columns + ")");
+ }
+}
diff --git a/pinot-common/src/thrift/query.thrift b/pinot-common/src/thrift/query.thrift
index 7b896a31da1..0b49b4a86af 100644
--- a/pinot-common/src/thrift/query.thrift
+++ b/pinot-common/src/thrift/query.thrift
@@ -32,6 +32,12 @@ struct PinotQuery {
11: optional map<string, string> queryOptions;
12: optional bool explain;
13: optional map<Expression, Expression> 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<i32> groupingSetMasks;
}
struct DataSource {
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/table/IndexedTable.java b/pinot-core/src/main/java/org/apache/pinot/core/data/table/IndexedTable.java
index 9ce27567d00..2bb526f82ae 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/data/table/IndexedTable.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/data/table/IndexedTable.java
@@ -87,7 +87,9 @@ public abstract class IndexedTable extends BaseTable {
List<ExpressionContext> groupByExpressions = queryContext.getGroupByExpressions();
assert groupByExpressions != null;
- _numKeyColumns = groupByExpressions.size();
+ /// Includes the synthetic $groupingId key column for GROUP BY GROUPING SETS / ROLLUP / CUBE queries, so
+ /// that rows from different grouping sets are keyed (and therefore merged) independently.
+ _numKeyColumns = queryContext.getNumGroupByKeyColumns();
_aggregationFunctions = queryContext.getAggregationFunctions();
_hasOrderBy = queryContext.getOrderByExpressions() != null;
_tableResizer = _hasOrderBy ? new TableResizer(dataSchema, hasFinalInput, queryContext) : null;
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/table/SortedRecordTable.java b/pinot-core/src/main/java/org/apache/pinot/core/data/table/SortedRecordTable.java
index a0bb60fe550..9db70fc92f6 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/data/table/SortedRecordTable.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/data/table/SortedRecordTable.java
@@ -53,7 +53,7 @@ public class SortedRecordTable extends BaseTable {
QueryContext queryContext, ExecutorService executorService) {
super(dataSchema);
assert queryContext.getGroupByExpressions() != null;
- _numKeyColumns = queryContext.getGroupByExpressions().size();
+ _numKeyColumns = queryContext.getNumGroupByKeyColumns();
_aggregationFunctions = queryContext.getAggregationFunctions();
_executorService = executorService;
_numThreadsExtractFinalResult = Math.min(queryContext.getNumThreadsExtractFinalResult(),
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/table/SortedRecordsMerger.java b/pinot-core/src/main/java/org/apache/pinot/core/data/table/SortedRecordsMerger.java
index 3363dc994b2..a167e1c93e7 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/data/table/SortedRecordsMerger.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/data/table/SortedRecordsMerger.java
@@ -42,7 +42,7 @@ public class SortedRecordsMerger {
public SortedRecordsMerger(QueryContext queryContext, int resultSize, Comparator<Record> comparator) {
assert queryContext.getGroupByExpressions() != null;
- _numKeyColumns = queryContext.getGroupByExpressions().size();
+ _numKeyColumns = queryContext.getNumGroupByKeyColumns();
_aggregationFunctions = queryContext.getAggregationFunctions();
_resultSize = resultSize;
_comparator = comparator;
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 60b601d2b91..3283c2aae36 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,10 +29,12 @@ import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import java.util.PriorityQueue;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.FilterContext;
import org.apache.pinot.common.request.context.FunctionContext;
+import org.apache.pinot.common.request.context.GroupingSets;
import org.apache.pinot.common.request.context.OrderByExpressionContext;
import org.apache.pinot.common.request.context.RequestContextUtils;
import org.apache.pinot.common.utils.DataSchema;
@@ -53,6 +55,9 @@ public class TableResizer {
private final DataSchema _dataSchema;
private final boolean _hasFinalInput;
private final int _numGroupByExpressions;
+ /// 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;
private final Map<ExpressionContext, Integer> _groupByExpressionIndexMap;
private final AggregationFunction[] _aggregationFunctions;
private final Map<Pair<FunctionContext, FilterContext>, Integer> _filteredAggregationIndexMap;
@@ -74,6 +79,7 @@ public class TableResizer {
List<ExpressionContext> groupByExpressions = queryContext.getGroupByExpressions();
assert groupByExpressions != null;
_numGroupByExpressions = groupByExpressions.size();
+ _numKeyColumns = queryContext.getNumGroupByKeyColumns();
_groupByExpressionIndexMap = new HashMap<>();
for (int i = 0; i < _numGroupByExpressions; i++) {
_groupByExpressionIndexMap.put(groupByExpressions.get(i), i);
@@ -96,7 +102,9 @@ public class TableResizer {
comparators[i] = orderByExpression.isAsc() ? Comparator.naturalOrder() : Comparator.reverseOrder();
nullComparisonResults[i] = orderByExpression.isNullsLast() ? -1 : 1;
}
- boolean nullHandlingEnabled = queryContext.isNullHandlingEnabled();
+ /// Grouping-set queries produce genuine null group-key values for rolled-up columns regardless of the
+ /// user's null-handling option, so ORDER BY over a grouping column must use the null-safe comparator.
+ boolean nullHandlingEnabled = queryContext.requiresNullAwareKeySerialization();
if (nullHandlingEnabled) {
_intermediateRecordComparator = (o1, o2) -> {
for (int i = 0; i < _numOrderByExpressions; i++) {
@@ -145,6 +153,13 @@ public class TableResizer {
FunctionContext function = expression.getFunction();
Preconditions.checkState(function != null, "Failed to find ORDER-BY expression: %s in the GROUP-BY clause",
expression);
+ String functionName = function.getFunctionName();
+ if (GroupingSets.isGroupingFunction(functionName)) {
+ /// GROUPING(col, ...) / GROUPING_ID(col, ...): computed from the synthetic $groupingId discriminator
+ /// column, mirroring PostAggregationHandler. Without this, the generic post-aggregation path below would
+ /// fail to resolve these context-dependent functions when building the ORDER BY comparator.
+ return new GroupingExtractor(function);
+ }
FunctionContext aggregation;
FilterContext filter;
if (function.getType() == FunctionContext.Type.AGGREGATION) {
@@ -163,7 +178,7 @@ public class TableResizer {
int index = _filteredAggregationIndexMap.get(Pair.of(aggregation, filter));
// For final aggregate result, we can handle it the same way as group key
- return _hasFinalInput ? new GroupByExpressionExtractor(_numGroupByExpressions + index)
+ return _hasFinalInput ? new GroupByExpressionExtractor(_numKeyColumns + index)
: new AggregationFunctionExtractor(index);
}
@@ -384,18 +399,56 @@ public class TableResizer {
return result;
}
+ /// Trims the per-segment aggregation results for a GROUP BY GROUPING SETS / ROLLUP / CUBE query, keeping the
+ /// top {@code perSetSize} groups WITHIN each grouping set, bucketed by the {@code discriminatorColumnIndex}
+ /// column (the $groupingId value). This keeps each grouping set's own top candidates so that a global
+ /// top-K cannot starve low-magnitude sets (e.g. the grand total). Records across all sets are returned
+ /// unsorted; the broker applies the final ORDER BY + LIMIT across sets.
+ ///
+ /// Like plain GROUP BY segment trim, this is opt-in (it only runs when {@code minSegmentGroupTrimSize > 0};
+ /// disabled by default, in which case the segment returns all groups and results are exact) and is an
+ /// approximation: a group ranking below {@code perSetSize} within its set ON THIS SEGMENT may still be
+ /// globally in the top-K once partial aggregates merge across segments/servers. The per-set bucketing only
+ /// bounds across-set starvation, not this within-set pre-merge approximation inherent to any segment trim.
+ public List<IntermediateRecord> trimInSegmentResultsByGroupingSet(GroupKeyGenerator groupKeyGenerator,
+ GroupByResultHolder[] groupByResultHolders, int perSetSize, int discriminatorColumnIndex) {
+ /// _intermediateRecordComparator orders the best records first. A bounded per-bucket min-heap keyed on
+ /// the reversed comparator keeps its WORST kept record at the head, so it retains the top 'perSetSize'.
+ Comparator<IntermediateRecord> worstFirst = _intermediateRecordComparator.reversed();
+ Map<Integer, PriorityQueue<IntermediateRecord>> bucketsByGroupingId = new HashMap<>();
+ Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator = groupKeyGenerator.getGroupKeys();
+ while (groupKeyIterator.hasNext()) {
+ IntermediateRecord record = getIntermediateRecord(groupKeyIterator.next(), groupByResultHolders);
+ int groupingId = ((Number) record._record.getValues()[discriminatorColumnIndex]).intValue();
+ PriorityQueue<IntermediateRecord> heap =
+ bucketsByGroupingId.computeIfAbsent(groupingId, k -> new PriorityQueue<>(worstFirst));
+ if (heap.size() < perSetSize) {
+ heap.offer(record);
+ } else if (_intermediateRecordComparator.compare(record, heap.peek()) < 0) {
+ /// 'record' ranks ahead of this set's current worst kept record: replace it.
+ heap.poll();
+ heap.offer(record);
+ }
+ }
+ List<IntermediateRecord> result = new ArrayList<>();
+ for (PriorityQueue<IntermediateRecord> heap : bucketsByGroupingId.values()) {
+ result.addAll(heap);
+ }
+ return result;
+ }
+
/**
* Constructs an IntermediateRecord for the given group.
*/
private IntermediateRecord getIntermediateRecord(GroupKeyGenerator.GroupKey groupKey,
GroupByResultHolder[] groupByResultHolders) {
int numAggregationFunctions = _aggregationFunctions.length;
- int numColumns = numAggregationFunctions + _numGroupByExpressions;
+ int numColumns = numAggregationFunctions + _numKeyColumns;
Object[] keys = groupKey._keys;
Object[] values = Arrays.copyOf(keys, numColumns);
int groupId = groupKey._groupId;
for (int i = 0; i < numAggregationFunctions; i++) {
- values[_numGroupByExpressions + i] =
+ values[_numKeyColumns + i] =
_aggregationFunctions[i].extractGroupByResult(groupByResultHolders[i], groupId);
}
return getIntermediateRecord(new Key(keys), new Record(values));
@@ -467,7 +520,9 @@ public class TableResizer {
final AggregationFunction _aggregationFunction;
AggregationFunctionExtractor(int aggregationFunctionIndex) {
- _index = aggregationFunctionIndex + _numGroupByExpressions;
+ /// Aggregation columns follow all key columns, which for grouping-set queries include the synthetic
+ /// $groupingId column (_numKeyColumns == _numGroupByExpressions for non-grouping-set queries).
+ _index = aggregationFunctionIndex + _numKeyColumns;
_aggregationFunction = _aggregationFunctions[aggregationFunctionIndex];
}
@@ -482,6 +537,44 @@ public class TableResizer {
}
}
+ /// 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).
+ private class GroupingExtractor implements OrderByValueExtractor {
+ final int[] _unionColumnIndexes;
+ final int _discriminatorIndex;
+
+ GroupingExtractor(FunctionContext function) {
+ List<ExpressionContext> arguments = function.getArguments();
+ Preconditions.checkState(!arguments.isEmpty(), "GROUPING requires at least one argument");
+ _unionColumnIndexes = 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;
+ }
+ /// The $groupingId column sits right after the union columns; it is absent for non-grouping-set queries.
+ _discriminatorIndex = _numKeyColumns > _numGroupByExpressions ? _numGroupByExpressions : -1;
+ }
+
+ @Override
+ public ColumnDataType getValueType() {
+ return ColumnDataType.INT;
+ }
+
+ @Override
+ public Comparable extract(Record record) {
+ if (_discriminatorIndex < 0) {
+ return 0;
+ }
+ int groupingId = ((Number) record.getValues()[_discriminatorIndex]).intValue();
+ return GroupingSets.groupingValue(groupingId, _unionColumnIndexes);
+ }
+ }
+
/**
* Extractor for a post-aggregation function.
*/
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/GroupByResultsBlock.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/GroupByResultsBlock.java
index 5baac68f099..75cfa97b6a5 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/GroupByResultsBlock.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/GroupByResultsBlock.java
@@ -207,10 +207,16 @@ public class GroupByResultsBlock extends BaseResultsBlock {
AggregationFunction[] aggregationFunctions = _queryContext.getAggregationFunctions();
List<ExpressionContext> groupByExpressions = _queryContext.getGroupByExpressions();
assert aggregationFunctions != null && groupByExpressions != null;
- int numKeyColumns = groupByExpressions.size();
+ /// Key columns (group-by union columns + the synthetic $groupingId discriminator for grouping sets)
+ /// precede the aggregation columns.
+ int numKeyColumns = _queryContext.getNumGroupByKeyColumns();
Iterator<Record> iterator = _table.iterator();
int numRowsAdded = 0;
- if (_queryContext.isNullHandlingEnabled()) {
+ /// Grouping sets produce NULL group keys (rolled-up columns) regardless of the user's null-handling
+ /// option, so they must be serialized through the null-aware path (null bitmaps) to survive the DataTable
+ /// round-trip -- the non-null path's setNull() writes to the variable buffer that the reducer does not
+ /// read back for fixed-width key columns.
+ if (_queryContext.requiresNullAwareKeySerialization()) {
RoaringBitmap[] nullBitmaps = new RoaringBitmap[numColumns];
Object[] nullPlaceholders = new Object[numColumns];
for (int colId = 0; colId < numColumns; colId++) {
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/ResultsBlockUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/ResultsBlockUtils.java
index 4864829eef0..e6c4617f0f4 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/ResultsBlockUtils.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/ResultsBlockUtils.java
@@ -25,6 +25,7 @@ import java.util.List;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.FilterContext;
+import org.apache.pinot.common.request.context.GroupingSets;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
@@ -86,7 +87,11 @@ public class ResultsBlockUtils {
queryContext.getFilteredAggregationFunctions();
List<ExpressionContext> groupByExpressions = queryContext.getGroupByExpressions();
assert filteredAggregationFunctions != null && groupByExpressions != null;
- int numColumns = groupByExpressions.size() + filteredAggregationFunctions.size();
+ /// Grouping-set queries append a synthetic $groupingId key column after the group-by columns; the empty
+ /// result schema must match the non-empty one (see GroupByOperator) so an empty match (e.g. fully-pruned
+ /// segments) does not produce a narrower DataTable that the reducer would reject as a schema mismatch.
+ int numExtraKeyColumns = queryContext.getNumExtraGroupByKeyColumns();
+ int numColumns = groupByExpressions.size() + numExtraKeyColumns + filteredAggregationFunctions.size();
String[] columnNames = new String[numColumns];
ColumnDataType[] columnDataTypes = new ColumnDataType[numColumns];
int index = 0;
@@ -96,6 +101,11 @@ public class ResultsBlockUtils {
columnDataTypes[index] = ColumnDataType.STRING;
index++;
}
+ if (numExtraKeyColumns > 0) {
+ columnNames[index] = GroupingSets.GROUPING_ID_COLUMN;
+ columnDataTypes[index] = ColumnDataType.INT;
+ index++;
+ }
for (Pair<AggregationFunction, FilterContext> pair : filteredAggregationFunctions) {
AggregationFunction aggregationFunction = pair.getLeft();
columnNames[index] = AggregationFunctionUtils.getResultColumnName(aggregationFunction, pair.getRight());
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java
index 3769ed41166..06835ce054c 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java
@@ -57,7 +57,9 @@ public class GroupByCombineOperator extends BaseSingleBlockCombineOperator<Group
private static final String EXPLAIN_NAME = "COMBINE_GROUP_BY";
private final int _numAggregationFunctions;
- private final int _numGroupByExpressions;
+ /// Number of key columns: union group-by columns plus the synthetic $groupingId column for grouping sets.
+ /// Key columns precede the aggregation columns in the record layout.
+ private final int _numKeyColumns;
private final int _numColumns;
// We use a CountDownLatch to track if all Futures are finished by the query timeout, and cancel the unfinished
// _futures (try to interrupt the execution if it already started).
@@ -75,8 +77,8 @@ public class GroupByCombineOperator extends BaseSingleBlockCombineOperator<Group
assert aggregationFunctions != null;
_numAggregationFunctions = aggregationFunctions.length;
assert _queryContext.getGroupByExpressions() != null;
- _numGroupByExpressions = _queryContext.getGroupByExpressions().size();
- _numColumns = _numGroupByExpressions + _numAggregationFunctions;
+ _numKeyColumns = _queryContext.getNumGroupByKeyColumns();
+ _numColumns = _numKeyColumns + _numAggregationFunctions;
_operatorLatch = new CountDownLatch(_numTasks);
}
@@ -150,7 +152,7 @@ public class GroupByCombineOperator extends BaseSingleBlockCombineOperator<Group
Object[] values = Arrays.copyOf(keys, _numColumns);
int groupId = groupKey._groupId;
for (int i = 0; i < _numAggregationFunctions; i++) {
- values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
+ values[_numKeyColumns + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
}
_indexedTable.upsert(new Key(keys), new Record(values));
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java
index 0cb56dfe5eb..96bc584ff8d 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java
@@ -29,6 +29,7 @@ import org.apache.pinot.common.metrics.ServerMeter;
import org.apache.pinot.common.metrics.ServerMetrics;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.FilterContext;
+import org.apache.pinot.common.request.context.GroupingSets;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.common.Operator;
import org.apache.pinot.core.data.table.IntermediateRecord;
@@ -48,6 +49,7 @@ import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder;
import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.core.startree.executor.StarTreeGroupByExecutor;
+import org.apache.pinot.core.util.GroupByUtils;
import org.apache.pinot.spi.query.QueryScanCostContext;
import org.apache.pinot.spi.trace.Tracing;
import org.slf4j.Logger;
@@ -86,7 +88,10 @@ public class FilteredGroupByOperator extends BaseOperator<GroupByResultsBlock> {
// NOTE: The indexedTable expects that the data schema will have group by columns before aggregation columns
int numGroupByExpressions = _groupByExpressions.length;
int numAggregationFunctions = _aggregationFunctions.length;
- int numColumns = numGroupByExpressions + numAggregationFunctions;
+ /// Grouping-set queries append a synthetic $groupingId key column after the union group-by columns.
+ int numExtraKeyColumns = queryContext.getNumExtraGroupByKeyColumns();
+ int numKeyColumns = numGroupByExpressions + numExtraKeyColumns;
+ int numColumns = numKeyColumns + numAggregationFunctions;
String[] columnNames = new String[numColumns];
DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[numColumns];
@@ -99,9 +104,15 @@ public class FilteredGroupByOperator extends BaseOperator<GroupByResultsBlock> {
projectOperator.getResultColumnContext(groupByExpression).getDataType());
}
+ /// Synthetic grouping-id discriminator column for GROUP BY GROUPING SETS / ROLLUP / CUBE
+ if (numExtraKeyColumns > 0) {
+ columnNames[numGroupByExpressions] = GroupingSets.GROUPING_ID_COLUMN;
+ columnDataTypes[numGroupByExpressions] = DataSchema.ColumnDataType.INT;
+ }
+
// Extract column names and data types for aggregation functions
for (int i = 0; i < numAggregationFunctions; i++) {
- int index = numGroupByExpressions + i;
+ int index = numKeyColumns + i;
Pair<AggregationFunction, FilterContext> pair = queryContext.getFilteredAggregationFunctions().get(i);
AggregationFunction aggregationFunction = pair.getLeft();
String columnName = AggregationFunctionUtils.getResultColumnName(aggregationFunction, pair.getRight());
@@ -204,6 +215,15 @@ public class FilteredGroupByOperator extends BaseOperator<GroupByResultsBlock> {
boolean unsafeTrim = _queryContext.isUnsafeTrim();
GroupByResultsBlock resultsBlock;
+ /// Grouping-set queries use a per-set bucketed segment trim (keyed on the $groupingId discriminator) so
+ /// that a global top-K cannot starve low-magnitude sets such as the grand total. The broker still applies
+ /// the final ORDER BY + LIMIT across all sets.
+ if (_queryContext.isGroupingSets()) {
+ /// The $groupingId discriminator is the key column immediately after the union group-by columns.
+ return GroupByUtils.buildGroupingSetsResultsBlock(_queryContext, _dataSchema, groupKeyGenerator,
+ groupByResultHolders, groupKeyGenerator.getNumKeys(), _groupByExpressions.length, numGroupsLimitReached,
+ numGroupsWarningLimitReached);
+ }
// sort and trim segment results if needed
if (trimSize > 0 && groupKeyGenerator.getNumKeys() > trimSize) {
TableResizer tableResizer = new TableResizer(_dataSchema, _queryContext);
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java
index 3465484b652..d03876e0a06 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java
@@ -26,6 +26,7 @@ import java.util.stream.Collectors;
import org.apache.pinot.common.metrics.ServerMeter;
import org.apache.pinot.common.metrics.ServerMetrics;
import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.GroupingSets;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.common.Operator;
import org.apache.pinot.core.data.table.IntermediateRecord;
@@ -42,6 +43,7 @@ import org.apache.pinot.core.query.aggregation.groupby.DefaultGroupByExecutor;
import org.apache.pinot.core.query.aggregation.groupby.GroupByExecutor;
import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.core.startree.executor.StarTreeGroupByExecutor;
+import org.apache.pinot.core.util.GroupByUtils;
import org.apache.pinot.spi.query.QueryScanCostContext;
import org.apache.pinot.spi.trace.Tracing;
import org.slf4j.Logger;
@@ -78,7 +80,11 @@ public class GroupByOperator extends BaseOperator<GroupByResultsBlock> {
// NOTE: The indexedTable expects that the data schema will have group by columns before aggregation columns
int numGroupByExpressions = _groupByExpressions.length;
int numAggregationFunctions = _aggregationFunctions.length;
- int numColumns = numGroupByExpressions + numAggregationFunctions;
+ /// Grouping-set queries append a synthetic $groupingId key column after the union group-by columns (the
+ /// per-set bitmask discriminator); the key columns thus precede the aggregation columns.
+ int numExtraKeyColumns = _queryContext.getNumExtraGroupByKeyColumns();
+ int numKeyColumns = numGroupByExpressions + numExtraKeyColumns;
+ int numColumns = numKeyColumns + numAggregationFunctions;
String[] columnNames = new String[numColumns];
DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[numColumns];
@@ -90,10 +96,16 @@ public class GroupByOperator extends BaseOperator<GroupByResultsBlock> {
_projectOperator.getResultColumnContext(groupByExpression).getDataType());
}
+ /// Synthetic grouping-id discriminator column for GROUP BY GROUPING SETS / ROLLUP / CUBE
+ if (numExtraKeyColumns > 0) {
+ columnNames[numGroupByExpressions] = GroupingSets.GROUPING_ID_COLUMN;
+ columnDataTypes[numGroupByExpressions] = DataSchema.ColumnDataType.INT;
+ }
+
// Extract column names and data types for aggregation functions
for (int i = 0; i < numAggregationFunctions; i++) {
AggregationFunction aggregationFunction = _aggregationFunctions[i];
- int index = numGroupByExpressions + i;
+ int index = numKeyColumns + i;
columnNames[index] = aggregationFunction.getResultColumnName();
columnDataTypes[index] = aggregationFunction.getIntermediateResultColumnType();
}
@@ -153,6 +165,16 @@ public class GroupByOperator extends BaseOperator<GroupByResultsBlock> {
boolean unsafeTrim = _queryContext.isUnsafeTrim();
GroupByResultsBlock resultsBlock;
+ /// Grouping-set queries use a per-set bucketed segment trim (keyed on the $groupingId discriminator) so
+ /// that a global top-K cannot starve low-magnitude sets such as the grand total. The broker still applies
+ /// the final ORDER BY + LIMIT across all sets.
+ if (_queryContext.isGroupingSets()) {
+ /// The $groupingId discriminator is the key column immediately after the union group-by columns.
+ return GroupByUtils.buildGroupingSetsResultsBlock(_queryContext, _dataSchema,
+ groupByExecutor.getGroupKeyGenerator(), groupByExecutor.getGroupByResultHolders(),
+ groupByExecutor.getNumGroups(), _groupByExpressions.length, numGroupsLimitReached,
+ numGroupsWarningLimitReached);
+ }
// sort and trim segment results if needed
if (trimSize > 0 && groupByExecutor.getNumGroups() > trimSize) {
TableResizer tableResizer = new TableResizer(_dataSchema, _queryContext);
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java
index a9e6f439320..bf9285b421d 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java
@@ -355,6 +355,11 @@ public class AggregationFunctionUtils {
public static AggregationInfo buildAggregationInfoWithStarTree(SegmentContext segmentContext,
QueryContext queryContext, AggregationFunction[] aggregationFunctions, @Nullable FilterContext filter,
BaseFilterOperator filterOperator, List<Pair<Predicate, PredicateEvaluator>> predicateEvaluators) {
+ /// Star-tree stores pre-aggregated values per group key and cannot expand a row across multiple grouping
+ /// sets, so it cannot serve GROUP BY GROUPING SETS / ROLLUP / CUBE queries. Fall back to the regular path.
+ if (queryContext.isGroupingSets()) {
+ return null;
+ }
if (!filterOperator.isResultEmpty()) {
BaseProjectOperator<?> projectOperator =
StarTreeUtils.createStarTreeBasedProjectOperator(segmentContext.getIndexSegment(), queryContext,
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java
index d5af6d9eae7..529f2b15f70 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java
@@ -92,7 +92,10 @@ public class DefaultGroupByExecutor implements GroupByExecutor {
// isDictionaryEncoded() flag rather than gating on dictionary nullness alone.
hasNoDictionaryGroupByExpression |= !columnContext.isDictionaryEncoded();
}
- _hasMVGroupByExpression = hasMVGroupByExpression;
+ /// Grouping-set queries expand each row into one group per grouping set, so they always use the
+ /// multi-value (int[][]) executor path even though the union group-by columns are single-valued.
+ boolean groupingSets = queryContext.isGroupingSets();
+ _hasMVGroupByExpression = hasMVGroupByExpression || groupingSets;
// Initialize group key generator
int numGroupsLimit = queryContext.getNumGroupsLimit();
@@ -103,6 +106,9 @@ public class DefaultGroupByExecutor implements GroupByExecutor {
}
if (groupKeyGenerator != null) {
_groupKeyGenerator = groupKeyGenerator;
+ } else if (groupingSets) {
+ _groupKeyGenerator = new GroupingSetsGroupKeyGenerator(projectOperator, groupByExpressions,
+ queryContext.getGroupingSets(), numGroupsLimit, _nullHandlingEnabled);
} else {
if (hasNoDictionaryGroupByExpression || _nullHandlingEnabled) {
if (groupByExpressions.length == 1) {
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
new file mode 100644
index 00000000000..a9d93563aef
--- /dev/null
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupingSetsGroupKeyGenerator.java
@@ -0,0 +1,423 @@
+/**
+ * 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.core.query.aggregation.groupby;
+
+import it.unimi.dsi.fastutil.ints.IntArrayList;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+import it.unimi.dsi.fastutil.objects.ObjectIterator;
+import java.util.Iterator;
+import java.util.List;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.core.common.BlockValSet;
+import org.apache.pinot.core.operator.BaseProjectOperator;
+import org.apache.pinot.core.operator.ColumnContext;
+import org.apache.pinot.core.operator.blocks.ValueBlock;
+import org.apache.pinot.core.query.aggregation.groupby.utils.ValueToIdMap;
+import org.apache.pinot.core.query.aggregation.groupby.utils.ValueToIdMapFactory;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.utils.ByteArray;
+import org.apache.pinot.spi.utils.FixedIntArray;
+import org.roaringbitmap.RoaringBitmap;
+
+
+/// {@link GroupKeyGenerator} for GROUP BY GROUPING SETS / ROLLUP / CUBE queries in the single-stage engine.
+///
+/// Given the union of all grouping columns `[c_0, ..., c_{k-1}]` and a list of grouping sets (each a subset
+/// of the union), this generator expands every input row into one group per grouping set in a single scan.
+/// For a grouping set `S`, the composite group key is
+///
+/// ```
+/// [ id(c_0) or ID_FOR_NULL, ..., id(c_{k-1}) or ID_FOR_NULL, groupingId ]
+/// ```
+///
+/// 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)`.
+///
+/// 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
+/// dictionaries (like {@link NoDictionaryMultiColumnGroupKeyGenerator}) so that NULL keys -- which grouping
+/// sets produce regardless of the query's null-handling option -- are representable and reconstructable.
+///
+/// A multi-value union column participating in a grouping set expands the row over its values (Cartesian
+/// product across participating MV columns), composed with the per-set expansion. Star-tree is not used for
+/// grouping-set queries (the planner falls back to the regular path).
+public class GroupingSetsGroupKeyGenerator implements GroupKeyGenerator {
+ /// Sentinel id for a NULL value or a rolled-up (non-participating) column. Distinct from real dictionary ids
+ /// (>= 0) and from GroupKeyGenerator.INVALID_ID (-1), so the three id spaces never collide in a composite key.
+ private static final int ID_FOR_NULL = INVALID_ID - 1;
+
+ private final ExpressionContext[] _groupByExpressions;
+ private final int _numGroupByExpressions;
+ private final DataType[] _storedTypes;
+ private final ValueToIdMap[] _onTheFlyDictionaries;
+ private final boolean[] _isSingleValue;
+ private final boolean _hasMvColumn;
+ 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().
+ private final boolean[][] _setContains;
+ private final int[] _setBitmasks;
+ private final int _numSets;
+
+ private final Object2IntOpenHashMap<FixedIntArray> _groupKeyMap;
+ /// Reusable single-element id list for a rolled-up (non-participating) column on the MV path. Per-instance
+ /// (segment-confined, never shared across query threads) and treated as read-only: expandGroupIds copies its
+ /// element into the composite key and never mutates it.
+ private final int[] _nullComponent = {ID_FOR_NULL};
+
+ public GroupingSetsGroupKeyGenerator(BaseProjectOperator<?> projectOperator,
+ ExpressionContext[] groupByExpressions, List<int[]> groupingSets, int numGroupsLimit,
+ boolean nullHandlingEnabled) {
+ _groupByExpressions = groupByExpressions;
+ _numGroupByExpressions = groupByExpressions.length;
+ _storedTypes = new DataType[_numGroupByExpressions];
+ _onTheFlyDictionaries = new ValueToIdMap[_numGroupByExpressions];
+ _isSingleValue = new boolean[_numGroupByExpressions];
+ _nullHandlingEnabled = nullHandlingEnabled;
+ _numGroupsLimit = numGroupsLimit;
+ boolean hasMvColumn = false;
+ for (int i = 0; i < _numGroupByExpressions; i++) {
+ ColumnContext columnContext = projectOperator.getResultColumnContext(groupByExpressions[i]);
+ _isSingleValue[i] = columnContext.isSingleValue();
+ hasMvColumn |= !_isSingleValue[i];
+ _storedTypes[i] = columnContext.getDataType().getStoredType();
+ _onTheFlyDictionaries[i] = ValueToIdMapFactory.get(_storedTypes[i]);
+ }
+ _hasMvColumn = hasMvColumn;
+
+ _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<>();
+ _groupKeyMap.defaultReturnValue(INVALID_ID);
+ }
+
+ @Override
+ public int getGlobalGroupKeyUpperBound() {
+ return _numGroupsLimit;
+ }
+
+ @Override
+ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) {
+ throw new UnsupportedOperationException(
+ "GroupingSetsGroupKeyGenerator only supports the multi-value group key path");
+ }
+
+ @Override
+ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) {
+ if (_hasMvColumn) {
+ generateKeysForBlockWithMv(valueBlock, groupKeys);
+ return;
+ }
+ int numDocs = valueBlock.getNumDocs();
+ /// Resolve each union column's value to an on-the-fly dictionary id per row (ID_FOR_NULL for nulls). The
+ /// id of a value is independent of which grouping set uses it, so it is computed once per row here.
+ int[][] columnIds = new int[_numGroupByExpressions][];
+ for (int col = 0; col < _numGroupByExpressions; col++) {
+ columnIds[col] = resolveColumnIds(valueBlock, col, numDocs);
+ }
+ /// For each row, emit one group id per grouping set. Reuse a single key buffer (flyweight) across rows and
+ /// sets, allocating a fresh one only when a new group is inserted (the map then owns that buffer). This
+ /// mirrors NoDictionaryMultiColumnGroupKeyGenerator and avoids a per-(row, set) allocation.
+ int[] keyValues = new int[_numGroupByExpressions + 1];
+ FixedIntArray flyweightKey = new FixedIntArray(keyValues);
+ for (int row = 0; row < numDocs; row++) {
+ int[] ids = new int[_numSets];
+ for (int s = 0; s < _numSets; s++) {
+ int numGroups = _groupKeyMap.size();
+ for (int col = 0; col < _numGroupByExpressions; col++) {
+ keyValues[col] = _setContains[s][col] ? columnIds[col][row] : ID_FOR_NULL;
+ }
+ keyValues[_numGroupByExpressions] = _setBitmasks[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.
+ keyValues = new int[_numGroupByExpressions + 1];
+ flyweightKey = new FixedIntArray(keyValues);
+ }
+ ids[s] = groupId;
+ }
+ groupKeys[row] = ids;
+ }
+ }
+
+ /// Multi-value path: a participating MV column contributes its full list of value ids, so a row expands
+ /// over the Cartesian product of participating MV columns' values, composed with the per-set expansion.
+ private void generateKeysForBlockWithMv(ValueBlock valueBlock, int[][] groupKeys) {
+ int numDocs = valueBlock.getNumDocs();
+ /// Per column, per row: the list of value ids (length 1 for a single-value column, length N for an MV
+ /// column with N values in that row).
+ int[][][] columnValueIds = new int[_numGroupByExpressions][][];
+ for (int col = 0; col < _numGroupByExpressions; col++) {
+ columnValueIds[col] = resolveColumnValueIds(valueBlock, col, numDocs);
+ }
+ int[][] keyComponents = new int[_numGroupByExpressions + 1][];
+ int[] bitmaskComponent = new int[1];
+ keyComponents[_numGroupByExpressions] = bitmaskComponent;
+ 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];
+ expandGroupIds(keyComponents, new int[_numGroupByExpressions + 1], 0, rowGroupIds);
+ }
+ groupKeys[row] = rowGroupIds.toIntArray();
+ }
+ }
+
+ /// Recursively builds composite keys from the per-column id lists (Cartesian product) and resolves each to
+ /// a group id appended to {@code out}.
+ private void expandGroupIds(int[][] keyComponents, int[] keyValues, int level, IntArrayList out) {
+ if (level == keyComponents.length) {
+ out.add(getGroupIdForKey(new FixedIntArray(keyValues.clone())));
+ return;
+ }
+ for (int id : keyComponents[level]) {
+ keyValues[level] = id;
+ expandGroupIds(keyComponents, keyValues, level + 1, out);
+ }
+ }
+
+ /// Resolves a column's value ids per row as a list (length 1 for single-value columns, the MV value list
+ /// for multi-value columns).
+ private int[][] resolveColumnValueIds(ValueBlock valueBlock, int col, int numDocs) {
+ int[][] ids = new int[numDocs][];
+ if (_isSingleValue[col]) {
+ int[] svIds = resolveColumnIds(valueBlock, col, numDocs);
+ for (int row = 0; row < numDocs; row++) {
+ ids[row] = new int[]{svIds[row]};
+ }
+ return ids;
+ }
+ BlockValSet blockValSet = valueBlock.getBlockValueSet(_groupByExpressions[col]);
+ ValueToIdMap dictionary = _onTheFlyDictionaries[col];
+ switch (_storedTypes[col]) {
+ case INT:
+ int[][] intValues = blockValSet.getIntValuesMV();
+ for (int row = 0; row < numDocs; row++) {
+ int[] values = intValues[row];
+ int[] rowIds = new int[values.length];
+ for (int k = 0; k < values.length; k++) {
+ rowIds[k] = dictionary.put(values[k]);
+ }
+ ids[row] = rowIds;
+ }
+ break;
+ case LONG:
+ long[][] longValues = blockValSet.getLongValuesMV();
+ for (int row = 0; row < numDocs; row++) {
+ long[] values = longValues[row];
+ int[] rowIds = new int[values.length];
+ for (int k = 0; k < values.length; k++) {
+ rowIds[k] = dictionary.put(values[k]);
+ }
+ ids[row] = rowIds;
+ }
+ break;
+ case FLOAT:
+ float[][] floatValues = blockValSet.getFloatValuesMV();
+ for (int row = 0; row < numDocs; row++) {
+ float[] values = floatValues[row];
+ int[] rowIds = new int[values.length];
+ for (int k = 0; k < values.length; k++) {
+ rowIds[k] = dictionary.put(values[k]);
+ }
+ ids[row] = rowIds;
+ }
+ break;
+ case DOUBLE:
+ double[][] doubleValues = blockValSet.getDoubleValuesMV();
+ for (int row = 0; row < numDocs; row++) {
+ double[] values = doubleValues[row];
+ int[] rowIds = new int[values.length];
+ for (int k = 0; k < values.length; k++) {
+ rowIds[k] = dictionary.put(values[k]);
+ }
+ ids[row] = rowIds;
+ }
+ break;
+ case STRING:
+ String[][] stringValues = blockValSet.getStringValuesMV();
+ for (int row = 0; row < numDocs; row++) {
+ String[] values = stringValues[row];
+ int[] rowIds = new int[values.length];
+ for (int k = 0; k < values.length; k++) {
+ rowIds[k] = dictionary.put(values[k]);
+ }
+ ids[row] = rowIds;
+ }
+ break;
+ default:
+ throw new IllegalArgumentException(
+ "Illegal multi-value data type for grouping-sets group key generator: " + _storedTypes[col]);
+ }
+ return ids;
+ }
+
+ /// Resolves the on-the-fly dictionary id for the given union column across all rows in the block, mapping
+ /// null values (when null handling is enabled) to {@link #ID_FOR_NULL}.
+ private int[] resolveColumnIds(ValueBlock valueBlock, int col, int numDocs) {
+ int[] ids = new int[numDocs];
+ BlockValSet blockValSet = valueBlock.getBlockValueSet(_groupByExpressions[col]);
+ ValueToIdMap dictionary = _onTheFlyDictionaries[col];
+ RoaringBitmap nullBitmap = _nullHandlingEnabled ? blockValSet.getNullBitmap() : null;
+ switch (_storedTypes[col]) {
+ case INT:
+ int[] intValues = blockValSet.getIntValuesSV();
+ for (int row = 0; row < numDocs; row++) {
+ ids[row] = isNull(nullBitmap, row) ? ID_FOR_NULL : dictionary.put(intValues[row]);
+ }
+ break;
+ case LONG:
+ long[] longValues = blockValSet.getLongValuesSV();
+ for (int row = 0; row < numDocs; row++) {
+ ids[row] = isNull(nullBitmap, row) ? ID_FOR_NULL : dictionary.put(longValues[row]);
+ }
+ break;
+ case FLOAT:
+ float[] floatValues = blockValSet.getFloatValuesSV();
+ for (int row = 0; row < numDocs; row++) {
+ ids[row] = isNull(nullBitmap, row) ? ID_FOR_NULL : dictionary.put(floatValues[row]);
+ }
+ break;
+ case DOUBLE:
+ double[] doubleValues = blockValSet.getDoubleValuesSV();
+ for (int row = 0; row < numDocs; row++) {
+ ids[row] = isNull(nullBitmap, row) ? ID_FOR_NULL : dictionary.put(doubleValues[row]);
+ }
+ break;
+ case BIG_DECIMAL:
+ Object[] bigDecimalValues = blockValSet.getBigDecimalValuesSV();
+ for (int row = 0; row < numDocs; row++) {
+ ids[row] = isNull(nullBitmap, row) ? ID_FOR_NULL : dictionary.put(bigDecimalValues[row]);
+ }
+ break;
+ case STRING:
+ String[] stringValues = blockValSet.getStringValuesSV();
+ for (int row = 0; row < numDocs; row++) {
+ ids[row] = isNull(nullBitmap, row) ? ID_FOR_NULL : dictionary.put(stringValues[row]);
+ }
+ break;
+ case BYTES:
+ byte[][] bytesValues = blockValSet.getBytesValuesSV();
+ for (int row = 0; row < numDocs; row++) {
+ ids[row] = isNull(nullBitmap, row) ? ID_FOR_NULL : dictionary.put(new ByteArray(bytesValues[row]));
+ }
+ break;
+ default:
+ throw new IllegalArgumentException(
+ "Illegal data type for grouping-sets group key generator: " + _storedTypes[col]);
+ }
+ return ids;
+ }
+
+ private static boolean isNull(RoaringBitmap nullBitmap, int row) {
+ return nullBitmap != null && nullBitmap.contains(row);
+ }
+
+ /// Returns the group id for the given composite key, creating a new group while the per-segment group
+ /// limit has not been reached. Once the limit is reached, only existing groups are returned and brand-new
+ /// keys map to {@link #INVALID_ID} (the aggregation result holders skip {@code INVALID_ID}).
+ private int getGroupIdForKey(FixedIntArray keyList) {
+ int numGroups = _groupKeyMap.size();
+ if (numGroups < _numGroupsLimit) {
+ return _groupKeyMap.computeIfAbsent(keyList, k -> numGroups);
+ } else {
+ return _groupKeyMap.getInt(keyList);
+ }
+ }
+
+ @Override
+ public int getCurrentGroupKeyUpperBound() {
+ return _groupKeyMap.size();
+ }
+
+ @Override
+ public int getNumKeys() {
+ return _groupKeyMap.size();
+ }
+
+ @Override
+ public Iterator<GroupKey> getGroupKeys() {
+ return new GroupKeyIterator();
+ }
+
+ /// Reconstructs the output row for a composite key: the union column values (NULL where rolled up) followed
+ /// by the integer grouping-id 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).
+ keys[_numGroupByExpressions] = ids[_numGroupByExpressions];
+ return keys;
+ }
+
+ private class GroupKeyIterator implements Iterator<GroupKey> {
+ private final ObjectIterator<Object2IntMap.Entry<FixedIntArray>> _iterator;
+ private final GroupKey _groupKey;
+
+ GroupKeyIterator() {
+ _iterator = _groupKeyMap.object2IntEntrySet().fastIterator();
+ _groupKey = new GroupKey();
+ }
+
+ @Override
+ public boolean hasNext() {
+ return _iterator.hasNext();
+ }
+
+ @Override
+ public GroupKey next() {
+ Object2IntMap.Entry<FixedIntArray> entry = _iterator.next();
+ _groupKey._groupId = entry.getIntValue();
+ _groupKey._keys = buildKeysFromIds(entry.getKey());
+ return _groupKey;
+ }
+
+ @Override
+ public void remove() {
+ throw new UnsupportedOperationException();
+ }
+ }
+}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java
index 839948f88be..574b4d94ca0 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java
@@ -78,7 +78,9 @@ public class GroupByDataTableReducer implements DataTableReducer {
private final QueryContext _queryContext;
private final AggregationFunction[] _aggregationFunctions;
private final int _numAggregationFunctions;
- private final int _numGroupByExpressions;
+ /// Number of key columns: union group-by columns plus the synthetic $groupingId column for grouping sets.
+ /// Key columns precede the aggregation columns in the merged row layout.
+ private final int _numKeyColumns;
private final int _numColumns;
public GroupByDataTableReducer(QueryContext queryContext) {
@@ -88,8 +90,8 @@ public class GroupByDataTableReducer implements DataTableReducer {
_numAggregationFunctions = _aggregationFunctions.length;
List<ExpressionContext> groupByExpressions = queryContext.getGroupByExpressions();
assert groupByExpressions != null;
- _numGroupByExpressions = groupByExpressions.size();
- _numColumns = _numAggregationFunctions + _numGroupByExpressions;
+ _numKeyColumns = queryContext.getNumGroupByKeyColumns();
+ _numColumns = _numAggregationFunctions + _numKeyColumns;
}
/**
@@ -202,8 +204,8 @@ public class GroupByDataTableReducer implements DataTableReducer {
FilterContext havingFilter = _queryContext.getHavingFilter();
if (havingFilter != null) {
rows = new ArrayList<>();
- HavingFilterHandler havingFilterHandler =
- new HavingFilterHandler(havingFilter, postAggregationHandler, _queryContext.isNullHandlingEnabled());
+ HavingFilterHandler havingFilterHandler = new HavingFilterHandler(havingFilter, postAggregationHandler,
+ _queryContext.requiresNullAwareKeySerialization());
int processedRows = 0;
while (rows.size() < limit && sortedIterator.hasNext()) {
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(processedRows++, "GroupByDataTableReducer");
@@ -248,9 +250,9 @@ public class GroupByDataTableReducer implements DataTableReducer {
private DataSchema getPrePostAggregationDataSchema(DataSchema dataSchema) {
String[] columnNames = dataSchema.getColumnNames();
ColumnDataType[] columnDataTypes = new ColumnDataType[_numColumns];
- System.arraycopy(dataSchema.getColumnDataTypes(), 0, columnDataTypes, 0, _numGroupByExpressions);
+ System.arraycopy(dataSchema.getColumnDataTypes(), 0, columnDataTypes, 0, _numKeyColumns);
for (int i = 0; i < _numAggregationFunctions; i++) {
- columnDataTypes[i + _numGroupByExpressions] = _aggregationFunctions[i].getFinalResultColumnType();
+ columnDataTypes[i + _numKeyColumns] = _aggregationFunctions[i].getFinalResultColumnType();
}
return new DataSchema(columnNames, columnDataTypes);
}
@@ -292,9 +294,11 @@ public class GroupByDataTableReducer implements DataTableReducer {
public void runJob() {
try {
for (DataTable dataTable : reduceGroup) {
- boolean nullHandlingEnabled = _queryContext.isNullHandlingEnabled();
+ /// Grouping-set queries serialize NULL group keys via null bitmaps regardless of the user's
+ /// null-handling option, so their key nulls must be restored here as well.
+ boolean restoreNulls = _queryContext.requiresNullAwareKeySerialization();
RoaringBitmap[] nullBitmaps = null;
- if (nullHandlingEnabled) {
+ if (restoreNulls) {
nullBitmaps = new RoaringBitmap[_numColumns];
for (int i = 0; i < _numColumns; i++) {
nullBitmaps[i] = dataTable.getNullRowIds(i);
@@ -358,7 +362,7 @@ public class GroupByDataTableReducer implements DataTableReducer {
if (customObject != null) {
assert _aggregationFunctions != null;
values[colId] =
- _aggregationFunctions[colId - _numGroupByExpressions].deserializeIntermediateResult(
+ _aggregationFunctions[colId - _numKeyColumns].deserializeIntermediateResult(
customObject);
}
break;
@@ -367,7 +371,7 @@ public class GroupByDataTableReducer implements DataTableReducer {
throw new IllegalStateException();
}
}
- if (nullHandlingEnabled) {
+ if (restoreNulls) {
for (int colId = 0; colId < _numColumns; colId++) {
if (nullBitmaps[colId] != null && nullBitmaps[colId].contains(rowId)) {
values[colId] = null;
@@ -448,8 +452,8 @@ public class GroupByDataTableReducer implements DataTableReducer {
FilterContext havingFilter = _queryContext.getHavingFilter();
if (havingFilter != null) {
rows = new ArrayList<>();
- HavingFilterHandler havingFilterHandler =
- new HavingFilterHandler(havingFilter, postAggregationHandler, _queryContext.isNullHandlingEnabled());
+ HavingFilterHandler havingFilterHandler = new HavingFilterHandler(havingFilter, postAggregationHandler,
+ _queryContext.requiresNullAwareKeySerialization());
for (int i = 0; i < numRows; i++) {
Object[] row = getConvertedRowWithFinalResult(dataTable, i);
if (havingFilterHandler.isMatch(row)) {
@@ -496,7 +500,7 @@ public class GroupByDataTableReducer implements DataTableReducer {
Object[] row = new Object[_numColumns];
ColumnDataType[] columnDataTypes = dataTable.getDataSchema().getColumnDataTypes();
for (int i = 0; i < _numColumns; i++) {
- if (i < _numGroupByExpressions) {
+ if (i < _numKeyColumns) {
row[i] = getConvertedKey(dataTable, columnDataTypes[i], rowId, i);
} else {
row[i] = AggregationFunctionUtils.getConvertedFinalResult(dataTable, columnDataTypes[i], rowId, i);
@@ -543,7 +547,9 @@ public class GroupByDataTableReducer implements DataTableReducer {
DataTableBuilder dataTableBuilder = DataTableBuilderFactory.getDataTableBuilder(dataSchema);
ColumnDataType[] storedColumnDataTypes = dataSchema.getStoredColumnDataTypes();
Iterator<Record> iterator = indexedTable.iterator();
- if (_queryContext.isNullHandlingEnabled()) {
+ /// Grouping-set queries carry NULL group keys that must be serialized via the null-aware path even when
+ /// the user did not enable null handling.
+ if (_queryContext.requiresNullAwareKeySerialization()) {
RoaringBitmap[] nullBitmaps = new RoaringBitmap[_numColumns];
Object[] nullPlaceholders = new Object[_numColumns];
for (int colId = 0; colId < _numColumns; colId++) {
@@ -562,7 +568,7 @@ public class GroupByDataTableReducer implements DataTableReducer {
dataTableBuilder.setNull(i);
} else {
dataTableBuilder.setColumn(i,
- _aggregationFunctions[i - _numGroupByExpressions].serializeIntermediateResult(value));
+ _aggregationFunctions[i - _numKeyColumns].serializeIntermediateResult(value));
}
} else {
if (value == null) {
@@ -590,7 +596,7 @@ public class GroupByDataTableReducer implements DataTableReducer {
dataTableBuilder.setNull(i);
} else if (storedColumnDataTypes[i] == ColumnDataType.OBJECT) {
dataTableBuilder.setColumn(i,
- _aggregationFunctions[i - _numGroupByExpressions].serializeIntermediateResult(value));
+ _aggregationFunctions[i - _numKeyColumns].serializeIntermediateResult(value));
} else {
DataTableBuilderUtils.setColumn(dataTableBuilder, storedColumnDataTypes[i], i, value);
}
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 4c4bcaf03c3..0f6d9b16f42 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
@@ -26,6 +26,7 @@ import org.apache.commons.lang3.tuple.Pair;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.FilterContext;
import org.apache.pinot.common.request.context.FunctionContext;
+import org.apache.pinot.common.request.context.GroupingSets;
import org.apache.pinot.common.request.context.RequestContextUtils;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
@@ -44,6 +45,12 @@ import org.apache.pinot.core.query.request.context.QueryContext;
public class PostAggregationHandler implements ValueExtractorFactory {
private final Map<Pair<FunctionContext, FilterContext>, Integer> _filteredAggregationsIndexMap;
private final int _numGroupByExpressions;
+ /// 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.
+ private final int _groupingIdColumnIndex;
private final Map<ExpressionContext, Integer> _groupByExpressionIndexMap;
private final DataSchema _dataSchema;
private final ValueExtractor[] _valueExtractors;
@@ -63,6 +70,8 @@ public class PostAggregationHandler implements ValueExtractorFactory {
_numGroupByExpressions = 0;
_groupByExpressionIndexMap = null;
}
+ _numKeyColumns = _numGroupByExpressions + queryContext.getNumExtraGroupByKeyColumns();
+ _groupingIdColumnIndex = queryContext.isGroupingSets() ? _numGroupByExpressions : -1;
// 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.
@@ -120,17 +129,25 @@ public class PostAggregationHandler implements ValueExtractorFactory {
FunctionContext function = expression.getFunction();
Preconditions
.checkState(function != null, "Failed to find SELECT expression: %s in the GROUP-BY clause", expression);
+ /// Function names reaching here are canonicalized (underscores removed, lower-cased), so GROUPING_ID
+ /// arrives as "groupingid".
+ String functionName = function.getFunctionName();
+ if (GroupingSets.isGroupingFunction(functionName)) {
+ /// GROUPING(col, ...) / GROUPING_ID(col, ...): computed from the synthetic $groupingId discriminator
+ /// column rather than from a server-side aggregation.
+ return new GroupingValueExtractor(function);
+ }
if (function.getType() == FunctionContext.Type.AGGREGATION) {
// Aggregation function
return new ColumnValueExtractor(
- _filteredAggregationsIndexMap.get(Pair.of(function, null)) + _numGroupByExpressions, _dataSchema);
+ _filteredAggregationsIndexMap.get(Pair.of(function, null)) + _numKeyColumns, _dataSchema);
} else if (function.getType() == FunctionContext.Type.TRANSFORM && function.getFunctionName()
.equalsIgnoreCase("filter")) {
FunctionContext aggregation = function.getArguments().get(0).getFunction();
ExpressionContext filterExpression = function.getArguments().get(1);
FilterContext filter = RequestContextUtils.getFilter(filterExpression);
return new ColumnValueExtractor(
- _filteredAggregationsIndexMap.get(Pair.of(aggregation, filter)) + _numGroupByExpressions, _dataSchema);
+ _filteredAggregationsIndexMap.get(Pair.of(aggregation, filter)) + _numKeyColumns, _dataSchema);
} else {
// Post-aggregation function
return new PostAggregationValueExtractor(function);
@@ -183,4 +200,47 @@ public class PostAggregationHandler implements ValueExtractorFactory {
return _postAggregationFunction.invoke(_arguments);
}
}
+
+ /// 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).
+ private class GroupingValueExtractor implements ValueExtractor {
+ private final FunctionContext _function;
+ private final int[] _unionColumnIndexes;
+
+ GroupingValueExtractor(FunctionContext function) {
+ _function = function;
+ List<ExpressionContext> arguments = function.getArguments();
+ Preconditions.checkState(!arguments.isEmpty(), "GROUPING requires at least one argument");
+ _unionColumnIndexes = 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;
+ }
+ }
+
+ @Override
+ public String getColumnName() {
+ return _function.toString();
+ }
+
+ @Override
+ public ColumnDataType getColumnDataType() {
+ return ColumnDataType.INT;
+ }
+
+ @Override
+ public Object extract(Object[] row) {
+ /// Without grouping sets no column is rolled up, so GROUPING is always 0.
+ if (_groupingIdColumnIndex < 0) {
+ return 0;
+ }
+ int groupingId = ((Number) row[_groupingIdColumnIndex]).intValue();
+ return GroupingSets.groupingValue(groupingId, _unionColumnIndexes);
+ }
+ }
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ReducerDataSchemaUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ReducerDataSchemaUtils.java
index 7bd751877dc..82a4a01afd0 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ReducerDataSchemaUtils.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ReducerDataSchemaUtils.java
@@ -23,6 +23,7 @@ import java.util.List;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.FilterContext;
+import org.apache.pinot.common.request.context.GroupingSets;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils;
@@ -69,19 +70,36 @@ public class ReducerDataSchemaUtils {
queryContext.getFilteredAggregationFunctions();
assert groupByExpressions != null && filteredAggregationFunctions != null;
int numGroupByExpression = groupByExpressions.size();
+ /// Grouping-set queries carry an extra synthetic $groupingId key column after the union group-by columns.
+ int numExtraKeyColumns = queryContext.getNumExtraGroupByKeyColumns();
+ int numKeyColumns = numGroupByExpression + numExtraKeyColumns;
int numAggregations = filteredAggregationFunctions.size();
- int numColumns = numGroupByExpression + numAggregations;
+ int numColumns = numKeyColumns + numAggregations;
String[] columnNames = new String[numColumns];
+ /// Rolling-upgrade safety: an older server (without grouping-set support) silently ignores the
+ /// groupingSets wire field and returns a plain GROUP BY result missing the synthetic $groupingId column.
+ /// Detect that specific shape and fail with an actionable message instead of an opaque column-count BUG.
+ /// NOTE: this is a column-count heuristic, not a version handshake; it should be replaced by a
+ /// server-reported capability signal once Pinot supports per-query capability negotiation.
+ if (numExtraKeyColumns > 0 && dataSchema.size() == numGroupByExpression + numAggregations) {
+ throw new IllegalStateException(
+ "GROUP BY GROUPING SETS / ROLLUP / CUBE requires all servers to support grouping sets, but a server "
+ + "returned a result without the grouping-id column. Upgrade all servers before issuing grouping-set "
+ + "queries.");
+ }
Preconditions.checkState(dataSchema.size() == numColumns,
- "BUG: Expect same number of group-by expressions, aggregations and columns in data schema, got %s group-by "
- + "expressions, %s aggregations, %s columns in data schema", numGroupByExpression, numAggregations,
+ "BUG: Expect same number of group-by key columns, aggregations and columns in data schema, got %s group-by "
+ + "key columns, %s aggregations, %s columns in data schema", numKeyColumns, numAggregations,
dataSchema.size());
for (int i = 0; i < numGroupByExpression; i++) {
columnNames[i] = groupByExpressions.get(i).toString();
}
+ if (numExtraKeyColumns > 0) {
+ columnNames[numGroupByExpression] = GroupingSets.GROUPING_ID_COLUMN;
+ }
for (int i = 0; i < numAggregations; i++) {
Pair<AggregationFunction, FilterContext> pair = filteredAggregationFunctions.get(i);
- columnNames[numGroupByExpression + i] =
+ columnNames[numKeyColumns + i] =
AggregationFunctionUtils.getResultColumnName(pair.getLeft(), pair.getRight());
}
return new DataSchema(columnNames, dataSchema.getColumnDataTypes());
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
index c0ea02e5f38..ab7a4cab57b 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
@@ -151,6 +151,10 @@ public class QueryContext {
private boolean _serverReturnFinalResultKeyUnpartitioned;
private boolean _accurateGroupByWithoutOrderBy;
private boolean _isUnsafeTrim;
+ /// Grouping sets for GROUP BY GROUPING SETS / ROLLUP / CUBE, as sorted column indexes into
+ /// _groupByExpressions (the union of all grouping columns). Null for a plain GROUP BY (or non-group-by)
+ /// query. Each entry is one grouping set; an empty array denotes the grand-total set ().
+ private List<int[]> _groupingSets;
// Collection of index types to skip per column
private Map<String, Set<FieldConfig.IndexType>> _skipIndexes;
@@ -246,6 +250,44 @@ public class QueryContext {
return _groupByExpressions;
}
+ /// Returns the grouping sets for a GROUP BY GROUPING SETS / ROLLUP / CUBE query, or {@code null} for a
+ /// plain GROUP BY (or non-group-by) query. Each entry is the sorted list of column indexes into
+ /// {@link #getGroupByExpressions()} (the union of all grouping columns) that participate in that set; an
+ /// empty array denotes the grand-total set ().
+ @Nullable
+ public List<int[]> getGroupingSets() {
+ return _groupingSets;
+ }
+
+ public void setGroupingSets(@Nullable List<int[]> groupingSets) {
+ _groupingSets = groupingSets;
+ }
+
+ /// Returns whether this is a GROUP BY GROUPING SETS / ROLLUP / CUBE query.
+ public boolean isGroupingSets() {
+ return _groupingSets != null;
+ }
+
+ /// Returns the number of synthetic key columns appended after the group-by (union) columns: 1 for a
+ /// grouping-set query (the {@code $groupingId} discriminator column), 0 otherwise.
+ public int getNumExtraGroupByKeyColumns() {
+ return isGroupingSets() ? 1 : 0;
+ }
+
+ /// Returns the total number of group-by key columns in the server result / reducer row layout: the union
+ /// group-by columns plus any synthetic key columns ({@link #getNumExtraGroupByKeyColumns()}).
+ public int getNumGroupByKeyColumns() {
+ int numGroupByExpressions = _groupByExpressions != null ? _groupByExpressions.size() : 0;
+ return numGroupByExpressions + getNumExtraGroupByKeyColumns();
+ }
+
+ /// Returns whether group-by key columns must be serialized/deserialized through the null-aware path (null
+ /// bitmaps). True when the user enabled null handling, or for grouping-set queries, which produce NULL keys
+ /// for rolled-up columns regardless of the user's null-handling option.
+ public boolean requiresNullAwareKeySerialization() {
+ return isNullHandlingEnabled() || isGroupingSets();
+ }
+
/**
* Returns the filter in the HAVING clause, or {@code null} if there is no HAVING clause.
*/
@@ -565,7 +607,31 @@ public class QueryContext {
return _effectiveSegmentGroupTrimSize;
}
+ /// Returns the PER-grouping-set segment group-trim size for a grouping-set query, or -1 to disable
+ /// per-segment trim (no ORDER BY, or trimming disabled). Unlike the global {@link
+ /// #getEffectiveSegmentGroupTrimSize()} (which is -1 for grouping sets), the per-segment trim for grouping
+ /// sets keeps up to this many groups WITHIN each grouping set (bucketed by the discriminator), so a global
+ /// top-K cannot starve low-magnitude sets such as the grand total. The broker still applies the final
+ /// ORDER BY + LIMIT across all sets.
+ public int getGroupingSetSegmentTrimSize() {
+ if (!isGroupingSets()) {
+ return -1;
+ }
+ int minGroupTrimSize = getMinSegmentGroupTrimSize();
+ if (getOrderByExpressions() != null && minGroupTrimSize > 0) {
+ return GroupByUtils.getTableCapacity(getLimit(), minGroupTrimSize);
+ }
+ return -1;
+ }
+
private int calculateEffectiveSegmentGroupTrimSize() {
+ /// Grouping sets expand each input row into one group per set; a global per-segment top-K could drop
+ /// groups belonging to low-magnitude sets (e.g. the grand-total set) before they are merged. Grouping-set
+ /// queries therefore use the per-set bucketed trim (getGroupingSetSegmentTrimSize()) instead of this
+ /// global trim, which is disabled for them here.
+ if (isGroupingSets()) {
+ return -1;
+ }
int minGroupTrimSize = getMinSegmentGroupTrimSize();
List<OrderByExpressionContext> orderByExpressions = getOrderByExpressions();
if (!isUnsafeTrim()) {
@@ -630,6 +696,7 @@ public class QueryContext {
private List<String> _aliasList;
private FilterContext _filter;
private List<ExpressionContext> _groupByExpressions;
+ private List<int[]> _groupingSets;
private FilterContext _havingFilter;
private List<OrderByExpressionContext> _orderByExpressions;
private int _limit;
@@ -678,6 +745,14 @@ public class QueryContext {
return this;
}
+ /// Sets the grouping sets for a GROUP BY GROUPING SETS / ROLLUP / CUBE query. Each entry is the sorted
+ /// list of column indexes into {@code groupByExpressions} (the union of all grouping columns); an empty
+ /// array denotes the grand-total set (). Leave unset (null) for a plain GROUP BY query.
+ public Builder setGroupingSets(@Nullable List<int[]> groupingSets) {
+ _groupingSets = groupingSets;
+ return this;
+ }
+
public Builder setHavingFilter(FilterContext havingFilter) {
_havingFilter = havingFilter;
return this;
@@ -736,6 +811,14 @@ public class QueryContext {
queryContext.setServerReturnFinalResult(QueryOptionsUtils.isServerReturnFinalResult(_queryOptions));
queryContext.setServerReturnFinalResultKeyUnpartitioned(
QueryOptionsUtils.isServerReturnFinalResultKeyUnpartitioned(_queryOptions));
+ queryContext.setGroupingSets(_groupingSets);
+ if (queryContext.isGroupingSets()) {
+ /// Phase 1: grouping-set queries always use the intermediate-result merge path. The final-result
+ /// reduce path (server.returnFinalResult) does not restore the NULL group keys that grouping sets
+ /// produce, so disable that optimization here.
+ queryContext.setServerReturnFinalResult(false);
+ queryContext.setServerReturnFinalResultKeyUnpartitioned(false);
+ }
// Pre-calculate the aggregation functions and columns for the query
generateAggregationFunctions(queryContext);
@@ -743,8 +826,10 @@ public class QueryContext {
// Pre-calculate group-by configs
if (queryContext.getGroupByExpressions() != null) {
- queryContext._isUnsafeTrim =
- !queryContext.isSameOrderAndGroupByColumns(queryContext) || queryContext.getHavingFilter() != null;
+ /// Grouping-set queries expand each row into multiple groups; force the unsafe-trim path so the
+ /// sort-aggregate fast path (which assumes a single flat grouping) is not taken.
+ queryContext._isUnsafeTrim = queryContext.isGroupingSets()
+ || !queryContext.isSameOrderAndGroupByColumns(queryContext) || queryContext.getHavingFilter() != null;
Integer sortAggregateLimitThreshold = QueryOptionsUtils.getSortAggregateLimitThreshold(_queryOptions);
if (sortAggregateLimitThreshold != null) {
queryContext.setSortAggregateLimitThreshold(sortAggregateLimitThreshold);
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 611ffccd5b5..ec165cf8f7e 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,6 +114,45 @@ public class QueryContextConverterUtils {
}
}
+ /// 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.
+ List<int[]> groupingSets = null;
+ List<Integer> 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");
+ }
+ 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 idx = 0;
+ for (int bit = 0; bit < Integer.SIZE && idx < indexes.length; bit++) {
+ if ((mask & (1 << bit)) != 0) {
+ indexes[idx++] = bit;
+ }
+ }
+ groupingSets.add(indexes);
+ }
+ }
+
// ORDER BY
List<OrderByExpressionContext> orderByExpressions = null;
List<Expression> orderByList = pinotQuery.getOrderByList();
@@ -174,6 +213,7 @@ public class QueryContextConverterUtils {
.setAliasList(aliasList)
.setFilter(filter)
.setGroupByExpressions(groupByExpressions)
+ .setGroupingSets(groupingSets)
.setOrderByExpressions(orderByExpressions)
.setHavingFilter(havingFilter)
.setLimit(pinotQuery.getLimit())
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java
index 086ba7e6a1e..36247442171 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java
@@ -23,6 +23,8 @@ import java.util.Comparator;
import java.util.List;
import java.util.concurrent.ExecutorService;
import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.metrics.ServerMeter;
+import org.apache.pinot.common.metrics.ServerMetrics;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.HashUtil;
import org.apache.pinot.core.data.table.ConcurrentIndexedTable;
@@ -33,8 +35,12 @@ import org.apache.pinot.core.data.table.Record;
import org.apache.pinot.core.data.table.SimpleIndexedTable;
import org.apache.pinot.core.data.table.SortedRecords;
import org.apache.pinot.core.data.table.SortedRecordsMerger;
+import org.apache.pinot.core.data.table.TableResizer;
import org.apache.pinot.core.data.table.UnboundedConcurrentIndexedTable;
import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
+import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult;
+import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder;
+import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
import org.apache.pinot.core.query.reduce.DataTableReducerContext;
import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.spi.query.QueryThreadContext;
@@ -47,6 +53,40 @@ public final class GroupByUtils {
public static final int DEFAULT_MIN_NUM_GROUPS = 5000;
public static final int MAX_TRIM_THRESHOLD = 1_000_000_000;
+ /// Builds the segment-level {@link GroupByResultsBlock} for a GROUP BY GROUPING SETS / ROLLUP / CUBE query,
+ /// shared by {@code GroupByOperator} and {@code FilteredGroupByOperator}. When the group count exceeds the
+ /// per-set budget ({@code perSetTrimSize * numGroupingSets}), a per-set bucketed trim (keyed on the
+ /// {@code $groupingId} discriminator at {@code discriminatorColumnIndex}) keeps each grouping set's own top
+ /// candidates so a global top-K cannot starve low-magnitude sets such as the grand total; otherwise the full
+ /// segment result is returned. The broker still applies the final ORDER BY + LIMIT across all sets.
+ ///
+ /// @param discriminatorColumnIndex index of the synthetic $groupingId column, i.e. the number of union
+ /// group-by columns
+ public static GroupByResultsBlock buildGroupingSetsResultsBlock(QueryContext queryContext, DataSchema dataSchema,
+ GroupKeyGenerator groupKeyGenerator, GroupByResultHolder[] groupByResultHolders, int numGroups,
+ int discriminatorColumnIndex, boolean numGroupsLimitReached, boolean numGroupsWarningLimitReached) {
+ GroupByResultsBlock resultsBlock;
+ int perSetTrimSize = queryContext.getGroupingSetSegmentTrimSize();
+ int numGroupingSets = queryContext.getGroupingSets().size();
+ if (perSetTrimSize > 0 && numGroups > (long) perSetTrimSize * numGroupingSets) {
+ TableResizer tableResizer = new TableResizer(dataSchema, queryContext);
+ List<IntermediateRecord> intermediateRecords =
+ tableResizer.trimInSegmentResultsByGroupingSet(groupKeyGenerator, groupByResultHolders, perSetTrimSize,
+ discriminatorColumnIndex);
+ groupKeyGenerator.close();
+ ServerMetrics.get().addMeteredGlobalValue(ServerMeter.AGGREGATE_TIMES_GROUPS_TRIMMED, 1);
+ resultsBlock = new GroupByResultsBlock(dataSchema, intermediateRecords, queryContext);
+ resultsBlock.setGroupsTrimmed(true);
+ } else {
+ AggregationGroupByResult aggregationGroupByResult =
+ new AggregationGroupByResult(groupKeyGenerator, queryContext.getAggregationFunctions(), groupByResultHolders);
+ resultsBlock = new GroupByResultsBlock(dataSchema, aggregationGroupByResult, queryContext);
+ }
+ resultsBlock.setNumGroupsLimitReached(numGroupsLimitReached);
+ resultsBlock.setNumGroupsWarningLimitReached(numGroupsWarningLimitReached);
+ return resultsBlock;
+ }
+
/**
* Returns the capacity of the table required by the given query. NOTE: It returns {@code max(limit * 5, 5000)} to
* ensure the result accuracy.
@@ -112,6 +152,17 @@ public final class GroupByUtils {
queryContext.getMinServerGroupTrimSize(); // it's minBrokerGroupTrimSize in broker
int minInitialIndexedTableCapacity = queryContext.getMinInitialIndexedTableCapacity();
+ /// Grouping-set queries must not trim per server: a global ORDER BY top-K here would drop a row that ranks
+ /// higher globally once partial aggregates are merged at the broker, and could starve entire grouping sets
+ /// (silently wrong results). Keep all groups (bounded by numGroupsLimit) and defer ORDER BY + LIMIT to the
+ /// broker. Per-set bucketed trim still happens at the segment level.
+ if (queryContext.isGroupingSets()) {
+ int resultSize = queryContext.getNumGroupsLimit();
+ int initialCapacity = getIndexedTableInitialCapacity(resultSize, numGroups, minInitialIndexedTableCapacity);
+ return getTrimDisabledIndexedTable(dataSchema, false, queryContext, resultSize, initialCapacity, numThreads,
+ executorService);
+ }
+
// Disable trim when min trim size is non-positive
int trimSize = minTrimSize > 0 ? getTableCapacity(limit, minTrimSize) : Integer.MAX_VALUE;
@@ -172,6 +223,16 @@ public final class GroupByUtils {
// TODO: Resolve the HAVING clause within the IndexedTable before returning the result
int resultSize = hasHaving ? trimSize : limit;
+ /// Grouping-set queries must not incrementally trim while merging server responses (a row could be dropped
+ /// before all its partial aggregates are merged) nor apply a per-server top-K. Force the trim-disabled path
+ /// so all groups are merged first; finish() then keeps the correct global top-K (resultSize) and the broker
+ /// applies the final ORDER BY + LIMIT over the fully-merged table.
+ if (queryContext.isGroupingSets()) {
+ int initialCapacity = getIndexedTableInitialCapacity(resultSize, numGroups, minInitialIndexedTableCapacity);
+ return getTrimDisabledIndexedTable(dataSchema, hasFinalInput, queryContext, resultSize, initialCapacity,
+ numThreads, executorService);
+ }
+
// When there is no ORDER BY, trim is not required because the indexed table stops accepting new groups once the
// result size is reached
if (!hasOrderBy) {
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/table/TableResizerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/table/TableResizerTest.java
index a41ead8ec82..a8d7a0ee2ac 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/data/table/TableResizerTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/data/table/TableResizerTest.java
@@ -375,4 +375,58 @@ public class TableResizerTest {
assertEquals(results.get(2)._record, _records.get(3));
}
}
+
+ @Test
+ public void testTrimInSegmentResultsByGroupingSet() {
+ /// ROLLUP(d1) yields grouping sets {d1} (grouping-id 0) and {} (grouping-id 1). The record layout is
+ /// [d1, $groupingId, sum(m1)], so the discriminator column index is 1 (after the single union column).
+ TableResizer tableResizer = new TableResizer(
+ new DataSchema(new String[]{"d1", "$groupingId", "sum(m1)"}, new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.STRING, DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.DOUBLE}),
+ QueryContextConverterUtils.getQueryContext(
+ "SELECT d1, SUM(m1) FROM testTable GROUP BY ROLLUP(d1) ORDER BY SUM(m1) DESC"));
+
+ /// Bucket {d1}: five detail groups with sums 10..50. Bucket {}: one grand-total group with sum 150.
+ String[] details = {"a", "b", "c", "d", "e"};
+ double[] sums = {10.0, 20.0, 30.0, 40.0, 50.0};
+ List<GroupKeyGenerator.GroupKey> groupKeys = new ArrayList<>();
+ DoubleGroupByResultHolder holder = new DoubleGroupByResultHolder(6, 6, 0.0);
+ for (int i = 0; i < details.length; i++) {
+ GroupKeyGenerator.GroupKey groupKey = new GroupKeyGenerator.GroupKey();
+ groupKey._groupId = i;
+ groupKey._keys = new Object[]{details[i], 0}; // grouping-id 0 => set {d1}
+ groupKeys.add(groupKey);
+ holder.setValueForKey(i, sums[i]);
+ }
+ GroupKeyGenerator.GroupKey grandTotal = new GroupKeyGenerator.GroupKey();
+ grandTotal._groupId = 5;
+ grandTotal._keys = new Object[]{null, 1}; // grouping-id 1 => set {} (d1 rolled up to NULL)
+ groupKeys.add(grandTotal);
+ holder.setValueForKey(5, 150.0);
+
+ GroupKeyGenerator groupKeyGenerator = mock(GroupKeyGenerator.class);
+ when(groupKeyGenerator.getGroupKeys()).then(invocation -> groupKeys.iterator());
+
+ /// Keep top-2 PER grouping set. A global top-2 by sum would keep only {grand total (150), e (50)} -- a
+ /// single detail row. The per-set trim must instead keep the grand total AND the top-2 details, so the
+ /// low-cardinality detail set is not starved by the dominant grand-total row.
+ List<IntermediateRecord> result =
+ tableResizer.trimInSegmentResultsByGroupingSet(groupKeyGenerator, new GroupByResultHolder[]{holder}, 2, 1);
+
+ int numDetailRows = 0;
+ int numGrandTotalRows = 0;
+ for (IntermediateRecord record : result) {
+ int groupingId = (int) record._record.getValues()[1];
+ if (groupingId == 0) {
+ numDetailRows++;
+ /// The kept details must be the two highest sums (e=50, d=40).
+ assertTrue((double) record._record.getValues()[2] >= 40.0);
+ } else {
+ numGrandTotalRows++;
+ }
+ }
+ assertEquals(result.size(), 3);
+ assertEquals(numDetailRows, 2);
+ assertEquals(numGrandTotalRows, 1);
+ }
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/reduce/ReducerDataSchemaUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/reduce/ReducerDataSchemaUtilsTest.java
index 4836199cbaf..d4bd72ad83f 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/query/reduce/ReducerDataSchemaUtilsTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/query/reduce/ReducerDataSchemaUtilsTest.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.core.query.reduce;
+import org.apache.pinot.common.request.context.GroupingSets;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.core.query.request.context.QueryContext;
@@ -25,6 +26,8 @@ import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUt
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
public class ReducerDataSchemaUtilsTest {
@@ -106,4 +109,30 @@ public class ReducerDataSchemaUtilsTest {
assertEquals(canonicalDataSchema, new DataSchema(new String[]{"col1", "plus(col2,col3)"},
new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.DOUBLE}));
}
+
+ /// A grouping-set query expects the synthetic $groupingId key column. During a rolling upgrade an older
+ /// server (without grouping-set support) silently ignores the wire field and returns a plain GROUP BY result
+ /// missing that column. The reducer must reject this with an actionable "upgrade all servers" message rather
+ /// than an opaque column-count BUG, and must still canonicalize the fully-upgraded (correct) shape.
+ @Test
+ public void testCanonicalizeDataSchemaForGroupByGroupingSetsMixedVersion() {
+ QueryContext queryContext =
+ QueryContextConverterUtils.getQueryContext("SELECT col3, COUNT(*) FROM testTable GROUP BY ROLLUP(col3)");
+ assertTrue(queryContext.isGroupingSets());
+
+ /// Old-server shape: [col3, count(*)] -- missing the $groupingId column the new broker expects.
+ DataSchema oldServerSchema =
+ new DataSchema(new String[]{"col3", "count(*)"}, new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.LONG});
+ IllegalStateException e = expectThrows(IllegalStateException.class,
+ () -> ReducerDataSchemaUtils.canonicalizeDataSchemaForGroupBy(queryContext, oldServerSchema));
+ assertTrue(e.getMessage().contains("Upgrade all servers"), "unexpected message: " + e.getMessage());
+
+ /// Fully-upgraded shape: [col3, $groupingId, count(*)] canonicalizes without throwing.
+ DataSchema newServerSchema = new DataSchema(new String[]{"col3", GroupingSets.GROUPING_ID_COLUMN, "count(*)"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.INT, ColumnDataType.LONG});
+ DataSchema canonicalDataSchema =
+ ReducerDataSchemaUtils.canonicalizeDataSchemaForGroupBy(queryContext, newServerSchema);
+ assertEquals(canonicalDataSchema.size(), 3);
+ assertEquals(canonicalDataSchema.getColumnName(1), GroupingSets.GROUPING_ID_COLUMN);
+ }
}
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 ef331ebf59c..61fa6295727 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
@@ -24,6 +24,7 @@ import java.util.HashSet;
import java.util.List;
import java.util.Map;
import org.apache.commons.lang3.tuple.Pair;
+import org.apache.pinot.common.request.PinotQuery;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.FilterContext;
import org.apache.pinot.common.request.context.FunctionContext;
@@ -38,6 +39,7 @@ import org.apache.pinot.core.query.aggregation.function.MinAggregationFunction;
import org.apache.pinot.core.query.aggregation.function.SumAggregationFunction;
import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.sql.parsers.CalciteSqlParser;
import org.testng.annotations.Test;
import static org.testng.Assert.*;
@@ -810,4 +812,38 @@ public class BrokerRequestToQueryContextConverterTest {
queryContext = QueryContextConverterUtils.getQueryContext(query);
assertNull(queryContext.getFilter());
}
+
+ @Test
+ public void testInvalidGroupingSetMasksRejected() {
+ /// 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}) {
+ PinotQuery corrupted =
+ CalciteSqlParser.compileToPinotQuery("SELECT COUNT(*) FROM testTable GROUP BY ROLLUP(a, b)");
+ corrupted.setGroupingSetMasks(List.of(3, badMask));
+ assertThrows(IllegalStateException.class, () -> QueryContextConverterUtils.getQueryContext(corrupted));
+ }
+
+ /// An empty set list must be rejected rather than silently producing zero groups.
+ PinotQuery emptyMasks =
+ CalciteSqlParser.compileToPinotQuery("SELECT COUNT(*) FROM testTable GROUP BY ROLLUP(a, b)");
+ emptyMasks.setGroupingSetMasks(List.of());
+ assertThrows(IllegalStateException.class, () -> QueryContextConverterUtils.getQueryContext(emptyMasks));
+
+ /// 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.
+ StringBuilder columns = new StringBuilder();
+ for (int i = 0; i < 32; 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));
+ }
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/util/GroupByUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/util/GroupByUtilsTest.java
index 99d5cdf82f9..30c2990cea9 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/util/GroupByUtilsTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/util/GroupByUtilsTest.java
@@ -18,10 +18,27 @@
*/
package org.apache.pinot.core.util;
+import java.util.Collections;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.request.context.GroupingSets;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.common.utils.HashUtil;
+import org.apache.pinot.core.common.datatable.DataTableBuilderFactory;
+import org.apache.pinot.core.data.table.ConcurrentIndexedTable;
+import org.apache.pinot.core.data.table.IndexedTable;
+import org.apache.pinot.core.data.table.UnboundedConcurrentIndexedTable;
+import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
+import org.apache.pinot.core.query.reduce.DataTableReducerContext;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.sql.parsers.CalciteSqlParser;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
public class GroupByUtilsTest {
@@ -71,4 +88,49 @@ public class GroupByUtilsTest {
assertEquals(GroupByUtils.getIndexedTableInitialCapacity(100, 100, 256), HashUtil.getHashMapCapacity(100));
assertEquals(GroupByUtils.getIndexedTableInitialCapacity(100, 1000, 256), HashUtil.getHashMapCapacity(100));
}
+
+ /// Grouping-set queries must keep all groups at the combine and reducer stages (no global ORDER BY trim), so
+ /// a row that ranks low on a single segment/server is never dropped before its partial aggregates are merged.
+ /// The IndexedTable must therefore be the trim-disabled [UnboundedConcurrentIndexedTable] even under
+ /// aggressive trim settings, whereas a plain GROUP BY with the same settings stays trim-enabled.
+ @Test
+ public void testGroupingSetQueriesDisableCombineAndReducerTrim() {
+ QueryContext groupingSets = QueryContextConverterUtils.getQueryContext(
+ CalciteSqlParser.compileToPinotQuery("SELECT a, COUNT(*) FROM t GROUP BY ROLLUP(a) ORDER BY COUNT(*) DESC"));
+ assertTrue(groupingSets.isGroupingSets());
+ DataSchema schema = new DataSchema(new String[]{"a", GroupingSets.GROUPING_ID_COLUMN, "count"},
+ new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.INT, ColumnDataType.LONG});
+ ExecutorService executorService = Executors.newCachedThreadPool();
+ try {
+ /// Aggressive trim settings (minGroupTrimSize=1, groupTrimThreshold=1) would enable trim for a plain query.
+ DataTableReducerContext aggressiveTrim = new DataTableReducerContext(executorService, 2, 10_000L, 1, 1, 128);
+ DataTable dataTable = DataTableBuilderFactory.getDataTableBuilder(schema).build();
+ IndexedTable reducerTable =
+ GroupByUtils.createIndexedTableForDataTableReducer(dataTable, groupingSets, aggressiveTrim, 2,
+ executorService);
+ assertTrue(reducerTable instanceof UnboundedConcurrentIndexedTable,
+ "grouping-set reducer table must be trim-disabled, got " + reducerTable.getClass().getSimpleName());
+
+ GroupByResultsBlock resultsBlock = new GroupByResultsBlock(schema, Collections.emptyList(), groupingSets);
+ IndexedTable combineTable =
+ GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, groupingSets, 2, executorService);
+ assertTrue(combineTable instanceof UnboundedConcurrentIndexedTable,
+ "grouping-set combine table must be trim-disabled, got " + combineTable.getClass().getSimpleName());
+
+ /// Contrast: a plain GROUP BY with the same aggressive trim settings still uses the trim-enabled table.
+ QueryContext plainGroupBy = QueryContextConverterUtils.getQueryContext(
+ CalciteSqlParser.compileToPinotQuery("SELECT a, COUNT(*) FROM t GROUP BY a ORDER BY COUNT(*) DESC"));
+ DataSchema plainSchema = new DataSchema(new String[]{"a", "count"},
+ new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.LONG});
+ DataTableReducerContext trimEnabled = new DataTableReducerContext(executorService, 2, 10_000L, 100, 1, 128);
+ DataTable plainDataTable = DataTableBuilderFactory.getDataTableBuilder(plainSchema).build();
+ IndexedTable plainReducerTable =
+ GroupByUtils.createIndexedTableForDataTableReducer(plainDataTable, plainGroupBy, trimEnabled, 2,
+ executorService);
+ assertEquals(plainReducerTable.getClass(), ConcurrentIndexedTable.class,
+ "plain GROUP BY with aggressive trim must remain trim-enabled");
+ } finally {
+ executorService.shutdownNow();
+ }
+ }
}
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
new file mode 100644
index 00000000000..dd774545eae
--- /dev/null
+++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GroupingSetsQueriesTest.java
@@ -0,0 +1,720 @@
+/**
+ * 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.integration.tests.custom;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import java.io.File;
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericData;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+/// 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.
+///
+/// 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
+/// discriminator test: a ROLLUP(d1, d2) produces BOTH a genuine {@code (a, NULL)} detail group (from the
+/// {d1, d2} set, GROUPING(d2)=0) and a rolled-up {@code (a, NULL)} subtotal group (from the {d1} set,
+/// 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)}.
+@Test(suiteName = "CustomClusterIntegrationTest")
+public class GroupingSetsQueriesTest extends CustomDataQueryClusterIntegrationTest {
+ private static final String DEFAULT_TABLE_NAME = "GroupingSetsQueriesTest";
+ private static final String D1 = "d1";
+ private static final String D2 = "d2";
+ /// LONG and DOUBLE grouping columns (functionally determined by d1) to cover the non-STRING key-resolution
+ /// branches of the generator. MV column to exercise the multi-value rejection guard.
+ private static final String LNG = "lng";
+ private static final String DBL = "dbl";
+ private static final String MV = "mv";
+ private static final String MET = "met";
+ /// 2 rows for each of (a,x), (a,null), (b,x), (b,null) => 8 docs.
+ private static final int NUM_DOCS = 8;
+
+ @Override
+ public String getTableName() {
+ return DEFAULT_TABLE_NAME;
+ }
+
+ @Override
+ protected long getCountStarResult() {
+ return NUM_DOCS;
+ }
+
+ @Override
+ public TableConfig createOfflineTableConfig() {
+ return new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName()).setNullHandlingEnabled(true).build();
+ }
+
+ @Override
+ public Schema createSchema() {
+ return new Schema.SchemaBuilder()
+ .setSchemaName(getTableName())
+ .addSingleValueDimension(D1, DataType.STRING)
+ .addSingleValueDimension(D2, DataType.STRING)
+ .addSingleValueDimension(LNG, DataType.LONG)
+ .addSingleValueDimension(DBL, DataType.DOUBLE)
+ .addMultiValueDimension(MV, DataType.STRING)
+ .addMetric(MET, DataType.INT)
+ .build();
+ }
+
+ @Override
+ public List<File> createAvroFiles()
+ throws IOException {
+ org.apache.avro.Schema stringSchema = org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING);
+ /// d2 is nullable: union [null, string].
+ org.apache.avro.Schema nullableString = org.apache.avro.Schema.createUnion(
+ List.of(org.apache.avro.Schema.create(org.apache.avro.Schema.Type.NULL), stringSchema));
+ org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
+ avroSchema.setFields(List.of(
+ new org.apache.avro.Schema.Field(D1, stringSchema, null, null),
+ new org.apache.avro.Schema.Field(D2, nullableString, null, null),
+ new org.apache.avro.Schema.Field(LNG, org.apache.avro.Schema.create(org.apache.avro.Schema.Type.LONG), null,
+ null),
+ new org.apache.avro.Schema.Field(DBL, org.apache.avro.Schema.create(org.apache.avro.Schema.Type.DOUBLE), null,
+ null),
+ new org.apache.avro.Schema.Field(MV, org.apache.avro.Schema.createArray(stringSchema), null, null),
+ new org.apache.avro.Schema.Field(MET, org.apache.avro.Schema.create(org.apache.avro.Schema.Type.INT), null,
+ null)));
+
+ try (AvroFilesAndWriters avroFilesAndWriters = createAvroFilesAndWriters(avroSchema)) {
+ List<DataFileWriter<GenericData.Record>> writers = avroFilesAndWriters.getWriters();
+ int doc = 0;
+ for (String d1 : new String[]{"a", "b"}) {
+ for (String d2 : new String[]{"x", null}) {
+ for (int i = 0; i < 2; i++) {
+ GenericData.Record record = new GenericData.Record(avroSchema);
+ record.put(D1, d1);
+ record.put(D2, d2);
+ /// lng/dbl are functionally determined by d1 so ROLLUP over them mirrors ROLLUP(d1).
+ record.put(LNG, d1.equals("a") ? 100L : 200L);
+ record.put(DBL, d1.equals("a") ? 1.5 : 2.5);
+ record.put(MV, List.of("t1", "t2"));
+ record.put(MET, 1);
+ writers.get(doc % getNumAvroFiles()).append(record);
+ doc++;
+ }
+ }
+ }
+ return avroFilesAndWriters.getAvroFiles();
+ }
+ }
+
+ /// Encodes a JSON cell value: literal {@code null} becomes the sentinel "NULL", everything else its text.
+ private static String cell(JsonNode row, int index) {
+ JsonNode value = row.get(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)
+ throws Exception {
+ setUseMultiStageQueryEngine(false);
+ String prefixed = nullHandling ? "SET enableNullHandling=true; " + query : query;
+ JsonNode response = postQuery(prefixed);
+ JsonNode exceptions = response.get("exceptions");
+ assertTrue(exceptions == null || exceptions.isEmpty(), "query failed: " + response.toPrettyString());
+ return response.get("resultTable").get("rows");
+ }
+
+ @Test
+ public void testRollupWithGenuineAndRolledUpNulls()
+ throws Exception {
+ /// 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);
+
+ /// Key: d1|d2|grouping(d1)|grouping(d2) -> count
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ String key = cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt() + "|" + row.get(4).asInt();
+ actual.put(key, row.get(2).asLong());
+ }
+
+ Map<String, Long> expected = new HashMap<>();
+ /// {d1, d2} set (GROUPING = 0,0): includes the genuine (a, NULL) and (b, NULL) detail groups.
+ expected.put("a|x|0|0", 2L);
+ expected.put("a|NULL|0|0", 2L);
+ expected.put("b|x|0|0", 2L);
+ expected.put("b|NULL|0|0", 2L);
+ /// {d1} set (d2 rolled up, GROUPING = 0,1): the (a, NULL) / (b, NULL) SUBTOTALS, distinct from the genuine
+ /// (a, NULL) / (b, NULL) detail rows above purely by the grouping-id discriminator.
+ expected.put("a|NULL|0|1", 4L);
+ expected.put("b|NULL|0|1", 4L);
+ /// {} grand total (GROUPING = 1,1)
+ expected.put("NULL|NULL|1|1", 8L);
+
+ assertEquals(actual, expected);
+ assertEquals(rows.size(), expected.size());
+ /// The crux: the genuine and rolled-up (a, NULL) rows did NOT merge.
+ assertTrue(actual.containsKey("a|NULL|0|0") && actual.containsKey("a|NULL|0|1"),
+ "genuine-NULL detail and rolled-up-NULL subtotal must both be present");
+ }
+
+ @Test
+ public void testRolledUpNullWithoutNullHandling()
+ throws Exception {
+ /// 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);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0), row.get(1).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("a", 4L);
+ expected.put("b", 4L);
+ expected.put("NULL", 8L);
+ assertEquals(actual, expected);
+ assertTrue(actual.containsKey("NULL"), "rolled-up grand-total row must have NULL key without null handling");
+ }
+
+ @Test
+ public void testCube()
+ throws Exception {
+ String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*), GROUPING(" + D1 + "), GROUPING(" + D2 + ") FROM "
+ + getTableName() + " GROUP BY CUBE(" + D1 + ", " + D2 + ")";
+ JsonNode rows = runRows(query, true);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ String key = cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt() + "|" + row.get(4).asInt();
+ actual.put(key, row.get(2).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ /// {d1, d2}
+ expected.put("a|x|0|0", 2L);
+ expected.put("a|NULL|0|0", 2L);
+ expected.put("b|x|0|0", 2L);
+ expected.put("b|NULL|0|0", 2L);
+ /// {d1}
+ expected.put("a|NULL|0|1", 4L);
+ expected.put("b|NULL|0|1", 4L);
+ /// {d2}: d1 rolled up; d2 values are x (4) and genuine NULL (4)
+ expected.put("NULL|x|1|0", 4L);
+ expected.put("NULL|NULL|1|0", 4L);
+ /// {}
+ expected.put("NULL|NULL|1|1", 8L);
+ assertEquals(actual, expected);
+ assertEquals(rows.size(), expected.size());
+ }
+
+ @Test
+ public void testGroupingSets()
+ throws Exception {
+ String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*) FROM " + getTableName()
+ + " GROUP BY GROUPING SETS ((" + D1 + "), (" + D2 + "))";
+ JsonNode rows = runRows(query, true);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0) + "|" + cell(row, 1), row.get(2).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ /// {d1}
+ expected.put("a|NULL", 4L);
+ expected.put("b|NULL", 4L);
+ /// {d2}
+ expected.put("NULL|x", 4L);
+ expected.put("NULL|NULL", 4L);
+ assertEquals(actual, expected);
+ assertEquals(rows.size(), 4);
+ }
+
+ @Test
+ public void testGroupingIdMultiArg()
+ throws Exception {
+ /// 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);
+ /// 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<Integer, Integer> rowsPerGroupingId = new HashMap<>();
+ Map<Integer, Long> docsPerGroupingId = new HashMap<>();
+ for (JsonNode row : rows) {
+ int groupingId = row.get(2).asInt();
+ rowsPerGroupingId.merge(groupingId, 1, Integer::sum);
+ docsPerGroupingId.merge(groupingId, row.get(3).asLong(), Long::sum);
+ }
+ assertEquals(rowsPerGroupingId, Map.of(0, 4, 1, 2, 3, 1));
+ assertEquals(docsPerGroupingId, Map.of(0, 8L, 1, 8L, 3, 8L));
+ assertEquals(rows.size(), 7);
+ }
+
+ @Test
+ public void testPlainGroupByRegression()
+ throws Exception {
+ /// 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;
+ JsonNode response = postQuery("SET enableNullHandling=true; " + query);
+ JsonNode rows = response.get("resultTable").get("rows");
+ JsonNode columnNames = response.get("resultTable").get("dataSchema").get("columnNames");
+ /// No $groupingId column should leak into the result schema.
+ assertEquals(columnNames.size(), 3);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0) + "|" + cell(row, 1), row.get(2).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("a|x", 2L);
+ expected.put("a|NULL", 2L);
+ expected.put("b|x", 2L);
+ expected.put("b|NULL", 2L);
+ assertEquals(actual, expected);
+ }
+
+ @Test
+ public void testHavingOnGrouping()
+ throws Exception {
+ /// 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);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0), row.get(1).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("a", 4L);
+ expected.put("b", 4L);
+ expected.put("NULL", 8L);
+ assertEquals(actual, expected);
+ }
+
+ @Test
+ public void testLongGroupingColumn()
+ throws Exception {
+ /// 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);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0), row.get(1).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("100", 4L);
+ expected.put("200", 4L);
+ expected.put("NULL", 8L);
+ assertEquals(actual, expected);
+ }
+
+ @Test
+ public void testDoubleGroupingColumn()
+ throws Exception {
+ /// 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);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0), row.get(1).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("1.5", 4L);
+ expected.put("2.5", 4L);
+ expected.put("NULL", 8L);
+ assertEquals(actual, expected);
+ }
+
+ @Test
+ public void testOrderByGroupingColumnWithoutNullHandling()
+ 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);
+ JsonNode response = postQuery(query);
+ assertEquals(response.get("exceptions").size(), 0, response.toPrettyString());
+ JsonNode rows = response.get("resultTable").get("rows");
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0), row.get(1).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("a", 4L);
+ expected.put("b", 4L);
+ expected.put("NULL", 8L);
+ assertEquals(actual, expected);
+ }
+
+ @Test
+ public void testGroupingOnNonGroupingColumnRejected()
+ 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);
+ 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()
+ throws Exception {
+ /// 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);
+ 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);
+ 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");
+ assertEquals(cell(rows.get(1), 0), "a");
+ assertEquals(cell(rows.get(2), 0), "b");
+ }
+
+ @Test
+ public void testAggregationFreeGroupingSetsRejected()
+ 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);
+ for (String query : new String[]{
+ "SELECT " + D1 + " FROM " + getTableName() + " GROUP BY ROLLUP(" + D1 + ")",
+ "SELECT " + D1 + ", " + D2 + " FROM " + getTableName() + " GROUP BY GROUPING SETS ((" + D1 + "), (" + D2
+ + "))"
+ }) {
+ JsonNode response = postQuery(query);
+ /// The broker surfaces either the specific compile error ("requires at least one aggregation function",
+ /// pinned in GroupingSetsParserTest) or, when the query compiles on the multi-stage engine, the generic
+ /// "retry using the multi-stage query engine" hint. Either way it must be an error, never rows.
+ assertTrue(response.get("exceptions").size() > 0,
+ "aggregation-free grouping-set query must be rejected: " + response.toPrettyString());
+ assertEquals(response.get("numRowsResultSet").asInt(), 0, response.toPrettyString());
+ }
+ }
+
+ @Test
+ public void testMultiValueColumnInGroupingSet()
+ 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.
+ String query = "SELECT " + MV + ", COUNT(*) FROM " + getTableName() + " GROUP BY ROLLUP(" + MV + ")";
+ JsonNode rows = runRows(query, true);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0), row.get(1).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("t1", 8L);
+ expected.put("t2", 8L);
+ expected.put("NULL", 8L);
+ assertEquals(actual, expected);
+ assertEquals(rows.size(), 3);
+ }
+
+ @Test
+ public void testFilteredAggregation()
+ throws Exception {
+ /// 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);
+ Map<String, long[]> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0), new long[]{row.get(1).asLong(), row.get(2).asLong()});
+ }
+ assertEquals(rows.size(), 3);
+ /// d1=a: 2 of 4 rows have d2='x'; d1=b likewise; grand total: 4 of 8.
+ assertEquals(actual.get("a"), new long[]{2L, 4L});
+ assertEquals(actual.get("b"), new long[]{2L, 4L});
+ assertEquals(actual.get("NULL"), new long[]{4L, 8L});
+ }
+
+ @Test
+ public void testOrderByAggregation()
+ throws Exception {
+ /// 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);
+ /// 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;
+ for (JsonNode row : rows) {
+ long count = row.get(2).asLong();
+ assertTrue(count <= previousCount, "results must be ordered by COUNT(*) DESC");
+ previousCount = count;
+ }
+ }
+
+ @Test
+ public void testOrderByGroupingFunction()
+ throws Exception {
+ /// 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);
+ 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;
+ for (JsonNode row : rows) {
+ int grouping = row.get(3).asInt();
+ assertTrue(grouping >= previous, "results must be ordered by GROUPING(d2) ascending");
+ previous = grouping;
+ }
+ assertEquals(rows.get(0).get(3).asInt(), 0);
+ assertEquals(rows.get(rows.size() - 1).get(3).asInt(), 1);
+
+ /// GROUPING_ID(d1, d2) in ORDER BY must also resolve. ROLLUP yields ids 0 ({d1,d2}), 1 ({d1}), 3 ({}).
+ 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);
+ 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;
+ for (JsonNode row : idRows) {
+ int groupingId = row.get(3).asInt();
+ assertTrue(groupingId <= previousId, "results must be ordered by GROUPING_ID(d1, d2) descending");
+ previousId = groupingId;
+ }
+ }
+
+ @Test
+ public void testEmptyMatchRollup()
+ throws Exception {
+ /// 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.
+ String query = "SELECT " + D1 + ", " + D2 + ", COUNT(*) FROM " + getTableName() + " WHERE " + D1
+ + " = 'no_such_value' GROUP BY ROLLUP(" + D1 + ", " + D2 + ")";
+ JsonNode rows = runRows(query, true);
+ assertEquals(rows.size(), 0);
+ }
+
+ @Test
+ public void testMixedPlainAndRollup()
+ throws Exception {
+ /// 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);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt(), row.get(2).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("a|x|0", 2L);
+ expected.put("a|NULL|0", 2L);
+ expected.put("a|NULL|1", 4L);
+ expected.put("b|x|0", 2L);
+ expected.put("b|NULL|0", 2L);
+ expected.put("b|NULL|1", 4L);
+ assertEquals(actual, expected);
+ assertEquals(rows.size(), 6);
+ }
+
+ @Test
+ public void testCompositeRollupLevel()
+ throws Exception {
+ /// 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);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt(), row.get(2).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("a|x|0", 2L);
+ expected.put("a|NULL|0", 2L);
+ expected.put("b|x|0", 2L);
+ expected.put("b|NULL|0", 2L);
+ expected.put("NULL|NULL|3", 8L);
+ assertEquals(actual, expected);
+ assertEquals(rows.size(), 5);
+ }
+
+ @Test
+ public void testNestedRollupInsideGroupingSets()
+ throws Exception {
+ /// 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);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt(), row.get(2).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("a|NULL|1", 4L);
+ expected.put("b|NULL|1", 4L);
+ expected.put("NULL|x|2", 4L);
+ expected.put("NULL|NULL|2", 4L);
+ expected.put("NULL|NULL|3", 8L);
+ assertEquals(actual, expected);
+ assertEquals(rows.size(), 5);
+ }
+
+ @Test
+ public void testExpressionGroupingColumn()
+ throws Exception {
+ /// 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);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0), row.get(1).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("A", 4L);
+ expected.put("B", 4L);
+ expected.put("NULL", 8L);
+ assertEquals(actual, expected);
+ assertEquals(rows.size(), 3);
+ }
+
+ @Test
+ public void testWhereFilterWithCube()
+ throws Exception {
+ /// 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);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0) + "|" + cell(row, 1) + "|" + row.get(3).asInt(), row.get(2).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("a|x|0", 2L);
+ expected.put("b|x|0", 2L);
+ expected.put("a|NULL|1", 2L);
+ expected.put("b|NULL|1", 2L);
+ expected.put("NULL|x|2", 4L);
+ expected.put("NULL|NULL|3", 4L);
+ assertEquals(actual, expected);
+ assertEquals(rows.size(), 6);
+ }
+
+ @Test
+ public void testCaseWhenGroupingRelabelsSubtotals()
+ throws Exception {
+ /// 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);
+ Map<String, Long> actual = new HashMap<>();
+ for (JsonNode row : rows) {
+ actual.put(cell(row, 0), row.get(1).asLong());
+ }
+ Map<String, Long> expected = new HashMap<>();
+ expected.put("a", 4L);
+ expected.put("b", 4L);
+ expected.put("ALL", 8L);
+ assertEquals(actual, expected);
+ assertEquals(rows.size(), 3);
+ }
+
+ @Test
+ public void testMultipleAggregationsWithOrderByAndLimit()
+ throws Exception {
+ /// 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);
+ assertEquals(rows.size(), 2);
+ assertEquals(cell(rows.get(0), 0), "NULL");
+ assertEquals(rows.get(0).get(1).asDouble(), 8.0);
+ assertEquals(rows.get(0).get(2).asDouble(), 200.0);
+ assertEquals(rows.get(0).get(3).asLong(), 8L);
+ assertEquals(cell(rows.get(1), 0), "a");
+ assertEquals(rows.get(1).get(1).asDouble(), 4.0);
+ assertEquals(rows.get(1).get(2).asDouble(), 100.0);
+ assertEquals(rows.get(1).get(3).asLong(), 4L);
+ }
+
+ @Test
+ public void testDistinctCountUnderRollup()
+ throws Exception {
+ /// 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);
+ Map<String, long[]> 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()});
+ }
+ assertEquals(rows.size(), 3);
+ assertEquals(actual.get("x|0"), new long[]{2L, 4L});
+ assertEquals(actual.get("NULL|0"), new long[]{2L, 4L});
+ assertEquals(actual.get("NULL|1"), new long[]{2L, 8L});
+ }
+
+ @Test
+ public void testHavingOnAggregationAndGrouping()
+ throws Exception {
+ /// 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);
+ assertEquals(rows.size(), 1);
+ assertEquals(cell(rows.get(0), 0), "NULL");
+ assertEquals(cell(rows.get(0), 1), "NULL");
+ assertEquals(rows.get(0).get(2).asLong(), 8L);
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org