diff --git a/native/Cargo.lock b/native/Cargo.lock index df5fd4ac14a..03b08363295 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2056,6 +2056,7 @@ name = "datafusion-comet-shuffle" version = "1.1.0" dependencies = [ "arrow", + "arrow-data", "arrow-select", "async-trait", "bytes", diff --git a/native/Cargo.toml b/native/Cargo.toml index 1805a185e9d..81ee8f53449 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -39,6 +39,7 @@ rust-version = "1.94.0" [workspace.dependencies] arrow = { version = "59.2.0", features = ["prettyprint", "ffi", "chrono-tz"] } +arrow-data = { version = "59.2.0" } arrow-select = { version = "59.2.0" } async-trait = { version = "0.1" } bytes = { version = "1.11.1" } diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 692b0d1bccf..510fb9016df 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -101,7 +101,7 @@ use tokio::sync::mpsc; use crate::execution::memory_pools::{create_memory_pool, parse_memory_pool_config}; use crate::execution::operators::{ScanExec, ShuffleScanExec}; use crate::execution::shuffle::{ - decode_remote_shuffle_batch, read_ipc_compressed, CompressionCodec, + decode_remote_shuffle_batch_with, CompressionCodec, ShuffleBlockDecoder, }; use crate::execution::spark_plan::SparkPlan; @@ -1317,11 +1317,31 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlock( ) -> jlong { try_unwrap_or_throw(&e, |env| { with_trace("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || { - decode_shuffle_block(env, byte_buffer, length, array_addrs, schema_addrs, None) + // This entry point carries no native handle, so the schema cache lives in a + // thread-local decoder. Spark runs one task per thread at a time and the decoder + // compares the schema message bytes before reusing a cached schema, so a decoder + // last used by a different task on this thread simply misses and re-parses. + LOCAL_SHUFFLE_BLOCK_DECODER.with_borrow_mut(|decoder| { + decode_shuffle_block( + env, + decoder, + byte_buffer, + length, + array_addrs, + schema_addrs, + None, + ) + }) }) }) } +thread_local! { + /// Schema-caching decoder for the handle-less local `decodeShuffleBlock` entry point. + static LOCAL_SHUFFLE_BLOCK_DECODER: std::cell::RefCell = + std::cell::RefCell::new(ShuffleBlockDecoder::new()); +} + #[no_mangle] /// Parse the expected schema once for a remote shuffle iterator. /// @@ -1338,14 +1358,19 @@ pub extern "system" fn Java_org_apache_comet_Native_createRemoteShuffleDecoder( })?; let decoder = RemoteShuffleDecoder { expected_types: schema.fields.iter().map(to_arrow_datatype).collect(), + decoder: ShuffleBlockDecoder::new(), }; Ok(Box::into_raw(Box::new(decoder)) as jlong) }) } -/// Immutable decoding state owned by one JVM remote shuffle iterator, not shared across tasks. +/// Decoding state owned by one JVM remote shuffle iterator, not shared across tasks. The JVM +/// side serializes decode calls on a handle, which is what the schema cache relies on. struct RemoteShuffleDecoder { expected_types: Vec, + /// Lives as long as the JVM-side reader holds the handle, so the IPC schema message is + /// parsed once per distinct schema encoding rather than once per block. + decoder: ShuffleBlockDecoder, } #[no_mangle] @@ -1384,17 +1409,21 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlockWit ) -> jlong { try_unwrap_or_throw(&e, |env| { with_trace("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || { - let decoder = unsafe { (decoder_handle as *const RemoteShuffleDecoder).as_ref() } + // SAFETY: the handle was returned by `createRemoteShuffleDecoder`, has not been + // released, and the JVM side does not decode concurrently on one handle, so this + // is the only live reference for the duration of the call. + let remote = unsafe { (decoder_handle as *mut RemoteShuffleDecoder).as_mut() } .ok_or_else(|| { CometError::Internal("Remote shuffle decoder is not initialized".to_owned()) })?; decode_shuffle_block( env, + &mut remote.decoder, byte_buffer, length, array_addrs, schema_addrs, - Some(&decoder.expected_types), + Some(&remote.expected_types), ) }) }) @@ -1402,6 +1431,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlockWit fn decode_shuffle_block( env: &mut Env, + decoder: &mut ShuffleBlockDecoder, byte_buffer: JByteBuffer, length: jint, array_addrs: JLongArray, @@ -1414,9 +1444,9 @@ fn decode_shuffle_block( let batch = if let Some(expected_types) = expected_types { // Reject incompatible logical types, then decode dictionaries before JVM import. The // JVM importer supports fewer dictionary key/value layouts than the shuffle writer. - decode_remote_shuffle_batch(slice, expected_types)? + decode_remote_shuffle_batch_with(decoder, slice, expected_types)? } else { - read_ipc_compressed(slice)? + decoder.decode(slice)? }; prepare_output(env, array_addrs, schema_addrs, batch, false) } diff --git a/native/core/src/execution/operators/shuffle_scan.rs b/native/core/src/execution/operators/shuffle_scan.rs index 8778ce36998..984bfa82161 100644 --- a/native/core/src/execution/operators/shuffle_scan.rs +++ b/native/core/src/execution/operators/shuffle_scan.rs @@ -20,7 +20,7 @@ use crate::{ execution::{ operators::ExecutionError, planner::TEST_EXEC_CONTEXT_ID, - shuffle::{decode_remote_shuffle_batch, read_ipc_compressed}, + shuffle::{decode_remote_shuffle_batch_with, ShuffleBlockDecoder}, }, jvm_bridge::{jni_call, JVMClasses}, }; @@ -73,6 +73,10 @@ pub struct ShuffleScanExec { decode_time: Time, /// Remote inputs require Arrow array and logical schema validation; queried once at construction. requires_validation: bool, + /// Block decoder held for the life of the scan so the IPC schema message, which every block + /// from the same writer repeats verbatim, is parsed once rather than per block. Behind a + /// mutex only because this exec is `Clone`; it is used from `get_next_batch` on one thread. + decoder: Arc>, } impl ShuffleScanExec { @@ -115,6 +119,7 @@ impl ShuffleScanExec { schema, decode_time, requires_validation, + decoder: Arc::new(Mutex::new(ShuffleBlockDecoder::new())), }) } @@ -134,12 +139,14 @@ impl ShuffleScanExec { let mut current_batch = self.batch.try_lock().unwrap(); if current_batch.is_none() { + let mut decoder = self.decoder.try_lock().unwrap(); let next_batch = Self::get_next( self.exec_context_id, self.input_source.as_ref().unwrap().as_obj(), &self.data_types, &self.decode_time, self.requires_validation, + &mut decoder, )?; *current_batch = Some(next_batch); } @@ -156,6 +163,7 @@ impl ShuffleScanExec { data_types: &[DataType], decode_time: &Time, requires_validation: bool, + decoder: &mut ShuffleBlockDecoder, ) -> Result { if exec_context_id == TEST_EXEC_CONTEXT_ID { return Ok(InputBatch::EOF); @@ -191,7 +199,8 @@ impl ShuffleScanExec { // Decode the compressed IPC data let mut timer = decode_time.timer(); - let batch = match decode_shuffle_batch(slice, data_types, requires_validation) { + let batch = match decode_shuffle_batch(decoder, slice, data_types, requires_validation) + { Ok(batch) => batch, Err(failure) => { // Remote inputs must invalidate the failed shuffle generation even when @@ -230,6 +239,7 @@ impl ShuffleScanExec { } fn decode_shuffle_batch( + decoder: &mut ShuffleBlockDecoder, bytes: &[u8], expected_types: &[DataType], requires_validation: bool, @@ -237,9 +247,9 @@ fn decode_shuffle_batch( if requires_validation { // Validate logical types before decoding dictionaries or normalizing nested fields. // Keep both validation and normalization failures inside get_next's recovery callback. - decode_remote_shuffle_batch(bytes, expected_types) + decode_remote_shuffle_batch_with(decoder, bytes, expected_types) } else { - check_column_count(read_ipc_compressed(bytes)?, expected_types.len()) + check_column_count(decoder.decode(bytes)?, expected_types.len()) } } @@ -474,19 +484,29 @@ mod tests { .values(), &[u32::MAX] ); - let error = super::decode_shuffle_batch(&payload, &[DataType::Int32], true) - .unwrap_err() - .to_string(); + let error = super::decode_shuffle_batch( + &mut crate::execution::shuffle::ShuffleBlockDecoder::new(), + &payload, + &[DataType::Int32], + true, + ) + .unwrap_err() + .to_string(); assert!(error.contains("type mismatch at column 0"), "{error}"); assert!(error.contains("UInt32"), "{error}"); assert!(error.contains("Int32"), "{error}"); // The new logical validation is confined to remote inputs. assert_eq!( - super::decode_shuffle_batch(&payload, &[DataType::Int32], false) - .unwrap() - .column(0) - .data_type(), + super::decode_shuffle_batch( + &mut crate::execution::shuffle::ShuffleBlockDecoder::new(), + &payload, + &[DataType::Int32], + false + ) + .unwrap() + .column(0) + .data_type(), &DataType::UInt32 ); } @@ -500,7 +520,13 @@ mod tests { ) .unwrap(); let payload = uncompressed_shuffle_payload(&batch); - let decoded = super::decode_shuffle_batch(&payload, &[], true).unwrap(); + let decoded = super::decode_shuffle_batch( + &mut crate::execution::shuffle::ShuffleBlockDecoder::new(), + &payload, + &[], + true, + ) + .unwrap(); assert_eq!(decoded.num_columns(), 0); assert_eq!(decoded.num_rows(), 3); } @@ -621,14 +647,24 @@ mod tests { // Local decoding preserves the wire encoding for get_next to unpack. Remote decoding // validates and unpacks first, so the same result can also be safely imported by the JVM. - let local = - super::decode_shuffle_batch(body, &[DataType::Int32, DataType::Utf8], false).unwrap(); + let local = super::decode_shuffle_batch( + &mut crate::execution::shuffle::ShuffleBlockDecoder::new(), + body, + &[DataType::Int32, DataType::Utf8], + false, + ) + .unwrap(); assert!(matches!( local.column(1).data_type(), DataType::Dictionary(_, _) )); - let decoded = - super::decode_shuffle_batch(body, &[DataType::Int32, DataType::Utf8], true).unwrap(); + let decoded = super::decode_shuffle_batch( + &mut crate::execution::shuffle::ShuffleBlockDecoder::new(), + body, + &[DataType::Int32, DataType::Utf8], + true, + ) + .unwrap(); assert_eq!(decoded.column(1).data_type(), &DataType::Utf8); // Create ShuffleScanExec with value types (Utf8, not Dictionary) — this is @@ -697,8 +733,13 @@ mod tests { let declared = list_of_struct_type(true); let block = RecordBatch::try_from_iter([("payload", block_column)]).unwrap(); let payload = uncompressed_shuffle_payload(&block); - let decoded = - super::decode_shuffle_batch(&payload, std::slice::from_ref(&declared), true).unwrap(); + let decoded = super::decode_shuffle_batch( + &mut crate::execution::shuffle::ShuffleBlockDecoder::new(), + &payload, + std::slice::from_ref(&declared), + true, + ) + .unwrap(); let mut scan = ShuffleScanExec::new( super::super::super::planner::TEST_EXEC_CONTEXT_ID, None, diff --git a/native/shuffle/Cargo.toml b/native/shuffle/Cargo.toml index 9504834ef4a..f0ed22ad730 100644 --- a/native/shuffle/Cargo.toml +++ b/native/shuffle/Cargo.toml @@ -30,6 +30,7 @@ publish = false [dependencies] arrow = { workspace = true } +arrow-data = { workspace = true } arrow-select = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } diff --git a/native/shuffle/benches/shuffle_reader.rs b/native/shuffle/benches/shuffle_reader.rs index 81efcda7aa1..f43ada2bc21 100644 --- a/native/shuffle/benches/shuffle_reader.rs +++ b/native/shuffle/benches/shuffle_reader.rs @@ -24,7 +24,9 @@ use arrow::ipc::reader::StreamReader; use arrow::ipc::writer::IpcWriteContext; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use datafusion::physical_plan::metrics::Time; -use datafusion_comet_shuffle::{read_ipc_compressed, CompressionCodec, ShuffleBlockWriter}; +use datafusion_comet_shuffle::{ + read_ipc_compressed, CompressionCodec, ShuffleBlockDecoder, ShuffleBlockWriter, +}; use std::hint::black_box; use std::io::Cursor; use std::sync::Arc; @@ -93,13 +95,24 @@ fn criterion_benchmark(c: &mut Criterion) { let id = format!("{num_columns}col_{num_rows}row"); - // full decode: schema parse plus record batch + // full decode with a throwaway decoder: schema parse plus record batch group.bench_with_input( BenchmarkId::new("decode_block", &id), &uncompressed, |b, block| b.iter(|| black_box(read_ipc_compressed(black_box(block)).unwrap())), ); + // full decode with a decoder held across blocks, the way a reader holds it: the + // schema message is byte-identical every time, so it is served from the cache + group.bench_with_input( + BenchmarkId::new("decode_block_cached_schema", &id), + &uncompressed, + |b, block| { + let mut decoder = ShuffleBlockDecoder::new(); + b.iter(|| black_box(decoder.decode(black_box(block)).unwrap())) + }, + ); + // schema parse alone: `try_new` stops before the record batch. Skips the codec tag. group.bench_with_input( BenchmarkId::new("parse_schema_only", &id), diff --git a/native/shuffle/src/ipc.rs b/native/shuffle/src/ipc.rs index 97890f50148..5de32e9f272 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -15,52 +15,312 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::RecordBatch; -use arrow::ipc::reader::StreamReader; +use arrow::array::{ArrayRef, RecordBatch}; +use arrow::buffer::MutableBuffer; +use arrow::datatypes::SchemaRef; +use arrow::ipc::convert::fb_to_schema; +use arrow::ipc::reader::{read_dictionary_impl, RecordBatchDecoder}; +use arrow::ipc::{root_as_message, MessageHeader}; +use arrow_data::UnsafeFlag; use datafusion::common::DataFusionError; use datafusion::error::Result; +use std::collections::HashMap; use std::io::{Error, ErrorKind, Read}; +use std::sync::Arc; /// Decode trusted local Comet output without revalidating every Arrow array value or offset. +/// +/// Convenience wrapper around a throwaway [`ShuffleBlockDecoder`]; callers that decode more than +/// one block should hold a decoder so the schema message is parsed once rather than per block. pub fn read_ipc_compressed(bytes: &[u8]) -> Result { - read_ipc_compressed_impl(bytes, false) + ShuffleBlockDecoder::new().decode(bytes) } /// Decode remotely fetched Comet output, including Arrow buffer and offset validation. +/// +/// See [`read_ipc_compressed`] for when to hold a [`ShuffleBlockDecoder`] instead. pub fn read_ipc_compressed_validated(bytes: &[u8]) -> Result { - read_ipc_compressed_impl(bytes, true) + ShuffleBlockDecoder::new().decode_validated(bytes) } -fn read_ipc_compressed_impl(bytes: &[u8], validate: bool) -> Result { - let codec = bytes.get(..4).ok_or_else(|| { - DataFusionError::Execution("Failed to decode batch: truncated compression codec".to_owned()) - })?; - let mut encoded = &bytes[4..]; - let batch = match codec { - b"SNAP" => read_single_batch(snap::read::FrameDecoder::new(&mut encoded), validate)?, - b"LZ4_" => read_single_batch( - lz4_flex::frame::FrameDecoder::new(RequireLz4EndMark(&mut encoded)), - validate, - )?, - // The slice already implements BufRead. Adding another BufReader would let read-ahead - // conceal compressed bytes left over after the decoder reaches its end marker. - b"ZSTD" => read_single_batch(zstd::Decoder::with_buffer(&mut encoded)?, validate)?, - b"NONE" => read_single_batch(&mut encoded, validate)?, - other => { +/// The largest message metadata length a block may carry. Arrow encodes it as an `i32`, so +/// anything beyond this is corruption rather than a large message. +const MAX_METADATA_LEN: usize = i32::MAX as usize; + +/// Decodes Comet shuffle blocks, caching the IPC schema across blocks. +/// +/// Every block is a complete Arrow IPC stream: a schema message, any dictionary batches, one +/// record batch, and the end-of-stream marker. `ShuffleBlockWriter` pre-encodes the schema +/// message once and writes it verbatim into every block, so consecutive blocks from the same +/// writer carry byte-identical schema messages. Parsing that message per block (a flatbuffer +/// verification plus one `Arc` and `String` per column) is a fixed cost that dominates the +/// decode of small blocks, which is exactly what high partition counts produce. +/// +/// The decoder keeps the raw bytes of the last schema message it parsed together with the parsed +/// [`SchemaRef`]. On each block it compares the incoming schema message bytes against the cached +/// ones and, on a match, reuses the schema without parsing. A mismatch (a different writer, a +/// different Comet version, or a block of a different shape) parses the new message and replaces +/// the cache, so correctness never depends on the cache: a block is decoded against the schema it +/// actually carries. +/// +/// A decoder is meant to be held for the life of a reader (one per `ShuffleScanExec`, one per JNI +/// decoder handle) and is not thread-safe. +#[derive(Debug, Default)] +pub struct ShuffleBlockDecoder { + /// Raw flatbuffer bytes of the last schema message and the schema parsed from them. + cached_schema: Option<(Vec, SchemaRef)>, + /// Scratch for message metadata so it is not reallocated per message. + metadata: Vec, + /// Number of schema messages actually parsed, i.e. cache misses. One per distinct schema + /// encoding seen; exposed so tests and metrics can confirm the cache is doing its job. + schema_parses: usize, +} + +impl ShuffleBlockDecoder { + pub fn new() -> Self { + Self::default() + } + + /// Decode a trusted local block without revalidating array values or offsets. + pub fn decode(&mut self, bytes: &[u8]) -> Result { + self.decode_impl(bytes, false) + } + + /// Decode a remotely fetched block, including Arrow buffer and offset validation. + pub fn decode_validated(&mut self, bytes: &[u8]) -> Result { + self.decode_impl(bytes, true) + } + + /// How many schema messages this decoder has parsed so far. Every block whose schema message + /// is byte-identical to the previous block's reuses the cached schema and does not count. + pub fn schema_parses(&self) -> usize { + self.schema_parses + } + + fn decode_impl(&mut self, bytes: &[u8], validate: bool) -> Result { + let codec = bytes.get(..4).ok_or_else(|| { + DataFusionError::Execution( + "Failed to decode batch: truncated compression codec".to_owned(), + ) + })?; + let mut encoded = &bytes[4..]; + let batch = match codec { + b"SNAP" => { + self.read_single_batch(snap::read::FrameDecoder::new(&mut encoded), validate)? + } + b"LZ4_" => self.read_single_batch( + lz4_flex::frame::FrameDecoder::new(RequireLz4EndMark(&mut encoded)), + validate, + )?, + // The slice already implements BufRead. Adding another BufReader would let + // read-ahead conceal compressed bytes left over after the decoder reaches its end + // marker. + b"ZSTD" => { + self.read_single_batch(zstd::Decoder::with_buffer(&mut encoded)?, validate)? + } + b"NONE" => self.read_single_batch(&mut encoded, validate)?, + other => { + return Err(DataFusionError::Execution(format!( + "Failed to decode batch: invalid compression codec: {other:?}" + ))) + } + }; + // LZ4 returns EOF at the end of one compressed frame without consuming the next one. + // Check the encoded source as well as the decoded IPC tail so an oversized outer frame + // cannot silently swallow another native frame's bytes. + if !encoded.is_empty() { + return Err(DataFusionError::Execution( + "Failed to decode batch: trailing data after compressed stream".to_owned(), + )); + } + Ok(batch) + } + + /// Reads one complete IPC stream holding exactly one record batch, mirroring what + /// `arrow::ipc::reader::StreamReader` does message by message but with the schema message + /// served from the cache when its bytes match. + fn read_single_batch(&mut self, mut input: R, validate: bool) -> Result { + let mut skip_validation = UnsafeFlag::new(); + if !validate { + // SAFETY: local blocks were written by this Comet version's ShuffleBlockWriter from + // arrays that were valid when encoded, which is the same trust the previous + // StreamReader-based path placed in them. Remote data keeps full validation. + unsafe { skip_validation.set(true) }; + } + + // Schema message: served from the cache on a byte match, parsed otherwise. + let schema = match self.read_metadata(&mut input)? { + None => { + return Err(DataFusionError::Execution( + "Failed to decode batch: empty IPC stream".to_owned(), + )) + } + Some(()) => self.schema_for_current_metadata()?, + }; + + let mut dictionaries_by_id: HashMap = HashMap::new(); + let mut decoded: Option = None; + while self.read_metadata(&mut input)?.is_some() { + let message = root_as_message(&self.metadata).map_err(|err| { + DataFusionError::Execution(format!( + "Failed to decode batch: unable to get root as message: {err:?}" + )) + })?; + let version = message.version(); + let body = read_body(&mut input, message.bodyLength())?; + match message.header_type() { + MessageHeader::DictionaryBatch => { + let dictionary = message.header_as_dictionary_batch().ok_or_else(|| { + DataFusionError::Execution( + "Failed to decode batch: unable to read dictionary batch".to_owned(), + ) + })?; + read_dictionary_impl( + &body.into(), + dictionary, + &schema, + &mut dictionaries_by_id, + &version, + false, + skip_validation.clone(), + )?; + } + MessageHeader::RecordBatch => { + // Each Comet frame contains one complete IPC stream with exactly one record + // batch. Stopping after that batch would skip codec footer/checksum + // validation and could silently discard further frames swallowed by a + // corrupt outer length prefix, so keep reading until the end-of-stream + // marker and reject a second batch. + if decoded.is_some() { + return Err(DataFusionError::Execution( + "Failed to decode batch: multiple record batches in one shuffle frame" + .to_owned(), + )); + } + let batch = message.header_as_record_batch().ok_or_else(|| { + DataFusionError::Execution( + "Failed to decode batch: unable to read record batch".to_owned(), + ) + })?; + let body = body.into(); + decoded = Some( + RecordBatchDecoder::try_new( + &body, + batch, + Arc::clone(&schema), + &dictionaries_by_id, + &version, + )? + .with_require_alignment(false) + .with_skip_validation(skip_validation.clone()) + .read_record_batch()?, + ); + } + MessageHeader::Schema => { + return Err(DataFusionError::Execution( + "Failed to decode batch: expected a record batch, but found a schema" + .to_owned(), + )); + } + other => { + return Err(DataFusionError::Execution(format!( + "Failed to decode batch: unsupported message header type in IPC \ + stream: '{other:?}'" + ))); + } + } + } + + let batch = decoded.ok_or_else(|| { + DataFusionError::Execution("Failed to decode batch: empty IPC stream".to_owned()) + })?; + if input.read(&mut [0])? != 0 { + return Err(DataFusionError::Execution( + "Failed to decode batch: trailing data after IPC stream".to_owned(), + )); + } + Ok(batch) + } + + /// Reads the next message's metadata length prefix and flatbuffer into `self.metadata`. + /// Returns `None` at the end of the stream, whether marked (a zero length, optionally after + /// a continuation marker) or a clean EOF before any length bytes. + fn read_metadata(&mut self, input: &mut R) -> Result> { + let mut prefix = [0u8; 4]; + match input.read_exact(&mut prefix) { + Ok(()) => {} + Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(None), + Err(e) => return Err(e.into()), + } + if prefix == [0xff; 4] { + input.read_exact(&mut prefix)?; + } + let len = i32::from_le_bytes(prefix); + if len == 0 { + return Ok(None); + } + let len = usize::try_from(len) + .ok() + .filter(|len| *len <= MAX_METADATA_LEN) + .ok_or_else(|| { + DataFusionError::Execution(format!( + "Failed to decode batch: invalid metadata length: {len}" + )) + })?; + self.metadata.resize(len, 0); + input.read_exact(&mut self.metadata)?; + Ok(Some(())) + } + + /// Resolves the schema for the schema message currently in `self.metadata`, reusing the + /// cached schema when the bytes match and parsing (and caching) otherwise. A schema message + /// has no body, which the parse path checks, so a byte-identical message needs nothing read. + fn schema_for_current_metadata(&mut self) -> Result { + if let Some((cached_bytes, cached_schema)) = &self.cached_schema { + if *cached_bytes == self.metadata { + return Ok(Arc::clone(cached_schema)); + } + } + + let message = root_as_message(&self.metadata).map_err(|err| { + DataFusionError::Execution(format!( + "Failed to decode batch: unable to get root as message: {err:?}" + )) + })?; + if message.header_type() != MessageHeader::Schema { return Err(DataFusionError::Execution(format!( - "Failed to decode batch: invalid compression codec: {other:?}" - ))) + "Failed to decode batch: expected a schema as the first message in the \ + stream, got: {:?}", + message.header_type() + ))); + } + if message.bodyLength() != 0 { + return Err(DataFusionError::Execution( + "Failed to decode batch: schema message with a non-empty body".to_owned(), + )); } - }; - // LZ4 returns EOF at the end of one compressed frame without consuming the next one. Check - // the encoded source as well as the decoded IPC tail so an oversized outer frame cannot - // silently swallow another native frame's bytes. - if !encoded.is_empty() { - return Err(DataFusionError::Execution( - "Failed to decode batch: trailing data after compressed stream".to_owned(), - )); + let schema = message.header_as_schema().ok_or_else(|| { + DataFusionError::Execution( + "Failed to decode batch: failed to parse schema from message header".to_owned(), + ) + })?; + let schema = Arc::new(fb_to_schema(schema)); + self.schema_parses += 1; + self.cached_schema = Some((self.metadata.clone(), Arc::clone(&schema))); + Ok(schema) } - Ok(batch) +} + +/// Reads a message body of `len` bytes into a fresh buffer, as `StreamReader` does. +fn read_body(input: &mut R, len: i64) -> Result { + let len = usize::try_from(len).map_err(|_| { + DataFusionError::Execution(format!( + "Failed to decode batch: invalid message body length: {len}" + )) + })?; + let mut body = MutableBuffer::from_len_zeroed(len); + input.read_exact(&mut body)?; + Ok(body) } // lz4_flex treats physical EOF (including a partial block header) as a clean end of frame. @@ -83,38 +343,9 @@ impl Read for RequireLz4EndMark { } } -fn read_single_batch(input: R, validate: bool) -> Result { - let reader = StreamReader::try_new(input, None)?; - let mut reader = if validate { - // Remote data must not escape as unchecked arrays and fail later in a native operator. - reader - } else { - // Preserve the existing local-shuffle fast path for trusted Comet-written arrays. - unsafe { reader.with_skip_validation(true) } - }; - let batch = reader.next().transpose()?.ok_or_else(|| { - DataFusionError::Execution("Failed to decode batch: empty IPC stream".to_owned()) - })?; - - // Each Comet frame contains one complete IPC stream with exactly one record batch. - // Stopping after that batch would skip codec footer/checksum validation and could silently - // discard further frames swallowed by a corrupt outer length prefix. - if reader.next().transpose()?.is_some() { - return Err(DataFusionError::Execution( - "Failed to decode batch: multiple record batches in one shuffle frame".to_owned(), - )); - } - if reader.get_mut().read(&mut [0])? != 0 { - return Err(DataFusionError::Execution( - "Failed to decode batch: trailing data after IPC stream".to_owned(), - )); - } - Ok(batch) -} - #[cfg(test)] mod tests { - use super::{read_ipc_compressed, read_ipc_compressed_validated}; + use super::{read_ipc_compressed, read_ipc_compressed_validated, ShuffleBlockDecoder}; use arrow::array::{Int32Array, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::ipc::writer::StreamWriter; @@ -161,6 +392,119 @@ mod tests { bytes } + /// Blocks that repeat the same schema message must be decoded against the cached schema + /// (one parse for the whole run), and a block carrying a different schema must be decoded + /// against its own schema and replace the cache, never against the stale one. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn schema_cache_hits_identical_messages_and_misses_different_ones() { + let int_stream = ipc_stream(1); + let utf8_stream = { + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(vec!["abc", "de"]))], + ) + .unwrap(); + let mut bytes = Vec::new(); + let mut writer = StreamWriter::try_new(&mut bytes, &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + bytes + }; + + for codec in [b"NONE", b"SNAP", b"LZ4_", b"ZSTD"] { + let int_frame = encode(codec, &int_stream); + let utf8_frame = encode(codec, &utf8_stream); + let fresh_int = read_ipc_compressed(&int_frame).unwrap(); + let fresh_utf8 = read_ipc_compressed(&utf8_frame).unwrap(); + + for validate in [false, true] { + let mut decoder = ShuffleBlockDecoder::new(); + let decode = |decoder: &mut ShuffleBlockDecoder, frame: &[u8]| { + if validate { + decoder.decode_validated(frame).unwrap() + } else { + decoder.decode(frame).unwrap() + } + }; + + for _ in 0..3 { + assert_eq!(decode(&mut decoder, &int_frame), fresh_int, "{codec:?}"); + } + assert_eq!( + decoder.schema_parses(), + 1, + "{codec:?}: repeats must hit the cache" + ); + + let utf8 = decode(&mut decoder, &utf8_frame); + assert_eq!(utf8, fresh_utf8, "{codec:?}"); + assert_eq!(utf8.schema().field(0).data_type(), &DataType::Utf8); + assert_eq!( + decoder.schema_parses(), + 2, + "{codec:?}: new schema must parse" + ); + + assert_eq!(decode(&mut decoder, &int_frame), fresh_int, "{codec:?}"); + assert_eq!( + decoder.schema_parses(), + 3, + "{codec:?}: switching back is a new encoding, not a stale hit" + ); + assert_eq!(decode(&mut decoder, &int_frame), fresh_int, "{codec:?}"); + assert_eq!(decoder.schema_parses(), 3, "{codec:?}"); + } + } + } + + /// Dictionary-encoded columns arrive as a dictionary batch before the record batch, the + /// layout the JVM columnar shuffle produces for strings. Both must decode through the + /// cached-schema path, with dictionaries scoped to their own block. + #[test] + fn dictionary_blocks_decode_with_cached_schema() { + use arrow::array::{DictionaryArray, Int32Array}; + use arrow::datatypes::Int32Type; + + let schema = Arc::new(Schema::new(vec![Field::new( + "d", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + )])); + let frame = |values: Vec<&str>| { + let keys = Int32Array::from((0..values.len() as i32).collect::>()); + let dictionary = + DictionaryArray::::try_new(keys, Arc::new(StringArray::from(values))) + .unwrap(); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(dictionary)]).unwrap(); + let mut bytes = Vec::new(); + let mut writer = StreamWriter::try_new(&mut bytes, &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + encode(b"NONE", &bytes) + }; + let first = frame(vec!["a", "b"]); + let second = frame(vec!["x", "y", "z"]); + + let mut decoder = ShuffleBlockDecoder::new(); + for validate in [false, true] { + for (block, expected) in [(&first, vec!["a", "b"]), (&second, vec!["x", "y", "z"])] { + let batch = if validate { + decoder.decode_validated(block).unwrap() + } else { + decoder.decode(block).unwrap() + }; + let values = arrow::compute::cast(batch.column(0), &DataType::Utf8).unwrap(); + let values = values.as_any().downcast_ref::().unwrap(); + let got: Vec<&str> = values.iter().map(|v| v.unwrap()).collect(); + assert_eq!(got, expected); + } + } + assert_eq!(decoder.schema_parses(), 1); + } + #[test] fn malformed_codec_prefix_returns_error() { for prefix in [&b""[..], b"N", b"NO", b"NON", b"BAD!"] { diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 766634eb71e..9fa4a66dad3 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -31,8 +31,10 @@ pub mod spark_unsafe; pub(crate) mod writers; pub use comet_partitioning::CometPartitioning; -pub use ipc::{read_ipc_compressed, read_ipc_compressed_validated}; -pub use remote_schema::{decode_remote_shuffle_batch, validate_remote_schema}; +pub use ipc::{read_ipc_compressed, read_ipc_compressed_validated, ShuffleBlockDecoder}; +pub use remote_schema::{ + decode_remote_shuffle_batch, decode_remote_shuffle_batch_with, validate_remote_schema, +}; pub use schema_align::SchemaAlignExec; pub use shuffle_writer::{ShuffleWriterDestination, ShuffleWriterExec}; pub use writers::{CompressionCodec, ShuffleBlockWriter}; diff --git a/native/shuffle/src/remote_schema.rs b/native/shuffle/src/remote_schema.rs index b205421aea4..c95974db345 100644 --- a/native/shuffle/src/remote_schema.rs +++ b/native/shuffle/src/remote_schema.rs @@ -25,15 +25,28 @@ use datafusion::error::Result; use datafusion_comet_common::cast_and_stamp_schema; use std::sync::Arc; +use crate::ShuffleBlockDecoder; + /// Decode a remote shuffle batch and reconcile its encoding with Spark's declared logical types. /// Validate buffers and logical types before casting, so a corrupt frame cannot be made to look /// compatible by a value-changing cast. Dictionary keys are an encoding detail, including inside /// containers: decode them here before either native execution or the JVM Arrow importer sees them. +/// Decode and validate one remote block with a throwaway decoder. Readers that decode many +/// blocks should use [`decode_remote_shuffle_batch_with`] and keep the decoder, so the IPC +/// schema message is parsed once rather than per block. pub fn decode_remote_shuffle_batch( bytes: &[u8], expected_types: &[DataType], ) -> Result { - let batch = crate::read_ipc_compressed_validated(bytes)?; + decode_remote_shuffle_batch_with(&mut ShuffleBlockDecoder::new(), bytes, expected_types) +} + +pub fn decode_remote_shuffle_batch_with( + decoder: &mut ShuffleBlockDecoder, + bytes: &[u8], + expected_types: &[DataType], +) -> Result { + let batch = decoder.decode_validated(bytes)?; validate_remote_schema(&batch, expected_types)?; if batch .columns()