[SPARK-46536][SQL] Support GROUP BY calendar_interval_type

common/unsafe/src/main/java/org/apache/spark/unsafe/types/CalendarInterval.java

@@ -44,7 +44,7 @@
  * @since 3.0.0
  */
 @Unstable
-public final class CalendarInterval implements Serializable {
+public final class CalendarInterval implements Serializable, Comparable<CalendarInterval> {
   // NOTE: If you're moving or renaming this file, you should also update Unidoc configuration
   // specified in 'SparkBuild.scala'.
   public final int months;
@@ -127,4 +127,27 @@ private void appendUnit(StringBuilder sb, long value, String unit) {
    * @throws ArithmeticException if a numeric overflow occurs
    */
   public Duration extractAsDuration() { return Duration.of(microseconds, ChronoUnit.MICROS); }
+
+  /**
+   * This method is not used to order CalendarInterval instances, as they are not orderable and
+   * can't be used in a ORDER BY statement.
+   * Instead, it is used to find identical interval instances for aggregation purposes.
+   * It compares the 'months', 'days', and 'microseconds' fields of this CalendarInterval
+   * with another instance. The comparison is done first on the 'months', then on the 'days',
+   * and finally on the 'microseconds'.
+   *
+   * @param o The CalendarInterval instance to compare with.
+   * @return A negative integer, zero, or a positive integer as this object is less than,
+   *         equal to, or greater than the specified object.
+   */
+  @Override
+  public int compareTo(CalendarInterval o) {
+    if (this.months != o.months) {
+      return Integer.compare(this.months, o.months);
+    } else if (this.days != o.days) {
+      return Integer.compare(this.days, o.days);
+    } else {
+      return Long.compare(this.microseconds, o.microseconds);
+    }
+  }
 }

common/unsafe/src/test/java/org/apache/spark/unsafe/types/CalendarIntervalSuite.java

@@ -76,6 +76,22 @@ public void toStringTest() {
       i.toString());
   }
 
+  @Test
+  public void compareToTest() {
+   CalendarInterval i = new CalendarInterval(0, 0, 0);
+
+   assertEquals(i.compareTo(new CalendarInterval(0, 0, 0)), 0);
+   assertEquals(i.compareTo(new CalendarInterval(0, 0, 1)), -1);
+   assertEquals(i.compareTo(new CalendarInterval(0, 1, 0)), -1);
+   assertEquals(i.compareTo(new CalendarInterval(0, 1, -1)), -1);
+   assertEquals(i.compareTo(new CalendarInterval(1, 0, 0)), -1);
+   assertEquals(i.compareTo(new CalendarInterval(1, 0, -1)), -1);
+   assertEquals(i.compareTo(new CalendarInterval(0, 0, -1)), 1);
+   assertEquals(i.compareTo(new CalendarInterval(0, -1, 0)), 1);
+   assertEquals(i.compareTo(new CalendarInterval(-1, 0, 0)), 1);
+   assertEquals(i.compareTo(new CalendarInterval(-1, 0, 1)), 1);
+  }
+
   @Test
   public void periodAndDurationTest() {
     CalendarInterval interval = new CalendarInterval(120, -40, 123456);

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExprUtils.scala

@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.plans.logical.Aggregate
 import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, CharVarcharUtils}
 import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase, QueryExecutionErrors}
-import org.apache.spark.sql.types.{DataType, MapType, StringType, StructType}
+import org.apache.spark.sql.types.{CalendarIntervalType, DataType, MapType, StringType, StructType}
 import org.apache.spark.unsafe.types.UTF8String
 
 object ExprUtils extends QueryErrorsBase {
@@ -193,8 +193,8 @@ object ExprUtils extends QueryErrorsBase {
           messageParameters = Map("sqlExpr" -> expr.sql))
       }
 
-      // Check if the data type of expr is orderable.
-      if (!RowOrdering.isOrderable(expr.dataType)) {
+      // Check if the data type of expr can be used in group by
+      if (!canBeUsedInGroupBy(expr.dataType)) {
         expr.failAnalysis(
           errorClass = "GROUP_EXPRESSION_TYPE_IS_NOT_ORDERABLE",
           messageParameters = Map(
@@ -217,4 +217,12 @@ object ExprUtils extends QueryErrorsBase {
     a.groupingExpressions.foreach(checkValidGroupingExprs)
     a.aggregateExpressions.foreach(checkValidAggregateExpression)
   }
+
+  /**
+   * Returns whether the data type can be used in group by
+   */
+  def canBeUsedInGroupBy(dt: DataType): Boolean = dt match {
+    case CalendarIntervalType => true
+    case _ => RowOrdering.isOrderable(dt)
+  }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala

@@ -627,6 +627,7 @@ class CodegenContext extends Logging {
     case array: ArrayType => genComp(array, c1, c2) + " == 0"
     case struct: StructType => genComp(struct, c1, c2) + " == 0"
     case udt: UserDefinedType[_] => genEqual(udt.sqlType, c1, c2)
+    case CalendarIntervalType => s"$c1.equals($c2)"
     case NullType => "false"
     case _ =>
       throw QueryExecutionErrors.cannotGenerateCodeForIncomparableTypeError(
@@ -652,6 +653,7 @@ class CodegenContext extends Logging {
     // use c1 - c2 may overflow
     case dt: DataType if isPrimitiveType(dt) => s"($c1 > $c2 ? 1 : $c1 < $c2 ? -1 : 0)"
     case BinaryType => s"org.apache.spark.unsafe.types.ByteArray.compareBinary($c1, $c2)"
+    case CalendarIntervalType => s"$c1.compareTo($c2)"
     case NullType => "0"
     case array: ArrayType =>
       val elementType = array.elementType

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/TypeUtils.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.util
 
 import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, TypeCoercion}
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch
-import org.apache.spark.sql.catalyst.expressions.{Expression, RowOrdering}
+import org.apache.spark.sql.catalyst.expressions.{Expression, ExprUtils}
 import org.apache.spark.sql.catalyst.types.{PhysicalDataType, PhysicalNumericType}
 import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase}
 import org.apache.spark.sql.types._
@@ -30,7 +30,7 @@ import org.apache.spark.sql.types._
 object TypeUtils extends QueryErrorsBase {
 
   def checkForOrderingExpr(dt: DataType, caller: String): TypeCheckResult = {
-    if (RowOrdering.isOrderable(dt)) {
+    if (ExprUtils.canBeUsedInGroupBy(dt)) {
       TypeCheckResult.TypeCheckSuccess
     } else {
       DataTypeMismatch(

sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashMapGenerator.scala

@@ -174,6 +174,7 @@ abstract class HashMapGenerator(
           """
         }
       case StringType => hashBytes(s"$input.getBytes()")
+      case CalendarIntervalType => hashInt(s"$input.hashCode()")
     }
   }
 }

sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala

@@ -37,6 +37,7 @@ import org.apache.spark.sql.test.SQLTestData.DecimalData
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.DayTimeIntervalType.{DAY, HOUR, MINUTE, SECOND}
 import org.apache.spark.sql.types.YearMonthIntervalType.{MONTH, YEAR}
+import org.apache.spark.unsafe.types.CalendarInterval
 
 case class Fact(date: Int, hour: Int, minute: Int, room_name: String, temp: Double)
 
@@ -2125,6 +2126,32 @@ class DataFrameAggregateSuite extends QueryTest
       Seq(Row(1))
     )
   }
+
+  test("SPARK-46536 Support GROUP BY CalendarIntervalType") {
+    // forces the use of sort aggregate by using min/max functions
+
+    val numRows = 50
+    val numRepeat = 25
+
+    val df = (0 to numRows)
+      .map(i => Tuple1(new CalendarInterval(i, i, i)))
+      .toDF("c0")
+
+    for (_ <- 0 until numRepeat) {
+      val shuffledDf = df.orderBy(rand())
+
+      checkAnswer(
+        shuffledDf.agg(max("c0")),
+        Row(new CalendarInterval(numRows, numRows, numRows))
+      )
+
+      checkAnswer(
+        shuffledDf.agg(min("c0")),
+        Row(new CalendarInterval(0, 0, 0))
+      )
+    }
+  }
+
 }
 
 case class B(c: Option[Double])

sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/ObjectHashAggregateSuite.scala

@@ -33,6 +33,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SQLTestUtils
 import org.apache.spark.sql.types._
 import org.apache.spark.tags.SlowHiveTest
+import org.apache.spark.unsafe.types.CalendarInterval
 
 @SlowHiveTest
 class ObjectHashAggregateSuite
@@ -457,4 +458,38 @@ class ObjectHashAggregateSuite
       )
     }
   }
+
+  test("SPARK-46536 Support GROUP BY CalendarIntervalType") {
+    withSQLConf(
+      SQLConf.USE_OBJECT_HASH_AGG.key -> "true",
+      SQLConf.OBJECT_AGG_SORT_BASED_FALLBACK_THRESHOLD.key -> "1"
+    ) {
+
+      val numRows = 50
+
+      val dfSame = createAggregate(_ => Tuple1(new CalendarInterval(1, 2, 3)), numRows)
+      assert(getPlanOutput(dfSame).contains("HashAggregate"))
+      assert(dfSame.count() == 1)
+
+      val dfDifferent = createAggregate(i => Tuple1(new CalendarInterval(i, i, i)), numRows)
+      assert(getPlanOutput(dfDifferent).contains("HashAggregate"))
+      assert(dfDifferent.count() == numRows)
+    }
+
+    def createAggregate(f: Int => Tuple1[CalendarInterval], numRows: Int): DataFrame = {
+      (1 to numRows)
+        .map(i => f(i))
+        .toDF("c0")
+        .groupBy("c0")
+        .agg(count("*"))
+    }
+
+    def getPlanOutput(df: DataFrame): String = {
+      val outputStream = new java.io.ByteArrayOutputStream()
+      Console.withOut(outputStream) {
+        df.explain(true)
+      }
+      outputStream.toString
+    }
+  }
 }