feat: add RSS partition writer and task-owned JNI callback (1/n) - #5473
feat: add RSS partition writer and task-owned JNI callback (1/n)#5473pingzh wants to merge 1 commit into
Conversation
| continue; | ||
| } | ||
|
|
||
| if encoded_size > self.max_frame_size { |
There was a problem hiding this comment.
[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).
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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" | ||
| )) | ||
| }) |
There was a problem hiding this comment.
in specification it Integer.MAX_VALUE, but HotSpot implementation details Integer.MAX_VALUE-8
so maybe make sense to test it as well
/**
* 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;
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?
ShufflePartitionPushercallback interface for task-owned shuffle output.RssPartitionWriter, which:How are these changes tested?