Comet native broadcast fails under spark.kryo.registrationRequired=true

[andygrove]

Describe the bug

CometBroadcastExchangeExec fails outright when spark.kryo.registrationRequired=true, which makes Kryo reject any class it has not been told about. This has nothing to do with the in-memory cache work; it reproduces on main with only Comet's broadcast exchange in play.

CometBroadcastExchangeExec.relationFuture broadcasts the result of Utils.coalesceBroadcastBatches, which is an Array[ChunkedByteBuffer]. Spark registers ChunkedByteBuffer in KryoSerializer.toRegister, but not an array of them, and Comet has no Kryo registrator of its own, so the broadcast throws:

java.util.concurrent.ExecutionException: java.lang.IllegalArgumentException:
  Class is not registered: org.apache.spark.util.io.ChunkedByteBuffer[]
Note: To register this class use: kryo.register(org.apache.spark.util.io.ChunkedByteBuffer[].class);
	at com.esotericsoftware.kryo.Kryo.getRegistration(Kryo.java:503)
	at com.esotericsoftware.kryo.util.DefaultClassResolver.writeClass(DefaultClassResolver.java:97)
	at com.esotericsoftware.kryo.Kryo.writeClassAndObject(Kryo.java:645)
	at org.apache.spark.serializer.KryoSerializationStream.writeObject(KryoSerializer.scala:285)
	...
	at org.apache.spark.sql.comet.CometBroadcastExchangeExec.doExecuteBroadcast(CometBroadcastExchangeExec.scala:232)

Because the message names a Spark utility class rather than anything Comet-shaped, it is hard to attribute to Comet at all.

Steps to reproduce

With spark.serializer=org.apache.spark.serializer.KryoSerializer and spark.kryo.registrationRequired=true, run any query that produces a CometBroadcastExchange:

SELECT /*+ BROADCAST(b) */ a._1, b._2 FROM tbl_a a JOIN tbl_b b ON a._1 = b._1

The plan contains CometBroadcastHashJoin over CometBroadcastExchange, and collect() fails as above. Disabling spark.comet.exec.broadcastExchange.enabled avoids it.

Expected behavior

A native broadcast should work under spark.kryo.registrationRequired=true, as Spark's own broadcast does.

Additional context

A fix is currently carried inside #5051, which needed a Kryo registrator anyway for its cached batch format and covers this case because the cache write path hands back the same ChunkedByteBuffer type. @viirya raised the cache half of that in review; this half fell out of checking whether the registrator needed anything beyond the cached batch.

This issue exists so that fix can be split out and land on its own, since it is a pre-existing bug rather than one that PR introduces, and it should not have to wait on an experimental default-off feature.

What a standalone fix needs:

• Register ChunkedByteBuffer, Array[ChunkedByteBuffer], Array[ByteBuffer] and the heap ByteBuffer class. In PR 5051 these live in Utils.arrowBytesKryoClasses, next to serializeBatches/serializeBatchColumns, which are what produce them.
• A KryoRegistrator implementation for users to point spark.kryo.registrator at.
• Note that Comet cannot install this itself. KryoSerializer reads spark.kryo.registrator into a val in its constructor, and SparkEnv.create builds it at SparkContext.scala:478, whereas PluginContainer is constructed at line 574. Unlike spark.sql.cache.serializer — a StaticSQLConf read lazily from the session, which is why CometDriverPlugin.maybeSetCacheSerializer can inject it — this one is already captured before any plugin runs. Setting it from the driver plugin would reach executors but not the driver's own SparkEnv, so it would appear to work in cluster mode and fail in local mode. So it has to be documented, with a startup warning when Kryo, registrationRequired, and a missing registrator are combined.
• A regression test. PR 5051 has one as Comet broadcast exchange survives Kryo with registration required in CometInMemoryCacheKryoSuite; standalone it belongs somewhere broadcast-shaped instead.

Worth checking as part of this whether any other Comet path hands Spark's serializer an unregistered class. getByteArrayRdd(child).collect() in the same file returns (Long, ChunkedByteBuffer) tuples through the task-result serializer; Tuple2 comes from Chill and ChunkedByteBuffer from Spark, so that one appears covered, but I have not swept the shuffle paths.