apache / apache/datafusion-comet
Comet native broadcast fails under spark.kryo.registrationRequired=true
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 198
Description
## 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`:
```sql
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 https://github.com/apache/datafusion-comet/pull/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.
Contributor guide
Research direction
Start in CometBroadcastExchangeExec.scala and inspect the serializer inputs from Utils.arrowBytesKryoClasses, then compare the proposed registrator setup with SparkContext.scala:478 and PluginContainer construction at line 574. Add the standalone registration and startup-warning coverage described in the issue, and run a broadcast-shaped regression test equivalent to “Comet broadcast exchange survives Kryo with registration required” in CometInMemoryCacheKryoSuite.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- scala
- Domain
- distributed-systems, testing-qa
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100