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

[andygrove]

Which issue does this PR close?

Closes #6256.

Rationale for this change

With spark.shuffle.checksum.enabled=false, every map task of a JVM columnar shuffle that uses the sort-based writer (CometUnsafeShuffleWriter, chosen when the partition count exceeds spark.shuffle.sort.bypassMergeThreshold) fails with ArrayIndexOutOfBoundsException: Index 0 out of bounds for length 0.

CometShuffleChecksumSupport.createPartitionChecksums returns an empty array when checksums are disabled. SpillSorter.writeSortedFileNative checks partitionChecksums.length > 0 before every access except the store that runs when the sorted records move on to the next partition, so the first partition switch indexes into the empty array.

Auditing the other JVM shuffle paths turned up a second problem in the same checksum-disabled path. SpillWriter uses checksum == -1 to mean "no checksum", but doSpilling always overwrites checksum with the value returned by the native writeSortedFileNative, which is Long.MIN_VALUE when no checksum was computed. From its second write on, a writer therefore passes checksumEnabled = true, and the native code computes an Adler32 checksum that nothing reads. SpillSorter is reused across partitions and spills, and CometDiskBlockWriter across the spills of a partition, so with checksums disabled the JVM shuffle still checksummed almost all of its output. Once the store above is guarded, this is the only remaining effect, so it is fixed here too.

The other uses of the checksum array are already safe. CometBypassMergeSortShuffleWriter guards setChecksum and collects checksums with a loop bounded by partitionChecksums.length. CometShuffleExternalSorter and CometUnsafeShuffleWriter only pass the (possibly empty) array to Spark's commitAllPartitions and transferMapSpillFile, which treat an empty array as checksums disabled, just as for Spark's own UnsafeShuffleWriter.

What changes are included in this PR?

• SpillSorter.writeSortedFileNative: guard the per-partition checksum store with partitionChecksums.length > 0, like the other accesses.
• SpillWriter.doSpilling: only keep the checksum returned by the native code when checksums are enabled, so a writer that starts with checksums disabled keeps them disabled for all of its writes.
• New tests (below). The new CometShuffleChecksumDisabledSuite is registered in the shuffle group of both PR workflows.

How are these changes tested?

• SpillSorterSuite: the new test "write sorted file across partitions with shuffle checksums disabled" writes a sorted file across three partitions with an empty checksum array and no checksum algorithm, which is what the sorter receives when checksums are disabled. The existing tests always passed a non-empty array and never wrote a file, which is why nothing caught this.
• CometDiskBlockWriterSuite: the new test "a writer computes no checksum when shuffle checksums are disabled" makes a writer with no checksum set spill twice before close(), then checks that getChecksum() still reports no checksum.
• New CometShuffleChecksumDisabledSuite in CometColumnarShuffleSuite.scala. spark.shuffle.checksum.enabled is a core setting, so the suite sets it in sparkConf, as CometShuffleEncryptionSuite does for encryption. It runs the query from the issue with spark.comet.shuffle.mode=jvm for 10 partitions (bypass merge sort writer) and 300 partitions (sort-based writer). Each runs with the default spill threshold and with a threshold of 2000 rows that makes every map task spill, so the sort-based writer also merges spill files. The test asserts that the plan has a Comet columnar shuffle exchange and that the result matches Spark.

I wrote the tests before the fix and ran them on the default Spark 4.1 profile. Without the fix, all three failed. The SpillSorterSuite test and the 300-partition case of the end-to-end test failed with the ArrayIndexOutOfBoundsException at SpillSorter.java:290 from the issue. The CometDiskBlockWriterSuite test failed with 2642559282 did not equal -1, which shows the writer computed a checksum. With the fix, on the same profile, SpillSorterSuite (11 tests), CometDiskBlockWriterSuite (5), CometShuffleChecksumDisabledSuite (1) and CometShuffleSuite (46, JVM columnar shuffle with checksums enabled and forced spills) all pass.

[sunchao]

Summary

• Prior state and problem: Disabling shuffle checksums caused partition-switch failures and unintentionally enabled checksum computation after the first write.
• Design approach: Guard the checksum-array store and preserve the writer’s disabled-checksum sentinel.
• Correctness / compatibility analysis: The production changes match Spark’s empty-checksum-array contract across supported versions 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. One introduced build blocker remains: the new test resolves SpillInfo incorrectly under Scala 2.12.
• Key design decisions: Reuses existing state and conventions without additional abstractions, allocations or JNI calls. Avoiding unintended checksum computation removes demonstrated unnecessary work.
• Implementation sketch: Two production guards, three regression tests covering partition transitions and repeated spills, and suite registration in both platform workflows.
• Behavioral changes worth calling out: JVM shuffle can retain disabled checksums across partitions and spills. Enabled checksum updates retain their existing behavior.
• Suggested improvements: Alias Comet’s SpillInfo in the new test and use that alias for both allocations, then rerun the affected builds.

Reviewed full SHA 207fab1faa05f7ab0cce7e07254a0e4a5a1e43d4. Covered all seven files in the PR diff against base 22a07067bda272ed68adc67733cbaa7462a13c85, using merge base 65b334bfdd25196091a42d773bc3e61d19212e52. Confirmed non-draft status and read the discussion snapshot and live comments. No existing review concerns were present.

Routed skills: review-comet-pr, review-comet-shuffle-pr, review-comet-memory-pr, and review-comet-ffi-pr.

Exact-head CI: Linux Spark 4.1 suites passed, including all 515 shuffle tests and the three new regressions. Native build, Rust tests, TPC-H/TPC-DS verification and Spark 4.0 build/lint checks passed. Spark 3.4/3.5 build checks and both Celeborn compatibility jobs failed at the new test’s type mismatch, leaving Required Checks red.

Validation limits: A bounded local compiler reproduction confirmed the failure and alias fix. No full local Comet build or runtime suite was run. macOS, Spark SQL and Iceberg suites were skipped in exact-head CI.

[sunchao]

[P1] Alias Comet’s SpillInfo to restore Scala 2.12 builds. When compiling the supported Spark 3.4 or 3.5 profile, Spark’s org.apache.spark.shuffle.sort.SpillInfo in this suite’s package hides the newly imported Comet class. This allocation therefore produces the wrong LinkedList type, and passing it to createSpillSorter fails compilation instead of running the regression test. The constructor at line 294 has the same collision. This blocks both profiles’ test compilation and both Celeborn compatibility jobs. Use an import alias such as SpillInfo => CometSpillInfo at both allocations, or fully qualify the Comet type.

Evidence: Exact-head CI reports found: java.util.LinkedList[org.apache.spark.shuffle.sort.SpillInfo], required: java.util.LinkedList[org.apache.spark.sql.comet.execution.shuffle.SpillInfo] at line 276, plus the hidden-import warning. See Spark 3.5 job https://github.com/apache/datafusion-comet/actions/runs/36499297541/job/109186916881 and Spark 3.4 job https://github.com/apache/datafusion-comet/actions/runs/36499297541/job/109186916902. A local JDK 17 compiler probe using Spark 3.5.9 and the checkout’s actual Comet SpillInfo reproduced both mismatches with Scala 2.12.18. Aliasing both uses compiled successfully. Original and aliased forms both compiled with Scala 2.13.17.

[comphead]

Thanks for the careful fix and for auditing the other checksum uses. Guarding on partitionChecksums.length > 0 matches Spark's ShuffleExternalSorter, and Spark's IndexShuffleBlockResolver treats an empty array as checksums disabled (checksums.nonEmpty), in 3.4.3 through 4.1.2. I didn't spot another unguarded use. The comments below are small test and cleanup suggestions.

[comphead]

Would it make sense to also run this scenario with checksums enabled? That is the default, and this PR changes the assignment in doSpilling that the enabled path relies on, but I couldn't find a test that checks a checksum value. With writer.setChecksum(Long.MIN_VALUE) and writer.setChecksumAlgo("adler32") before the inserts, I'd expect getChecksum to equal a one-shot java.util.zip.Adler32 over the partition file (I haven't run this). Parameterizing this test over both settings would keep the test count flat.

[comphead]

The comment says 300 partitions use the sort-based writer. Would it be worth asserting that, so the test can't silently stop covering it if a threshold changes? checkCometExchange returns the exchanges, so something like assert(exchanges.head.shuffleDependency.shuffleHandle.getClass.getName.contains("CometSerializedShuffleHandle") == (numPartitions > 200)) should work (I haven't run it). The similar checks in Columnar shuffle for large shuffle partition number discard the boolean, so as written they don't assert anything.

[comphead]

Nit: the partition-switch block here and the final-partition block below are near copies, and the final one already had this guard, so they drifted apart. Would a small private helper that sets the checksum, calls doSpilling, records the length and stores the checksum keep the guard in one place? Happy to leave that for a follow-up if you'd rather keep this fix small for backporting.

[sunchao]

Summary

• Prior state and problem: Disabling shuffle checksums caused partition-switch failures and unintentionally enabled checksum computation after the first write.
• Design approach: Guard the checksum-array store and preserve the writer’s disabled-checksum sentinel.
• Correctness / compatibility analysis: The production guards match Spark’s empty-array contract across supported versions 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. The previously reported Scala 2.12 compilation blocker remains unresolved.
• Key design decisions: Reuses existing state without additional abstractions, allocations or JNI calls. Prevents unnecessary checksum computation without changing enabled-checksum updates.
• Implementation sketch: Two production guards, three regression tests covering partition transitions and repeated spills, and registration in both platform workflows.
• Behavioral changes worth calling out: JVM shuffle retains disabled checksums across partitions and spills. Partitioning, output format and memory ownership remain unchanged.
• Suggested improvements: Address the existing SpillInfo import collision using an alias at both allocations, then rerun the affected builds. No additional introduced P1/P2 issues found within this review.

Reviewed full SHA 207fab1faa05f7ab0cce7e07254a0e4a5a1e43d4, covering all seven files in the full PR diff against base 22a07067bda272ed68adc67733cbaa7462a13c85, using merge base 65b334bfdd25196091a42d773bc3e61d19212e52. Confirmed non-draft status and read existing reviews, comments and threads.

Routed skills: review-comet-pr, review-comet-shuffle-pr, review-comet-memory-pr and review-comet-ffi-pr.

Existing blocker: The reported P1 remains reproducible. Scala 2.12 resolves the new test’s SpillInfo references to Spark’s class instead of Comet’s. A bounded local compiler reproduction with JDK 17, Scala 2.12.18, Spark 3.5.9 and the checkout’s Comet class reproduced both mismatches. Aliasing both uses compiled successfully. This is not duplicated as a new finding.

Exact-head CI: All 38 checks reference the requested head. The build jobs tested GitHub’s merge with the requested base. Linux Spark 4.1 passed all 515 shuffle tests, including the three new regressions. Other Linux 4.1 suites, native/Rust checks, TPC-H/TPC-DS verification and Spark 4.0 build checks passed. Spark 3.4/3.5 and both Celeborn jobs failed on the existing type collision, leaving Required Checks red.

Validation limits: No full local Comet build or runtime suite was run. Runtime CI coverage was limited to the default Spark profile. macOS, Spark SQL and Iceberg suites were skipped.