From b4295d73a8eaa15e6a16c973bbd854b6d5fcd22d Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Mon, 20 Jul 2026 12:37:40 +0000 Subject: [PATCH 1/4] [SPARK-58069][SQL][FOLLOWUP] Handle empty approx_top_k combine buffers --- .../aggregate/ApproxTopKAggregates.scala | 31 +++++++++++++++---- .../apache/spark/sql/ApproxTopKSuite.scala | 17 ++++++++++ 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala index 2690a52b2d31d..b96535e0c31c8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala @@ -697,9 +697,20 @@ class CombineInternal[T]( * boundaries (SPARK-58069). */ def serialize(): Array[Byte] = { - val sketchWithNullCountBytes = sketchWithNullCount.serialize( - ApproxTopK.genSketchSerDe(itemDataType).asInstanceOf[ArrayOfItemsSerDe[T]]) - val typeBytes: Array[Byte] = itemDataType.json.getBytes(StandardCharsets.UTF_8) + // An empty partition has a placeholder buffer whose itemDataType has not been initialized. + // It can still be serialized between aggregation stages, so use a default serde for its empty + // sketch and encode the missing type as a zero-length section. + val serDe: ArrayOfItemsSerDe[T] = if (itemDataType == null) { + new ArrayOfStringsSerDe().asInstanceOf[ArrayOfItemsSerDe[T]] + } else { + ApproxTopK.genSketchSerDe(itemDataType).asInstanceOf[ArrayOfItemsSerDe[T]] + } + val sketchWithNullCountBytes = sketchWithNullCount.serialize(serDe) + val typeBytes: Array[Byte] = if (itemDataType == null) { + Array.emptyByteArray + } else { + itemDataType.json.getBytes(StandardCharsets.UTF_8) + } val byteArray = new Array[Byte]( sketchWithNullCountBytes.length + Integer.BYTES + Integer.BYTES + typeBytes.length) @@ -729,12 +740,20 @@ object CombineInternal { val typeLength = byteBuffer.getInt val typeBytes = new Array[Byte](typeLength) byteBuffer.get(typeBytes) - val itemDataType = DataType.fromJson(new String(typeBytes, StandardCharsets.UTF_8)) + val itemDataType = if (typeLength == 0) { + null + } else { + DataType.fromJson(new String(typeBytes, StandardCharsets.UTF_8)) + } // read sketchBytes val sketchBytes = new Array[Byte](buffer.length - Integer.BYTES - Integer.BYTES - typeLength) byteBuffer.get(sketchBytes) - val sketchWithNullCount = ApproxTopKAggregateBuffer.deserialize( - sketchBytes, ApproxTopK.genSketchSerDe(itemDataType)) + val serDe = if (itemDataType == null) { + new ArrayOfStringsSerDe().asInstanceOf[ArrayOfItemsSerDe[Any]] + } else { + ApproxTopK.genSketchSerDe(itemDataType) + } + val sketchWithNullCount = ApproxTopKAggregateBuffer.deserialize(sketchBytes, serDe) new CombineInternal[Any](sketchWithNullCount, itemDataType, maxItemsTracked) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala index 98aba80682e68..663091dbffb6b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala @@ -20,8 +20,11 @@ package org.apache.spark.sql import java.sql.{Date, Timestamp} import java.time.{LocalDateTime, LocalTime} +import org.apache.datasketches.frequencies.ItemsSketch + import org.apache.spark.{SparkArithmeticException, SparkRuntimeException} import org.apache.spark.sql.catalyst.ExtendedAnalysisException +import org.apache.spark.sql.catalyst.expressions.aggregate.{ApproxTopKAggregateBuffer, CombineInternal} import org.apache.spark.sql.errors.DataTypeErrors.toSQLType import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, TimestampNTZType, TimestampType, TimeType} @@ -575,6 +578,20 @@ class ApproxTopKSuite extends SharedSparkSession { "CAST('13:00:00.123' AS TIME(3))")) ) + test("SPARK-58069: serialize an empty approx_top_k_combine buffer") { + val maxItemsTracked = 100 + val buffer = new CombineInternal[Any]( + new ApproxTopKAggregateBuffer[Any](new ItemsSketch[Any](128), 0L), + null, + maxItemsTracked) + + val restored = CombineInternal.deserialize(buffer.serialize()) + + assert(restored.getItemDataType == null) + assert(restored.getMaxItemsTracked == maxItemsTracked) + assert(restored.getSketchWithNullCount.sketch.isEmpty) + } + // positive tests for approx_top_k_combine on every types gridTest("SPARK-52798: same type, same size, specified combine size - success")(itemsWithTopK) { case (input, expected) => From 41eecbac20e4530291820698fc020ecd91bc3c5c Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Wed, 22 Jul 2026 12:23:19 +0000 Subject: [PATCH 2/4] [SPARK-58069][SQL][TEST] Cover empty combine partitions --- .../apache/spark/sql/ApproxTopKSuite.scala | 34 +++++++++++++------ 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala index 663091dbffb6b..dda1f687146ce 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala @@ -24,7 +24,8 @@ import org.apache.datasketches.frequencies.ItemsSketch import org.apache.spark.{SparkArithmeticException, SparkRuntimeException} import org.apache.spark.sql.catalyst.ExtendedAnalysisException -import org.apache.spark.sql.catalyst.expressions.aggregate.{ApproxTopKAggregateBuffer, CombineInternal} +import org.apache.spark.sql.catalyst.expressions.aggregate.{ApproxTopK, ApproxTopKAggregateBuffer, + CombineInternal} import org.apache.spark.sql.errors.DataTypeErrors.toSQLType import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, TimestampNTZType, TimestampType, TimeType} @@ -579,17 +580,30 @@ class ApproxTopKSuite extends SharedSparkSession { ) test("SPARK-58069: serialize an empty approx_top_k_combine buffer") { - val maxItemsTracked = 100 - val buffer = new CombineInternal[Any]( - new ApproxTopKAggregateBuffer[Any](new ItemsSketch[Any](128), 0L), - null, - maxItemsTracked) + Seq(100, ApproxTopK.VOID_MAX_ITEMS_TRACKED).foreach { maxItemsTracked => + val buffer = new CombineInternal[Any]( + new ApproxTopKAggregateBuffer[Any](new ItemsSketch[Any](128), 0L), + null, + maxItemsTracked) + + val restored = CombineInternal.deserialize(buffer.serialize()) + + assert(restored.getItemDataType == null) + assert(restored.getMaxItemsTracked == maxItemsTracked) + assert(restored.getSketchWithNullCount.sketch.isEmpty) + assert(restored.getSketchWithNullCount.nullCount == 0L) + } + } - val restored = CombineInternal.deserialize(buffer.serialize()) + test("SPARK-58069: combine sketches with empty shuffle partitions") { + val sketches = sql( + "SELECT approx_top_k_accumulate(id) AS sketch FROM range(10) GROUP BY id") + .repartition(20) + val combined = sketches.selectExpr("approx_top_k_combine(sketch) AS sketch") - assert(restored.getItemDataType == null) - assert(restored.getMaxItemsTracked == maxItemsTracked) - assert(restored.getSketchWithNullCount.sketch.isEmpty) + checkAnswer( + combined.selectExpr("inline(approx_top_k_estimate(sketch))"), + (0L until 10L).map(Row(_, 1L))) } // positive tests for approx_top_k_combine on every types From d778721566576d1ab0a3989427b85fb820353813 Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Thu, 23 Jul 2026 07:02:33 +0000 Subject: [PATCH 3/4] [SPARK-58069][SQL][FOLLOWUP] Fix ApproxTopKSuite compilation --- .../src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala index dda1f687146ce..2ec8c584f2033 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala @@ -591,7 +591,7 @@ class ApproxTopKSuite extends SharedSparkSession { assert(restored.getItemDataType == null) assert(restored.getMaxItemsTracked == maxItemsTracked) assert(restored.getSketchWithNullCount.sketch.isEmpty) - assert(restored.getSketchWithNullCount.nullCount == 0L) + assert(restored.getSketchWithNullCount.eval(1, StringType).numElements() == 0) } } From 5478360e3c3a269f4ed351c275bebd5d1837a552 Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Thu, 23 Jul 2026 12:41:31 +0000 Subject: [PATCH 4/4] [SPARK-58069][SQL][FOLLOWUP] Ignore empty combine sketch size --- .../expressions/aggregate/ApproxTopKAggregates.scala | 2 +- .../scala/org/apache/spark/sql/ApproxTopKSuite.scala | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala index b96535e0c31c8..c69ecf8759735 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala @@ -646,7 +646,7 @@ class CombineInternal[T]( def getMaxItemsTracked: Int = maxItemsTracked def updateMaxItemsTracked(combineSizeSpecified: Boolean, newMaxItemsTracked: Int): Unit = { - if (!combineSizeSpecified) { + if (!combineSizeSpecified && newMaxItemsTracked != ApproxTopK.VOID_MAX_ITEMS_TRACKED) { // check size if (this.maxItemsTracked == ApproxTopK.VOID_MAX_ITEMS_TRACKED) { // If buffer's maxItemsTracked VOID_MAX_ITEMS_TRACKED, it means the buffer is a placeholder diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala index 2ec8c584f2033..46acffbaffe0b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala @@ -595,6 +595,16 @@ class ApproxTopKSuite extends SharedSparkSession { } } + test("SPARK-58069: merge an empty approx_top_k_combine buffer") { + val initializedBuffer = new CombineInternal[Any]( + new ApproxTopKAggregateBuffer[Any](new ItemsSketch[Any](128), 0L), + StringType, + 100) + initializedBuffer.updateMaxItemsTracked( + combineSizeSpecified = false, ApproxTopK.VOID_MAX_ITEMS_TRACKED) + assert(initializedBuffer.getMaxItemsTracked == 100) + } + test("SPARK-58069: combine sketches with empty shuffle partitions") { val sketches = sql( "SELECT approx_top_k_accumulate(id) AS sketch FROM range(10) GROUP BY id")