[SPARK-58044][SQL] Support the TIME data type in the aggregate function avg

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala

@@ -52,10 +52,10 @@ case class Average(
   override def prettyName: String = getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("avg")
 
   override def inputTypes: Seq[AbstractDataType] =
-    Seq(TypeCollection(NumericType, YearMonthIntervalType, DayTimeIntervalType))
+    Seq(TypeCollection(NumericType, YearMonthIntervalType, DayTimeIntervalType, AnyTimeType))
 
   override def checkInputDataTypes(): TypeCheckResult =
-    TypeUtils.checkForAnsiIntervalOrNumericType(child)
+    TypeUtils.checkForAnsiIntervalOrNumericOrTimeType(child)
 
   override def nullable: Boolean = true
 
@@ -69,13 +69,17 @@ case class Average(
       DecimalType.bounded(p + 4, s + 4)
     case _: YearMonthIntervalType => YearMonthIntervalType()
     case _: DayTimeIntervalType => DayTimeIntervalType()
+    // The average of a TIME column is a TIME value with the same precision.
+    case t: TimeType => t
     case _ => DoubleType
   }
 
   lazy val sumDataType = child.dataType match {
     case _ @ DecimalType.Fixed(p, s) => DecimalType.bounded(p + 10, s)
     case _: YearMonthIntervalType => YearMonthIntervalType()
     case _: DayTimeIntervalType => DayTimeIntervalType()
+    // TIME is stored as a `Long` number of nanoseconds, so accumulate the sum as `LongType`.
+    case _: TimeType => LongType
     case _ => DoubleType
   }
 
@@ -117,15 +121,26 @@ case class Average(
     case _: DayTimeIntervalType =>
       If(EqualTo(count, Literal(0L)),
         Literal(null, DayTimeIntervalType()), DivideDTInterval(sum, count))
+    case t: TimeType =>
+      If(EqualTo(count, Literal(0L)),
+        Literal(null, t), TimeDivide(sum, count, t))
     case _ =>
       Divide(sum.cast(resultType), count.cast(resultType), EvalMode.LEGACY)
   }
 
+  // Convert the child value to the sum's data type before accumulating. For TIME, a regular cast
+  // to `LongType` is lossy (it truncates to whole seconds), so reinterpret the underlying `Long`
+  // nanoseconds without changing the value instead.
+  private def childToSum: Expression = child.dataType match {
+    case _: TimeType => PreciseTimestampConversion(child, child.dataType, LongType)
+    case _ => child.cast(sumDataType)
+  }
+
   override lazy val updateExpressions: Seq[Expression] = Seq(
     /* sum = */
     add(
       sum,
-      coalesce(child.cast(sumDataType), Literal.default(sumDataType))),
+      coalesce(childToSum, Literal.default(sumDataType))),
     /* count = */ If(child.isNull, count, count + 1L)
   )
 

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

@@ -902,3 +902,39 @@ case class DivideDTInterval(
       newLeft: Expression, newRight: Expression): DivideDTInterval =
     copy(interval = newLeft, num = newRight)
 }
+
+// Divide a TIME value (stored as a `Long` number of nanoseconds since midnight) by an integral
+// number, rounding to the nearest nanosecond. Used by `Average` to compute the mean of a TIME
+// column. The result preserves the input `TimeType`'s precision. `num` is expected to be a
+// positive count, so no divide-by-zero or `Long.MinValue` overflow checks are needed (the caller
+// guards the zero-count case).
+case class TimeDivide(
+    time: Expression,
+    num: Expression,
+    targetType: TimeType)
+  extends BinaryExpression with ImplicitCastInputTypes with Serializable {
+  override def nullIntolerant: Boolean = true
+
+  override def left: Expression = time
+  override def right: Expression = num
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(LongType, LongType)
+  override def dataType: DataType = targetType
+
+  override def nullSafeEval(nanos: Any, num: Any): Any = {
+    LongMath.divide(nanos.asInstanceOf[Long], num.asInstanceOf[Long], RoundingMode.HALF_UP)
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val math = classOf[LongMath].getName
+    nullSafeCodeGen(ctx, ev, (t, n) =>
+      s"${ev.value} = $math.divide($t, $n, java.math.RoundingMode.HALF_UP);")
+  }
+
+  override def toString: String = s"($left / $right)"
+  override def sql: String = s"(${left.sql} / ${right.sql})"
+
+  override protected def withNewChildrenInternal(
+      newLeft: Expression, newRight: Expression): TimeDivide =
+    copy(time = newLeft, num = newRight)
+}

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

@@ -96,6 +96,22 @@ object TypeUtils extends QueryErrorsBase {
           "inputType" -> toSQLType(other)))
   }
 
+  def checkForAnsiIntervalOrNumericOrTimeType(input: Expression): TypeCheckResult =
+    input.dataType match {
+      case _: AnsiIntervalType | _: TimeType | NullType =>
+        TypeCheckResult.TypeCheckSuccess
+      case dt if dt.isInstanceOf[NumericType] => TypeCheckResult.TypeCheckSuccess
+      case other =>
+        DataTypeMismatch(
+          errorSubClass = "UNEXPECTED_INPUT_TYPE",
+          messageParameters = Map(
+            "paramIndex" -> ordinalNumber(0),
+            "requiredType" ->
+              Seq(NumericType, AnsiIntervalType, AnyTimeType).map(toSQLType).mkString(" or "),
+            "inputSql" -> toSQLExpr(input),
+            "inputType" -> toSQLType(other)))
+    }
+
   def getNumeric(t: DataType, exactNumericRequired: Boolean = false): Numeric[Any] = {
     if (exactNumericRequired) {
       PhysicalNumericType.exactNumeric(t.asInstanceOf[NumericType])

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala

@@ -1406,7 +1406,7 @@ class AnalysisSuite extends AnalysisTest with Matchers {
         "paramIndex" -> "first",
         "inputSql" -> "\"c\"",
         "inputType" -> "\"BOOLEAN\"",
-        "requiredType" -> "\"NUMERIC\" or \"ANSI INTERVAL\""),
+        "requiredType" -> "\"NUMERIC\" or \"ANSI INTERVAL\" or \"TIME\""),
       queryContext = Array(ExpectedContext("mean(t.c)", 65, 73)),
       caseSensitive = false
     )
@@ -1425,7 +1425,7 @@ class AnalysisSuite extends AnalysisTest with Matchers {
         "paramIndex" -> "first",
         "inputSql" -> "\"c\"",
         "inputType" -> "\"BOOLEAN\"",
-        "requiredType" -> "\"NUMERIC\" or \"ANSI INTERVAL\""),
+        "requiredType" -> "\"NUMERIC\" or \"ANSI INTERVAL\" or \"TIME\""),
       queryContext = Array(ExpectedContext("mean(c)", 91, 97)),
       caseSensitive = false)
 

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala

@@ -471,7 +471,7 @@ class ExpressionTypeCheckingSuite extends SparkFunSuite with SQLHelper with Quer
         "paramIndex" -> ordinalNumber(0),
         "inputSql" -> "\"booleanField\"",
         "inputType" -> "\"BOOLEAN\"",
-        "requiredType" -> "\"NUMERIC\" or \"ANSI INTERVAL\""))
+        "requiredType" -> "\"NUMERIC\" or \"ANSI INTERVAL\" or \"TIME\""))
   }
 
   test("check types for others") {

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

@@ -3798,6 +3798,46 @@ class DataFrameAggregateSuite extends SharedSparkSession
       Row(LocalTime.of(22, 1, 0), LocalTime.of(3, 0, 0)))
   }
 
+  test("SPARK-58044: Support the TIME data type in the aggregate function `avg`") {
+    // Basic average of TIME values.
+    val df = Seq("10:00:00", "12:00:00", "14:00:00").map(LocalTime.parse).toDF("t")
+    val avgDF = df.select(avg($"t"))
+    checkAnswer(avgDF, Row(LocalTime.of(12, 0, 0)))
+    assert(find(avgDF.queryExecution.executedPlan)(_.isInstanceOf[HashAggregateExec]).isDefined)
+    // The average of a TIME column is a TIME value.
+    assert(avgDF.schema == StructType(Seq(StructField("avg(t)", TimeType()))))
+
+    // Average that requires rounding to the nearest microsecond/nanosecond.
+    val df2 = Seq("10:00:00", "10:00:01").map(LocalTime.parse).toDF("t")
+    checkAnswer(df2.select(avg($"t")), Row(LocalTime.of(10, 0, 0, 500000000)))
+
+    // NULL values are ignored; all-null input yields NULL.
+    val df3 = Seq(Some("10:00:00"), None, Some("14:00:00"))
+      .map(_.map(LocalTime.parse)).toDF("t")
+    checkAnswer(df3.select(avg($"t")), Row(LocalTime.of(12, 0, 0)))
+    val emptyDF = Seq.empty[LocalTime].toDF("t")
+    checkAnswer(emptyDF.select(avg($"t")), Row(null))
+
+    // The output preserves the input precision.
+    val df4 = Seq("10:00:00", "12:00:00").map(LocalTime.parse).toDF("t")
+      .selectExpr("CAST(t AS TIME(3)) AS t")
+    assert(df4.select(avg($"t")).schema ==
+      StructType(Seq(StructField("avg(t)", TimeType(3)))))
+
+    // Grouped average.
+    val df5 = Seq(
+      ("a", "10:00:00"), ("a", "12:00:00"), ("b", "20:00:00"), ("b", "22:00:00"))
+      .map { case (c, t) => (c, LocalTime.parse(t)) }.toDF("class", "t")
+    checkAnswer(
+      df5.groupBy($"class").agg(avg($"t")).orderBy($"class"),
+      Seq(Row("a", LocalTime.of(11, 0, 0)), Row("b", LocalTime.of(21, 0, 0))))
+
+    // The SQL `avg` surface works over TIME literals as well.
+    checkAnswer(
+      sql("SELECT avg(t) FROM VALUES (TIME'10:00:00'), (TIME'14:00:00') AS tab(t)"),
+      Row(LocalTime.of(12, 0, 0)))
+  }
+
   test("SPARK-53155: global lower aggregation should not be removed") {
     val df = emptyTestData
       .groupBy().agg(lit(1).as("col1"), lit(2).as("col2"), lit(3).as("col3"))