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..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 @@ -20,8 +20,12 @@ 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.{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} @@ -575,6 +579,33 @@ class ApproxTopKSuite extends SharedSparkSession { "CAST('13:00:00.123' AS TIME(3))")) ) + test("SPARK-58069: serialize an empty approx_top_k_combine buffer") { + 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) + } + } + + 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") + + 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 gridTest("SPARK-52798: same type, same size, specified combine size - success")(itemsWithTopK) { case (input, expected) =>