[SPARK-59440][SQL] Support the TIME data type in the hll_sketch_agg function

docs/sql-ref-sketch-aggregates.md

@@ -98,7 +98,7 @@ hll_sketch_agg(expr [, lgConfigK])
 
 | Argument | Type | Description |
 |----------|------|-------------|
-| `expr` | INT, BIGINT, STRING, or BINARY | The expression whose distinct values will be counted |
+| `expr` | INT, BIGINT, TIME, STRING, or BINARY | The expression whose distinct values will be counted |
 | `lgConfigK` | INT (optional) | Log-base-2 of K, where K is the number of buckets. Range: 4-21. Default: 12. Higher values provide more accuracy but use more memory. |
 
 Returns a BINARY containing the HLL sketch in updatable binary representation.

python/pyspark/sql/functions/builtin.py

@@ -30008,10 +30008,13 @@ def hll_sketch_agg(
 
     .. versionadded:: 3.5.0
 
+    .. versionchanged:: 4.4.0
+        Supports the TIME type for the ``col`` argument.
+
     Parameters
     ----------
     col : :class:`~pyspark.sql.Column` or column name
-        A column that evaluates to an integer, long, string, or binary.
+        A column that evaluates to an integer, long, time, string, or binary.
     lgConfigK : :class:`~pyspark.sql.Column` or int, optional
         The log-base-2 of K, where K is the number of buckets or slots for the HllSketch.
         A column that evaluates to an integer.

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

@@ -898,8 +898,8 @@ object functions {
    * configured with lgConfigK arg.
    *
    * @param e
-   *   the column to compute the sketch on. A column that evaluates to an integral, a string or a
-   *   binary.
+   *   the column to compute the sketch on. A column that evaluates to an integral, a time, a
+   *   string or a binary.
    * @param lgConfigK
    *   the log-base-2 of K, where K is the number of buckets or slots for the HllSketch. A column
    *   that evaluates to an integral. Must be a constant.
@@ -916,8 +916,8 @@ object functions {
    * configured with lgConfigK arg.
    *
    * @param e
-   *   the column to compute the sketch on. A column that evaluates to an integral, a string or a
-   *   binary.
+   *   the column to compute the sketch on. A column that evaluates to an integral, a time, a
+   *   string or a binary.
    * @param lgConfigK
    *   the log-base-2 of K, where K is the number of buckets or slots for the HllSketch. A column
    *   that evaluates to an integral. Must be a constant.
@@ -935,7 +935,7 @@ object functions {
    *
    * @param columnName
    *   the name of the column to compute the sketch on. A column that evaluates to an integral, a
-   *   string or a binary.
+   *   time, a string or a binary.
    * @param lgConfigK
    *   the log-base-2 of K, where K is the number of buckets or slots for the HllSketch. A column
    *   that evaluates to an integral. Must be a constant.
@@ -953,8 +953,8 @@ object functions {
    * configured with default lgConfigK value.
    *
    * @param e
-   *   the column to compute the sketch on. A column that evaluates to an integral, a string or a
-   *   binary.
+   *   the column to compute the sketch on. A column that evaluates to an integral, a time, a
+   *   string or a binary.
    * @group agg_funcs
    * @since 3.5.0
    * @return
@@ -969,7 +969,7 @@ object functions {
    *
    * @param columnName
    *   the name of the column to compute the sketch on. A column that evaluates to an integral, a
-   *   string or a binary.
+   *   time, a string or a binary.
    * @group agg_funcs
    * @since 3.5.0
    * @return

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

@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.trees.BinaryLike
 import org.apache.spark.sql.catalyst.util.CollationFactory
 import org.apache.spark.sql.errors.QueryExecutionErrors
 import org.apache.spark.sql.internal.types.StringTypeWithCollation
-import org.apache.spark.sql.types.{AbstractDataType, BinaryType, BooleanType, DataType, IntegerType, LongType, StringType, TypeCollection}
+import org.apache.spark.sql.types.{AbstractDataType, AnyTimeType, BinaryType, BooleanType, DataType, IntegerType, LongType, StringType, TimeType, TypeCollection}
 import org.apache.spark.unsafe.types.UTF8String
 
 
@@ -51,7 +51,7 @@ import org.apache.spark.unsafe.types.UTF8String
   arguments = """
     Arguments:
       * expr - The expression to aggregate into the HLL sketch.
-        An expression that evaluates to an integer, long, string, or binary.
+        An expression that evaluates to an integer, long, time, string, or binary.
       * lgConfigK - The log-base-2 of K, where K is the number of buckets for the sketch.
         An expression that evaluates to an integer.
   """,
@@ -119,6 +119,7 @@ case class HllSketchAgg(
       TypeCollection(
         IntegerType,
         LongType,
+        AnyTimeType,
         StringTypeWithCollation(supportsTrimCollation = true),
         BinaryType),
       IntegerType)
@@ -156,9 +157,12 @@ case class HllSketchAgg(
         // Spark SQL doesn't have equivalent types for ByteBuffer or char[] so leave those out.
         // We leave out support for Array types, as unique counting these aren't a common use case.
         // We leave out support for floating point types (such as DoubleType) due to imprecision.
-        // TODO: implement support for decimal/datetime/interval types
+        // TODO: implement support for decimal/date/timestamp/interval types
         case IntegerType => sketch.update(v.asInstanceOf[Int])
         case LongType => sketch.update(v.asInstanceOf[Long])
+        // TIME is physically stored as a long (nanoseconds since midnight), so it hashes exactly
+        // like LongType: equal times share the same nanos and therefore the same sketch entry.
+        case _: TimeType => sketch.update(v.asInstanceOf[Long])
         case st: StringType =>
           val collation = CollationFactory.fetchCollation(st.collationId)
           val str = v.asInstanceOf[UTF8String]

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/DatasketchesHllSketchSuite.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
+import java.time.LocalTime
+
 import scala.collection.immutable.NumericRange
 import scala.util.Random
 
@@ -26,7 +28,7 @@ import org.apache.datasketches.memory.Memory
 import org.apache.spark.{SparkFunSuite, SparkRuntimeException}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{BoundReference, HllSketchEstimate, HllUnion, Literal}
-import org.apache.spark.sql.types.{BinaryType, DataType, IntegerType, LongType, StringType}
+import org.apache.spark.sql.types.{BinaryType, DataType, IntegerType, LongType, StringType, TimeType}
 import org.apache.spark.unsafe.types.UTF8String
 
 
@@ -88,6 +90,44 @@ class DatasketchesHllSketchSuite extends SparkFunSuite {
       binaryEstimateRange.contains(binaryRange.size.toLong))
   }
 
+  test("Test hll_sketch_agg and hll_union_agg over the TIME type") {
+    // The analyzer admits TIME as a value to sketch.
+    assert(
+      new HllSketchAgg(BoundReference(0, TimeType(6), nullable = true), 12)
+        .checkInputDataTypes().isSuccess)
+
+    // TIME is physically a long of nanos-of-day, so distinct-counting a TIME column behaves
+    // exactly like counting the underlying longs.
+    val timeRange = (0 until 1000).map(_.toLong * 1000000000L) // 0s..999s of the day, in nanos
+    val (estimate, estimateRange) = simulateUpdateMerge(TimeType(), timeRange)
+    assert(estimate == timeRange.size || estimateRange.contains(timeRange.size.toLong))
+
+    val nineAm = LocalTime.of(9, 0, 0).toNanoOfDay
+    val noon = LocalTime.of(12, 0, 0).toNanoOfDay
+    val fivePm = LocalTime.of(17, 0, 0).toNanoOfDay
+
+    def timeSketch(precision: Int, values: Seq[Long]): Array[Byte] = {
+      val agg = new HllSketchAgg(BoundReference(0, TimeType(precision), nullable = true), 12)
+      val buf = values.foldLeft(agg.createAggregationBuffer())((b, v) =>
+        agg.update(b, InternalRow(v)))
+      agg.eval(buf).asInstanceOf[Array[Byte]]
+    }
+
+    // Repeated values are counted once (deduplication).
+    assert(estimateOf(timeSketch(9, Seq(noon, noon, noon, nineAm, nineAm))) == 2L)
+
+    // The full nanos-of-day is hashed: times that differ only in sub-microsecond digits are
+    // distinct. A regression that truncated to micros before hashing would under-count these.
+    assert(estimateOf(timeSketch(9, Seq(noon, noon + 1L, noon + 2L))) == 3L)
+
+    // Sketches built from TIME columns of different precisions round-trip through hll_union_agg,
+    // which only ever sees the serialized BINARY sketch and so needs no TIME-specific handling.
+    val merged = unionAgg(
+      Seq[Any](timeSketch(3, Seq(nineAm, noon)), timeSketch(9, Seq(noon, fivePm))),
+      allowDifferentLgConfigK = false)
+    assert(estimateOf(merged) == 3L) // distinct {09:00, 12:00, 17:00}
+  }
+
   test("Test lgMaxK results in downsampling sketches with larger lgConfigK") {
     val aggFunc1 = new HllSketchAgg(BoundReference(0, IntegerType, nullable = true), 12)
     val sketch1 = aggFunc1.createAggregationBuffer()

sql/core/src/test/resources/sql-tests/analyzer-results/hll.sql.out

@@ -190,6 +190,34 @@ Aggregate [hll_sketch_estimate(hll_union_agg(sketch#x, true, 0, 0)) AS hll_sketc
             +- LocalRelation [col#x]
 
 
+-- !query
+SELECT hll_sketch_estimate(hll_sketch_agg(col))
+FROM VALUES (TIME'12:00:00'), (TIME'12:00:00'), (TIME'09:00:00'), (TIME'17:00:00') tab(col)
+-- !query analysis
+Aggregate [hll_sketch_estimate(hll_sketch_agg(col#x, 12, 0, 0)) AS hll_sketch_estimate(hll_sketch_agg(col, 12))#xL]
++- SubqueryAlias tab
+   +- LocalRelation [col#x]
+
+
+-- !query
+SELECT hll_sketch_estimate(hll_union_agg(sketch, true))
+    FROM (SELECT hll_sketch_agg(col) as sketch
+            FROM VALUES (TIME'12:00:00'), (TIME'09:00:00') AS tab(col)
+          UNION ALL
+          SELECT hll_sketch_agg(col) as sketch
+            FROM VALUES (TIME'12:00:00'), (TIME'17:00:00') AS tab(col))
+-- !query analysis
+Aggregate [hll_sketch_estimate(hll_union_agg(sketch#x, true, 0, 0)) AS hll_sketch_estimate(hll_union_agg(sketch, true))#xL]
++- SubqueryAlias __auto_generated_subquery_name
+   +- Union false, false
+      :- Aggregate [hll_sketch_agg(col#x, 12, 0, 0) AS sketch#x]
+      :  +- SubqueryAlias tab
+      :     +- LocalRelation [col#x]
+      +- Aggregate [hll_sketch_agg(col#x, 12, 0, 0) AS sketch#x]
+         +- SubqueryAlias tab
+            +- LocalRelation [col#x]
+
+
 -- !query
 SELECT hll_sketch_agg(col)
 FROM VALUES (ARRAY(1, 2)), (ARRAY(3, 4)) tab(col)
@@ -202,7 +230,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException
     "inputSql" : "\"col\"",
     "inputType" : "\"ARRAY<INT>\"",
     "paramIndex" : "first",
-    "requiredType" : "(\"INT\" or \"BIGINT\" or \"STRING\" or \"BINARY\")",
+    "requiredType" : "(\"INT\" or \"BIGINT\" or \"TIME\" or \"STRING\" or \"BINARY\")",
     "sqlExpr" : "\"hll_sketch_agg(col, 12)\""
   },
   "queryContext" : [ {

sql/core/src/test/resources/sql-tests/inputs/hll.sql

@@ -59,6 +59,18 @@ SELECT hll_sketch_estimate(hll_union_agg(sketch, true))
           SELECT hll_sketch_agg(col, 20) as sketch
             FROM VALUES (1) AS tab(col));
 
+-- TIME type: hll_sketch_agg counts distinct times, and the resulting sketches merge via
+-- hll_union_agg (which only sees the serialized binary sketch).
+SELECT hll_sketch_estimate(hll_sketch_agg(col))
+FROM VALUES (TIME'12:00:00'), (TIME'12:00:00'), (TIME'09:00:00'), (TIME'17:00:00') tab(col);
+
+SELECT hll_sketch_estimate(hll_union_agg(sketch, true))
+    FROM (SELECT hll_sketch_agg(col) as sketch
+            FROM VALUES (TIME'12:00:00'), (TIME'09:00:00') AS tab(col)
+          UNION ALL
+          SELECT hll_sketch_agg(col) as sketch
+            FROM VALUES (TIME'12:00:00'), (TIME'17:00:00') AS tab(col));
+
 -- Negative test cases
 SELECT hll_sketch_agg(col)
 FROM VALUES (ARRAY(1, 2)), (ARRAY(3, 4)) tab(col);

sql/core/src/test/resources/sql-tests/results/hll.sql.out

@@ -191,6 +191,28 @@ struct<hll_sketch_estimate(hll_union_agg(sketch, true)):bigint>
 1
 
 
+-- !query
+SELECT hll_sketch_estimate(hll_sketch_agg(col))
+FROM VALUES (TIME'12:00:00'), (TIME'12:00:00'), (TIME'09:00:00'), (TIME'17:00:00') tab(col)
+-- !query schema
+struct<hll_sketch_estimate(hll_sketch_agg(col, 12)):bigint>
+-- !query output
+3
+
+
+-- !query
+SELECT hll_sketch_estimate(hll_union_agg(sketch, true))
+    FROM (SELECT hll_sketch_agg(col) as sketch
+            FROM VALUES (TIME'12:00:00'), (TIME'09:00:00') AS tab(col)
+          UNION ALL
+          SELECT hll_sketch_agg(col) as sketch
+            FROM VALUES (TIME'12:00:00'), (TIME'17:00:00') AS tab(col))
+-- !query schema
+struct<hll_sketch_estimate(hll_union_agg(sketch, true)):bigint>
+-- !query output
+3
+
+
 -- !query
 SELECT hll_sketch_agg(col)
 FROM VALUES (ARRAY(1, 2)), (ARRAY(3, 4)) tab(col)
@@ -205,7 +227,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException
     "inputSql" : "\"col\"",
     "inputType" : "\"ARRAY<INT>\"",
     "paramIndex" : "first",
-    "requiredType" : "(\"INT\" or \"BIGINT\" or \"STRING\" or \"BINARY\")",
+    "requiredType" : "(\"INT\" or \"BIGINT\" or \"TIME\" or \"STRING\" or \"BINARY\")",
     "sqlExpr" : "\"hll_sketch_agg(col, 12)\""
   },
   "queryContext" : [ {