DRILL-3962: Add Support For ROLLUP, CUBE, GROUPING SETS, GROUPING, GROUPING_ID, GROUP_ID (#3026)

Author: cgivre

exec/java-exec/src/main/codegen/data/DateIntervalFunc.tdd

@@ -22,6 +22,8 @@
     {truncInputTypes: ["Date", "TimeStamp", "Time", "Interval", "IntervalDay", "IntervalYear"] },
     {truncUnits : ["Second", "Minute", "Hour", "Day", "Month", "Year", "Week", "Quarter", "Decade", "Century", "Millennium" ] },
     {timestampDiffUnits : ["Nanosecond", "Microsecond", "Millisecond", "Second", "Minute", "Hour", "Day", "Month", "Year", "Week", "Quarter"] },
+    {timestampAddUnits : ["Nanosecond", "Microsecond", "Millisecond", "Second", "Minute", "Hour", "Day", "Month", "Year", "Week", "Quarter"] },
+    {timestampAddInputTypes : ["Date", "TimeStamp", "Time"] },
 
     {
         varCharToDate: [

exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/impl/GroupingFunctions.java

@@ -0,0 +1,79 @@
+/*
+ * 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.drill.exec.expr.fn.impl;
+
+import org.apache.drill.exec.expr.DrillSimpleFunc;
+import org.apache.drill.exec.expr.annotations.FunctionTemplate;
+import org.apache.drill.exec.expr.annotations.Output;
+import org.apache.drill.exec.expr.annotations.Param;
+import org.apache.drill.exec.expr.holders.IntHolder;
+
+/**
+ * Functions for working with GROUPING SETS, ROLLUP, and CUBE.
+ *
+ * Note: These are internal helper functions. The actual GROUPING() and GROUPING_ID()
+ * SQL functions need special query rewriting to work correctly with GROUPING SETS.
+ */
+public class GroupingFunctions {
+
+  /**
+   * GROUPING_ID_INTERNAL - Returns the grouping ID bitmap.
+   * This is an internal function that will be called with the $g column value.
+   */
+  @FunctionTemplate(name = "grouping_id_internal",
+      scope = FunctionTemplate.FunctionScope.SIMPLE,
+      nulls = FunctionTemplate.NullHandling.NULL_IF_NULL)
+  public static class GroupingIdInternal implements DrillSimpleFunc {
+
+    @Param IntHolder groupingId;
+    @Output IntHolder out;
+
+    public void setup() {
+    }
+
+    public void eval() {
+      out.value = groupingId.value;
+    }
+  }
+
+  /**
+   * GROUPING_INTERNAL - Returns 1 if the specified bit in the grouping ID is set, 0 otherwise.
+   * This is an internal function that extracts a specific bit from the grouping ID.
+   *
+   * @param groupingId The grouping ID bitmap ($g column value)
+   * @param bitPosition The bit position to check (0-based)
+   */
+  @FunctionTemplate(name = "grouping_internal",
+      scope = FunctionTemplate.FunctionScope.SIMPLE,
+      nulls = FunctionTemplate.NullHandling.NULL_IF_NULL)
+  public static class GroupingInternal implements DrillSimpleFunc {
+
+    @Param IntHolder groupingId;
+    @Param IntHolder bitPosition;
+    @Output IntHolder out;
+
+    public void setup() {
+    }
+
+    public void eval() {
+      // Extract the bit at bitPosition from groupingId
+      // Bit is 1 if column is NOT in the grouping set (i.e., it's a grouping NULL)
+      out.value = (groupingId.value >> bitPosition.value) & 1;
+    }
+  }
+}

exec/java-exec/src/main/java/org/apache/drill/exec/physical/config/UnionAll.java

@@ -33,9 +33,22 @@ public class UnionAll extends AbstractMultiple {
 
   public static final String OPERATOR_TYPE = "UNION";
 
+  private final boolean isGroupingSetsExpansion;
+
   @JsonCreator
-  public UnionAll(@JsonProperty("children") List<PhysicalOperator> children) {
+  public UnionAll(@JsonProperty("children") List<PhysicalOperator> children,
+                  @JsonProperty("isGroupingSetsExpansion") Boolean isGroupingSetsExpansion) {
     super(children);
+    this.isGroupingSetsExpansion = isGroupingSetsExpansion != null ? isGroupingSetsExpansion : false;
+  }
+
+  public UnionAll(List<PhysicalOperator> children) {
+    this(children, false);
+  }
+
+  @JsonProperty("isGroupingSetsExpansion")
+  public boolean isGroupingSetsExpansion() {
+    return isGroupingSetsExpansion;
   }
 
   @Override
@@ -45,7 +58,7 @@ public <T, X, E extends Throwable> T accept(PhysicalVisitor<T, X, E> physicalVis
 
   @Override
   public PhysicalOperator getNewWithChildren(List<PhysicalOperator> children) {
-    return new UnionAll(children);
+    return new UnionAll(children, isGroupingSetsExpansion);
   }
 
   @Override

exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/union/UnionAllRecordBatch.java

@@ -17,13 +17,7 @@
  */
 package org.apache.drill.exec.physical.impl.union;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.Iterator;
-import java.util.List;
-import java.util.NoSuchElementException;
-import java.util.Stack;
-
+import com.google.common.base.Preconditions;
 import org.apache.calcite.util.Pair;
 import org.apache.drill.common.exceptions.DrillRuntimeException;
 import org.apache.drill.common.expression.ErrorCollector;
@@ -59,10 +53,16 @@
 import org.apache.drill.exec.vector.FixedWidthVector;
 import org.apache.drill.exec.vector.SchemaChangeCallBack;
 import org.apache.drill.exec.vector.ValueVector;
-import com.google.common.base.Preconditions;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+import java.util.Stack;
+
 public class UnionAllRecordBatch extends AbstractBinaryRecordBatch<UnionAll> {
   private static final Logger logger = LoggerFactory.getLogger(UnionAllRecordBatch.class);
 
@@ -278,10 +278,14 @@ private void inferOutputFieldsBothSide(final BatchSchema leftSchema, final Batch
     final Iterator<MaterializedField> leftIter = leftSchema.iterator();
     final Iterator<MaterializedField> rightIter = rightSchema.iterator();
 
+    logger.debug("UnionAll inferring schema: isGroupingSetsExpansion={}", popConfig.isGroupingSetsExpansion());
     int index = 1;
     while (leftIter.hasNext() && rightIter.hasNext()) {
       MaterializedField leftField  = leftIter.next();
       MaterializedField rightField = rightIter.next();
+      logger.debug("Column {}: left='{}' type={}, right='{}' type={}",
+          index, leftField.getName(), leftField.getType().getMinorType(),
+          rightField.getName(), rightField.getType().getMinorType());
 
       if (Types.isSameTypeAndMode(leftField.getType(), rightField.getType())) {
         MajorType.Builder builder = MajorType.newBuilder()
@@ -301,15 +305,7 @@ private void inferOutputFieldsBothSide(final BatchSchema leftSchema, final Batch
           builder.setMinorType(leftField.getType().getMinorType());
           builder = Types.calculateTypePrecisionAndScale(leftField.getType(), rightField.getType(), builder);
         } else {
-          TypeProtos.MinorType outputMinorType = TypeCastRules.getLeastRestrictiveType(
-            leftField.getType().getMinorType(),
-            rightField.getType().getMinorType()
-          );
-          if (outputMinorType == null) {
-            throw new DrillRuntimeException("Type mismatch between " + leftField.getType().getMinorType().toString() +
-                " on the left side and " + rightField.getType().getMinorType().toString() +
-                " on the right side in column " + index + " of UNION ALL");
-          }
+          TypeProtos.MinorType outputMinorType = resolveUnionColumnType(leftField, rightField, index);
           builder.setMinorType(outputMinorType);
         }
 
@@ -328,6 +324,46 @@ private void inferOutputFieldsBothSide(final BatchSchema leftSchema, final Batch
       "Mismatch of column count should have been detected when validating sqlNode at planning";
   }
 
+  /**
+   * Determines the output type for a UNION ALL column when combining two types.
+   * <p>
+   * Special handling is applied for GROUPING SETS expansion:
+   * - Drill represents NULL columns as INT during grouping sets expansion.
+   * - If one side is INT (likely a NULL placeholder) and the other is not, prefer the non-INT type.
+   * <p>
+   * For all other cases, the least restrictive type according to Drill's type cast rules is returned.
+   *
+   * @param leftField  The type of the left column
+   * @param rightField The type of the right column
+   * @param index     The column index (for logging)
+   * @return The resolved output type
+   * @throws DrillRuntimeException if types are incompatible
+   */
+  private TypeProtos.MinorType resolveUnionColumnType(MaterializedField leftField,
+      MaterializedField rightField,
+      int index) {
+    TypeProtos.MinorType leftType = leftField.getType().getMinorType();
+    TypeProtos.MinorType rightType = rightField.getType().getMinorType();
+
+    boolean isGroupingSets = popConfig.isGroupingSetsExpansion();
+    boolean leftIsPlaceholder = leftType == TypeProtos.MinorType.INT && rightType != TypeProtos.MinorType.INT;
+    boolean rightIsPlaceholder = rightType == TypeProtos.MinorType.INT && leftType != TypeProtos.MinorType.INT;
+
+    if (isGroupingSets && (leftIsPlaceholder || rightIsPlaceholder)) {
+      TypeProtos.MinorType outputType = leftIsPlaceholder ? rightType : leftType;
+      logger.debug("GROUPING SETS: Preferring {} over INT for column {}", outputType, index);
+      return outputType;
+    }
+
+    TypeProtos.MinorType outputType = TypeCastRules.getLeastRestrictiveType(leftType, rightType);
+    if (outputType == null) {
+      throw new DrillRuntimeException("Type mismatch between " + leftType +
+          " and " + rightType + " in column " + index + " of UNION ALL");
+    }
+    logger.debug("Using standard type rules: {} + {} -> {}", leftType, rightType, outputType);
+    return outputType;
+  }
+
   private void inferOutputFieldsOneSide(final BatchSchema schema) {
     for (MaterializedField field : schema) {
       container.addOrGet(field, callBack);

exec/java-exec/src/main/java/org/apache/drill/exec/planner/PlannerPhase.java

@@ -356,6 +356,11 @@ static RuleSet getDrillUserConfigurableLogicalRules(OptimizerRulesContext optimi
        Convert from Calcite Logical to Drill Logical Rules.
        */
       RuleInstance.EXPAND_CONVERSION_RULE,
+
+      // Expand GROUPING SETS, ROLLUP, and CUBE BEFORE converting aggregates to Drill logical operators
+      // This prevents multi-grouping-set aggregates from being converted to DrillAggregateRel
+      RuleInstance.AGGREGATE_EXPAND_GROUPING_SETS_RULE,
+
       DrillScanRule.INSTANCE,
       DrillFilterRule.INSTANCE,
       DrillProjectRule.INSTANCE,

exec/java-exec/src/main/java/org/apache/drill/exec/planner/RuleInstance.java

@@ -39,6 +39,7 @@
 import org.apache.calcite.rel.rules.SortRemoveRule;
 import org.apache.calcite.rel.rules.SubQueryRemoveRule;
 import org.apache.calcite.rel.rules.UnionToDistinctRule;
+import org.apache.drill.exec.planner.logical.DrillAggregateExpandGroupingSetsRule;
 import org.apache.drill.exec.planner.logical.DrillConditions;
 import org.apache.drill.exec.planner.logical.DrillRelFactories;
 import com.google.common.base.Preconditions;
@@ -107,6 +108,9 @@ public boolean matches(RelOptRuleCall call) {
       .withRelBuilderFactory(DrillRelFactories.LOGICAL_BUILDER)
       .toRule();
 
+  RelOptRule AGGREGATE_EXPAND_GROUPING_SETS_RULE =
+    DrillAggregateExpandGroupingSetsRule.INSTANCE;
+
   /**
    * Instance of the rule that works on logical joins only, and pushes to the
    * right.