Skip to content

fix: cap JVM shuffle batch size at read time instead of during CometConf init - #6367

Closed
comphead wants to merge 1 commit into
apache:mainfrom
comphead:jvm-shuffle-batch-size-6286
Closed

comphead wants to merge 1 commit into
apache:mainfrom
comphead:jvm-shuffle-batch-size-6286

Conversation

@comphead

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #6286.

Rationale for this change

createWithDefault runs an entry's checks on its default while CometConf initializes. The check on spark.comet.shuffle.jvm.batchSize compared the value with COMET_BATCH_SIZE.get(), which reads whatever SQLConf.get returns at that moment. An executor first loads CometConf inside a task, where that is the session's conf, so a spark.comet.batchSize below 8192 failed the default and left CometConf unable to initialize for the rest of the executor's life. The driver hits the same failure when Comet is loaded only through spark.sql.extensions.

The existing tests missed it for two reasons:

  • The suites run in local mode. CometConf is initialized once, on the driver, during suite setup, while SQLConf.get still returns the defaults. The tests that set a smaller batch size through withSQLConf never run the initializer again, and no suite starts separate executor JVMs.
  • When a config is read, ConfigEntryWithDefault.get returns the default without running the checks. So those tests never compared the default with their smaller batch size, and JVM shuffle kept writing batches of up to 8192 rows under them. The limit the check was meant to enforce never applied to the default.

What changes are included in this PR?

  1. The check on spark.comet.shuffle.jvm.batchSize is removed. The new CometConf.jvmShuffleBatchSize returns that value capped at spark.comet.batchSize. CometDiskBlockWriter and SpillWriter, the only two readers, call it.
  2. The scaladoc of checkValue now says that a check also runs on the default while CometConf initializes, so it must not read other configs.
  3. The config description, the memory tuning guide and the JVM shuffle contributor guide describe the cap.

Capping the value where it is read, rather than failing there, is what lets the default work with a smaller spark.comet.batchSize. Two things behave differently as a result:

  • A spark.comet.shuffle.jvm.batchSize larger than spark.comet.batchSize is now capped. Since 0.14.0 it failed the task.
  • Where a spark.comet.batchSize below 8192 already worked, as in local mode, JVM shuffle now writes batches of at most that many rows instead of 8192.

This overlaps with #6287, which adds a positivity check to the same entry on the line above the one removed here. Whichever PR merges second should keep both.

How are these changes tested?

Two new tests in CometConfSuite:

  • JVM shuffle batch size is capped at spark.comet.batchSize covers the default, the issue's case of a 4096 batch size with the default JVM shuffle batch size, and explicit JVM shuffle batch sizes above and below the batch size.
  • config checks do not read other configs runs every registered entry's checks on its default, as createWithDefault does, under SQLConf.withExistingConf with a conf that fails the test on any getConfString. That is the rule spark.comet.batchSize below 8192 makes CometConf fail to initialize on every executor #6286 broke, and the test applies it to every entry, not only this one. The check this PR removes read spark.comet.batchSize through that conf.

There is no end-to-end test. The executor case needs CometConf to be loaded for the first time inside a task, which takes separate executor JVMs, as in the issue's local-cluster reproduction. In Comet's local-mode suites, CometConf is already loaded by the time a test sets a smaller batch size, so neither the executor case nor the driver case can be reproduced there.

…onf init

The check on spark.comet.shuffle.jvm.batchSize read spark.comet.batchSize while CometConf initialized. An executor first loads CometConf inside a task, so a batch size below 8192 failed the default and left CometConf unable to initialize there. Cap the JVM shuffle batch size where it is read instead, and test that no config check reads another config.
@github-actions github-actions Bot added bug Something isn't working area:shuffle Shuffle (JVM and native) labels Sep 29, 2026
@comphead
comphead requested a review from andygrove September 29, 2026 00:59
@andygrove

andygrove commented Sep 29, 2026 •

Copy link
Copy Markdown
Member

@comphead Is this a duplicate of #6286?

I meant #6366

@comphead

Copy link
Copy Markdown
Contributor Author

correct, we made PRs at almost the same time

@comphead comphead closed this Sep 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:shuffle Shuffle (JVM and native) bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

2 participants