diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala index 988b1f73a6e66..0aebdc0223144 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala @@ -2090,11 +2090,11 @@ object WidthBucket { arguments = """ Arguments: * value - The value to assign to a bucket. - An expression that evaluates to a double or interval. + An expression that evaluates to a double, interval, or time. * min_value - The minimum value of the histogram range. - An expression that evaluates to a double or interval. + An expression that evaluates to a double, interval, or time. * max_value - The maximum value of the histogram range. - An expression that evaluates to a double or interval. + An expression that evaluates to a double, interval, or time. * num_bucket - The number of equiwidth buckets in the histogram. An expression that evaluates to a long. """, @@ -2116,6 +2116,8 @@ object WidthBucket { 1 > SELECT _FUNC_(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10); 2 + > SELECT _FUNC_(TIME'12:00:00', TIME'09:00:00', TIME'17:00:00', 8); + 4 """, since = "3.1.0", group = "math_funcs") @@ -2128,9 +2130,9 @@ case class WidthBucket( override def nullIntolerant: Boolean = true override def inputTypes: Seq[AbstractDataType] = Seq( - TypeCollection(DoubleType, YearMonthIntervalType, DayTimeIntervalType), - TypeCollection(DoubleType, YearMonthIntervalType, DayTimeIntervalType), - TypeCollection(DoubleType, YearMonthIntervalType, DayTimeIntervalType), + TypeCollection(DoubleType, YearMonthIntervalType, DayTimeIntervalType, AnyTimeType), + TypeCollection(DoubleType, YearMonthIntervalType, DayTimeIntervalType, AnyTimeType), + TypeCollection(DoubleType, YearMonthIntervalType, DayTimeIntervalType, AnyTimeType), LongType) override def checkInputDataTypes(): TypeCheckResult = { @@ -2141,6 +2143,8 @@ case class WidthBucket( TypeCheckSuccess case (_: DayTimeIntervalType, _: DayTimeIntervalType, _: DayTimeIntervalType) => TypeCheckSuccess + case (_: TimeType, _: TimeType, _: TimeType) => + TypeCheckSuccess case _ => val types = Seq(value.dataType, minValue.dataType, maxValue.dataType) TypeUtils.checkForSameTypeInputExpr(types, prettyName) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala index 90148932237c5..761f5268d4562 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala @@ -18,7 +18,7 @@ package org.apache.spark.sql.catalyst.expressions import java.nio.charset.StandardCharsets -import java.time.{Duration, Period} +import java.time.{Duration, LocalTime, Period} import java.time.temporal.ChronoUnit import com.google.common.math.LongMath @@ -1044,6 +1044,38 @@ class MathExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { checkEvaluation(WidthBucket(5.35, 0.024, Double.PositiveInfinity, 5L), null) } + test("width_bucket with the TIME type") { + def time(t: LocalTime, precision: Int = TimeType.MICROS_PRECISION): Literal = + Literal.create(t.toNanoOfDay, TimeType(precision)) + val t08 = time(LocalTime.of(8, 0, 0)) + val t09 = time(LocalTime.of(9, 0, 0)) + val t12 = time(LocalTime.of(12, 0, 0)) + val t17 = time(LocalTime.of(17, 0, 0)) + val n8 = Literal(8L) + + // 09:00:00..17:00:00 split into 8 one-hour buckets. + checkEvaluation(WidthBucket(t12, t09, t17, n8), 4L) + checkEvaluation(WidthBucket(t09, t09, t17, n8), 1L) // lower edge + checkEvaluation(WidthBucket(t17, t09, t17, n8), 9L) // >= max -> numBucket + 1 + checkEvaluation(WidthBucket(t08, t09, t17, n8), 0L) // < min -> 0 + // Reversed range (min > max). + checkEvaluation(WidthBucket(t12, t17, t09, n8), 6L) + + // Mixed precisions are allowed; comparison uses the shared nanos-of-day value. + checkEvaluation( + WidthBucket( + time(LocalTime.of(12, 0, 0), 0), + time(LocalTime.of(9, 0, 0), 3), + time(LocalTime.of(17, 0, 0), 9), + n8), + 4L) + + // Null and degenerate inputs. + checkEvaluation(WidthBucket(Literal.create(null, TimeType()), t09, t17, n8), null) + checkEvaluation(WidthBucket(t12, t09, t17, Literal.create(null, LongType)), null) + checkEvaluation(WidthBucket(t12, t09, t09, n8), null) // min == max + } + test("context independent foldable math expressions") { // Create some base literals val intLit = Literal(5) diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/time.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/time.sql.out index 3fb5379900209..e3767ca4e4d5c 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/time.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/time.sql.out @@ -2527,3 +2527,59 @@ DROP TABLE time_narrow_tbl -- !query analysis DropTable false, false +- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.time_narrow_tbl + + +-- !query +select width_bucket(TIME'12:00:00', TIME'09:00:00', TIME'17:00:00', 8) +-- !query analysis +Project [width_bucket(12:00:00, 09:00:00, 17:00:00, cast(8 as bigint)) AS width_bucket(TIME '12:00:00', TIME '09:00:00', TIME '17:00:00', 8)#xL] ++- OneRowRelation + + +-- !query +select width_bucket(TIME'09:00:00', TIME'09:00:00', TIME'17:00:00', 8) +-- !query analysis +Project [width_bucket(09:00:00, 09:00:00, 17:00:00, cast(8 as bigint)) AS width_bucket(TIME '09:00:00', TIME '09:00:00', TIME '17:00:00', 8)#xL] ++- OneRowRelation + + +-- !query +select width_bucket(TIME'17:00:00', TIME'09:00:00', TIME'17:00:00', 8) +-- !query analysis +Project [width_bucket(17:00:00, 09:00:00, 17:00:00, cast(8 as bigint)) AS width_bucket(TIME '17:00:00', TIME '09:00:00', TIME '17:00:00', 8)#xL] ++- OneRowRelation + + +-- !query +select width_bucket(TIME'08:00:00', TIME'09:00:00', TIME'17:00:00', 8) +-- !query analysis +Project [width_bucket(08:00:00, 09:00:00, 17:00:00, cast(8 as bigint)) AS width_bucket(TIME '08:00:00', TIME '09:00:00', TIME '17:00:00', 8)#xL] ++- OneRowRelation + + +-- !query +select width_bucket(TIME'12:00:00', TIME'17:00:00', TIME'09:00:00', 8) +-- !query analysis +Project [width_bucket(12:00:00, 17:00:00, 09:00:00, cast(8 as bigint)) AS width_bucket(TIME '12:00:00', TIME '17:00:00', TIME '09:00:00', 8)#xL] ++- OneRowRelation + + +-- !query +select width_bucket(TIME'12:00:00', CAST(TIME'09:00:00' AS TIME(3)), CAST(TIME'17:00:00' AS TIME(9)), 8) +-- !query analysis +Project [width_bucket(12:00:00, cast(09:00:00 as time(3)), cast(17:00:00 as time(9)), cast(8 as bigint)) AS width_bucket(TIME '12:00:00', CAST(TIME '09:00:00' AS TIME(3)), CAST(TIME '17:00:00' AS TIME(9)), 8)#xL] ++- OneRowRelation + + +-- !query +select width_bucket(CAST(null AS TIME), TIME'09:00:00', TIME'17:00:00', 8) +-- !query analysis +Project [width_bucket(cast(null as time(6)), 09:00:00, 17:00:00, cast(8 as bigint)) AS width_bucket(CAST(NULL AS TIME(6)), TIME '09:00:00', TIME '17:00:00', 8)#xL] ++- OneRowRelation + + +-- !query +select width_bucket(TIME'12:00:00', TIME'09:00:00', TIME'09:00:00', 8) +-- !query analysis +Project [width_bucket(12:00:00, 09:00:00, 09:00:00, cast(8 as bigint)) AS width_bucket(TIME '12:00:00', TIME '09:00:00', TIME '09:00:00', 8)#xL] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/inputs/time.sql b/sql/core/src/test/resources/sql-tests/inputs/time.sql index ee27d2acc1f9f..e7dd8f57aca72 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/time.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/time.sql @@ -430,3 +430,16 @@ INSERT INTO time_narrow_tbl SELECT '01:02:03.456789' :: TIME(6); INSERT INTO time_narrow_tbl SELECT CAST('01:02:03.456789' :: TIME(6) AS TIME(3)); SELECT typeof(t3), t3 FROM time_narrow_tbl; DROP TABLE time_narrow_tbl; + +-- width_bucket over the TIME type: 09:00:00 .. 17:00:00 split into 8 one-hour buckets +select width_bucket(TIME'12:00:00', TIME'09:00:00', TIME'17:00:00', 8); +select width_bucket(TIME'09:00:00', TIME'09:00:00', TIME'17:00:00', 8); +select width_bucket(TIME'17:00:00', TIME'09:00:00', TIME'17:00:00', 8); +select width_bucket(TIME'08:00:00', TIME'09:00:00', TIME'17:00:00', 8); +-- reversed range (min > max) +select width_bucket(TIME'12:00:00', TIME'17:00:00', TIME'09:00:00', 8); +-- mixed precisions are allowed +select width_bucket(TIME'12:00:00', CAST(TIME'09:00:00' AS TIME(3)), CAST(TIME'17:00:00' AS TIME(9)), 8); +-- null and degenerate inputs +select width_bucket(CAST(null AS TIME), TIME'09:00:00', TIME'17:00:00', 8); +select width_bucket(TIME'12:00:00', TIME'09:00:00', TIME'09:00:00', 8); diff --git a/sql/core/src/test/resources/sql-tests/results/time.sql.out b/sql/core/src/test/resources/sql-tests/results/time.sql.out index 5bb6fb64c7485..0e73fd366eefe 100644 --- a/sql/core/src/test/resources/sql-tests/results/time.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/time.sql.out @@ -3034,3 +3034,67 @@ DROP TABLE time_narrow_tbl struct<> -- !query output + + +-- !query +select width_bucket(TIME'12:00:00', TIME'09:00:00', TIME'17:00:00', 8) +-- !query schema +struct +-- !query output +4 + + +-- !query +select width_bucket(TIME'09:00:00', TIME'09:00:00', TIME'17:00:00', 8) +-- !query schema +struct +-- !query output +1 + + +-- !query +select width_bucket(TIME'17:00:00', TIME'09:00:00', TIME'17:00:00', 8) +-- !query schema +struct +-- !query output +9 + + +-- !query +select width_bucket(TIME'08:00:00', TIME'09:00:00', TIME'17:00:00', 8) +-- !query schema +struct +-- !query output +0 + + +-- !query +select width_bucket(TIME'12:00:00', TIME'17:00:00', TIME'09:00:00', 8) +-- !query schema +struct +-- !query output +6 + + +-- !query +select width_bucket(TIME'12:00:00', CAST(TIME'09:00:00' AS TIME(3)), CAST(TIME'17:00:00' AS TIME(9)), 8) +-- !query schema +struct +-- !query output +4 + + +-- !query +select width_bucket(CAST(null AS TIME), TIME'09:00:00', TIME'17:00:00', 8) +-- !query schema +struct +-- !query output +NULL + + +-- !query +select width_bucket(TIME'12:00:00', TIME'09:00:00', TIME'09:00:00', 8) +-- !query schema +struct +-- !query output +NULL