Skip to content

fix: [branch-1.0] read shuffle write buffer, spill limit and off-heap sizes in bytes (#6191) - #6209

Merged
andygrove merged 1 commit into
apache:branch-1.0from
andygrove:backport-6191-branch-1.0
Sep 24, 2026
Merged

andygrove merged 1 commit into
apache:branch-1.0from
andygrove:backport-6191-branch-1.0

Conversation

@andygrove

Copy link
Copy Markdown
Member

Backport of #6191 to branch-1.0.

Cherry-picked from 02e84047a423a08debe3c6b6daab19fd9233427c. The fixes themselves are unchanged. Five adaptations were needed, described under "What changes are included" below. Two of them touch production code, and neither changes behaviour: a visibility widening and an unused import.

Which issue does this PR close?

Closes #6183, #6184 and #6185 on branch-1.0. Listed in #6201.

Rationale for this change

All three bugs ship in 1.0.0, through the same code as on main before #6191:

  • spark.comet.shuffle.native.writeBufferSize is declared in MiB but sent to native code as a byte count, so the native shuffle writer runs with a 1-byte write buffer by default. On branch-1.0 the setting sizes the multi-partition data file's BufWriter and the flush threshold of every BufBatchWriter, so each encoded block goes straight to its own write call, for shuffle and spill files alike.
  • spark.comet.maxTempDirectorySize reaches native code as the raw string, and native code parses only a bare integer, so a value with a unit silently falls back to 100 GiB. The boolean flags native code reads have the same problem with anything other than lowercase true. Native code on branch-1.0 reads the same configs as on main.
  • spark.memory.offHeap.size is read as MiB when sizing the memory pool, while Spark reads a bare number as bytes, so a bare byte count makes the fair_unified per-task cap effectively unlimited.

What changes are included in this PR?

The fix is the original one, so see #6191 for the details. The adaptations:

How are these changes tested?

Same tests as the original PR, verified locally on branch-1.0 with Spark 4.1.3, Scala 2.13 and JDK 17 unless noted:

  • CometNativeShuffleWriterSuite and CometExecSuite pass, 146 tests including the three new ones.
  • The bugs are present on branch-1.0, and the tests catch them. With the CometConf.scala and CometExecIterator.scala changes reverted and the tests kept, exactly the three new tests fail and the other 143 pass. The plan sent to native code carries a write buffer of 1 rather than 1048576, the serialized configs have no maxTempDirectorySize entry, and a bare 4294967296 off-heap size gives a memory limit of 4503599627370496, which is 4 PiB.
  • The same two suites pass on Spark 3.4.3, Scala 2.12 and JDK 11, the oldest profile branch-1.0 CI runs: 143 passed, and 3 were cancelled by existing Spark version guards.
  • CometNativeShuffleSuite, CometShuffleSuite and CometShuffleEncryptionSuite pass, 73 tests, now that the native shuffle writer uses a 1 MiB buffer instead of 1 byte.
  • scalafix in CHECK mode on Spark 3.5 / Scala 2.12, Spotless, Scalastyle, apache-rat:check, cargo fmt --all -- --check, cargo clippy --all-targets --workspace -- -D warnings, dev/ci/check-suites.py, dev/ci/check-ci-config.py, actionlint, and prettier --check on the two changed docs all pass.

Are there any user-facing changes?

Yes, the same ones as #6191, which are worth weighing for a patch release:

  • The native shuffle writer now uses the 1 MiB write buffer the setting documents. On branch-1.0 each shuffle-writing task holds up to about 2 MiB of buffers that no memory pool tracks: the data file's BufWriter and one BufBatchWriter staging buffer at a time. That is about half of the 4 MiB on main, where perf: reuse per-partition scratch in the shuffle write path #5568 and perf: spill every shuffle partition of a task into one file #5916 added buffers of the same size.
  • writeBufferSize values with a unit now mean what they say, and a bare number keeps its meaning. A value of 2 GiB or more, which the native int32 field cannot hold, now fails with a config error. Previously 4g, for example, became a 4096-byte buffer.
  • maxTempDirectorySize values with a unit, and upper-case booleans for the flags native code reads, now take effect. An invalid value for any of them now fails with the usual config error instead of silently becoming the native default.
  • A bare byte count for spark.memory.offHeap.size now sizes the Comet memory pool in bytes, as Spark does.

…es (apache#6191)

* fix: read shuffle write buffer, spill limit and off-heap sizes in bytes

spark.comet.shuffle.native.writeBufferSize was declared in MiB but its
value was sent to native code as a byte count, so the native shuffle
writer ran with a 1-byte write buffer by default. Declare it in bytes
with a 1 MiB default; a bare number keeps its meaning.

Native code parsed spark.comet.maxTempDirectorySize and the boolean
flags it reads from the raw session strings, so a size with a unit or
an upper-case boolean silently fell back to the native default. Resolve
every config that native code reads on the JVM before it crosses JNI,
and document that the spill limit applies per native plan rather than
per task.

getMemoryConfig read spark.memory.offHeap.size as MiB, while Spark reads
a bare number as bytes.

Closes apache#6183, apache#6184, apache#6185.

* docs: give a real example of a task running several native plans

A Spark operator that Comet does not support does not split a stage into
two native plans: Comet does not resume native execution above it. The
native operators on either side of a union or a coalesce do run as
separate plans in the same task.

* Update spark/src/main/scala/org/apache/comet/CometConf.scala

Co-authored-by: Matt Butrovich <mbutrovich@users.noreply.github.com>

---------

Co-authored-by: Matt Butrovich <mbutrovich@users.noreply.github.com>
(cherry picked from commit 02e8404)

Adapted for branch-1.0:
- CI suite lists: added only CometNativeShuffleWriterSuite. The
  CometNativePositionalRoundRobinSuite and CometDiskBlockWriterSuite
  neighbours on main come from later commits and do not exist here.
- The getMemoryConfig test moved from CometExecIteratorLifecycleSuite,
  which does not exist on branch-1.0 (it comes from apache#5494), into
  CometExecSuite next to the PR's other CometExecIterator test. Its body
  is unchanged.
- CometNativeShuffleWriter.buildUnifiedPlan widened from private to
  private[shuffle], the same change apache#5513 made on main, so the new suite
  can call it. It still takes separate data and index paths here, so the
  suite passes both.
- Removed the ByteUnit import from CometExecIterator, which the
  offHeap.size fix leaves unused on branch-1.0.
- Updated the spark.comet.maxTempDirectorySize row of the frozen
  configs.md. branch-1.0 commits the generated config reference; main
  generates it at publish time.
@github-actions github-actions Bot added bug Something isn't working area:shuffle Shuffle (JVM and native) labels Sep 24, 2026

@comphead comphead left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @andygrove

@andygrove
andygrove merged commit 634c885 into apache:branch-1.0 Sep 24, 2026
70 checks passed
@andygrove
andygrove deleted the backport-6191-branch-1.0 branch September 24, 2026 22:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:shuffle Shuffle (JVM and native) bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants