diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala index fe30e2ea6f3ff..f175f4c5337e8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala +++ b/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,6 +69,8 @@ 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 } @@ -76,6 +78,8 @@ case class Average( 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) ) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/intervalExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/intervalExpressions.scala index 3e4d6772c4fc4..fdba3f2ea25a7 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/intervalExpressions.scala +++ b/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) +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/TypeUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/TypeUtils.scala index c6792af9f0ea9..bc1517d658b01 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/TypeUtils.scala +++ b/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]) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala index 4a7adcc050a00..a3094dd47dc3d 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala +++ b/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) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala index b7f360424e14f..e4a4085c14735 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala +++ b/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") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index 36c2a95401090..7806ff729aa56 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/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"))