This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new aca0501f2932 [SPARK-58731][SQL] Exempt empty sketches from the
hll_union lgConfigK check
aca0501f2932 is described below
commit aca0501f2932c1ad634044e81d9abe6eab89ea2c
Author: Joel Robin P <[email protected]>
AuthorDate: Thu Aug 13 00:37:59 2026 +0800
[SPARK-58731][SQL] Exempt empty sketches from the hll_union lgConfigK check
### What changes were proposed in this pull request?
`hll_union_agg` has no `lgConfigK` parameter, so it infers precision from
the sketches it reads. For a group whose sketches are all NULL it never sees
one, and `eval` falls back to a `Union` at the DataSketches default
(`lgConfigK=12`). That empty sketch is observable -- it gets stored -- and
could not afterwards be merged with sketches of any other precision, so the
function produced output it could not itself consume.
This PR makes empty sketches exempt from the precision check, on the
grounds that an empty sketch holds no coupons and therefore carries no
precision to reconcile:
- `HllUnionAgg` gains a `mergeSketch` helper used by both `update` and
`merge`. It skips `compareLgConfigK` unless both sides are non-empty, and
re-seeds an empty `Union` at the first non-empty sketch's `lgConfigK`. Routing
both paths through one helper matters because the failure reproduced through
`merge` as well, so fixing only `update` would leave the partial-to-final
aggregation path broken.
- Scalar `HllUnion` gets the same exemption, and picks the non-empty side's
`lgConfigK` instead of an unconditional `Math.min` -- otherwise an empty sketch
at 12 would silently downsample a populated sketch at 15.
Merging an empty sketch is now a no-op with respect to precision, which
also makes the result independent of the order rows reach the aggregate.
### Why are the changes needed?
Rolling up stored sketches fails under default settings when any group had
only NULL input:
```
[HLL_UNION_DIFFERENT_LG_K] Sketches have different `lgConfigK` values: 12
and 15.
Set the `allowDifferentLgConfigK` parameter to true to enable unions of
different lgConfigK values.
```
The failure is self-referential: `hll_union_agg` emits a sketch that
`hll_union_agg` then refuses. Reproduction:
```sql
WITH sketches AS (
SELECT 'has_data' AS grp, hll_sketch_agg(CAST(id AS STRING), 15) AS sketch
FROM (SELECT explode(sequence(1, 1000)) AS id)
UNION ALL
SELECT 'all_null' AS grp, CAST(NULL AS BINARY) AS sketch
),
stored AS (SELECT grp, hll_union_agg(sketch) AS sketch FROM sketches GROUP
BY grp)
SELECT hll_sketch_estimate(hll_union_agg(sketch)) FROM stored;
```
The workaround -- setting `allowDifferentLgConfigK = true` -- is a poor
fit, since it also permits genuine precision loss between two populated
sketches, which is the case the check exists to catch. An all-NULL group is
common in practice (a partition with no matching rows), so a user can hit this
without ever having mixed precisions deliberately.
### Does this PR introduce _any_ user-facing change?
Yes, a bug fix. Queries that previously failed with
`HLL_UNION_DIFFERENT_LG_K` when unioning an empty sketch with a
differing-precision sketch now succeed, and the populated sketch's `lgConfigK`
is preserved.
No change for previously-succeeding queries: sketches of equal precision
are unaffected, and two **non-empty** sketches with differing `lgConfigK` still
raise `HLL_UNION_DIFFERENT_LG_K` unless `allowDifferentLgConfigK` is set. One
narrow behavior change worth noting for review: `hll_union(empty12,
populated15)` previously downsampled to 12 when `allowDifferentLgConfigK =
true`; it now returns 15, since the empty side has no precision to impose.
### How was this patch tested?
New tests, all passing locally:
- `DatasketchesHllSketchSuite` -- 4 tests covering the aggregate and scalar
`hll_union`, in both merge orders; plus assertions that mismatched
**non-empty** sketches still throw `HLL_UNION_DIFFERENT_LG_K`, and that
`allowDifferentLgConfigK = true` still downsamples.
- `DataFrameAggregateSuite` -- 2 end-to-end SQL tests reproducing the JIRA
scenario, run at 1 and 2 partitions so both `update` and the partial-to-final
`merge` path are exercised.
Existing `hll_sketch_agg` / `hll_union_agg` tests and the `hll.sql` golden
file pass unmodified. The full `sql/catalyst` suite also passes (10531 tests,
380 suites, 0 failures).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
Closes #57958 from joelrobin18/SPARK-58731-hll-union-empty-sketch-lgk.
Authored-by: Joel Robin P <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../aggregate/datasketchesAggregates.scala | 35 ++++--
.../expressions/datasketchesExpressions.scala | 12 +-
.../aggregate/DatasketchesHllSketchSuite.scala | 130 ++++++++++++++++++++-
.../apache/spark/sql/DataFrameAggregateSuite.scala | 64 ++++++++++
4 files changed, 229 insertions(+), 12 deletions(-)
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 4470c415f5c9..0a06c54409c0 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
@@ -335,6 +335,30 @@ case class HllUnionAgg(
}
}
+ /**
+ * Merges `sketch` into the Union acting as the aggregation buffer,
instantiating it if absent.
+ *
+ * An empty sketch holds no coupons, so it carries no precision: it is
exempt from the lgConfigK
+ * check, and an empty Union is re-seeded at the lgConfigK of the first
non-empty sketch. This
+ * keeps the empty sketch that `eval` emits for an all-NULL group mergeable,
and makes the
+ * result independent of the order in which rows reach the aggregate.
+ *
+ * @param unionOption A previously initialized Union instance, or None
+ * @param sketch The sketch to merge in
+ */
+ private def mergeSketch(unionOption: Option[Union], sketch: HllSketch):
Option[Union] = {
+ val union = unionOption match {
+ case Some(buffer) if buffer.isEmpty && !sketch.isEmpty => new
Union(sketch.getLgConfigK)
+ case Some(buffer) => buffer
+ case None => new Union(sketch.getLgConfigK)
+ }
+ if (!union.isEmpty && !sketch.isEmpty) {
+ compareLgConfigK(union.getLgConfigK, sketch.getLgConfigK)
+ }
+ union.update(sketch)
+ Some(union)
+ }
+
/**
* Update the Union instance with the HllSketch byte array obtained from the
row.
*
@@ -348,10 +372,7 @@ case class HllUnionAgg(
case BinaryType =>
try {
val sketch =
HllSketch.wrap(Memory.wrap(v.asInstanceOf[Array[Byte]]))
- val union = unionOption.getOrElse(new Union(sketch.getLgConfigK))
- compareLgConfigK(union.getLgConfigK, sketch.getLgConfigK)
- union.update(sketch)
- Some(union)
+ mergeSketch(unionOption, sketch)
} catch {
case _: SketchesArgumentException | _: java.lang.Error
| _: ArrayIndexOutOfBoundsException =>
@@ -373,10 +394,8 @@ case class HllUnionAgg(
*/
override def merge(unionOption: Option[Union], inputOption: Option[Union]):
Option[Union] = {
(unionOption, inputOption) match {
- case (Some(union), Some(input)) =>
- compareLgConfigK(union.getLgConfigK, input.getLgConfigK)
- union.update(input.getResult(targetType))
- Some(union)
+ case (Some(_), Some(input)) =>
+ mergeSketch(unionOption, input.getResult(targetType))
// unclear if these scenarios can ever occur
case (Some(_), None) =>
unionOption
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datasketchesExpressions.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datasketchesExpressions.scala
index 743dd27c9240..46f9e494d2a8 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datasketchesExpressions.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datasketchesExpressions.scala
@@ -135,11 +135,19 @@ case class HllUnion(first: Expression, second:
Expression, third: Expression)
throw QueryExecutionErrors.hllInvalidInputSketchBuffer(prettyName)
}
val allowDifferentLgConfigK = value3.asInstanceOf[Boolean]
- if (!allowDifferentLgConfigK && sketch1.getLgConfigK !=
sketch2.getLgConfigK) {
+ if (!allowDifferentLgConfigK && !sketch1.isEmpty && !sketch2.isEmpty &&
+ sketch1.getLgConfigK != sketch2.getLgConfigK) {
throw QueryExecutionErrors.hllUnionDifferentLgK(
sketch1.getLgConfigK, sketch2.getLgConfigK, function = prettyName)
}
- val union = new Union(Math.min(sketch1.getLgConfigK, sketch2.getLgConfigK))
+ // An empty sketch holds no coupons, so it carries no precision: it is
exempt from the
+ // lgConfigK check and does not drag the result down to its own lgConfigK.
+ val lgConfigK = (sketch1.isEmpty, sketch2.isEmpty) match {
+ case (true, false) => sketch2.getLgConfigK
+ case (false, true) => sketch1.getLgConfigK
+ case _ => Math.min(sketch1.getLgConfigK, sketch2.getLgConfigK)
+ }
+ val union = new Union(lgConfigK)
union.update(sketch1)
union.update(sketch2)
union.getResult(targetType).toUpdatableByteArray
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 0f7f5ca54be0..4dcc3bf28d24 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
@@ -23,9 +23,9 @@ import scala.util.Random
import org.apache.datasketches.hll.HllSketch
import org.apache.datasketches.memory.Memory
-import org.apache.spark.SparkFunSuite
+import org.apache.spark.{SparkFunSuite, SparkRuntimeException}
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{BoundReference,
HllSketchEstimate}
+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.unsafe.types.UTF8String
@@ -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)) == 12)
+ }
+
+ test("hll_union merges an empty sketch of a different lgConfigK without an
error") {
+ val emptyAtDefaultLgK = unionAgg(Seq(null), allowDifferentLgConfigK =
false)
+ val sketchAt15 = sketchAgg(stringValues(1000), 15)
+
+ Seq(true, false).foreach { allowDifferentLgConfigK =>
+ Seq(
+ ("empty sketch first", emptyAtDefaultLgK, sketchAt15),
+ ("empty sketch last", sketchAt15, emptyAtDefaultLgK)
+ ).foreach { case (order, left, right) =>
+ val merged = scalarUnion(left, right, allowDifferentLgConfigK)
+ 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 still rejects non-empty sketches with different lgConfigK") {
+ val sketchAt12 = sketchAgg(stringValues(1000), 12)
+ val sketchAt15 = sketchAgg(stringValues(1000), 15)
+
+ Seq((sketchAt12, sketchAt15), (sketchAt15, sketchAt12)).foreach { case
(left, right) =>
+ val exception = intercept[SparkRuntimeException] {
+ scalarUnion(left, right, allowDifferentLgConfigK = false)
+ }
+ assert(exception.getCondition == "HLL_UNION_DIFFERENT_LG_K")
+ }
+
+ // Two populated sketches still downsample to the smaller lgConfigK when
opted in.
+ assert(lgConfigKOf(scalarUnion(sketchAt15, sketchAt12,
allowDifferentLgConfigK = true)) == 12)
+ }
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
index 3bc766860372..1918acab6c4e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
@@ -2807,6 +2807,70 @@ class DataFrameAggregateSuite extends SharedSparkSession
}
}
+ // Byte 3 of the serialized sketch holds lgConfigK, so SUBSTRING(hex(...),
7, 2) reads it as a
+ // hex pair: 0x0F = 15, 0x0C = 12.
+ private val allNullGroupSketches =
+ """
+ |with sketches as (
+ | select 'has_data' as grp, hll_sketch_agg(cast(id as string), 15) as
sketch
+ | from (select explode(sequence(1, 1000)) as id)
+ | union all
+ | select 'all_null' as grp, cast(null as binary) as sketch
+ |)
+ |""".stripMargin
+
+ test("hll_union_agg reports the default lgConfigK for a group with no
non-NULL sketch") {
+ // hll_union_agg has no lgConfigK parameter, and the all-NULL group holds
no sketch to take one
+ // from, so it cannot report the 15 the other group was built at. The
estimate is correct either
+ // way; the mismatched precision is what the next test has to cope with.
+ checkAnswer(
+ sql(allNullGroupSketches +
+ """
+ |select
+ | grp,
+ | substring(hex(hll_union_agg(sketch)), 7, 2) as lg_config_k,
+ | hll_sketch_estimate(hll_union_agg(sketch)) as estimate
+ |from sketches
+ |group by grp
+ |""".stripMargin),
+ Seq(Row("all_null", "0C", 0L), Row("has_data", "0F", 1000L)))
+ }
+
+ test("hll_union_agg and hll_union re-merge the empty sketch produced for an
all-NULL group") {
+ // The stored sketches are now an empty lgConfigK=12 one beside a
populated lgConfigK=15 one.
+ // Rolling them back up must not fail under default settings, because an
empty sketch holds no
+ // coupons and so cannot cost precision at any lgConfigK.
+ val stored = sql(allNullGroupSketches +
+ "select grp, hll_union_agg(sketch) as sketch from sketches group by grp")
+ .collect()
+ .map(row => row.getString(0) -> row.getAs[Array[Byte]]("sketch"))
+ .toMap
+
+ Seq(
+ "empty sketch first" -> Seq(stored("all_null"), stored("has_data")),
+ "empty sketch last" -> Seq(stored("has_data"), stored("all_null"))
+ ).foreach { case (order, sketches) =>
+ // One partition exercises update(); two exercise the partial-to-final
merge() as well.
+ Seq(1, 2).foreach { numPartitions =>
+ withClue(s"hll_union_agg, $order, $numPartitions partition(s): ") {
+ checkAnswer(
+ sketches.toDF("sketch").repartition(numPartitions).selectExpr(
+ "substring(hex(hll_union_agg(sketch)), 7, 2) as lg_config_k",
+ "hll_sketch_estimate(hll_union_agg(sketch)) as estimate"),
+ Seq(Row("0F", 1000L)))
+ }
+ }
+
+ withClue(s"hll_union, $order: ") {
+ checkAnswer(
+ Seq((sketches.head, sketches.last)).toDF("left", "right").selectExpr(
+ "substring(hex(hll_union(left, right)), 7, 2) as lg_config_k",
+ "hll_sketch_estimate(hll_union(left, right)) as estimate"),
+ Seq(Row("0F", 1000L)))
+ }
+ }
+ }
+
test("hll_sketch_agg") {
val df = Seq(1, 1, 2, 2, 3).toDF("col")
checkAnswer(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
