spark.comet.batchSize below 8192 makes CometConf fail to initialize on every executor

[andygrove]

Describe the bug

Setting spark.comet.batchSize below 8192 leaves CometConf unable to initialize on every executor, so every task that runs Comet code there fails.

createWithDefault passes the default through the entry's validators (CometConf.scala#L1335-L1336), and it does that inside CometConf's static initializer. The validator on spark.comet.shuffle.jvm.batchSize reads another config: the value must not be larger than COMET_BATCH_SIZE.get() (CometConf.scala#L688-L699). That reads SQLConf.get at whatever moment CometConf is loaded, so the default of 8192 is checked against the spark.comet.batchSize of whichever conf is current at that point.

On the driver, CometDriverPlugin.init loads CometConf before any session exists, so the check sees the default batch size and passes. Nothing on an executor loads CometConf at startup, so it is first loaded inside a task (CometBypassMergeSortShuffleWriter.java#L148 in the repro below), where SQLConf.get reads the session's confs from the task's local properties. The first task fails with ExceptionInInitializerError, and every later one on that executor with NoClassDefFoundError: Could not initialize class org.apache.comet.CometConf$. Setting spark.comet.shuffle.jvm.batchSize=4096 as well doesn't help, because it is the default that fails the check.

An application that registers CometSparkSessionExtensions through spark.sql.extensions without the plugin hits the same thing on the driver. The first query fails with ExceptionInInitializerError, and every later query in the session fails with NoClassDefFoundError, whether or not Comet would accelerate it.

The check was added in #3540, so this affects 0.14.0 onwards.

Steps to reproduce

Run on local-cluster[1,1,2048] with spark.plugins=org.apache.spark.CometPlugin, spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager, off-heap memory enabled, spark.comet.shuffle.mode=jvm and spark.comet.batchSize=4096:

spark.range(0, 1000, 1, 2).selectExpr("id", "cast(id as string) s").repartition(2).collect()

The job aborts after four task failures. The executor log has:

java.lang.ExceptionInInitializerError: null
	at org.apache.spark.sql.comet.execution.shuffle.CometBypassMergeSortShuffleWriter.<init>(CometBypassMergeSortShuffleWriter.java:148)
	at org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager.getWriter(CometShuffleManager.scala:258)
	at org.apache.spark.shuffle.ShuffleWriteProcessor.write(ShuffleWriteProcessor.scala:56)
	at org.apache.spark.scheduler.ShuffleMapTask.runTask(ShuffleMapTask.scala:111)
	...
Caused by: java.lang.IllegalArgumentException: '8192' in spark.comet.shuffle.jvm.batchSize is invalid. Should not be larger than batch size `spark.comet.batchSize`

The same query succeeds with spark.comet.batchSize unset. Local mode doesn't show the executor case, because the driver has already loaded CometConf in the same JVM. It does show the driver case: use local[1] with spark.sql.extensions=org.apache.comet.CometSparkSessionExtensions, no plugin and spark.comet.batchSize=4096, then run spark.range(100).selectExpr("id + 1").collect().

I reproduced both on main at 31b3819 with Spark 4.1.3.

Expected behavior

CometConf initializes whatever conf is current when it loads, and a spark.comet.batchSize below 8192 works. A constraint between two configs belongs where the values are read, not in a validator that runs during class initialization.

Additional context

Found while fixing #6259, whose PR adds a check that the same value is positive but leaves this one as it is.

[comphead]

hm, so this only for jvm shuffle, how come the tests didn't capture it