[SPARK-58069][SQL][FOLLOWUP] Handle empty approx_top_k combine buffers

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala

@@ -723,9 +723,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)
 
@@ -755,12 +766,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)
   }
 }

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,43 @@ 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.serialize().sameElements(buffer.serialize()))
+    }
+  }
+
+  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")
+      .repartition(20)
+    val combined = sketches.selectExpr("approx_top_k_combine(sketch) AS sketch")
+
+    checkAnswer(
+      combined.selectExpr("inline(approx_top_k_estimate(sketch, 10))"),
+      (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) =>