[SPARK-57784][SQL] Support the TIME data type in cost-based optimizer statistics estimation

Author: yadavay-amzn

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/interface.scala

@@ -1047,6 +1047,7 @@ object CatalogColumnStat extends Logging {
       case TimestampType => getTimestampFormatter(isParsing = true).parse(s)
       case TimestampNTZType =>
         getTimestampFormatter(isParsing = true, forTimestampNTZ = true).parse(s)
+      case _: TimeType => TimeFormatter(isParsing = true).parse(s)
       case ByteType => s.toByte
       case ShortType => s.toShort
       case IntegerType => s.toInt
@@ -1073,6 +1074,7 @@ object CatalogColumnStat extends Logging {
       case TimestampNTZType =>
         getTimestampFormatter(isParsing = false, forTimestampNTZ = true)
           .format(v.asInstanceOf[Long])
+      case _: TimeType => TimeFormatter(isParsing = false).format(v.asInstanceOf[Long])
       case BooleanType | _: IntegralType | FloatType | DoubleType => v
       case _: DecimalType => v.asInstanceOf[Decimal].toJavaBigDecimal
       // This version of Spark does not use min/max for binary/string types so we ignore it.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/statsEstimation/EstimationUtils.scala

@@ -135,7 +135,7 @@ object EstimationUtils {
    */
   def toDouble(value: Any, dataType: DataType): Double = {
     dataType match {
-      case _: NumericType | DateType | TimestampType => value.toString.toDouble
+      case _: NumericType | DateType | TimestampType | _: TimeType => value.toString.toDouble
       case BooleanType => if (value.asInstanceOf[Boolean]) 1 else 0
     }
   }
@@ -144,7 +144,7 @@ object EstimationUtils {
     dataType match {
       case BooleanType => double.toInt == 1
       case DateType => double.toInt
-      case TimestampType => double.toLong
+      case TimestampType | _: TimeType => double.toLong
       case ByteType => double.toByte
       case ShortType => double.toShort
       case IntegerType => double.toInt

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/statsEstimation/FilterEstimation.scala

@@ -291,7 +291,7 @@ case class FilterEstimation(plan: Filter) extends Logging {
     }
 
     attr.dataType match {
-      case _: NumericType | DateType | TimestampType | BooleanType =>
+      case _: NumericType | DateType | TimestampType | BooleanType | _: TimeType =>
         evaluateBinaryForNumeric(op, attr, literal, update)
       case StringType | BinaryType =>
         // TODO: It is difficult to support other binary comparisons for String/Binary
@@ -413,7 +413,7 @@ case class FilterEstimation(plan: Filter) extends Logging {
 
     // use [min, max] to filter the original hSet
     dataType match {
-      case _: NumericType | BooleanType | DateType | TimestampType =>
+      case _: NumericType | BooleanType | DateType | TimestampType | _: TimeType =>
         if (ndv.toDouble == 0) {
           return Some(0.0)
         }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/statsEstimation/UnionEstimation.scala

@@ -38,7 +38,7 @@ object UnionEstimation {
   private def isTypeSupported(dt: DataType): Boolean = dt match {
     case ByteType | IntegerType | ShortType | FloatType | LongType |
          DoubleType | DateType | _: DecimalType | TimestampType | TimestampNTZType |
-         _: AnsiIntervalType => true
+         _: AnsiIntervalType | _: TimeType => true
     case _ => false
   }
 

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/statsEstimation/FilterEstimationSuite.scala

@@ -54,6 +54,15 @@ class FilterEstimationSuite extends StatsEstimationTestBase {
     min = Some(dMin), max = Some(dMax),
     nullCount = Some(0), avgLen = Some(4), maxLen = Some(4))
 
+  // column ctime has 10 values from 08:00:00 through 17:00:00 (nanos of day).
+  // 08:00 = 28800000000000L nanos, 17:00 = 61200000000000L nanos
+  val tMin = DateTimeUtils.localTimeToNanos(java.time.LocalTime.of(8, 0, 0))
+  val tMax = DateTimeUtils.localTimeToNanos(java.time.LocalTime.of(17, 0, 0))
+  val attrTime = AttributeReference("ctime", TimeType())()
+  val colStatTime = ColumnStat(distinctCount = Some(10),
+    min = Some(tMin), max = Some(tMax),
+    nullCount = Some(0), avgLen = Some(8), maxLen = Some(8))
+
   // column cdecimal has 4 values from 0.20 through 0.80 at increment of 0.20.
   val decMin = Decimal("0.200000000000000000")
   val decMax = Decimal("0.800000000000000000")
@@ -118,6 +127,7 @@ class FilterEstimationSuite extends StatsEstimationTestBase {
     attrInt -> colStatInt,
     attrBool -> colStatBool,
     attrDate -> colStatDate,
+    attrTime -> colStatTime,
     attrDecimal -> colStatDecimal,
     attrDouble -> colStatDouble,
     attrString -> colStatString,
@@ -523,6 +533,44 @@ class FilterEstimationSuite extends StatsEstimationTestBase {
       expectedRowCount = 3)
   }
 
+  test("ctime = cast('10:00:00' AS TIME)") {
+    val t10 = DateTimeUtils.localTimeToNanos(java.time.LocalTime.of(10, 0, 0))
+    validateEstimatedStats(
+      Filter(EqualTo(attrTime, Literal(t10, TimeType())),
+        childStatsTestPlan(Seq(attrTime), 10L)),
+      Seq(attrTime -> ColumnStat(distinctCount = Some(1),
+        min = Some(t10), max = Some(t10),
+        nullCount = Some(0), avgLen = Some(8), maxLen = Some(8))),
+      expectedRowCount = 1)
+  }
+
+  test("ctime < cast('12:00:00' AS TIME)") {
+    // 12:00 is 43200000000000L nanos. Range is [08:00, 17:00] = 10 distinct values.
+    // Fraction: (12:00 - 08:00) / (17:00 - 08:00) = 4/9 hours => ~4.44 => rounded to 5
+    val t12 = DateTimeUtils.localTimeToNanos(java.time.LocalTime.of(12, 0, 0))
+    validateEstimatedStats(
+      Filter(LessThan(attrTime, Literal(t12, TimeType())),
+        childStatsTestPlan(Seq(attrTime), 10L)),
+      Seq(attrTime -> ColumnStat(distinctCount = Some(5),
+        min = Some(tMin), max = Some(t12),
+        nullCount = Some(0), avgLen = Some(8), maxLen = Some(8))),
+      expectedRowCount = 5)
+  }
+
+  test("""ctime IN ( cast('09:00:00' AS TIME),
+      cast('10:00:00' AS TIME), cast('11:00:00' AS TIME) )""") {
+    val t09 = DateTimeUtils.localTimeToNanos(java.time.LocalTime.of(9, 0, 0))
+    val t10 = DateTimeUtils.localTimeToNanos(java.time.LocalTime.of(10, 0, 0))
+    val t11 = DateTimeUtils.localTimeToNanos(java.time.LocalTime.of(11, 0, 0))
+    validateEstimatedStats(
+      Filter(In(attrTime, Seq(Literal(t09, TimeType()), Literal(t10, TimeType()),
+        Literal(t11, TimeType()))), childStatsTestPlan(Seq(attrTime), 10L)),
+      Seq(attrTime -> ColumnStat(distinctCount = Some(3),
+        min = Some(t09), max = Some(t11),
+        nullCount = Some(0), avgLen = Some(8), maxLen = Some(8))),
+      expectedRowCount = 3)
+  }
+
   test("cdecimal = 0.400000000000000000") {
     val dec_0_40 = Decimal("0.400000000000000000")
     validateEstimatedStats(

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/statsEstimation/JoinEstimationSuite.scala

@@ -472,6 +472,7 @@ class JoinEstimationSuite extends StatsEstimationTestBase {
       val dec = Decimal("1.000000000000000000")
       val date = DateTimeUtils.fromJavaDate(Date.valueOf("2016-05-08"))
       val timestamp = DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2016-05-08 00:00:01"))
+      val time = DateTimeUtils.localTimeToNanos(java.time.LocalTime.of(10, 30, 0))
       mutable.LinkedHashMap[Attribute, ColumnStat](
         AttributeReference("cbool", BooleanType)() -> ColumnStat(distinctCount = Some(1),
           min = Some(false), max = Some(false),
@@ -506,6 +507,9 @@ class JoinEstimationSuite extends StatsEstimationTestBase {
           nullCount = Some(0), avgLen = Some(4), maxLen = Some(4)),
         AttributeReference("ctimestamp", TimestampType)() -> ColumnStat(distinctCount = Some(1),
           min = Some(timestamp), max = Some(timestamp),
+          nullCount = Some(0), avgLen = Some(8), maxLen = Some(8)),
+        AttributeReference("ctime", TimeType())() -> ColumnStat(distinctCount = Some(1),
+          min = Some(time), max = Some(time),
           nullCount = Some(0), avgLen = Some(8), maxLen = Some(8))
       )
     }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/statsEstimation/UnionEstimationSuite.scala

@@ -60,6 +60,7 @@ class UnionEstimationSuite extends StatsEstimationTestBase {
     val attrTimestampNTZ = AttributeReference("ctimestamp_ntz", TimestampNTZType)()
     val attrYMInterval = AttributeReference("cyminterval", YearMonthIntervalType())()
     val attrDTInterval = AttributeReference("cdtinterval", DayTimeIntervalType())()
+    val attrTime = AttributeReference("ctime", TimeType())()
 
     val s1 = 1.toShort
     val s2 = 4.toShort
@@ -90,7 +91,8 @@ class UnionEstimationSuite extends StatsEstimationTestBase {
         attrTimestamp -> ColumnStat(min = Some(1L), max = Some(4L)),
         attrTimestampNTZ -> ColumnStat(min = Some(1L), max = Some(4L)),
         attrYMInterval -> ColumnStat(min = Some(2), max = Some(5)),
-        attrDTInterval -> ColumnStat(min = Some(2L), max = Some(5L))))
+        attrDTInterval -> ColumnStat(min = Some(2L), max = Some(5L)),
+        attrTime -> ColumnStat(min = Some(1000L), max = Some(4000L))))
 
     val s3 = 2.toShort
     val s4 = 6.toShort
@@ -133,7 +135,10 @@ class UnionEstimationSuite extends StatsEstimationTestBase {
           max = Some(8)),
         AttributeReference("cdttimestamp1", DayTimeIntervalType())() -> ColumnStat(
           min = Some(4L),
-          max = Some(8L))))
+          max = Some(8L)),
+        AttributeReference("ctime1", TimeType())() -> ColumnStat(
+          min = Some(3000L),
+          max = Some(6000L))))
 
     val child1 = StatsTestPlan(
       outputList = columnInfo.keys.toSeq.sortWith(_.exprId.id < _.exprId.id),
@@ -167,7 +172,8 @@ class UnionEstimationSuite extends StatsEstimationTestBase {
           attrTimestamp -> ColumnStat(min = Some(1L), max = Some(6L)),
           attrTimestampNTZ -> ColumnStat(min = Some(1L), max = Some(6L)),
           attrYMInterval -> ColumnStat(min = Some(2), max = Some(8)),
-          attrDTInterval -> ColumnStat(min = Some(2L), max = Some(8L)))))
+          attrDTInterval -> ColumnStat(min = Some(2L), max = Some(8L)),
+          attrTime -> ColumnStat(min = Some(1000L), max = Some(6000L)))))
     assert(union.stats === expectedStats)
   }