joelrobin18 commented on code in PR #57958:
URL: https://github.com/apache/spark/pull/57958#discussion_r3766802276
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/DatasketchesHllSketchSuite.scala:
##########
@@ -153,4 +153,130 @@ class DatasketchesHllSketchSuite extends SparkFunSuite {
s"but got: ${exception.getClass.getName}: ${exception.getMessage}"
)
}
+ /** Runs HllUnionAgg over `inputs` (a NULL entry stands for a NULL sketch)
for a single group. */
+ private def unionAgg(inputs: Seq[Any], allowDifferentLgConfigK: Boolean):
Array[Byte] = {
+ val aggFunc = new HllUnionAgg(
+ BoundReference(0, BinaryType, nullable = true), allowDifferentLgConfigK)
+ val buffer = inputs.foldLeft(aggFunc.createAggregationBuffer()) { (buf,
input) =>
+ aggFunc.update(buf, InternalRow(input))
+ }
+ aggFunc.eval(buffer).asInstanceOf[Array[Byte]]
+ }
+
+ /** Runs HllSketchAgg at `lgConfigK` over `values` (a NULL entry stands for
a NULL value). */
+ private def sketchAgg(values: Seq[Any], lgConfigK: Int): Array[Byte] = {
+ val aggFunc = new HllSketchAgg(BoundReference(0, StringType, nullable =
true), lgConfigK)
+ val buffer = values.foldLeft(aggFunc.createAggregationBuffer()) { (buf,
value) =>
+ aggFunc.update(buf, InternalRow(value))
+ }
+ aggFunc.eval(buffer).asInstanceOf[Array[Byte]]
+ }
+
+ /** Evaluates the scalar hll_union over two serialized sketches. */
+ private def scalarUnion(
+ left: Array[Byte], right: Array[Byte], allowDifferentLgConfigK:
Boolean): Array[Byte] =
+ HllUnion(
+ Literal(left, BinaryType),
+ Literal(right, BinaryType),
+
Literal(allowDifferentLgConfigK)).eval(InternalRow.empty).asInstanceOf[Array[Byte]]
+
+ private def lgConfigKOf(sketch: Array[Byte]): Int =
+ HllSketch.heapify(Memory.wrap(sketch)).getLgConfigK
+
+ private def estimateOf(sketch: Array[Byte]): Long =
+ HllSketchEstimate(BoundReference(0, BinaryType, nullable = true))
+ .eval(InternalRow(sketch)).asInstanceOf[Long]
+
+ private def stringValues(n: Int): Seq[Any] =
+ Seq.tabulate(n)(i => UTF8String.fromString(i.toString))
+
+ test("hll_union_agg on a group with no non-NULL sketch yields an empty
default-lgConfigK " +
+ "sketch") {
+ // The aggregate has no lgConfigK parameter and never saw a sketch, so it
has no precision to
+ // report and falls back to the Datasketches default. Documented here
because the resulting
+ // sketch is observable, and must stay harmless to later unions (see the
tests below).
+ val allNull = unionAgg(Seq(null, null), allowDifferentLgConfigK = false)
+ assert(estimateOf(allNull) == 0L)
+ assert(lgConfigKOf(allNull) == HllSketch.DEFAULT_LG_K)
+
+ // hll_sketch_agg does not share the problem only because it has the
parameter: it builds its
+ // buffer eagerly at the requested lgConfigK. Without the argument it
defaults to 12 as well.
+ val emptyAt15 = sketchAgg(Seq(null, null), 15)
+ assert(estimateOf(emptyAt15) == 0L)
+ assert(lgConfigKOf(emptyAt15) == 15)
+ }
+
+ test("hll_union_agg merges an empty sketch of a different lgConfigK without
an error") {
+ // An empty sketch holds no coupons, so unioning it cannot lose
information at any lgConfigK.
+ // This is the shape produced by the test above, i.e. what a persisted
table ends up holding
+ // for a group whose sketches were all NULL.
+ val emptyAtDefaultLgK = unionAgg(Seq(null), allowDifferentLgConfigK =
false)
+ val sketchAt15 = sketchAgg(stringValues(1000), 15)
+
+ Seq(true, false).foreach { allowDifferentLgConfigK =>
+ Seq(
+ ("empty sketch first", Seq[Any](emptyAtDefaultLgK, sketchAt15)),
+ ("empty sketch last", Seq[Any](sketchAt15, emptyAtDefaultLgK))
+ ).foreach { case (order, inputs) =>
+ val merged = unionAgg(inputs, allowDifferentLgConfigK)
+ // The non-empty sketch decides the precision, whichever order the
rows arrive in: the
+ // result must not depend on which row the aggregate happens to see
first.
+ assert(lgConfigKOf(merged) == 15,
+ s"$order (allowDifferentLgConfigK=$allowDifferentLgConfigK) changed
the lgConfigK")
+ assert(estimateOf(merged) == estimateOf(sketchAt15),
+ s"$order (allowDifferentLgConfigK=$allowDifferentLgConfigK) changed
the estimate")
+ }
+ }
+ }
+
+ test("hll_union_agg still rejects non-empty sketches with different
lgConfigK") {
+ val sketchAt12 = sketchAgg(stringValues(1000), 12)
+ val sketchAt15 = sketchAgg(stringValues(1000), 15)
+
+ Seq(
+ Seq[Any](sketchAt12, sketchAt15),
+ Seq[Any](sketchAt15, sketchAt12)
+ ).foreach { inputs =>
+ val exception = intercept[SparkRuntimeException] {
+ unionAgg(inputs, allowDifferentLgConfigK = false)
+ }
+ assert(exception.getCondition == "HLL_UNION_DIFFERENT_LG_K")
+ }
+
+ // And still downsamples rather than erroring when the caller opts in.
+ assert(lgConfigKOf(unionAgg(Seq(sketchAt15, sketchAt12),
allowDifferentLgConfigK = true)) <= 15)
Review Comment:
Updated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
