Skip to content

feat: add RSS partition writer and task-owned JNI callback (1/n) - #5473

Open
pingzh wants to merge 1 commit into
apache:mainfrom
pingzh:pingzh/rss-partition-writer-jni-callback
Open

feat: add RSS partition writer and task-owned JNI callback (1/n)#5473
pingzh wants to merge 1 commit into
apache:mainfrom
pingzh:pingzh/rss-partition-writer-jni-callback

Conversation

@pingzh

@pingzh pingzh commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Part of #5352. This is the first foundational PR and does not close the issue.

Rationale for this change

Comet’s native shuffle currently writes partition data to local files. Supporting remote shuffle services requires a transport-independent mechanism for delivering complete encoded shuffle blocks to task-owned JVM callbacks.

This PR introduces that foundation without changing shuffle planning, adding Celeborn integration, or modifying existing shuffle behavior.

What changes are included in this PR?

  • Add a Java ShufflePartitionPusher callback interface for task-owned shuffle output.
  • Add a Rust callback abstraction and JNI adapter that:
    • Safely invokes callbacks from native execution threads.
    • Preserves Java exceptions across the JNI boundary.
    • Validates partition identifiers and payload lengths.
  • Add RssPartitionWriter, which:
    • Encodes complete, length-prefixed Arrow IPC shuffle blocks.
    • Supports existing shuffle compression codecs and dictionary-encoded batches.
    • Enforces configurable frame limits and partition lifecycle ordering.
    • Preserves shuffle metrics and propagates callback failures.
  • Expose the writer for integration in subsequent PRs.

How are these changes tested?

  • 15 new Rust tests cover frame boundaries, partition routing, all compression codecs, dictionary batches, error propagation, frame limits, and lifecycle validation.
  • Three new JNI tests cover partition and payload validation.
  • All 50 native shuffle unit tests pass.
  • All 28 existing JVM native-shuffle regression tests pass.
  • Full-workspace Rust Clippy, Rust formatting, ScalaStyle, Spotless, CI preflight, and the complete four-module Maven build all pass.

continue;
}

if encoded_size > self.max_frame_size {

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

[P2] Split oversized batches before encoding instead of failing valid RSS writes. max_frame_size is checked only after ShuffleBlockWriter::write_batch has already encoded the entire Arrow batch into an unbounded Vec. Consequently a normal multi-row batch whose IPC output exceeds the transport frame limit fails the whole shuffle task even when every row would fit in its own complete frame, and the configured limit does not bound the allocation that can OOM an executor. The current oversized-frame test codifies this failure. Estimate/bound allocation before encoding and split batches at row boundaries into independently decodable complete frames; reject only a genuinely unsplittable row. AI-assisted review (Codex).

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Split oversized batches before encoding instead of failing valid RSS writes.

it important and need to introduce splitting before encoding,
but not exclude requirement to check size after compression
because depends from data and codec, final size can be bigger than initial data

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

LGTM for this foundational scope.

Verified 50 native shuffle tests, 3 JNI validation tests, 9 additional adversarial cases, and JNI runtime probes for native-thread callbacks and exception preservation. The full Spark/Maven regression suite was not rerun.

The frame-size concern is a follow-up before RSS integration: the current contract caps callback payloads and intentionally rejects oversized frames; it does not promise automatic splitting or bounded serialization memory.

AI-assisted review (Codex).

DataFusionError::Execution(format!(
"Remote shuffle payload size {payload_length} exceeds the JVM array limit"
))
})

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

in specification it Integer.MAX_VALUE, but HotSpot implementation details Integer.MAX_VALUE-8
so maybe make sense to test it as well

https://github.com/openjdk/jdk/blob/master/src/java.base/share/classes/jdk/internal/util/ArraysSupport.java#L854-L866

    /**
     * A soft maximum array length imposed by array growth computations.
     * Some JVMs (such as HotSpot) have an implementation limit that will cause
     *
     *     OutOfMemoryError("Requested array size exceeds VM limit")
     *
     * to be thrown if a request is made to allocate an array of some length near
     * Integer.MAX_VALUE, even if there is sufficient heap available. The actual
     * limit might depend on some JVM implementation-specific characteristics such
     * as the object header size. The soft maximum value is chosen conservatively so
     * as to be smaller than any implementation limit that is likely to be encountered.
     */
    public static final int SOFT_MAX_ARRAY_LENGTH = Integer.MAX_VALUE - 8;

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants