You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/30 13:20:22 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6366-c7aa38bd0b1e7dde39cf650da43a4452a7a5ea4c created (now 9f68a4144f)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6366-c7aa38bd0b1e7dde39cf650da43a4452a7a5ea4c
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 9f68a4144f fix: cap spark.comet.shuffle.jvm.batchSize at spark.comet.batchSize where it is read (#6366)
This branch includes the following new commits:
new 9f68a4144f fix: cap spark.comet.shuffle.jvm.batchSize at spark.comet.batchSize where it is read (#6366)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: fix: cap spark.comet.shuffle.jvm.batchSize at spark.comet.batchSize where it is read (#6366)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6366-c7aa38bd0b1e7dde39cf650da43a4452a7a5ea4c
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 9f68a4144fdefc70ec26f3d80b0e4eafe72a8df4
Author: Andy Grove <ag...@apache.org>
AuthorDate: Wed Sep 30 13:20:00 2026 +0000
fix: cap spark.comet.shuffle.jvm.batchSize at spark.comet.batchSize where it is read (#6366)
* fix: validate spark.comet.shuffle.jvm.batchSize and spark.comet.exec.memoryPool
spark.comet.shuffle.jvm.batchSize was only checked to be no larger than
spark.comet.batchSize, so 0 was accepted. The JVM shuffle writer then
passed 0 to process_sorted_row_partition, whose loop never advanced, and
because the loop runs inside a JNI call the task hung and could not be
killed. The value must now be positive.
spark.comet.exec.memoryPool had no check, so a misspelled value, or one
that differed only in case, reached native code, and in off-heap mode
every task failed with "Unsupported memory pool type" when it created a
native plan. The value is now lowercased and must be fair_unified or
greedy_unified, the same way spark.comet.shuffle.mode is handled.
Closes #6259
* fix: cap spark.comet.shuffle.jvm.batchSize at spark.comet.batchSize where it is read
The entry for spark.comet.shuffle.jvm.batchSize had a validator that
compared the value with COMET_BATCH_SIZE.get(). createWithDefault runs
the validators on the default inside CometConf's static initializer, so
the check read SQLConf.get at whatever moment CometConf was loaded. On
an executor that is inside the first task, where SQLConf.get holds the
session's confs, so with spark.comet.batchSize below 8192 the default
failed the check, CometConf threw ExceptionInInitializerError, and every
later task on that executor failed with NoClassDefFoundError. A session
that registered CometSparkSessionExtensions without the plugin hit the
same failure on the driver.
The validator is removed, and the JVM shuffle writers now read the batch
size through CometConf.jvmShuffleBatchSize, which caps
spark.comet.shuffle.jvm.batchSize at spark.comet.batchSize. Capping
rather than failing keeps a job that only lowers spark.comet.batchSize
working, and it keeps what the check was for: the JVM shuffle never
writes batches larger than spark.comet.batchSize. No other CometConf
entry reads another config while CometConf initializes.
Closes #6286
---
docs/source/contributor-guide/jvm_shuffle.md | 3 +-
docs/source/user-guide/latest/tuning/memory.md | 2 +-
.../execution/shuffle/CometDiskBlockWriter.java | 2 +-
.../sql/comet/execution/shuffle/SpillWriter.java | 2 +-
.../main/scala/org/apache/comet/CometConf.scala | 22 +++++---
.../scala/org/apache/comet/CometConfSuite.scala | 59 ++++++++++++++++++++++
.../comet/exec/CometColumnarShuffleSuite.scala | 17 +++++++
7 files changed, 97 insertions(+), 10 deletions(-)
diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md
index 1226b4c82b..11014f3490 100644
--- a/docs/source/contributor-guide/jvm_shuffle.md
+++ b/docs/source/contributor-guide/jvm_shuffle.md
@@ -201,7 +201,8 @@ writes the same Arrow IPC block format as native shuffle, so direct read applies
That mode exists only so the Spark SQL tests can run against Comet. See
[Memory Management](memory_management.md).
- Row count still triggers spilling in either mode. `CometDiskBlockWriter` spills at
- `min(spark.comet.shuffle.jvm.spillThreshold, spark.comet.shuffle.jvm.batchSize)`.
+ `min(spark.comet.shuffle.jvm.spillThreshold, batch size)`, where the batch size is
+ `spark.comet.shuffle.jvm.batchSize` capped at `spark.comet.batchSize`.
`CometShuffleExternalSorter` spills at `spark.comet.shuffle.jvm.spillThreshold` alone, which
defaults to `Int.MaxValue`, so on the sort path that trigger is effectively off by default.
- `CometDiskBlockWriter` coordinates spilling across all partition writers (largest first)
diff --git a/docs/source/user-guide/latest/tuning/memory.md b/docs/source/user-guide/latest/tuning/memory.md
index 9be6c72ae7..9b38348c42 100644
--- a/docs/source/user-guide/latest/tuning/memory.md
+++ b/docs/source/user-guide/latest/tuning/memory.md
@@ -208,7 +208,7 @@ if you see frequent spilling or out-of-memory errors on wide tables; increase it
narrow tables when memory is plentiful.
`spark.comet.shuffle.jvm.batchSize` controls the batch size used when the JVM columnar shuffle writer
-flushes sorted spill files. It must not exceed `spark.comet.batchSize`.
+flushes sorted spill files. A value larger than `spark.comet.batchSize` is capped at `spark.comet.batchSize`.
## Limiting Spill Disk Usage
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java
index 105452a434..839a69bd76 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java
@@ -143,7 +143,7 @@ public final class CometDiskBlockWriter {
this.file = file;
this.tracingEnabled = tracingEnabled;
- this.columnarBatchSize = (int) CometConf$.MODULE$.COMET_SHUFFLE_JVM_BATCH_SIZE().get();
+ this.columnarBatchSize = CometConf$.MODULE$.jvmShuffleBatchSize();
this.compressionCodec = CometConf$.MODULE$.COMET_SHUFFLE_COMPRESSION_CODEC().get();
this.compressionLevel = (int) CometConf$.MODULE$.COMET_SHUFFLE_COMPRESSION_ZSTD_LEVEL().get();
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java
index 1683af3e35..a28f89e1a0 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java
@@ -196,7 +196,7 @@ public abstract class SpillWriter {
long currentChecksum = checksumEnabled ? checksum : 0L;
long start = System.nanoTime();
- int batchSize = (int) CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE().get();
+ int batchSize = CometConf.jvmShuffleBatchSize();
long[] results =
nativeLib.writeSortedFileNative(
addresses,
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 71dc902baf..14dbd3d78b 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -692,16 +692,26 @@ object CometConf extends ShimCometConf {
conf("spark.comet.shuffle.jvm.batchSize")
.withAlternative("spark.comet.columnar.shuffle.batch.size")
.category(CATEGORY_SHUFFLE)
- .doc("Batch size when writing out sorted spill files on the native side. Note that " +
- "this should not be larger than batch size (i.e., `spark.comet.batchSize`). Otherwise " +
- "it will produce larger batches than expected in the native operator after shuffle.")
+ .doc(
+ "Batch size when writing out sorted spill files on the native side. A value larger " +
+ "than the batch size (i.e., `spark.comet.batchSize`) is capped at the batch size, so " +
+ "that the native operators after the shuffle do not receive larger batches than " +
+ "expected.")
.intConf
.checkValue(v => v > 0, "Batch size must be positive")
- .checkValue(
- v => v <= COMET_BATCH_SIZE.get(),
- "Should not be larger than batch size `spark.comet.batchSize`")
.createWithDefault(8192)
+ /**
+ * The batch size that the JVM columnar shuffle writes with: `spark.comet.shuffle.jvm.batchSize`
+ * capped at `spark.comet.batchSize`. The cap is applied where the values are read rather than
+ * in a validator on the entry, because validators also check the default while `CometConf`
+ * initializes. On an executor that happens inside a task, under the session's confs.
+ */
+ def jvmShuffleBatchSize(): Int = jvmShuffleBatchSize(SQLConf.get)
+
+ def jvmShuffleBatchSize(conf: SQLConf): Int =
+ math.min(COMET_SHUFFLE_JVM_BATCH_SIZE.get(conf), COMET_BATCH_SIZE.get(conf))
+
val COMET_SHUFFLE_NATIVE_WRITE_BUFFER_SIZE: ConfigEntry[Long] =
conf("spark.comet.shuffle.native.writeBufferSize")
.withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.writeBufferSize")
diff --git a/spark/src/test/scala/org/apache/comet/CometConfSuite.scala b/spark/src/test/scala/org/apache/comet/CometConfSuite.scala
index 679d7519bc..a5314fc39a 100644
--- a/spark/src/test/scala/org/apache/comet/CometConfSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometConfSuite.scala
@@ -206,6 +206,65 @@ class CometConfSuite extends AnyFunSuite {
assert(entry.get(conf) == 1)
}
+ test("CometConf initializes when spark.comet.batchSize is below the JVM shuffle batch size") {
+ // Defines its own copy of every org.apache.comet class, so that loading CometConf through it
+ // runs CometConf's initializer again. Everything else comes from the parent.
+ val loader = new ClassLoader(getClass.getClassLoader) {
+ override def loadClass(name: String, resolve: Boolean): Class[_] = {
+ if (!name.startsWith("org.apache.comet.")) {
+ super.loadClass(name, resolve)
+ } else {
+ getClassLoadingLock(name).synchronized {
+ Option(findLoadedClass(name)).getOrElse {
+ val in = getParent.getResourceAsStream(name.replace('.', '/') + ".class")
+ if (in == null) throw new ClassNotFoundException(name)
+ val bytes =
+ try in.readAllBytes()
+ finally in.close()
+ defineClass(name, bytes, 0, bytes.length)
+ }
+ }
+ }
+ }
+ }
+
+ // An executor first loads CometConf inside a task, where SQLConf.get holds the session's
+ // confs. The initializer must not depend on them.
+ val conf = new SQLConf
+ conf.setConfString(CometConf.COMET_BATCH_SIZE.key, "4096")
+ val cometConfClass =
+ try {
+ SQLConf.withExistingConf(conf) {
+ // scalastyle:off classforname
+ Class.forName(CometConf.getClass.getName, true, loader)
+ // scalastyle:on classforname
+ }
+ } catch {
+ // ScalaTest aborts the whole run on this error, so report it as a failure instead.
+ case e: ExceptionInInitializerError => fail("CometConf failed to initialize", e.getCause)
+ }
+ assert(cometConfClass ne CometConf.getClass)
+ }
+
+ test("JVM shuffle batch size is capped at spark.comet.batchSize where it is read") {
+ val conf = new SQLConf
+ val entry = CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE
+ conf.setConfString(CometConf.COMET_BATCH_SIZE.key, "4096")
+
+ // The JVM shuffle writers read the conf that is current on the task thread. The default is
+ // capped there, while the entry itself keeps it.
+ SQLConf.withExistingConf(conf) {
+ assert(entry.get() == 8192)
+ assert(CometConf.jvmShuffleBatchSize() == 4096)
+ }
+
+ // A larger value set explicitly is capped too, and a smaller one is used as it is.
+ conf.setConfString(entry.key, "16384")
+ assert(CometConf.jvmShuffleBatchSize(conf) == 4096)
+ conf.setConfString(entry.key, "1024")
+ assert(CometConf.jvmShuffleBatchSize(conf) == 1024)
+ }
+
test("memory pool type is accepted in any case and lowercased") {
val conf = new SQLConf
val entry = CometConf.COMET_OFFHEAP_MEMORY_POOL_TYPE
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
index 157cf972d6..c6a6805c72 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
@@ -30,6 +30,7 @@ import org.apache.hadoop.fs.Path
import org.apache.spark.{Partitioner, SparkConf}
import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
import org.apache.spark.sql.comet.execution.shuffle.{CometShuffleDependency, CometShuffleExchangeExec, CometShuffleManager}
+import org.apache.spark.sql.execution.SQLExecution
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEShuffleReadExec, ShuffleQueryStageExec}
import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
import org.apache.spark.sql.execution.joins.SortMergeJoinExec
@@ -922,6 +923,22 @@ class DisableAQECometShuffleSuite extends CometColumnarShuffleSuite {
}
}
}
+
+ test("JVM shuffle writes batches no larger than spark.comet.batchSize") {
+ // spark.comet.shuffle.jvm.batchSize keeps its default of 8192, which is capped at 4096.
+ withSQLConf(
+ CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> Int.MaxValue.toString,
+ CometConf.COMET_BATCH_SIZE.key -> "4096") {
+ val df = spark.range(0, 20000, 1, 1).toDF().repartition(2, $"id")
+ val shuffle = checkCometExchange(df, 1, native = false).head
+ // Propagating the SQL conf lets the tasks see the batch size, as they would under an action.
+ val batchRows = SQLExecution.withSQLConfPropagated(spark) {
+ shuffle.executeColumnar().map(_.numRows()).collect()
+ }
+ assert(batchRows.sum == 20000)
+ assert(batchRows.max <= 4096)
+ }
+ }
}
class CometShuffleEncryptionSuite extends CometTestBase {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org