Skip to content

perf: cache expected schemas for remote shuffle decoding #5722

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pingzh wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf: cache expected schemas for remote shuffle decoding #5722

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 53 additions & 8 deletions native/core/src/execution/jni_api.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1252,10 +1252,56 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlock(
})
}

#[no_mangle]
/// Parse the expected schema once for a remote shuffle iterator.
///
/// The iterator owns the returned decoder and releases it when the input is closed.
pub extern "system" fn Java_org_apache_comet_Native_createRemoteShuffleDecoder(
e: EnvUnowned,
_class: JClass,
expected_schema: JByteArray,
) -> jlong {
try_unwrap_or_throw(&e, |env| {
let bytes = env.convert_byte_array(expected_schema)?;
let schema = ShuffleScan::decode(bytes.as_slice()).map_err(|error| {
CometError::Internal(format!("Invalid expected remote shuffle schema: {error}"))
})?;
let decoder = RemoteShuffleDecoder {
expected_types: schema.fields.iter().map(to_arrow_datatype).collect(),
};
Ok(Box::into_raw(Box::new(decoder)) as jlong)
})
}

/// Immutable decoding state owned by one JVM remote shuffle iterator, not shared across tasks.
struct RemoteShuffleDecoder {
expected_types: Vec<ArrowDataType>,
}

#[no_mangle]
/// Release a remote shuffle iterator's decoder.
///
/// # Safety
/// A nonzero handle must have been returned by `createRemoteShuffleDecoder`, must not have
/// been released, and must not be in use by a concurrent decode call.
pub unsafe extern "system" fn Java_org_apache_comet_Native_releaseRemoteShuffleDecoder(
e: EnvUnowned,
_class: JClass,
decoder_handle: jlong,
) {
try_unwrap_or_throw(&e, |_| {
if decoder_handle != 0 {
drop(unsafe { Box::from_raw(decoder_handle as *mut RemoteShuffleDecoder) });
}
Ok(())
})
}

#[no_mangle]
/// Decode a remote native shuffle block with Arrow array and logical type validation enabled.
/// # Safety
/// This function is inherently unsafe since it deals with raw pointers passed from JNI.
/// Buffer and output pointers must be valid. The decoder handle must have been returned by
/// `createRemoteShuffleDecoder` and must remain alive for the duration of this call.
pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlockWithValidation(
e: EnvUnowned,
_class: JClass,
Expand All @@ -1264,22 +1310,21 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlockWit
array_addrs: JLongArray,
schema_addrs: JLongArray,
tracing_enabled: jboolean,
expected_schema: JByteArray,
decoder_handle: jlong,
) -> jlong {
try_unwrap_or_throw(&e, |env| {
with_trace("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || {
let bytes = env.convert_byte_array(expected_schema)?;
let schema = ShuffleScan::decode(bytes.as_slice()).map_err(|error| {
CometError::Internal(format!("Invalid expected remote shuffle schema: {error}"))
})?;
let expected_types: Vec<_> = schema.fields.iter().map(to_arrow_datatype).collect();
let decoder = unsafe { (decoder_handle as *const RemoteShuffleDecoder).as_ref() }
.ok_or_else(|| {
CometError::Internal("Remote shuffle decoder is not initialized".to_owned())
})?;
decode_shuffle_block(
env,
byte_buffer,
length,
array_addrs,
schema_addrs,
Some(&expected_types),
Some(&decoder.expected_types),
)
})
})
Expand Down
16 changes: 12 additions & 4 deletions spark/src/main/scala/org/apache/comet/Native.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -201,17 +201,25 @@ class Native extends NativeBase {
tracingEnabled: Boolean): Long

/**
* Decode a remote shuffle block with Arrow buffer/offset and logical type validation. The
* expected schema is serialized as a ShuffleScan protobuf. Keep the existing local decoder
* entry point unchanged so trusted local shuffle reads retain their fast path.
* Create a remote shuffle decoder that retains the expected Spark types for one iterator. The
* expected schema is serialized as a ShuffleScan protobuf.
*/
@native def createRemoteShuffleDecoder(expectedSchema: Array[Byte]): Long

/** Release a remote shuffle decoder after its iterator finishes reading or closes early. */
@native def releaseRemoteShuffleDecoder(decoderHandle: Long): Unit

/**
* Decode a remote shuffle block with Arrow buffer/offset and logical type validation, using the
* expected Spark types retained by the decoder.
*/
@native def decodeShuffleBlockWithValidation(
shuffleBlock: ByteBuffer,
length: Int,
arrayAddrs: Array[Long],
schemaAddrs: Array[Long],
tracingEnabled: Boolean,
expectedSchema: Array[Byte]): Long
decoderHandle: Long): Long

/**
* Log the beginning of an event.
Expand Down
144 changes: 84 additions & 60 deletions .../main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,11 +45,14 @@ case class NativeBatchDecoderIterator(
expectedSchema: Option[Array[Byte]] = None)
extends Iterator[ColumnarBatch] {

// One consumer reads this iterator, while task completion may close it from another thread.
// The monitor protects decoder and batch ownership; transport reads stay outside it.
private var isClosed = false
private val longBuf = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN)
private var currentBatch: ColumnarBatch = null
private var batch: Option[ColumnarBatch] = None
private val validateRemoteFrames = in.isInstanceOf[CometShuffleReadFailureHandler]
private var remoteDecoderHandle = 0L

require(
!validateRemoteFrames || expectedSchema.exists(_ != null),
Expand All @@ -64,78 +67,100 @@ case class NativeBatchDecoderIterator(
}

override def hasNext: Boolean = {
if (channel == null || isClosed) {
return false
}
if (batch.isDefined) {
return true
}
synchronized {
if (channel == null || isClosed) {
return false
}
if (batch.isDefined) {
return true
}

// Release the previous batch.
if (currentBatch != null) {
currentBatch.close()
currentBatch = null
// Clear ownership before cleanup so a failed close cannot release this batch twice.
if (currentBatch != null) {
val previous = currentBatch
currentBatch = null
previous.close()
}
}

batch = fetchNext()
if (batch.isEmpty) {
close()
return false
}
true
fetchNext()
}

def next(): ColumnarBatch = {
if (!hasNext) {
throw new NoSuchElementException
}

val nextBatch = batch.get

currentBatch = nextBatch
batch = None
currentBatch
synchronized {
// Completion may have closed the iterator after hasNext() returned.
if (isClosed) {
throw new NoSuchElementException
}
currentBatch = batch.get
batch = None
currentBatch
}
}

private def fetchNext(): Option[ColumnarBatch] = {
private def fetchNext(): Boolean = {
// The remote input owns metadata, transport, and frame-boundary failure classification. Do
// not turn deliberately unreported metadata timeouts or stream-close failures into corruption.
// Read outside the monitor so close() can unblock the underlying stream.
val block = readNextBlock()
block.flatMap { case (fieldCount, dataBuf, bytesToRead) =>
// Allocation and Arrow import failures are not evidence of a corrupt remote shuffle. Only
// forward failures from native codec/IPC decoding and logical type validation.
val startTime = System.nanoTime()
val decoded = nativeUtil.getNextBatch(
fieldCount,
(arrayAddrs, schemaAddrs) => {
handleReadFailure {
if (validateRemoteFrames) {
nativeLib.decodeShuffleBlockWithValidation(
dataBuf,
bytesToRead,
arrayAddrs,
schemaAddrs,
tracingEnabled,
expectedSchema.get)
} else {
nativeLib.decodeShuffleBlock(
dataBuf,
bytesToRead,
arrayAddrs,
schemaAddrs,
tracingEnabled)
}
var nativeFailure: Throwable = null
try {
synchronized {
// Cleanup may have finished during the read. Do not create a decoder or import a batch
// after close(), and publish each decoded batch before cleanup can inspect ownership.
if (isClosed) {
return false
}
batch = block.flatMap { case (fieldCount, dataBuf, bytesToRead) =>
val startTime = System.nanoTime()
// Invalid expected schemas are setup failures, not evidence of corrupt persisted data.
if (validateRemoteFrames && remoteDecoderHandle == 0L) {
remoteDecoderHandle = nativeLib.createRemoteShuffleDecoder(expectedSchema.get)
}
})
decodeTime.add(System.nanoTime() - startTime)
decoded
}
}

private def handleReadFailure[T](read: => T): T = {
try read
catch {
case NonFatal(failure) =>
val decoded = nativeUtil.getNextBatch(
fieldCount,
(arrayAddrs, schemaAddrs) => {
try {
if (validateRemoteFrames) {
nativeLib.decodeShuffleBlockWithValidation(
dataBuf,
bytesToRead,
arrayAddrs,
schemaAddrs,
tracingEnabled,
remoteDecoderHandle)
} else {
nativeLib.decodeShuffleBlock(
dataBuf,
bytesToRead,
arrayAddrs,
schemaAddrs,
tracingEnabled)
}
} catch {
case NonFatal(failure) =>
// Record only native decode failures; creation, allocation and import failures
// must not be reported as corrupt remote data. Let NativeUtil clean up first.
nativeFailure = failure
throw failure
}
})
decodeTime.add(System.nanoTime() - startTime)
decoded
}
if (batch.isEmpty) {
close()
}
batch.isDefined
}
} catch {
case NonFatal(failure) if failure eq nativeFailure =>
// Reporting can make an RPC. It must run after Arrow cleanup and outside the monitor so
// task completion can release the decoder while that reporting is in progress.
in match {
case handler: CometShuffleReadFailureHandler => handler.onShuffleReadFailure(failure)
case _ =>
Expand All @@ -145,10 +170,6 @@ case class NativeBatchDecoderIterator(
}

private def readNextBlock(): Option[(Int, ByteBuffer, Int)] = {
if (channel == null || isClosed) {
return None
}

// read compressed batch size from header
longBuf.clear()
while (longBuf.hasRemaining && channel.read(longBuf) >= 0) {}
Expand Down Expand Up @@ -212,6 +233,8 @@ case class NativeBatchDecoderIterator(
currentBatch = null
val prefetched = batch
batch = None
val decoderHandle = remoteDecoderHandle
remoteDecoderHandle = 0L

var failure: Throwable = null
def release(resource: => Unit): Unit = {
Expand All @@ -225,6 +248,7 @@ case class NativeBatchDecoderIterator(

if (previous != null) release(previous.close())
prefetched.filterNot(_ eq previous).foreach(pending => release(pending.close()))
if (decoderHandle != 0L) release(nativeLib.releaseRemoteShuffleDecoder(decoderHandle))

@jayanth86 jayanth86 Sep 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

[P1] Keep the decoder alive while a read is in flight

Could creation and decoding be coordinated with this release? On supported Spark 3.4.3, a Python UDF downstream of a Celeborn exchange consumes this iterator on PythonRunner's writer thread. Its completion listener interrupts and joins that writer, but task cancellation with interruptThread=true can interrupt the join. TaskContextImpl catches the listener exception and continues cleanup, so Comet can reach this release while the writer is still in native decompression.

fetchNext() does not take the close() lock. This release therefore frees the boxed decoder while the active JNI call still borrows its expected_types, which are read after decompression in decode_remote_shuffle_batch. That introduces a use-after-free capable of crashing the executor. Previously, the types were owned locally by the JNI invocation and could not be freed by task cleanup.

Please protect the handle across creation, active decoding, and release, and add coverage for interrupted cleanup overlapping a read. A latch-based test using this exact iterator and mocked JNI reproduces release while decoding is still in flight. This is source-verified reachability plus a lifecycle reproduction, not a reproduced native crash.

@andygrove andygrove Sep 6, 2026

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.

The P1 lifetime concern looks right to me. I reproduced the JVM half of it locally with a latch-based probe on this iterator and a mocked Native: releaseRemoteShuffleDecoder runs to completion while decodeShuffleBlockWithValidation is still inside the JNI call, so nothing in the iterator serializes the release against an active decode. remoteDecoderHandle is published from fetchNext() without the monitor that close() takes, and decode_remote_shuffle_batch reads expected_types after read_ipc_compressed_validated has already decompressed the block, so the borrow outlives the free in that ordering. The mirror-image ordering is a leak rather than a crash, since close() reads 0L and the box is never freed.

What I do not think is true is that this PR introduces the hazard. CometExecIterator has had the same shape for a long time, with a wider window on a much hotter path. It registers close() as a task completion listener at construction (CometExecIterator.scala:189), close() is synchronized and calls nativeLib.releasePlan(plan) (:275 and :305), and getNextBatch calls nativeLib.executePlan(..., plan, ...) outside that monitor (:205). On the native side releasePlan is a Box::from_raw drop of the context (jni_api.rs:1051) while executePlan holds get_execution_context(exec_context), an unbounded &'a mut ExecutionContext derived from the raw handle, for the duration of an entire plan execution (jni_api.rs:848). So if the threaded Python consumer plus an interrupted cleanup wait reaches the shuffle decoder, it already reaches every Comet native plan today.

Given that, I would rather not hold this PR for it. Could we file one issue covering JVM-owned native handle lifetime for both NativeBatchDecoderIterator and CometExecIterator, and fix them the same way? Fixing only the new handle here would leave the wider window untouched and leave us with two different ownership conventions for the same problem.

For what it is worth, the minimal fix for this iterator looks like taking the monitor around just the nativeUtil.getNextBatch call and leaving readNextBlock() outside it, so in.close() can still unblock a transport read. Decoding one block is bounded work, so cleanup cannot hang waiting on the monitor. Whatever we settle on should apply to releasePlan as well.

@pingzh pingzh Sep 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Codex acting for the user: Fixed the NativeBatchDecoderIterator lifetime race in 92fa68c5b.

Decoder creation, native decoding, Arrow import, and batch publication now share the monitor used by close(). next() also transfers batch ownership under that monitor. A transport read that finishes after cleanup rechecks isClosed before creating or using a decoder.

readNextBlock() stays outside the monitor so in.close() can unblock it. Native decode failures unwind Arrow cleanup before reporting outside the monitor, since reporting may perform an RPC.

Added five latch-based regressions covering creation, active decoding (including interrupted cleanup), import/publication, blocked transport reads, and failure reporting. All five failed before the fix; all 69 reader-suite tests now pass on both Spark 4.1 and Spark 3.4.3. Spotless and Scalastyle also pass.

The broader CometExecIterator lifetime concern remains a separate follow-up.

if (in != null) release(in.close())
release(resetDataBuf())
if (failure != null) throw failure
Expand Down
Loading
Loading