fix: support JVM columnar shuffle with spark.shuffle.checksum.enabled=false

.github/workflows/pr_build_linux.yml

@@ -542,6 +542,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.DisableAQECometShuffleSuite
               org.apache.spark.shuffle.comet.CometUnboundedShuffleMemoryAllocatorSuite

.github/workflows/pr_build_macos.yml

@@ -208,6 +208,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.DisableAQECometShuffleSuite
               org.apache.spark.shuffle.comet.CometUnboundedShuffleMemoryAllocatorSuite

spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java

@@ -292,7 +292,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;
       }

spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java

@@ -215,7 +215,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]);

spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala

@@ -1069,6 +1069,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") {

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()
@@ -93,7 +99,7 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach {
       "zstd", // compressionCodec
       1, // compressionLevel
       true, // offsetsAsLengths
-      "adler32", // checksumAlgorithm
+      checksumAlgorithm,
       partitionChecksums,
       writeMetrics,
       taskContext,
@@ -267,6 +273,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.

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}
@@ -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)))
@@ -216,6 +218,49 @@ class CometDiskBlockWriterSuite extends AnyFunSuite {
     }
   }
 
+  private def newWriter(
+      file: File,
+      allocator: CometShuffleMemoryAllocatorTrait,
+      taskContext: TaskContextImpl,
+      conf: SparkConf): CometDiskBlockWriter = {
+    new CometDiskBlockWriter(
+      file,
+      allocator,
+      taskContext,
+      new UnsafeRowSerializer(1).newInstance(),
+      schema,
+      new ShuffleWriteMetrics,
+      conf,
+      false,
+      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