fix: keep a JVM shuffle writer's checksum disabled across spills, and test shuffles with checksums disabled by andygrove · Pull Request #6361 · apache/datafusion-comet

andygrove · 2026-09-28T23:41:07Z

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 used the sort-based writer (CometUnsafeShuffleWriter, chosen when the partition count exceeds spark.shuffle.sort.bypassMergeThreshold) failed with ArrayIndexOutOfBoundsException: Index 0 out of bounds for length 0. CometShuffleChecksumSupport.createPartitionChecksums returns an empty array when checksums are disabled, and SpillSorter indexed into it when the sorted records moved on to the next partition. #7065, which writes every partition of a spill in one native call, replaced that code, and SpillSorter now passes the empty array to native code with no checksum algorithm. This PR adds the regression tests for it, so the end-to-end case from the issue is pinned.

Auditing the checksum-disabled path turned up a second problem that is still there. SpillWriter uses checksum == -1 to mean "no checksum", but doSpilling always overwrites checksum with the value returned by the native write, 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. CometDiskBlockWriter is reused across the spills of a partition, so with checksums disabled the JVM shuffle still checksummed almost all of its output.

The other uses of the checksum array are 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?

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.
Tests (below). The new CometShuffleChecksumDisabledSuite is registered in the shuffle group of both PR workflows.

How are these changes tested?

CometDiskBlockWriterSuite: "a writer's checksum follows shuffle checksums being disabled across spills" makes a writer with no checksum set spill several times before close(), then checks that getChecksum() still reports no checksum. Its enabled twin sets the checksum to start from the sentinel with adler32 and checks that getChecksum() equals a one-shot Adler32 over the partition file, so a checksum that stops following the bytes of later writes fails it. With the SpillWriter change reverted, the disabled test fails.
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 with the default spill threshold and with a threshold of 2000 rows that makes every map task spill. It asserts that the plan has a Comet columnar shuffle exchange, that the shuffle uses the sort-based writer exactly when the partition count exceeds 200, and that the result matches Spark.
SpillSorterSuite on main covers the sorter without checksums.

On the default Spark 4.1 profile, CometDiskBlockWriterSuite, CometShuffleChecksumDisabledSuite and SpillSorterSuite pass, 21 tests in all. test-compile -Pspark-3.5 -Pstrict-warnings and dev/ci/check-suites.py pass.

…=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.

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 · 2026-09-29T02:00:15Z

+  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]()


[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.

Moot now. I merged main, where #7065 replaced the sorter code this change guarded and added its own SpillSorterSuite tests, including one with checksums disabled. I took main's SpillSorterSuite, so the SpillInfo collision is gone with the test that had it, and main's version already imports it as CometSpillInfo (7904068).

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 · 2026-09-30T16:22:39Z

+      (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)


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.

Done. The test runs with checksums disabled and enabled. The enabled run sets the checksum to start from the sentinel with adler32, and checks that getChecksum() equals a one-shot Adler32 over the partition file after several spills. The disabled run fails when I revert the SpillWriter change.

comphead · 2026-09-30T16:22:39Z

+            .range(0, 100000, 1, 4)
+            .selectExpr("id", "cast(id as string) as s")
+            .repartition(numPartitions, col("id"))
+          checkCometExchange(df, 1, false)


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.

Added. The test asserts that the shuffle handle is the CometSerializedShuffleHandle exactly when the partition count is over 200, with your assertion, and it passes for 10 and 300 partitions.

comphead · 2026-09-30T16:22:39Z


          // Store the checksum for the current partition.
-          partitionChecksums[currentPartition] = getChecksum();
+          if (partitionChecksums.length > 0) {


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.

No longer applies. #7065 rewrote writeSortedFileNative to write every partition of a spill in one native call, so the two near-copies of that block are gone, and the checksum store now happens in native code.

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.

andygrove · 2026-10-11T16:04:03Z

@Visorgood you worked on JVM columnar shuffle in #6110. If you have time, would you be able to review this one? Thanks.

sunchao

Summary

Prior state and problem: Disabling shuffle checksums caused the sort-based JVM writer to index an empty checksum array. Repeated writes also unintentionally enabled checksum computation after native code returned its unavailable-checksum sentinel.
Design approach: Guard the partition-switch checksum store and preserve SpillWriter’s disabled state after native writes.
Correctness: The guards match Spark’s empty-array contract. I traced partition transitions, repeated spills, final writes and commit paths. Enabled checksum updates remain unchanged. No additional correctness issue meeting the P1/P2 bar was identified.
Compatibility analysis: Checked Spark sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3 and experimental 4.2.0. Their checksum creation and commit paths consistently treat empty arrays as disabled. The previously reported Scala 2.12 test-compilation blocker remains unresolved.
Key design decisions: Retains the distinction between -1 for disabled checksums and Long.MIN_VALUE for an enabled checksum without an initial value. The JNI contract stays unchanged.
Implementation sketch: Two Java guards, regression tests for partition transitions and repeated writes, and an end-to-end suite covering 10/300 partitions with default/forced spilling. The new suite is registered in both platform workflows.
Performance: Preserving the disabled sentinel prevents unnecessary checksum creation and traversal of serialized bytes on subsequent writes. The change adds no allocations or JNI calls. No throughput improvement was measured.
Design: The focused state fix fits the existing shuffle implementation and is straightforward to reason about. No broader redesign is needed to address this failure.
Abstraction & complexity: Reusing SpillWriter applies the sentinel correction to both JVM writer variants without another abstraction. Buffer ownership, memory accounting and spill thresholds remain unchanged.
Behavioral changes worth calling out: Checksum-disabled JVM shuffle can cross partition boundaries and retain disabled checksums across spills. Partitioning, output format and enabled-checksum behavior are preserved.
Suggested improvements: Address the existing SpillInfo collision by aliasing Comet’s type at both new allocations, then rerun the affected Spark and Celeborn builds. No new introduced P1/P2 issues found within this review beyond that already reported blocker.

Reviewed full SHA 207fab1faa05f7ab0cce7e07254a0e4a5a1e43d4. Covered all seven files in the full PR diff against base 22a07067bda272ed68adc67733cbaa7462a13c85, using merge base 65b334bfdd25196091a42d773bc3e61d19212e52. Confirmed non-draft status and read existing reviews, conversation comments, inline 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 substantiated. Scala 2.12 resolves the new test’s SpillInfo references to Spark’s same-package class. Exact-head CI fails at SpillSorterSuite.scala:276 with incompatible Spark/Comet LinkedList types. A bounded local compiler probe reproduced both allocation mismatches with Scala 2.12.18. Aliasing both uses compiled successfully. Both forms compiled with Scala 2.13.17. This concern is not duplicated as a new finding.

Exact-head CI: All 38 checks reference the requested head. Downloaded logs confirm builds used its merge with the requested base. Linux Spark 4.1 passed all 515 shuffle tests, including all 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 compatibility jobs failed on the existing type collision, leaving Required Checks red.

Validation limits: No full local Comet build, runtime suite or benchmark was run. The local compiler probe used actual Spark/Comet SpillInfo sources with a placeholder TempShuffleBlockId, so it validates name resolution only. Suite-registration and diff checks passed. Runtime CI covered the default Spark profile. macOS, Spark SQL and Iceberg suites were skipped.

…s-disabled

sunchao

Reviewed commit 2a438eb9b25d37d1d1061ec572b94979334a5f29. New P1/P2 findings are detailed inline.

sunchao · 2026-10-19T15:42:05Z

+      // 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)


[P2] Make this offset explicitly Long, for example Platform.BYTE_ARRAY_OFFSET + 4L. Spark declares BYTE_ARRAY_OFFSET as int, so + 4 produces an Int that is implicitly widened for UnsafeRow.pointTo(Object, long, int). The required strict Scala job enables -Ywarn-numeric-widen and -Xfatal-warnings, causing test compilation to fail here even after the separately reported SpillInfo collision is fixed. An explicit long expression preserves the intended address and removes this independent build blocker.

Evidence: Exact-head job https://github.com/apache/datafusion-comet/actions/runs/39169798179/job/117800950319 reports SpillSorterSuite.scala:291: implicit numeric widening. A bounded JDK 17 / Scala 2.12.18 probe using Spark’s method signatures and the current pom.xml strict test compiler arguments reproduced this error with SpillInfo already aliased. The original offset exited 1 with the numeric-widening warning promoted to an error. Changing only + 4 to + 4L exited 0.

Gone with the same test, since main's SpillSorterSuite replaced it. test-compile -Pspark-3.5 -Pstrict-warnings passes on the merged tree.

…s-disabled apache#7065 replaced the SpillSorter code this change guarded, so take main's SpillSorter and SpillSorterSuite and keep the SpillWriter fix. The disk writer test runs with checksums disabled and enabled, and the end-to-end test asserts which writer the shuffle uses.

andygrove · 2026-11-12T13:23:58Z

I merged main (7904068). #7065 replaced the SpillSorter code this PR guarded, and the end-to-end suite passes without the guard, so the ArrayIndexOutOfBoundsException of #6256 no longer happens on main. The PR is now the SpillWriter fix, which main still needs, and the regression tests. I rewrote the title and the description to match.

github-actions Bot added bug Something isn't working area:shuffle Shuffle (JVM and native) labels Sep 28, 2026

sunchao reviewed Sep 29, 2026

View reviewed changes

comphead reviewed Sep 30, 2026

View reviewed changes

sunchao reviewed Oct 1, 2026

View reviewed changes

sunchao reviewed Oct 11, 2026

View reviewed changes

andygrove requested review from comphead and parthchandra October 15, 2026 14:58

andygrove mentioned this pull request Oct 16, 2026

perf: write every partition of a JVM shuffle spill in one native call #7065

Merged

andygrove added 2 commits October 18, 2026 22:27

Merge remote-tracking branch 'apache/main' into spill-sorter-checksum…

f5f5800

…s-disabled

test: import CometShuffleMemoryAllocatorTrait after merging main

2a438eb

sunchao reviewed Oct 19, 2026

View reviewed changes

andygrove changed the title ~~fix: support JVM columnar shuffle with spark.shuffle.checksum.enabled=false~~ fix: keep a JVM shuffle writer's checksum disabled across spills, and test shuffles with checksums disabled Nov 12, 2026

Conversation

andygrove commented Sep 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

comphead left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

andygrove commented Oct 11, 2026

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andygrove commented Nov 12, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

andygrove commented Sep 28, 2026 •

edited

Loading