diff --git a/docs/sql-ref-sketch-aggregates.md b/docs/sql-ref-sketch-aggregates.md index fd07cdf928bdf..a761be1d495d3 100644 --- a/docs/sql-ref-sketch-aggregates.md +++ b/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. diff --git a/python/pyspark/sql/functions/builtin.py b/python/pyspark/sql/functions/builtin.py index 528073a105158..b9adcaa4c4300 100644 --- a/python/pyspark/sql/functions/builtin.py +++ b/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. diff --git a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala index 2845c670a7691..193018c953ca0 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala +++ b/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 diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala index 0a06c54409c08..7ead098ee4ca1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala +++ b/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] diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/DatasketchesHllSketchSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/DatasketchesHllSketchSuite.scala index 4dcc3bf28d24e..ed037cbc19b11 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/DatasketchesHllSketchSuite.scala +++ b/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() diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/hll.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/hll.sql.out index 291f071ef06c2..22608c0e08475 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/hll.sql.out +++ b/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\"", "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" : [ { diff --git a/sql/core/src/test/resources/sql-tests/inputs/hll.sql b/sql/core/src/test/resources/sql-tests/inputs/hll.sql index 35128da97fd61..552c30f2e0be8 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/hll.sql +++ b/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); diff --git a/sql/core/src/test/resources/sql-tests/results/hll.sql.out b/sql/core/src/test/resources/sql-tests/results/hll.sql.out index 908221f0e7c40..46f59198fa06f 100644 --- a/sql/core/src/test/resources/sql-tests/results/hll.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/hll.sql.out @@ -191,6 +191,28 @@ struct 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 +-- !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 +-- !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\"", "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" : [ {