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/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..3771d9fcc3d 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, reset_schema_cache, CompressionCodec, ShuffleBlockWriter, +}; use std::hint::black_box; use std::io::Cursor; use std::sync::Arc; @@ -100,6 +102,18 @@ fn criterion_benchmark(c: &mut Criterion) { |b, block| b.iter(|| black_box(read_ipc_compressed(black_box(block)).unwrap())), ); + // same decode with the cache cleared each iteration, so drift moves both arms together + group.bench_with_input( + BenchmarkId::new("decode_block_uncached", &id), + &uncompressed, + |b, block| { + b.iter(|| { + reset_schema_cache(); + black_box(read_ipc_compressed(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..3389f9d9df5 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -15,11 +15,17 @@ // 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::Buffer; +use arrow::datatypes::SchemaRef; +use arrow::ipc::reader::{RecordBatchDecoder, StreamReader}; +use arrow::ipc::{root_as_message, MessageHeader}; use datafusion::common::DataFusionError; use datafusion::error::Result; -use std::io::{Error, ErrorKind, Read}; +use std::cell::RefCell; +use std::collections::HashMap; +use std::io::{Cursor, Error, ErrorKind, Read}; +use std::sync::Arc; /// Decode trusted local Comet output without revalidating every Arrow array value or offset. pub fn read_ipc_compressed(bytes: &[u8]) -> Result { @@ -31,21 +37,260 @@ pub fn read_ipc_compressed_validated(bytes: &[u8]) -> Result { read_ipc_compressed_impl(bytes, true) } +/// Arrow IPC continuation marker introducing a message length. +const CONTINUATION_MARKER: [u8; 4] = [0xff, 0xff, 0xff, 0xff]; + +/// Distinct schemas cached per thread. More than one because a reduce task can interleave blocks +/// from several shuffles. Keyed on the raw schema message, so a hit costs one memcmp. +const SCHEMA_CACHE_CAPACITY: usize = 4; + +thread_local! { + static SCHEMA_CACHE: RefCell, SchemaRef)>> = + const { RefCell::new(Vec::new()) }; + /// Empty dictionary map; the fast path only runs for blocks with no dictionary messages. + static NO_DICTIONARIES: HashMap = HashMap::new(); +} + +fn cached_schema(schema_message: &[u8]) -> Option { + SCHEMA_CACHE.with(|cache| { + let mut cache = cache.borrow_mut(); + let hit = cache + .iter() + .position(|(message, _)| message.as_ref() == schema_message)?; + // most recently used first, so an alternating pair stays resident + if hit != 0 { + cache.swap(0, hit); + } + Some(Arc::clone(&cache[0].1)) + }) +} + +fn cache_schema(schema_message: &[u8], schema: SchemaRef) { + SCHEMA_CACHE.with(|cache| { + let mut cache = cache.borrow_mut(); + if cache + .iter() + .any(|(message, _)| message.as_ref() == schema_message) + { + return; + } + if cache.len() == SCHEMA_CACHE_CAPACITY { + cache.pop(); + } + cache.insert(0, (schema_message.into(), schema)); + }); +} + +/// Empties this thread's schema cache, so the next decode re-parses its schema. For benchmarks +/// comparing the cached and uncached paths; not part of the decode contract. +#[doc(hidden)] +pub fn reset_schema_cache() { + SCHEMA_CACHE.with(|cache| cache.borrow_mut().clear()); +} + +/// One Arrow IPC message located inside a decoded block. +struct IpcMessage<'a> { + /// The flatbuffer metadata, without the continuation marker or length prefix. + metadata: &'a [u8], + /// Offset of the message body within the block. + body_start: usize, + /// Offset just past this message, where the next one begins. + end: usize, +} + +/// Reads the message at `offset`. `Ok(None)` at a well-formed end, an end-of-stream marker or a +/// clean message boundary; anything truncated or inconsistent is an error. +fn read_message(block: &[u8], offset: usize) -> Result>> { + fn corrupt(what: &str) -> DataFusionError { + DataFusionError::Execution(format!("Failed to decode batch: {what}")) + } + + // ending on a message boundary is the legacy stream ending, and is valid + if offset == block.len() { + return Ok(None); + } + + let mut cursor = offset; + let first = block + .get(cursor..cursor + 4) + .ok_or_else(|| corrupt("truncated IPC message length"))?; + cursor += 4; + + let length_bytes = if first == CONTINUATION_MARKER { + let bytes = block + .get(cursor..cursor + 4) + .ok_or_else(|| corrupt("truncated IPC message length"))?; + cursor += 4; + bytes + } else { + first + }; + + let metadata_len = i32::from_le_bytes(length_bytes.try_into().expect("four bytes")); + if metadata_len == 0 { + // End-of-stream marker. + return Ok(None); + } + let metadata_len = + usize::try_from(metadata_len).map_err(|_| corrupt("negative IPC metadata length"))?; + + let metadata_end = cursor + .checked_add(metadata_len) + .ok_or_else(|| corrupt("IPC metadata length overflows the block"))?; + let metadata = block + .get(cursor..metadata_end) + .ok_or_else(|| corrupt("truncated IPC metadata"))?; + + let message = root_as_message(metadata) + .map_err(|error| corrupt(&format!("invalid IPC metadata: {error}")))?; + let body_len = + usize::try_from(message.bodyLength()).map_err(|_| corrupt("negative IPC body length"))?; + + let body_start = metadata_end; + let end = body_start + .checked_add(body_len) + .ok_or_else(|| corrupt("IPC body length overflows the block"))?; + if end > block.len() { + return Err(corrupt("truncated IPC body")); + } + + Ok(Some(IpcMessage { + metadata, + body_start, + end, + })) +} + +/// Confirms nothing follows the record batch but a well-formed end of stream. `read_message` +/// alone would not catch trailing bytes after an end-of-stream marker. +fn expect_end_of_stream(block: &[u8], offset: usize) -> Result<()> { + let trailing = || { + DataFusionError::Execution( + "Failed to decode batch: trailing data after IPC stream".to_owned(), + ) + }; + + if offset == block.len() { + return Ok(()); + } + + let mut cursor = offset; + let first = block.get(cursor..cursor + 4).ok_or_else(trailing)?; + cursor += 4; + let length_bytes = if first == CONTINUATION_MARKER { + let bytes = block.get(cursor..cursor + 4).ok_or_else(trailing)?; + cursor += 4; + bytes + } else { + first + }; + + if i32::from_le_bytes(length_bytes.try_into().expect("four bytes")) != 0 { + return Err(trailing()); + } + if cursor != block.len() { + return Err(trailing()); + } + Ok(()) +} + +/// Decodes a block whose schema is already known. `Ok(None)` if the block is not the simple +/// `[schema][record batch][end]` shape, leaving it to the general decoder. +fn decode_with_known_schema( + block: &Buffer, + schema: SchemaRef, + batch_message: &IpcMessage<'_>, + validate: bool, +) -> Result> { + let message = root_as_message(batch_message.metadata).map_err(|error| { + DataFusionError::Execution(format!( + "Failed to decode batch: invalid IPC metadata: {error}" + )) + })?; + let Some(record_batch) = message.header_as_record_batch() else { + return Ok(None); + }; + + let body = block.slice_with_length( + batch_message.body_start, + batch_message.end - batch_message.body_start, + ); + + let version = message.version(); + let batch = NO_DICTIONARIES.with(|dictionaries| { + let decoder = + RecordBatchDecoder::try_new(&body, record_batch, schema, dictionaries, &version)?; + let decoder = if validate { + decoder + } else { + // matches the trusted-local path the general decoder takes + let mut flag = arrow_data::UnsafeFlag::new(); + unsafe { flag.set(true) }; + decoder.with_skip_validation(flag) + }; + decoder.read_record_batch() + })?; + + Ok(Some(batch)) +} + +/// Decodes one decompressed block, reusing a cached schema when its schema message is known. +fn decode_block(block: Buffer, validate: bool) -> Result { + if let Some(batch) = try_decode_with_cached_schema(&block, validate) { + return Ok(batch); + } + + // general path: the only one that parses a schema, and it caches what it parsed + let (batch, schema, schema_message) = read_single_batch_cached(block.as_slice(), validate)?; + if let Some(schema_message) = schema_message { + cache_schema(schema_message, schema); + } + Ok(batch) +} + +/// Decodes a block against an already-parsed schema, or `None` if it cannot. +/// +/// Never reports an error of its own: anything it does not handle yields `None` and the general +/// decoder runs instead, so validation and error messages are unchanged. +fn try_decode_with_cached_schema(block: &Buffer, validate: bool) -> Option { + let bytes = block.as_slice(); + + let schema_message = read_message(bytes, 0).ok()??; + let is_schema = root_as_message(schema_message.metadata) + .map(|message| message.header_type() == MessageHeader::Schema) + .unwrap_or(false); + if !is_schema { + return None; + } + + let schema = cached_schema(schema_message.metadata)?; + + // the record batch must follow the schema directly; a dictionary message lands here instead + let batch_message = read_message(bytes, schema_message.end).ok()??; + expect_end_of_stream(bytes, batch_message.end).ok()?; + + decode_with_known_schema(block, schema, &batch_message, validate).ok()? +} + 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, - )?, + // materialized so messages can be walked in place; the decoded arrays borrow this buffer + let block = match codec { + b"SNAP" => decompress(snap::read::FrameDecoder::new(&mut encoded))?, + b"LZ4_" => decompress(lz4_flex::frame::FrameDecoder::new(RequireLz4EndMark( + &mut encoded, + )))?, // 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)?, + b"ZSTD" => decompress(zstd::Decoder::with_buffer(&mut encoded)?)?, + b"NONE" => { + let block = Buffer::from(encoded); + encoded = &[]; + block + } other => { return Err(DataFusionError::Execution(format!( "Failed to decode batch: invalid compression codec: {other:?}" @@ -60,7 +305,14 @@ fn read_ipc_compressed_impl(bytes: &[u8], validate: bool) -> Result "Failed to decode batch: trailing data after compressed stream".to_owned(), )); } - Ok(batch) + decode_block(block, validate) +} + +/// Reads a decompressor to the end, yielding the decoded block. +fn decompress(mut reader: R) -> Result { + let mut decoded = Vec::new(); + reader.read_to_end(&mut decoded)?; + Ok(Buffer::from_vec(decoded)) } // lz4_flex treats physical EOF (including a partial block header) as a clean end of frame. @@ -83,8 +335,14 @@ impl Read for RequireLz4EndMark { } } -fn read_single_batch(input: R, validate: bool) -> Result { - let reader = StreamReader::try_new(input, None)?; +/// General decoder: the original `StreamReader` path. Also returns the parsed schema and the raw +/// schema message it came from, for the caller to cache. +fn read_single_batch_cached( + block: &[u8], + validate: bool, +) -> Result<(RecordBatch, SchemaRef, Option<&[u8]>)> { + let mut input = Cursor::new(block); + let reader = StreamReader::try_new(&mut input, None)?; let mut reader = if validate { // Remote data must not escape as unchecked arrays and fail later in a native operator. reader @@ -92,6 +350,7 @@ fn read_single_batch(input: R, validate: bool) -> Result { // Preserve the existing local-shuffle fast path for trusted Comet-written arrays. unsafe { reader.with_skip_validation(true) } }; + let schema = reader.schema(); let batch = reader.next().transpose()?.ok_or_else(|| { DataFusionError::Execution("Failed to decode batch: empty IPC stream".to_owned()) })?; @@ -109,13 +368,22 @@ fn read_single_batch(input: R, validate: bool) -> Result { "Failed to decode batch: trailing data after IPC stream".to_owned(), )); } - Ok(batch) + + // only a leading schema message is a key the fast path can match + let schema_message = read_message(block, 0)?.and_then(|message| { + let is_schema = root_as_message(message.metadata) + .map(|parsed| parsed.header_type() == MessageHeader::Schema) + .unwrap_or(false); + is_schema.then_some(message.metadata) + }); + + Ok((batch, schema, schema_message)) } #[cfg(test)] mod tests { use super::{read_ipc_compressed, read_ipc_compressed_validated}; - use arrow::array::{Int32Array, RecordBatch, StringArray}; + use arrow::array::{Array, Int32Array, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::ipc::writer::StreamWriter; use std::io::Write; @@ -161,6 +429,128 @@ mod tests { bytes } + /// One encoded block, without the 16-byte Comet header. + fn block_for(batch: &RecordBatch, codec: &[u8; 4]) -> Vec { + let mut payload = Vec::new(); + let mut writer = StreamWriter::try_new(&mut payload, batch.schema_ref()).unwrap(); + writer.write(batch).unwrap(); + writer.finish().unwrap(); + encode(codec, &payload) + } + + fn mixed_batch() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, true), + Field::new("s", DataType::Utf8, true), + Field::new("f", DataType::Float64, false), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])), + Arc::new(StringArray::from(vec![Some("a"), Some(""), None])), + Arc::new(arrow::array::Float64Array::from(vec![1.5, -0.0, 2.25])), + ], + ) + .unwrap() + } + + fn dictionary_batch() -> RecordBatch { + let values = StringArray::from(vec!["x", "y"]); + let keys = Int32Array::from(vec![0, 1, 0]); + let dictionary = arrow::array::DictionaryArray::try_new( + keys, + Arc::new(values) as arrow::array::ArrayRef, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "d", + dictionary.data_type().clone(), + false, + )])); + RecordBatch::try_new(schema, vec![Arc::new(dictionary)]).unwrap() + } + + /// A warm decode must equal a cold one, on every codec and both entry points. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn cached_schema_decode_matches_the_first_decode() { + for batch in [mixed_batch(), dictionary_batch()] { + for codec in [b"NONE", b"LZ4_", b"ZSTD", b"SNAP"] { + let block = block_for(&batch, codec); + + let cold = read_ipc_compressed(&block).unwrap(); + let warm = read_ipc_compressed(&block).unwrap(); + assert_eq!(cold, batch, "cold decode differs, codec {codec:?}"); + assert_eq!(warm, batch, "warm decode differs, codec {codec:?}"); + assert_eq!(warm.schema(), batch.schema()); + + let validated = read_ipc_compressed_validated(&block).unwrap(); + assert_eq!( + validated, batch, + "validated decode differs, codec {codec:?}" + ); + } + } + } + + /// A dictionary block never takes the fast path, but must decode with a warm cache. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn dictionary_blocks_keep_decoding_with_a_warm_cache() { + let batch = dictionary_batch(); + let block = block_for(&batch, b"ZSTD"); + for _ in 0..3 { + assert_eq!(read_ipc_compressed(&block).unwrap(), batch); + } + } + + /// Trailing bytes after the end-of-stream marker must stay an error with a warm cache. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn trailing_data_still_fails_with_a_warm_cache() { + let batch = mixed_batch(); + let mut payload = Vec::new(); + let mut writer = StreamWriter::try_new(&mut payload, batch.schema_ref()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + + // warm the cache with the well-formed block first + let good = encode(b"NONE", &payload); + assert_eq!(read_ipc_compressed(&good).unwrap(), batch); + + let mut corrupted = payload.clone(); + corrupted.extend_from_slice(&[0u8; 8]); + let error = read_ipc_compressed(&encode(b"NONE", &corrupted)).unwrap_err(); + assert!( + error.to_string().contains("trailing data"), + "unexpected error: {error}" + ); + } + + /// A block truncated inside its body must fail cold and warm. Dropping only the + /// end-of-stream marker is not truncation: a stream ending on a message boundary is valid. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn truncated_block_fails_with_a_warm_cache() { + let batch = mixed_batch(); + let block = block_for(&batch, b"NONE"); + + // cold, before anything is cached + let cut_into_body = &block[..block.len() - 24]; + assert!(read_ipc_compressed(cut_into_body).is_err()); + + // warm, and the same truncation must still fail + assert_eq!(read_ipc_compressed(&block).unwrap(), batch); + assert!(read_ipc_compressed(cut_into_body).is_err()); + + // dropping just the end-of-stream marker stays valid + assert_eq!( + read_ipc_compressed(&block[..block.len() - 8]).unwrap(), + batch + ); + } + #[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..a9bb905c97e 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -31,7 +31,7 @@ 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 ipc::{read_ipc_compressed, read_ipc_compressed_validated, reset_schema_cache}; pub use remote_schema::{decode_remote_shuffle_batch, validate_remote_schema}; pub use schema_align::SchemaAlignExec; pub use shuffle_writer::{ShuffleWriterDestination, ShuffleWriterExec};