andygrove opened a new issue, #5510:
URL: https://github.com/apache/datafusion-comet/issues/5510
## 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
