From 207fab1faa05f7ab0cce7e07254a0e4a5a1e43d4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 28 Sep 2026 17:40:49 -0600 Subject: [PATCH 1/2] fix: support JVM columnar shuffle with spark.shuffle.checksum.enabled=false With shuffle checksums disabled, the partition checksum array is empty. SpillSorter.writeSortedFileNative stored a checksum into it whenever the sorted records moved to the next partition, so every task of the sort-based JVM columnar shuffle writer failed with an ArrayIndexOutOfBoundsException. Guard that store like the other accesses. SpillWriter.doSpilling also replaced its -1 "no checksum" sentinel with the Long.MIN_VALUE the native writer returns when it computes no checksum, so every later write by the same writer computed an unused Adler32 checksum. Only keep the returned checksum when checksums are enabled. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../spark/shuffle/sort/SpillSorter.java | 4 +- .../comet/execution/shuffle/SpillWriter.java | 6 ++- .../exec/CometColumnarShuffleSuite.scala | 28 +++++++++++ .../spark/shuffle/sort/SpillSorterSuite.scala | 47 ++++++++++++++++++- .../shuffle/CometDiskBlockWriterSuite.scala | 28 +++++++++++ 7 files changed, 111 insertions(+), 4 deletions(-) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index ca2b3e38cb1..2593ca89fb4 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -514,6 +514,7 @@ jobs: org.apache.spark.sql.comet.execution.shuffle.CometNativePositionalRoundRobinSuite org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite org.apache.comet.exec.CometShuffleEncryptionSuite + org.apache.comet.exec.CometShuffleChecksumDisabledSuite org.apache.comet.exec.CometShuffleManagerSuite org.apache.comet.exec.CometAsyncShuffleSuite org.apache.comet.exec.DisableAQECometShuffleSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 74fa0128bf9..7166a2b3b18 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -162,6 +162,7 @@ jobs: org.apache.spark.sql.comet.execution.shuffle.CometNativePositionalRoundRobinSuite org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite org.apache.comet.exec.CometShuffleEncryptionSuite + org.apache.comet.exec.CometShuffleChecksumDisabledSuite org.apache.comet.exec.CometShuffleManagerSuite org.apache.comet.exec.CometAsyncShuffleSuite org.apache.comet.exec.DisableAQECometShuffleSuite diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java b/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java index f695a235f41..830b6ec063a 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java +++ b/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java @@ -287,7 +287,9 @@ public void writeSortedFileNative(boolean isLastFile, boolean tracingEnabled) th spillInfo.partitionLengths[currentPartition] = written; // Store the checksum for the current partition. - partitionChecksums[currentPartition] = getChecksum(); + if (partitionChecksums.length > 0) { + partitionChecksums[currentPartition] = getChecksum(); + } } currentPartition = partition; } 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 1683af3e35a..a96f4311664 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 @@ -213,7 +213,11 @@ protected long doSpilling( tracingEnabled); long written = results[0]; - checksum = results[1]; + // Without a checksum the native code returns Long.MIN_VALUE here. Keep -1 in that case, or the + // next write by this writer would start computing a checksum that nothing reads. + if (checksumEnabled) { + checksum = results[1]; + } // results[2] = encode + compression time in nanoseconds (measured by native code) if (encodeNanosAcc != null) { encodeNanosAcc.addAndGet(results[2]); 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 157cf972d67..ca725b09c38 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala @@ -959,6 +959,34 @@ class CometShuffleEncryptionSuite extends CometTestBase { } } +class CometShuffleChecksumDisabledSuite extends CometTestBase { + + override protected def sparkConf: SparkConf = { + val conf = super.sparkConf + conf.set("spark.shuffle.checksum.enabled", "false") + } + + test("comet columnar shuffle with shuffle checksums disabled") { + // 10 partitions use the bypass merge sort writer, while 300 partitions exceed + // spark.shuffle.sort.bypassMergeThreshold and use the sort-based writer. The lower spill + // threshold makes every map task spill, so the sort-based writer also merges spill files. + Seq(10, 300).foreach { numPartitions => + Seq(Int.MaxValue, 2000).foreach { spillThreshold => + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> spillThreshold.toString) { + val df = spark + .range(0, 100000, 1, 4) + .selectExpr("id", "cast(id as string) as s") + .repartition(numPartitions, col("id")) + checkCometExchange(df, 1, false) + checkSparkAnswer(df) + } + } + } + } +} + class CometShuffleManagerSuite extends CometTestBase { test("should not bypass merge sort if executor cores are too high") { diff --git a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala index 42ecc7611e8..aa7234bd07a 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala @@ -19,6 +19,8 @@ package org.apache.spark.shuffle.sort +import java.io.File +import java.util.UUID import java.util.concurrent.atomic.AtomicInteger import org.scalatest.BeforeAndAfterEach @@ -28,7 +30,10 @@ import org.apache.spark.{SparkConf, TaskContext} import org.apache.spark.executor.ShuffleWriteMetrics import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} import org.apache.spark.shuffle.comet.{CometShuffleMemoryAllocator, CometShuffleMemoryAllocatorTrait} +import org.apache.spark.sql.catalyst.expressions.UnsafeRow +import org.apache.spark.sql.comet.execution.shuffle.SpillInfo import org.apache.spark.sql.types._ +import org.apache.spark.storage.TempShuffleBlockId import org.apache.spark.unsafe.Platform import org.apache.spark.unsafe.UnsafeAlignedOffset @@ -79,7 +84,8 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { spillCallback: SpillSorter.SpillCallback = () => {}, spills: java.util.LinkedList[org.apache.spark.sql.comet.execution.shuffle.SpillInfo] = new java.util.LinkedList[org.apache.spark.sql.comet.execution.shuffle.SpillInfo](), - partitionChecksums: Array[Long] = new Array[Long](10)): SpillSorter = { + partitionChecksums: Array[Long] = new Array[Long](10), + checksumAlgorithm: String = "adler32"): SpillSorter = { val schema = createTestSchema() val writeMetrics = new ShuffleWriteMetrics() val taskContext = TaskContext.empty() @@ -92,7 +98,7 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { 0.5, // preferDictionaryRatio "zstd", // compressionCodec 1, // compressionLevel - "adler32", // checksumAlgorithm + checksumAlgorithm, partitionChecksums, writeMetrics, taskContext, @@ -262,6 +268,43 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { } } + test("write sorted file across partitions with shuffle checksums disabled") { + // With spark.shuffle.checksum.enabled=false the sorter gets an empty checksum array and no + // checksum algorithm. + val spills = new java.util.LinkedList[SpillInfo]() + val sorter = createSpillSorter( + spills = spills, + partitionChecksums = Array.emptyLongArray, + checksumAlgorithm = null) + val file = File.createTempFile("spill-sorter", ".data") + try { + val numPartitions = 3 + val rowSize = UnsafeRow.calculateBitSetWidthInBytes(1) + 8 + // Each record is a 4-byte key followed by an UnsafeRow with one int column. + val record = new Array[Byte](4 + rowSize) + val row = new UnsafeRow(1) + row.pointTo(record, Platform.BYTE_ARRAY_OFFSET + 4, rowSize) + sorter.initialCurrentPage(numPartitions * 2 * (record.length + UAO_SIZE)) + for (p <- 0 until numPartitions; i <- 0 until 2) { + row.setInt(0, p * 2 + i) + sorter.insertRecord(record, Platform.BYTE_ARRAY_OFFSET, record.length, p) + } + + sorter.setSpillInfo( + new SpillInfo(numPartitions, file, TempShuffleBlockId(UUID.randomUUID()))) + sorter.writeSortedFileNative(true, false) + + assert(spills.size() === 1) + val partitionLengths = spills.getFirst.partitionLengths + assert(partitionLengths.forall(_ > 0)) + assert(partitionLengths.sum === file.length()) + } finally { + sorter.freeMemory() + sorter.freeArray() + file.delete() + } + } + test("pointer array growth is sized from the array, not from the data pages") { // Use the unified (off-heap) allocator so that every allocation made by this sorter is // visible through `getUsed` on an allocator that is private to this test. diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala index 7a9a4cb0637..2359875db32 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala @@ -40,6 +40,8 @@ import org.apache.spark.sql.types.{BinaryType, IntegerType, MetadataBuilder, Str import org.apache.spark.unsafe.UnsafeAlignedOffset import org.apache.spark.util.Utils +import org.apache.comet.CometConf + class CometDiskBlockWriterSuite extends AnyFunSuite { private val schema = StructType(Seq(StructField("a", BinaryType))) @@ -212,6 +214,32 @@ class CometDiskBlockWriterSuite extends AnyFunSuite { new JLinkedList[CometDiskBlockWriter]()) } + test("a writer computes no checksum when shuffle checksums are disabled") { + // With spark.shuffle.checksum.enabled=false the bypass merge sort writer never calls + // setChecksum, so every write, including the ones after the first spill, must skip it. + val conf = new SparkConf() + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "1g") + val tmm = new TaskMemoryManager(new TestMemoryManager(conf), 0L) + val allocator = CometShuffleMemoryAllocator.getInstance(tmm, pageSize) + val tempDir = Utils.createTempDir() + // Spill every two rows so that the writer writes to its file several times. + SQLConf.get.setConfString(CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key, "2") + try { + val writer = + newWriter(new File(tempDir, "partition0"), allocator, newTaskContext(tmm, 0L), conf) + val toUnsafe = UnsafeProjection.create(schema) + (0 until 5).foreach(_ => writer.insertRow(toUnsafe(InternalRow(new Array[Byte](16))), 0)) + assert(writer.close().length > 0) + assert(writer.getOutputRecords == 5) + assert(writer.getChecksum == -1) + } finally { + SQLConf.get.unsetConf(CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key) + Utils.deleteRecursively(tempDir) + tmm.cleanUpAllAllocatedMemory() + } + } + test("a fatal error during write() frees the task's buffered pages") { // Spark's ShuffleWriteProcessor only calls stop(false) when write() throws an Exception, so a // fatal error such as SparkOutOfMemoryError skips it. write() itself must therefore free the From 2a438eb9b25d37d1d1061ec572b94979334a5f29 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 18 Oct 2026 22:52:48 +0100 Subject: [PATCH 2/2] test: import CometShuffleMemoryAllocatorTrait after merging main --- .../sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala index dc4ca78ef69..1c75926771b 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala @@ -29,7 +29,7 @@ import org.apache.spark.executor.{ShuffleWriteMetrics, TaskMetrics} import org.apache.spark.memory.{SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager} import org.apache.spark.shuffle.api.{ShuffleExecutorComponents, ShuffleMapOutputWriter, ShufflePartitionWriter} import org.apache.spark.shuffle.api.metadata.MapOutputCommitMessage -import org.apache.spark.shuffle.comet.CometShuffleMemoryAllocator +import org.apache.spark.shuffle.comet.{CometShuffleMemoryAllocator, CometShuffleMemoryAllocatorTrait} import org.apache.spark.shuffle.sort.SpillSorter import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow}