[SPARK-37203][SQL] Fix NotSerializableException when observe with TypedImperativeAggregate

core/src/main/scala/org/apache/spark/util/AccumulatorV2.scala

@@ -157,6 +157,9 @@ abstract class AccumulatorV2[IN, OUT] extends Serializable {
    */
   def value: OUT
 
+  // We assume that serialization of AccumulatorV2 runs on executor is not necessary.
+  protected def withBufferSerialized(): AccumulatorV2[IN, OUT] = this
+
   // Called by Java when serializing an object
   final protected def writeReplace(): Any = {
     if (atDriverSide) {
@@ -179,7 +182,7 @@ abstract class AccumulatorV2[IN, OUT] extends Serializable {
       }
       copyAcc
     } else {
-      this
+      withBufferSerialized()
     }
   }
 

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

@@ -276,7 +276,7 @@ object ApproximatePercentile {
   }
 
   /**
-   * Serializer  for class [[PercentileDigest]]
+   * Serializer for class [[PercentileDigest]]
    *
    * This class is thread safe.
    */

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

@@ -621,6 +621,17 @@ abstract class TypedImperativeAggregate[T] extends ImperativeAggregate {
     buffer(mutableAggBufferOffset) = serialize(getBufferObject(buffer))
   }
 
+  /**
+   * In-place replaces SparkSQL internally supported underlying storage format (BinaryType),
+   * with the aggregation buffer object stored at buffer's index `mutableAggBufferOffset`.
+   *
+   * This is only called when AggregatingAccumulator running on driver, after the framework
+   * shuffle in aggregate buffers.
+   */
+  final def deserializeAggregateBufferInPlace(buffer: InternalRow): Unit = {
+    buffer(mutableAggBufferOffset) = deserialize(buffer.getBinary(inputAggBufferOffset))
+  }
+
   /**
    * Merge an input buffer into the aggregation buffer, where both buffers contain the deserialized
    * java object. This function is used by aggregating accumulators.

sql/core/src/main/scala/org/apache/spark/sql/execution/AggregatingAccumulator.scala

@@ -156,6 +156,15 @@ class AggregatingAccumulator private(
         case agg: AggregatingAccumulator =>
           val buffer = getOrCreateBuffer()
           val otherBuffer = agg.buffer
+          // If AggregatingAccumulator runs on driver,
+          // we should deserialize all TypedImperativeAggregate.
+          if (isAtDriverSide) {
+            var i = 0
+            while (i < typedImperatives.length) {
+              typedImperatives(i).deserializeAggregateBufferInPlace(otherBuffer)
+              i += 1
+            }
+          }
           mergeProjection.target(buffer)(joinedRow.withRight(otherBuffer))
           var i = 0
           while (i < imperatives.length) {
@@ -188,6 +197,18 @@ class AggregatingAccumulator private(
     resultProjection(input)
   }
 
+  override def withBufferSerialized(): AggregatingAccumulator = {
+    if (!isAtDriverSide) {
+      var i = 0
+      // AggregatingAccumulator runs on executor, we should serialize all TypedImperativeAggregate.
+      while (i < typedImperatives.length) {
+        typedImperatives(i).serializeAggregateBufferInPlace(buffer)
+        i += 1
+      }
+    }
+    this
+  }
+
   /**
    * Get the output schema of the aggregating accumulator.
    */

sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala

@@ -754,6 +754,17 @@ class DatasetSuite extends QueryTest
     assert(err2.getMessage.contains("Name must not be empty"))
   }
 
+  test("SPARK-37203: Fix NotSerializableException when observe with percentile_approx") {
+    val namedObservation = Observation("named")
+
+    val df = spark.range(100)
+    val observed_df = df.observe(
+      namedObservation, percentile_approx($"id", lit(0.5), lit(100)).as("percentile_approx_val"))
+
+    observed_df.collect()
+    assert(namedObservation.get === Map("percentile_approx_val" -> 49))
+  }
+
   test("sample with replacement") {
     val n = 100
     val data = sparkContext.parallelize(1 to n, 2).toDS()

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryListenerSuite.scala

@@ -417,7 +417,8 @@ class StreamingQueryListenerSuite extends StreamTest with BeforeAndAfter {
         min($"value").as("min_val"),
         max($"value").as("max_val"),
         sum($"value").as("sum_val"),
-        count(when($"value" % 2 === 0, 1)).as("num_even"))
+        count(when($"value" % 2 === 0, 1)).as("num_even"),
+        percentile_approx($"value", lit(0.5), lit(100)).as("percentile_approx_val"))
       .observe(
         name = "other_event",
         avg($"value").cast("int").as("avg_val"))
@@ -444,15 +445,15 @@ class StreamingQueryListenerSuite extends StreamTest with BeforeAndAfter {
         AddData(inputData, 1, 2),
         AdvanceManualClock(100),
         checkMetrics { metrics =>
-          assert(metrics.get("my_event") === Row(1, 2, 3L, 1L))
+          assert(metrics.get("my_event") === Row(1, 2, 3L, 1L, 1))
           assert(metrics.get("other_event") === Row(1))
         },
 
         // Batch 2
         AddData(inputData, 10, 30, -10, 5),
         AdvanceManualClock(100),
         checkMetrics { metrics =>
-          assert(metrics.get("my_event") === Row(-10, 30, 35L, 3L))
+          assert(metrics.get("my_event") === Row(-10, 30, 35L, 3L, 5))
           assert(metrics.get("other_event") === Row(8))
         },