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

[beliefer]

What changes were proposed in this pull request?

Currently,

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()
namedObservation.get

throws exception as follows:

15:16:27.994 ERROR org.apache.spark.util.Utils: Exception encountered
java.io.NotSerializableException: org.apache.spark.sql.catalyst.expressions.aggregate.ApproximatePercentile$PercentileDigest
	at java.io.ObjectOutputStream.writeObject0(ObjectOutputStream.java:1184)
	at java.io.ObjectOutputStream.defaultWriteFields(ObjectOutputStream.java:1548)
	at java.io.ObjectOutputStream.writeSerialData(ObjectOutputStream.java:1509)
	at java.io.ObjectOutputStream.writeOrdinaryObject(ObjectOutputStream.java:1432)
	at java.io.ObjectOutputStream.writeObject0(ObjectOutputStream.java:1178)
	at java.io.ObjectOutputStream.writeArray(ObjectOutputStream.java:1378)
	at java.io.ObjectOutputStream.writeObject0(ObjectOutputStream.java:1174)
	at java.io.ObjectOutputStream.defaultWriteFields(ObjectOutputStream.java:1548)
	at java.io.ObjectOutputStream.writeSerialData(ObjectOutputStream.java:1509)
	at java.io.ObjectOutputStream.writeOrdinaryObject(ObjectOutputStream.java:1432)
	at java.io.ObjectOutputStream.writeObject0(ObjectOutputStream.java:1178)
	at java.io.ObjectOutputStream.defaultWriteFields(ObjectOutputStream.java:1548)
	at java.io.ObjectOutputStream.writeSerialData(ObjectOutputStream.java:1509)
	at java.io.ObjectOutputStream.writeOrdinaryObject(ObjectOutputStream.java:1432)
	at java.io.ObjectOutputStream.writeObject0(ObjectOutputStream.java:1178)
	at java.io.ObjectOutputStream.writeObject(ObjectOutputStream.java:348)
	at org.apache.spark.scheduler.DirectTaskResult.$anonfun$writeExternal$2(TaskResult.scala:55)
	at org.apache.spark.scheduler.DirectTaskResult.$anonfun$writeExternal$2$adapted(TaskResult.scala:55)
	at scala.collection.Iterator.foreach(Iterator.scala:943)
	at scala.collection.Iterator.foreach$(Iterator.scala:943)
	at scala.collection.AbstractIterator.foreach(Iterator.scala:1431)
	at scala.collection.IterableLike.foreach(IterableLike.scala:74)
	at scala.collection.IterableLike.foreach$(IterableLike.scala:73)
	at scala.collection.AbstractIterable.foreach(Iterable.scala:56)
	at org.apache.spark.scheduler.DirectTaskResult.$anonfun$writeExternal$1(TaskResult.scala:55)
	at scala.runtime.java8.JFunction0$mcV$sp.apply(JFunction0$mcV$sp.java:23)
	at org.apache.spark.util.Utils$.tryOrIOException(Utils.scala:1434)
	at org.apache.spark.scheduler.DirectTaskResult.writeExternal(TaskResult.scala:51)
	at java.io.ObjectOutputStream.writeExternalData(ObjectOutputStream.java:1459)
	at java.io.ObjectOutputStream.writeOrdinaryObject(ObjectOutputStream.java:1430)
	at java.io.ObjectOutputStream.writeObject0(ObjectOutputStream.java:1178)
	at java.io.ObjectOutputStream.writeObject(ObjectOutputStream.java:348)
	at org.apache.spark.serializer.JavaSerializationStream.writeObject(JavaSerializer.scala:44)
	at org.apache.spark.serializer.JavaSerializerInstance.serialize(JavaSerializer.scala:101)
	at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:616)
	at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149)
	at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624)
	at java.lang.Thread.run(Thread.java:748)

This PR will fix the issue. After the change,
assert(namedObservation.get === Map("percentile_approx_val" -> 49))
java.io.NotSerializableException will not happen.

Why are the changes needed?

Fix NotSerializableException when observe with TypedImperativeAggregate.

Does this PR introduce any user-facing change?

No. This PR change the implement of AggregatingAccumulator who uses serialize and deserialize of TypedImperativeAggregate now.

How was this patch tested?

New tests.

[SparkQA]

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49343/

[SparkQA]

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49343/

[SparkQA]

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49345/

[SparkQA]

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49345/

[SparkQA]

Test build #144873 has finished for PR 34474 at commit 1eacf82.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[MaxGekk]

hmm, there are special methods for (de-)serialization of PercentileDigest:

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

Lines 212 to 218 in 08123a3

	override def serialize(obj: PercentileDigest): Array[Byte] = {
	ApproximatePercentile.serializer.serialize(obj)
	}
	override def deserialize(bytes: Array[Byte]): PercentileDigest = {
	ApproximatePercentile.serializer.deserialize(bytes)
	}

Do you know why the methods are used instead of extending of Serializable?

cc @cloud-fan

[hvanhovell]

Can we just add the proper logic to observe? Adding serialization will fix this particular problem, but it won't work for other aggregates.

[srowen]

I agree if that's possible - I couldn't tell if this was something observe() should never have to pull back and does accidentally, or whether it is necessary to return metrics about percentile_approx somehow

[hvanhovell]

I probably forgot to add it...

[beliefer]

@MaxGekk @hvanhovell @srowen I updated the code and it will fix the issue for AggregatingAccumulator takes any TypedImperativeAggregate.

[SparkQA]

Test build #144875 has finished for PR 34474 at commit a2f1fbb.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[srowen]

Looking at the implementation, it's never meant to be serialized by Java or Kryo, but is serialized 'manually' and internally by Spark. I think it must be related to observe() here, that somehow this object gets caught up in what is returned? I don't know a lot about observe().

It's possible the answer is that observe() should do something different, but not sure.

We can make this Serializable, but then need to implement methods for the Java serializer that would reuse the existing serializer code. That seems plausible.

[cloud-fan]

why do we need to create temp buffer here?

[SparkQA]

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49368/

[cloud-fan]

// Serialize the buffer of this accumulator before sending back this accumulator to the driver.
// By default this method does nothing.

[cloud-fan]

I think we can do assert(!isAtDriverSide)

[cloud-fan]

stored at buffer's index mutableAggBufferOffset describes the binary, should be put before with the aggregation buffer object ...

[cloud-fan]

This is nothing to do with shuffle

[cloud-fan]

We don't need to add a new code. We can just change the existing code

while (i < typedImperatives.length) {
  typedImperatives(i).mergeBuffersObjects(buffer, otherBuffer)
  i += 1
}

to

if (atDriverSide) {
  while (i < typedImperatives.length) {
    // The input buffer stores serialized data
    typedImperatives(i).merge(buffer, otherBuffer)
    i += 1
  }
} else {
  while (i < typedImperatives.length) {
    // The input buffer stores deserialized object
    typedImperatives(i). mergeBuffersObjects(buffer, otherBuffer)
    i += 1
  }
}

[SparkQA]

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49368/

[SparkQA]

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49369/

[SparkQA]

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49369/

[SparkQA]

Test build #144898 has finished for PR 34474 at commit 6c2f71c.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[cloud-fan]

can we test a DataFrame with no data?

[beliefer]

OK

[SparkQA]

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49373/

[SparkQA]

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49373/

[SparkQA]

Test build #144899 has finished for PR 34474 at commit 84f9657.

• This patch fails from timeout after a configured wait of 500m.
• This patch merges cleanly.
• This patch adds no public classes.

[SparkQA]

Test build #144903 has finished for PR 34474 at commit 5e7afac.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[SparkQA]

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49384/

[SparkQA]

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/49384/

[SparkQA]

Test build #144913 has finished for PR 34474 at commit acddbf3.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[cloud-fan]

@beliefer can you update the PR description? Yes. PercentileDigest extends Serializable. this is wrong now.

[beliefer]

@beliefer can you update the PR description? Yes. PercentileDigest extends Serializable. this is wrong now.

OK

[cloud-fan]

Yes. This PR let AggregatingAccumulator uses serialize and deserialize of TypedImperativeAggregate.

Is it a user-facing change?

[beliefer]

Yes. This PR let AggregatingAccumulator uses serialize and deserialize of TypedImperativeAggregate.

Is it a user-facing change?

No

[beliefer]

@cloud-fan Thanks for your hard work！@srowen @MaxGekk @hvanhovell ， thank you for review.

[cloud-fan]

thanks, merging to master/3.2/3.1!

[MaxGekk]

I have an example in which we get NPE here because buffer is null, see SPARK-40535

[beliefer]

I will investigate it.