This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5809-32fae06d24a1489a1b1d729372585211b4f64185 in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 587e387681859121f12eb9bc6c8bdce9ffe5074e Author: Peter Lee <[email protected]> AuthorDate: Sun Sep 20 20:11:32 2026 +0000 perf: decode shuffle blocks against a cached schema instead of re-parsing per block (#5809) * bench: add a shuffle read benchmark covering the per-block schema parse Every shuffle block is a self-contained Arrow IPC stream, so read_single_batch builds a fresh StreamReader per block and parses the schema flatbuffer once per block, even though every block in a shuffle carries the same schema. The write side already avoids the mirror image of this, encoding the schema once in ShuffleBlockWriter::try_new and writing the pre-encoded bytes verbatim, but there was no read-side benchmark to say whether the reader's half is worth removing. This adds one, parameterized by column count and rows per block, measuring the schema parse separately from the full block decode. On an M-series laptop: shape decode schema parse share 5 col x 64 row 1.93 us 1.14 us 59% 5 col x 512 row 2.38 us 0.91 us 38% 5 col x 8192 row 10.99 us 0.86 us 8% 50 col x 64 row 12.77 us 6.03 us 47% 50 col x 512 row 17.89 us 6.05 us 34% 50 col x 8192 row 218 us 6.05 us 3% The parse cost is constant per block and independent of row count, so its share is set by how many rows land in a block. That is largest exactly where the issue predicted: wide shuffles, where rows per partition are few, and repeated spilling, where each spill round emits its own block per partition. Co-Authored-By: Claude Opus 5 <[email protected]> * perf: decode shuffle blocks against a cached schema instead of re-parsing per block Every shuffle block is a self-contained Arrow IPC stream, so read_single_batch built a fresh StreamReader per block and parsed the schema flatbuffer once per block, even though every block in a shuffle carries the same schema. The write side already avoids the mirror image of this, encoding the schema once in ShuffleBlockWriter::try_new and writing the pre-encoded bytes verbatim. Blocks are now decoded against a per-thread cache keyed on the raw schema message, so a hit costs one memcmp. On a hit the block is decoded in place with RecordBatchDecoder; on a miss the original StreamReader path runs unchanged and its parsed schema is cached for later blocks. The cache holds four schemas, since a reduce task can interleave blocks from more than one shuffle and a single entry would thrash. The fast path never reports an error of its own. A cache miss, a dictionary message, more than one record batch, trailing bytes after the end-of-stream marker, or a block that simply fails to decode all fall back to the general decoder, so validation behaviour and every error message are unchanged and the fast path is always safe to skip. The measured win is not where #5792 predicted. Comparing this commit against its parent back to back, with the parse_schema_only arm as a control that this change does not touch (it drifted within 5% between the runs): shape before after change 5 col x 64 row 1.663 us 1.775 us +6.7% 5 col x 512 row 2.120 us 1.913 us -9.8% 5 col x 8192 row 11.098 us 7.841 us -29.3% 50 col x 64 row 13.479 us 12.849 us -4.7% 50 col x 512 row 18.606 us 16.090 us -13.5% 50 col x 8192 row 159.49 us 77.03 us -51.7% The issue expected the gain at small blocks, where the constant per-block parse is the largest share of decode. It is the other way round: the parse is worth under a microsecond, while decoding in place avoids the per-body MutableBuffer that StreamReader allocates and zero-fills before copying into it, and that cost scales with body size. Small blocks are marginally slower, since materializing the block and walking its messages is not repaid when the body is tiny. Co-Authored-By: Claude Opus 5 <[email protected]> * review: trim comments to what they need to say Co-Authored-By: Claude Opus 5 <[email protected]> * review: trim comments to what they need to say Co-Authored-By: Claude Opus 5 <[email protected]> * review: stream each block message by message and serve the cached schema without parsing it Replaces the materialize-then-probe fast path with one message loop that mirrors StreamReader. A cached schema is matched on its raw bytes and never verified or parsed again; the record batch and any dictionary batches are parsed once each. Bodies are read into exactly sized buffers, so a decoded batch reports the same memory as before, and dictionary blocks decode from the cache with dictionaries scoped to their own block. A #[cfg(test)] hit/miss counter proves which path each decode took; the tests reset the cache so cold and warm phases are explicit. The benchmark adds Lz4Frame, the default codec, and a dictionary-encoded string column. Co-Authored-By: Claude Fable 5.1 <[email protected]> * review: probe allocations against StreamReader, validate corrupt arrays on a warm cache, bench the validated entry point The RSS tests' allocation observer is shared with the reader tests, which compare one warm decode against the StreamReader path this change replaced: no more allocations, bytes or peak live memory on any codec. The corrupt offsets test now also fails validation with the schema served from the cache, and the benchmark times read_ipc_compressed_validated as well. Co-Authored-By: Claude Fable 5.1 <[email protected]> --------- Co-authored-by: Claude Opus 5 <[email protected]> --- native/Cargo.lock | 1 + native/Cargo.toml | 1 + native/shuffle/Cargo.toml | 1 + native/shuffle/benches/shuffle_reader.rs | 140 +++-- native/shuffle/src/ipc.rs | 873 +++++++++++++++++++++++++++++-- native/shuffle/src/lib.rs | 2 +- native/shuffle/src/writers/mod.rs | 2 +- native/shuffle/src/writers/rss/mod.rs | 9 +- 8 files changed, 934 insertions(+), 95 deletions(-) diff --git a/native/Cargo.lock b/native/Cargo.lock index d4c592e1f8..fe01bd97ec 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2066,6 +2066,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 6cf35f8676..a21658884b 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -38,6 +38,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 9504834ef4..f0ed22ad73 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 47903d002f..43d4d44b99 100644 --- a/native/shuffle/benches/shuffle_reader.rs +++ b/native/shuffle/benches/shuffle_reader.rs @@ -16,15 +16,17 @@ // under the License. //! Shuffle read benchmarks: the per-block schema parse measured against a full block decode, -//! across column counts and rows per block. +//! across column counts, rows per block, the default codec and no codec, and a dictionary-encoded +//! string column. -use arrow::array::{Int64Array, RecordBatch, StringArray}; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::array::{ArrayRef, DictionaryArray, Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; use arrow::ipc::reader::StreamReader; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use datafusion::physical_plan::metrics::Time; use datafusion_comet_shuffle::{ - read_ipc_compressed, CompressionCodec, ShuffleBlockWriter, ShuffleCodecContext, + read_ipc_compressed, read_ipc_compressed_validated, reset_schema_cache, CompressionCodec, + ShuffleBlockWriter, ShuffleCodecContext, }; use std::hint::black_box; use std::io::Cursor; @@ -33,15 +35,30 @@ use std::sync::Arc; /// 8-byte compressed length plus 8-byte field count; `read_ipc_compressed` expects what follows. const BLOCK_HEADER_LEN: usize = 16; -/// Alternating `Int64` and `Utf8`. -fn schema_of(num_columns: usize) -> SchemaRef { +/// How the odd columns hold their strings. +#[derive(Clone, Copy)] +enum Strings { + Plain, + /// `Dictionary(Int32, Utf8)`: the block carries a dictionary batch before its record batch, + /// as the JVM columnar shuffle writes for strings. + Dictionary, +} + +/// Alternating `Int64` and string columns. +fn schema_of(num_columns: usize, strings: Strings) -> SchemaRef { Arc::new(Schema::new( (0..num_columns) .map(|i| { let data_type = if i % 2 == 0 { DataType::Int64 } else { - DataType::Utf8 + match strings { + Strings::Plain => DataType::Utf8, + Strings::Dictionary => DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8), + ), + } }; Field::new(format!("column_{i}"), data_type, false) }) @@ -49,8 +66,8 @@ fn schema_of(num_columns: usize) -> SchemaRef { )) } -fn batch_of(num_columns: usize, num_rows: usize) -> RecordBatch { - let schema = schema_of(num_columns); +fn batch_of(num_columns: usize, num_rows: usize, strings: Strings) -> RecordBatch { + let schema = schema_of(num_columns, strings); let columns = (0..num_columns) .map(|i| { if i % 2 == 0 { @@ -58,13 +75,26 @@ fn batch_of(num_columns: usize, num_rows: usize) -> RecordBatch { (0..num_rows) .map(|r| Some(r as i64)) .collect::<Int64Array>(), - ) as arrow::array::ArrayRef + ) as ArrayRef } else { - Arc::new( - (0..num_rows) - .map(|r| Some(format!("value_{r}"))) - .collect::<StringArray>(), - ) as arrow::array::ArrayRef + match strings { + Strings::Plain => Arc::new( + (0..num_rows) + .map(|r| Some(format!("value_{r}"))) + .collect::<StringArray>(), + ) as ArrayRef, + // a small dictionary that every row's key points into + Strings::Dictionary => { + let values: Vec<String> = + (0..num_rows).map(|r| format!("value_{}", r % 16)).collect(); + Arc::new( + values + .iter() + .map(String::as_str) + .collect::<DictionaryArray<Int32Type>>(), + ) as ArrayRef + } + } } }) .collect::<Vec<_>>(); @@ -86,25 +116,43 @@ fn encode_block(batch: &RecordBatch, codec: CompressionCodec) -> Vec<u8> { fn criterion_benchmark(c: &mut Criterion) { let mut group = c.benchmark_group("shuffle_reader"); - // rows per block shrink as partition count rises, so the small cases stand in for wide shuffles - for num_columns in [5usize, 50] { - for num_rows in [64usize, 512, 8192] { - let batch = batch_of(num_columns, num_rows); - let uncompressed = encode_block(&batch, CompressionCodec::None); - - let id = format!("{num_columns}col_{num_rows}row"); + // Lz4Frame is the default codec; None isolates the decode from decompression. + for (codec_name, codec) in [ + ("none", CompressionCodec::None), + ("lz4", CompressionCodec::Lz4Frame), + ] { + // rows per block shrink as partition count rises, so the small cases stand in for wide + // shuffles + for num_columns in [5usize, 50] { + for num_rows in [64usize, 512, 8192] { + let batch = batch_of(num_columns, num_rows, Strings::Plain); + let block = encode_block(&batch, codec.clone()); + let id = format!("{codec_name}/{num_columns}col_{num_rows}row"); + bench_block(&mut group, &id, &block); + } + } - // full decode: 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())), - ); + // the dictionary batch before every record batch, at a narrow and a wide block + for num_rows in [64usize, 8192] { + let batch = batch_of(5, num_rows, Strings::Dictionary); + let block = encode_block(&batch, codec.clone()); + let id = format!("{codec_name}/5col_{num_rows}row_dict"); + bench_block(&mut group, &id, &block); + } + } - // schema parse alone: `try_new` stops before the record batch. Skips the codec tag. + // schema parse alone: `try_new` stops before the record batch. Skips the codec tag, so it + // only applies to uncompressed blocks. A control arm: this change does not touch it. + for num_columns in [5usize, 50] { + for num_rows in [64usize, 512, 8192] { + let batch = batch_of(num_columns, num_rows, Strings::Plain); + let block = encode_block(&batch, CompressionCodec::None); group.bench_with_input( - BenchmarkId::new("parse_schema_only", &id), - &uncompressed, + BenchmarkId::new( + "parse_schema_only", + format!("none/{num_columns}col_{num_rows}row"), + ), + &block, |b, block| { b.iter(|| { let mut ipc = &black_box(block)[4..]; @@ -118,5 +166,35 @@ fn criterion_benchmark(c: &mut Criterion) { group.finish(); } +fn bench_block( + group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>, + id: &str, + block: &[u8], +) { + // full decode with the schema served from the cache after the first iteration + group.bench_with_input(BenchmarkId::new("decode_block", id), block, |b, block| { + b.iter(|| black_box(read_ipc_compressed(black_box(block)).unwrap())) + }); + + // the remote entry point: the same decode with array validation on + group.bench_with_input( + BenchmarkId::new("decode_block_validated", id), + block, + |b, block| b.iter(|| black_box(read_ipc_compressed_validated(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), + block, + |b, block| { + b.iter(|| { + reset_schema_cache(); + black_box(read_ipc_compressed(black_box(block)).unwrap()) + }) + }, + ); +} + criterion_group!(benches, criterion_benchmark); criterion_main!(benches); diff --git a/native/shuffle/src/ipc.rs b/native/shuffle/src/ipc.rs index 97890f5014..7e54367a33 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -15,11 +15,19 @@ // 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, 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, Message, MessageHeader}; +use arrow_data::UnsafeFlag; use datafusion::common::DataFusionError; use datafusion::error::Result; +use std::cell::RefCell; +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. pub fn read_ipc_compressed(bytes: &[u8]) -> Result<RecordBatch> { @@ -31,24 +39,142 @@ pub fn read_ipc_compressed_validated(bytes: &[u8]) -> Result<RecordBatch> { read_ipc_compressed_impl(bytes, true) } +/// Arrow IPC continuation marker introducing a message length. +const CONTINUATION_MARKER: [u8; 4] = [0xff; 4]; + +/// Distinct schemas cached per thread. More than one because a reduce task can interleave blocks +/// from several shuffles, and a single entry would thrash. +const SCHEMA_CACHE_CAPACITY: usize = 4; + +/// Metadata scratch larger than this is released after the block rather than kept for the thread. +/// Real metadata is a few KiB even for wide schemas; only a corrupt length gets anywhere near. +const SCRATCH_RETAIN_LIMIT: usize = 1 << 20; + +/// Per-thread decoder state. +/// +/// Every block is a complete IPC stream that opens with a schema message. `ShuffleBlockWriter` +/// encodes that message once and writes it verbatim into every block, so consecutive blocks carry +/// byte-identical schema messages. The cache is keyed on those bytes: a hit is one memcmp, and +/// the schema message is neither verified nor parsed. +#[derive(Default)] +struct DecoderState { + /// Parsed schemas keyed on the raw schema message, most recently used first. + schemas: Vec<(Box<[u8]>, SchemaRef)>, + /// Message metadata read from a decompressor lands here, so it is not reallocated per block. + scratch: Vec<u8>, + #[cfg(test)] + stats: SchemaCacheStats, +} + +thread_local! { + static STATE: RefCell<DecoderState> = RefCell::new(DecoderState::default()); +} + +/// Empties this thread's schema cache, so the next decode re-parses its schema. For benchmarks +/// and tests comparing the cold and warm paths; not part of the decode contract. +#[doc(hidden)] +pub fn reset_schema_cache() { + STATE.with_borrow_mut(|state| { + state.schemas.clear(); + #[cfg(test)] + { + state.stats = SchemaCacheStats::default(); + } + }); +} + +/// Schema cache hits and misses on this thread since the last [`reset_schema_cache`]. +#[cfg(test)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +struct SchemaCacheStats { + hits: usize, + misses: usize, +} + +#[cfg(test)] +fn schema_cache_stats() -> SchemaCacheStats { + STATE.with_borrow(|state| state.stats) +} + +#[cfg(test)] +fn scratch_capacity() -> usize { + STATE.with_borrow(|state| state.scratch.capacity()) +} + +fn cached_schema( + schemas: &mut [(Box<[u8]>, SchemaRef)], + schema_message: &[u8], +) -> Option<SchemaRef> { + let hit = schemas + .iter() + .position(|(message, _)| message.as_ref() == schema_message)?; + // most recently used first, so an alternating pair stays resident + if hit != 0 { + schemas.swap(0, hit); + } + Some(Arc::clone(&schemas[0].1)) +} + +fn cache_schema( + schemas: &mut Vec<(Box<[u8]>, SchemaRef)>, + schema_message: &[u8], + schema: SchemaRef, +) { + if schemas.len() == SCHEMA_CACHE_CAPACITY { + schemas.pop(); + } + schemas.insert(0, (schema_message.into(), schema)); +} + +fn decode_error(what: &str) -> DataFusionError { + DataFusionError::Execution(format!("Failed to decode batch: {what}")) +} + +fn parse_message(metadata: &[u8]) -> Result<Message<'_>> { + root_as_message(metadata) + .map_err(|error| decode_error(&format!("unable to get root as message: {error:?}"))) +} + +fn body_length(message: &Message<'_>) -> Result<usize> { + usize::try_from(message.bodyLength()).map_err(|_| { + decode_error(&format!( + "invalid message body length: {}", + message.bodyLength() + )) + }) +} + fn read_ipc_compressed_impl(bytes: &[u8], validate: bool) -> Result<RecordBatch> { - let codec = bytes.get(..4).ok_or_else(|| { - DataFusionError::Execution("Failed to decode batch: truncated compression codec".to_owned()) - })?; + let codec = bytes + .get(..4) + .ok_or_else(|| decode_error("truncated compression codec"))?; 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)), + b"SNAP" => decode( + Streamed(snap::read::FrameDecoder::new(&mut encoded)), + validate, + )?, + b"LZ4_" => decode( + Streamed(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)?, + b"ZSTD" => decode( + Streamed(zstd::Decoder::with_buffer(&mut encoded)?), + validate, + )?, + // Uncompressed messages are located in place, so only bodies are copied. + b"NONE" => { + let batch = decode(Sliced::new(encoded), validate)?; + encoded = &[]; + batch + } other => { - return Err(DataFusionError::Execution(format!( - "Failed to decode batch: invalid compression codec: {other:?}" + return Err(decode_error(&format!( + "invalid compression codec: {other:?}" ))) } }; @@ -56,13 +182,291 @@ fn read_ipc_compressed_impl(bytes: &[u8], validate: bool) -> Result<RecordBatch> // 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(), - )); + return Err(decode_error("trailing data after compressed stream")); + } + Ok(batch) +} + +fn decode<'b, S: BlockSource<'b>>(source: S, validate: bool) -> Result<RecordBatch> { + STATE.with_borrow_mut(|state| { + let batch = read_single_batch(state, source, validate); + // a corrupt length can grow the scratch arbitrarily; do not pin that for the thread's life + if state.scratch.capacity() > SCRATCH_RETAIN_LIMIT { + state.scratch = Vec::new(); + } + batch + }) +} + +/// Reads one complete IPC stream holding exactly one record batch. Mirrors what +/// `arrow::ipc::reader::StreamReader` does message by message, except that the schema message is +/// served from the cache when its bytes match one already parsed. +fn read_single_batch<'b, S: BlockSource<'b>>( + state: &mut DecoderState, + mut source: S, + validate: bool, +) -> Result<RecordBatch> { + let DecoderState { + schemas, scratch, .. + } = state; + + 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, the same trust the StreamReader path placed in them. + // Remote blocks keep full validation. + unsafe { skip_validation.set(true) }; + } + + let Some(metadata) = source.next_metadata(scratch)? else { + return Err(decode_error("empty IPC stream")); + }; + let schema = match cached_schema(schemas, metadata) { + Some(schema) => { + #[cfg(test)] + { + state.stats.hits += 1; + } + schema + } + None => { + #[cfg(test)] + { + state.stats.misses += 1; + } + let message = parse_message(metadata)?; + if message.header_type() != MessageHeader::Schema { + return Err(decode_error(&format!( + "expected a schema as the first message in the stream, got: {:?}", + message.header_type() + ))); + } + let schema = message + .header_as_schema() + .ok_or_else(|| decode_error("failed to parse schema from message header"))?; + let schema = Arc::new(fb_to_schema(schema)); + // A schema message has no body. Only bodiless ones are cached, so a hit never has a + // body to skip; anything else is read past as StreamReader does, without caching. + match body_length(&message)? { + 0 => cache_schema(schemas, metadata, Arc::clone(&schema)), + len => { + source.body(len)?; + } + } + schema + } + }; + + // dictionaries belong to the block that carries them, never to the cached schema + let mut dictionaries: HashMap<i64, ArrayRef> = HashMap::new(); + let mut batch = None; + while let Some(metadata) = source.next_metadata(scratch)? { + let message = parse_message(metadata)?; + let version = message.version(); + let body_len = body_length(&message)?; + match message.header_type() { + MessageHeader::DictionaryBatch => { + let dictionary = message + .header_as_dictionary_batch() + .ok_or_else(|| decode_error("unable to read dictionary batch"))?; + let body = source.body(body_len)?; + read_dictionary_impl( + &body, + dictionary, + &schema, + &mut dictionaries, + &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 to the end-of-stream marker and reject a second batch. + if batch.is_some() { + return Err(decode_error("multiple record batches in one shuffle frame")); + } + let record_batch = message + .header_as_record_batch() + .ok_or_else(|| decode_error("unable to read record batch"))?; + let body = source.body(body_len)?; + batch = Some( + RecordBatchDecoder::try_new( + &body, + record_batch, + Arc::clone(&schema), + &dictionaries, + &version, + )? + .with_require_alignment(false) + .with_skip_validation(skip_validation.clone()) + .read_record_batch()?, + ); + } + MessageHeader::Schema => { + return Err(decode_error("expected a record batch, but found a schema")); + } + other => { + return Err(decode_error(&format!( + "unsupported message header type in IPC stream: '{other:?}'" + ))); + } + } } + + let batch = batch.ok_or_else(|| decode_error("empty IPC stream"))?; + source.expect_exhausted()?; Ok(batch) } +/// Where a block's IPC messages come from. Metadata is borrowed one message at a time; bodies +/// become exactly sized buffers that the decoded arrays keep. +/// +/// `'b` is the lifetime of an in-memory block, so [`Sliced`] can hand out metadata without +/// copying it; a streamed source uses `'static` and copies metadata into the caller's scratch. +trait BlockSource<'b> { + /// The next message's metadata, or `None` at the end of the stream: an explicit + /// end-of-stream marker, or a clean EOF on a message boundary, which is the legacy ending. + fn next_metadata<'a>(&mut self, scratch: &'a mut Vec<u8>) -> Result<Option<&'a [u8]>> + where + 'b: 'a; + + /// The next message's body, `len` bytes long. + fn body(&mut self, len: usize) -> Result<Buffer>; + + /// Errors unless every byte of the block has been consumed. + fn expect_exhausted(&mut self) -> Result<()>; +} + +/// Decodes the metadata length a message starts with, from its first four bytes and a reader for +/// four more should those be the continuation marker. `None` is the end-of-stream marker. +fn metadata_length( + first: [u8; 4], + next: impl FnOnce() -> Result<[u8; 4]>, +) -> Result<Option<usize>> { + let length_bytes = if first == CONTINUATION_MARKER { + next()? + } else { + first + }; + match i32::from_le_bytes(length_bytes) { + 0 => Ok(None), + len => usize::try_from(len) + .map(Some) + .map_err(|_| decode_error(&format!("invalid metadata length: {len}"))), + } +} + +/// A block read through a decompressor. +struct Streamed<R>(R); + +impl<R: Read> Streamed<R> { + fn read_exact(&mut self, buffer: &mut [u8], what: &str) -> Result<()> { + self.0.read_exact(buffer).map_err(|error| { + if error.kind() == ErrorKind::UnexpectedEof { + decode_error(what) + } else { + error.into() + } + }) + } +} + +impl<R: Read> BlockSource<'static> for Streamed<R> { + fn next_metadata<'a>(&mut self, scratch: &'a mut Vec<u8>) -> Result<Option<&'a [u8]>> + where + 'static: 'a, + { + let mut prefix = [0u8; 4]; + // EOF on a message boundary ends the stream; a partial length prefix does not + if self.0.read(&mut prefix[..1])? == 0 { + return Ok(None); + } + self.read_exact(&mut prefix[1..], "truncated IPC message length")?; + let Some(len) = metadata_length(prefix, || { + let mut bytes = [0u8; 4]; + self.read_exact(&mut bytes, "truncated IPC message length")?; + Ok(bytes) + })? + else { + return Ok(None); + }; + scratch.resize(len, 0); + self.read_exact(scratch, "truncated IPC metadata")?; + Ok(Some(scratch.as_slice())) + } + + fn body(&mut self, len: usize) -> Result<Buffer> { + let mut body = MutableBuffer::from_len_zeroed(len); + self.read_exact(&mut body, "truncated IPC body")?; + Ok(body.into()) + } + + fn expect_exhausted(&mut self) -> Result<()> { + if self.0.read(&mut [0])? != 0 { + return Err(decode_error("trailing data after IPC stream")); + } + Ok(()) + } +} + +/// An uncompressed block, walked in place. +struct Sliced<'b> { + block: &'b [u8], + offset: usize, +} + +impl<'b> Sliced<'b> { + fn new(block: &'b [u8]) -> Self { + Self { block, offset: 0 } + } + + fn take(&mut self, len: usize, what: &str) -> Result<&'b [u8]> { + let end = self + .offset + .checked_add(len) + .filter(|end| *end <= self.block.len()) + .ok_or_else(|| decode_error(what))?; + let bytes = &self.block[self.offset..end]; + self.offset = end; + Ok(bytes) + } +} + +impl<'b> BlockSource<'b> for Sliced<'b> { + fn next_metadata<'a>(&mut self, _scratch: &'a mut Vec<u8>) -> Result<Option<&'a [u8]>> + where + 'b: 'a, + { + if self.offset == self.block.len() { + return Ok(None); + } + let first = self.take(4, "truncated IPC message length")?; + let Some(len) = metadata_length(first.try_into().expect("four bytes"), || { + let bytes = self.take(4, "truncated IPC message length")?; + Ok(bytes.try_into().expect("four bytes")) + })? + else { + return Ok(None); + }; + Ok(Some(self.take(len, "truncated IPC metadata")?)) + } + + fn body(&mut self, len: usize) -> Result<Buffer> { + // an exactly sized copy, with no zero fill before it + Ok(Buffer::from(self.take(len, "truncated IPC body")?)) + } + + fn expect_exhausted(&mut self) -> Result<()> { + if self.offset != self.block.len() { + return Err(decode_error("trailing data after IPC stream")); + } + Ok(()) + } +} + // lz4_flex treats physical EOF (including a partial block header) as a clean end of frame. // Comet always writes an explicit LZ4 EndMark, so a decoder trying to read past the supplied // bytes has encountered a truncated frame. InvalidData is deliberate: UnexpectedEof is swallowed @@ -83,44 +487,27 @@ impl<R: Read> Read for RequireLz4EndMark<R> { } } -fn read_single_batch<R: Read>(input: R, validate: bool) -> Result<RecordBatch> { - 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 arrow::array::{Int32Array, RecordBatch, StringArray}; - use arrow::datatypes::{DataType, Field, Schema}; + use super::{ + read_ipc_compressed, read_ipc_compressed_validated, reset_schema_cache, schema_cache_stats, + scratch_capacity, RequireLz4EndMark, SchemaCacheStats, SCHEMA_CACHE_CAPACITY, + SCRATCH_RETAIN_LIMIT, + }; + use crate::writers::rss::tests::allocations; + use arrow::array::{Array, DictionaryArray, Int32Array, RecordBatch, StringArray}; + use arrow::datatypes::{DataType, Field, Int32Type, Schema}; + use arrow::ipc::reader::StreamReader; use arrow::ipc::writer::StreamWriter; - use std::io::Write; + use std::io::{Cursor, Read, Write}; use std::sync::Arc; + const CODECS: [&[u8; 4]; 4] = [b"NONE", b"LZ4_", b"ZSTD", b"SNAP"]; + + fn stats(hits: usize, misses: usize) -> SchemaCacheStats { + SchemaCacheStats { hits, misses } + } + fn ipc_stream(batch_count: usize) -> Vec<u8> { let schema = Arc::new(Schema::new(vec![Field::new("n", DataType::Int32, false)])); let batch = RecordBatch::try_new( @@ -161,6 +548,366 @@ mod tests { bytes } + /// One batch as a complete IPC stream. + fn ipc_bytes(batch: &RecordBatch) -> Vec<u8> { + let mut payload = Vec::new(); + let mut writer = StreamWriter::try_new(&mut payload, batch.schema_ref()).unwrap(); + writer.write(batch).unwrap(); + writer.finish().unwrap(); + payload + } + + /// One encoded block, without the 16-byte Comet header. + fn block_for(batch: &RecordBatch, codec: &[u8; 4]) -> Vec<u8> { + encode(codec, &ipc_bytes(batch)) + } + + 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() + } + + /// One dictionary-encoded string column; every call shares the same schema, so blocks built + /// from different values share a schema message but carry their own dictionary batch. + fn dictionary_batch(values: &[&str]) -> RecordBatch { + let dictionary: DictionaryArray<Int32Type> = values.iter().copied().collect(); + let schema = Arc::new(Schema::new(vec![Field::new( + "d", + dictionary.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![Arc::new(dictionary)]).unwrap() + } + + fn strings(batch: &RecordBatch) -> Vec<String> { + let values = arrow::compute::cast(batch.column(0), &DataType::Utf8).unwrap(); + let values = values.as_any().downcast_ref::<StringArray>().unwrap(); + values.iter().map(|v| v.unwrap().to_owned()).collect() + } + + fn n_column_batch(num_columns: usize) -> RecordBatch { + let fields = (0..num_columns) + .map(|i| Field::new(format!("c{i}"), DataType::Int32, false)) + .collect::<Vec<_>>(); + let columns = (0..num_columns) + .map(|_| Arc::new(Int32Array::from(vec![1, 2])) as arrow::array::ArrayRef) + .collect(); + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() + } + + /// After a cold decode, the same schema is served from the cache by both entry points, and + /// the warm decodes equal the cold one on every codec. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn warm_decodes_hit_the_cache_and_match_the_cold_one() { + for batch in [mixed_batch(), dictionary_batch(&["x", "y", "x"])] { + for codec in CODECS { + let block = block_for(&batch, codec); + reset_schema_cache(); + + let cold = read_ipc_compressed(&block).unwrap(); + assert_eq!(schema_cache_stats(), stats(0, 1), "codec {codec:?}"); + let warm = read_ipc_compressed(&block).unwrap(); + assert_eq!(schema_cache_stats(), stats(1, 1), "codec {codec:?}"); + let validated = read_ipc_compressed_validated(&block).unwrap(); + assert_eq!(schema_cache_stats(), stats(2, 1), "codec {codec:?}"); + + for decoded in [&cold, &warm, &validated] { + assert_eq!(decoded, &batch, "codec {codec:?}"); + assert_eq!(decoded.schema(), batch.schema(), "codec {codec:?}"); + } + } + } + } + + /// Blocks that share a schema each carry their own dictionary batch. With the schema served + /// from the cache, a record batch must still be decoded against the dictionary in its own + /// block, never against a previous block's. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn dictionaries_are_scoped_to_their_block_under_a_cached_schema() { + let first = dictionary_batch(&["a", "b", "a"]); + let second = dictionary_batch(&["x", "y", "z"]); + assert_eq!(first.schema(), second.schema()); + + for codec in CODECS { + for validate in [false, true] { + let decode = |block: &[u8]| { + if validate { + read_ipc_compressed_validated(block).unwrap() + } else { + read_ipc_compressed(block).unwrap() + } + }; + reset_schema_cache(); + assert_eq!(strings(&decode(&block_for(&first, codec))), ["a", "b", "a"]); + assert_eq!( + strings(&decode(&block_for(&second, codec))), + ["x", "y", "z"] + ); + assert_eq!(strings(&decode(&block_for(&first, codec))), ["a", "b", "a"]); + assert_eq!( + schema_cache_stats(), + stats(2, 1), + "codec {codec:?}, validate {validate}" + ); + } + } + } + + /// Each distinct schema misses once. The cache keeps several, so blocks from two shuffles + /// can alternate without evicting each other, and only the least recently used one goes + /// when the capacity is exceeded. + #[test] + fn distinct_schemas_miss_once_and_recent_ones_stay_cached() { + let blocks: Vec<Vec<u8>> = (1..=SCHEMA_CACHE_CAPACITY + 1) + .map(|num_columns| block_for(&n_column_batch(num_columns), b"NONE")) + .collect(); + let decode = |block: &[u8]| read_ipc_compressed(block).unwrap(); + + reset_schema_cache(); + decode(&blocks[0]); + decode(&blocks[1]); + decode(&blocks[0]); + decode(&blocks[1]); + assert_eq!(schema_cache_stats(), stats(2, 2)); + + // one more schema than the capacity evicts the least recently used one + for block in &blocks { + decode(block); + } + assert_eq!(schema_cache_stats(), stats(4, 5)); + decode(&blocks[0]); + assert_eq!(schema_cache_stats(), stats(4, 6), "evicted"); + decode(&blocks[SCHEMA_CACHE_CAPACITY]); + assert_eq!(schema_cache_stats(), stats(5, 6), "most recent stays"); + } + + /// An `Int32` and a `Utf8` column, `num_rows` long. + fn wide_batch(num_rows: i32) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new((0..num_rows).collect::<Int32Array>()), + Arc::new( + (0..num_rows) + .map(|i| Some(format!("value_{i}"))) + .collect::<StringArray>(), + ), + ], + ) + .unwrap() + } + + /// The reader this change replaced: a `StreamReader` per block over the decompressor, + /// exactly one batch, then the end of the stream. + fn stream_reader_decode(block: &[u8]) -> RecordBatch { + fn read<R: Read>(input: R) -> RecordBatch { + let mut reader = unsafe { + StreamReader::try_new(input, None) + .unwrap() + .with_skip_validation(true) + }; + let batch = reader.next().unwrap().unwrap(); + assert!(reader.next().is_none()); + batch + } + let mut encoded = &block[4..]; + match &block[..4] { + b"NONE" => read(&mut encoded), + b"LZ4_" => read(lz4_flex::frame::FrameDecoder::new(RequireLz4EndMark( + &mut encoded, + ))), + b"ZSTD" => read(zstd::Decoder::with_buffer(&mut encoded).unwrap()), + b"SNAP" => read(snap::read::FrameDecoder::new(&mut encoded)), + _ => unreachable!(), + } + } + + /// With the schema cached, a decode allocates no more than the `StreamReader` path did: + /// no more allocations, no more bytes, and no higher peak, on every codec, for a tiny block + /// and a typical one. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn warm_decode_allocates_no_more_than_stream_reader() { + /// (allocations, bytes requested, peak live bytes) of one decode + fn probe( + decode: impl FnOnce() -> RecordBatch, + expected: &RecordBatch, + ) -> (usize, usize, usize) { + let ((batch, (allocations, bytes)), peak) = allocations::measure(|| { + let batch = decode(); + (batch, allocations::totals()) + }); + assert_eq!(&batch, expected); + (allocations, bytes, peak) + } + + for (shape, batch) in [("3 rows", mixed_batch()), ("8192 rows", wide_batch(8192))] { + for codec in CODECS { + let block = block_for(&batch, codec); + reset_schema_cache(); + assert_eq!(read_ipc_compressed(&block).unwrap(), batch); + + let old = probe(|| stream_reader_decode(&block), &batch); + let new = probe(|| read_ipc_compressed(&block).unwrap(), &batch); + assert_eq!(schema_cache_stats(), stats(1, 1)); + + let codec = std::str::from_utf8(codec).unwrap(); + println!( + "{shape} {codec}: stream reader (allocations, bytes, peak) {old:?}, \ + cached {new:?}" + ); + assert!( + new.0 <= old.0 && new.1 <= old.1 && new.2 <= old.2, + "{shape} {codec}: {old:?} -> {new:?}" + ); + } + } + } + + /// Bodies read from a decompressor are allocated at exactly their length, as `StreamReader` + /// allocates them, so the arrays carry no growth slack and report the same memory size. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn decoded_arrays_report_the_same_memory_size_as_stream_reader() { + let batch = wide_batch(100_000); + let ipc = ipc_bytes(&batch); + let via_stream_reader = StreamReader::try_new(Cursor::new(&ipc), None) + .unwrap() + .next() + .unwrap() + .unwrap(); + + for codec in CODECS { + reset_schema_cache(); + // cold, then warm + for _ in 0..2 { + let decoded = read_ipc_compressed(&encode(codec, &ipc)).unwrap(); + assert_eq!(decoded, batch); + assert_eq!( + decoded.get_array_memory_size(), + via_stream_reader.get_array_memory_size(), + "codec {codec:?}" + ); + } + } + } + + /// Trailing bytes after the end-of-stream marker must stay an error with a warm cache. + #[test] + fn trailing_data_still_fails_with_a_warm_cache() { + let batch = mixed_batch(); + let payload = ipc_bytes(&batch); + + reset_schema_cache(); + assert_eq!( + read_ipc_compressed(&encode(b"NONE", &payload)).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}" + ); + assert_eq!(schema_cache_stats(), stats(1, 1), "failed on the warm path"); + } + + /// 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] + fn truncated_block_fails_with_a_warm_cache() { + let batch = mixed_batch(); + let block = block_for(&batch, b"NONE"); + reset_schema_cache(); + + // cold: the schema parses and is cached before the truncation is reached + let cut_into_body = &block[..block.len() - 24]; + assert!(read_ipc_compressed(cut_into_body).is_err()); + assert_eq!(schema_cache_stats(), stats(0, 1)); + + // 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()); + assert_eq!(schema_cache_stats(), stats(2, 1)); + + // dropping just the end-of-stream marker stays valid + assert_eq!( + read_ipc_compressed(&block[..block.len() - 8]).unwrap(), + batch + ); + } + + /// A partial message length after the record batch is an error on every codec, whether it + /// follows the end-of-stream marker or stands in for it. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn partial_length_prefix_is_an_error() { + let payload = ipc_stream(1); + for codec in CODECS { + let mut after_marker = payload.clone(); + after_marker.extend_from_slice(&[0, 0]); + let error = read_ipc_compressed(&encode(codec, &after_marker)) + .unwrap_err() + .to_string(); + assert!(error.contains("trailing data"), "{codec:?}: {error}"); + + let mut instead_of_marker = payload[..payload.len() - 8].to_vec(); + instead_of_marker.extend_from_slice(&[0, 0]); + let error = read_ipc_compressed(&encode(codec, &instead_of_marker)) + .unwrap_err() + .to_string(); + assert!( + error.contains("truncated IPC message length"), + "{codec:?}: {error}" + ); + } + } + + /// A corrupt metadata length makes the streamed reader grow its scratch before the read + /// fails. That growth must not stay pinned in the thread-local state afterwards. + #[test] + fn oversized_metadata_length_is_an_error_and_releases_the_scratch() { + let mut payload = ipc_stream(1); + // the record batch message follows the schema message: continuation marker, length, body + let schema_len = i32::from_le_bytes(payload[4..8].try_into().unwrap()) as usize; + let batch_message = 8 + schema_len; + assert_eq!(payload[batch_message..batch_message + 4], [0xff; 4]); + let forged = (2 * SCRATCH_RETAIN_LIMIT) as i32; + payload[batch_message + 4..batch_message + 8].copy_from_slice(&forged.to_le_bytes()); + + let error = read_ipc_compressed(&encode(b"LZ4_", &payload)) + .unwrap_err() + .to_string(); + assert!(error.contains("truncated IPC metadata"), "{error}"); + assert!(scratch_capacity() <= SCRATCH_RETAIN_LIMIT); + + // the in-place reader rejects the same length without allocating anything + let error = read_ipc_compressed(&encode(b"NONE", &payload)) + .unwrap_err() + .to_string(); + assert!(error.contains("truncated IPC metadata"), "{error}"); + } + #[test] fn malformed_codec_prefix_returns_error() { for prefix in [&b""[..], b"N", b"NO", b"NON", b"BAD!"] { @@ -172,7 +919,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. fn empty_or_multiple_batch_stream_returns_error() { - for codec in [b"NONE", b"SNAP", b"LZ4_", b"ZSTD"] { + for codec in CODECS { for batch_count in [0, 2] { let error = read_ipc_compressed(&encode(codec, &ipc_stream(batch_count))) .unwrap_err() @@ -194,7 +941,7 @@ mod tests { fn trailing_data_after_ipc_stream_returns_error() { let mut payload = ipc_stream(1); payload.extend_from_slice(b"another shuffle frame"); - for codec in [b"NONE", b"SNAP", b"LZ4_", b"ZSTD"] { + for codec in CODECS { let error = read_ipc_compressed(&encode(codec, &payload)) .unwrap_err() .to_string(); @@ -205,7 +952,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. fn trailing_data_after_compressed_stream_returns_error() { - for codec in [b"NONE", b"SNAP", b"LZ4_", b"ZSTD"] { + for codec in CODECS { let mut frame = encode(codec, &ipc_stream(1)); frame.extend_from_slice(&20_u64.to_le_bytes()); frame.extend_from_slice(b"another native frame"); @@ -224,19 +971,18 @@ mod tests { } } + /// Validation must reject a corrupt array whether the schema is parsed for this block or + /// served from the cache by an earlier valid block of the same schema. #[test] #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. - fn invalid_array_offsets_return_error() { + fn invalid_array_offsets_fail_validation_cold_and_warm() { 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", "def"]))], ) .unwrap(); - let mut payload = Vec::new(); - let mut writer = StreamWriter::try_new(&mut payload, &schema).unwrap(); - writer.write(&batch).unwrap(); - writer.finish().unwrap(); + let mut payload = ipc_bytes(&batch); let offsets: Vec<u8> = [0_i32, 3, 6] .into_iter() @@ -250,15 +996,26 @@ mod tests { assert_eq!(positions.len(), 1); // Change [0, 3, 6] to [0, 3, 2]: the second string now has decreasing offsets. payload[positions[0] + 8..positions[0] + 12].copy_from_slice(&2_i32.to_le_bytes()); - for codec in [b"NONE", b"SNAP", b"LZ4_", b"ZSTD"] { + let valid = ipc_bytes(&batch); + for codec in CODECS { + reset_schema_cache(); assert!(read_ipc_compressed_validated(&encode(codec, &payload)).is_err()); + assert_eq!( + read_ipc_compressed_validated(&encode(codec, &valid)).unwrap(), + batch + ); + assert!( + read_ipc_compressed_validated(&encode(codec, &payload)).is_err(), + "{codec:?}: warm" + ); + assert_eq!(schema_cache_stats(), stats(2, 1), "{codec:?}"); } } #[test] #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. fn valid_single_batch_frames_decode_with_all_codecs() { - for codec in [b"NONE", b"SNAP", b"LZ4_", b"ZSTD"] { + for codec in CODECS { let frame = encode(codec, &ipc_stream(1)); let batch = read_ipc_compressed(&frame).unwrap(); let validated = read_ipc_compressed_validated(&frame).unwrap(); diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 0eb18b517f..1158a2b1e2 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -33,7 +33,7 @@ pub(crate) mod writers; pub use codec_context::ShuffleCodecContext; 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::{PartitionOffsets, ShuffleWriterDestination, ShuffleWriterExec}; diff --git a/native/shuffle/src/writers/mod.rs b/native/shuffle/src/writers/mod.rs index fb3af2c991..4586e46c25 100644 --- a/native/shuffle/src/writers/mod.rs +++ b/native/shuffle/src/writers/mod.rs @@ -19,7 +19,7 @@ mod buf_batch_writer; mod checksum; mod local; mod partition_writer; -mod rss; +pub(crate) mod rss; mod shuffle_block_writer; pub(crate) use buf_batch_writer::BufBatchWriter; diff --git a/native/shuffle/src/writers/rss/mod.rs b/native/shuffle/src/writers/rss/mod.rs index 164680d4a6..061c6b53a1 100644 --- a/native/shuffle/src/writers/rss/mod.rs +++ b/native/shuffle/src/writers/rss/mod.rs @@ -18,7 +18,7 @@ pub(crate) mod rss_partition_writer; #[cfg(test)] -mod tests { +pub(crate) mod tests { use super::rss_partition_writer::RssPartitionWriter; use crate::metrics::ShufflePartitionerMetrics; use crate::writers::PartitionWriter; @@ -44,7 +44,8 @@ mod tests { /// Test-only allocation observation on a synchronous encoder thread. Production execution /// does not use thread-local state. Zstd's C allocations are covered separately by its public /// streaming-workspace estimate; this observes Rust buffers and their realloc overlap. - mod allocations { + /// Shared with the reader tests in `ipc.rs`, since a crate has one global allocator. + pub(crate) mod allocations { use std::alloc::{GlobalAlloc, Layout, System}; use std::cell::Cell; @@ -125,7 +126,7 @@ mod tests { } // Allocation/reallocation requests and requested bytes, not retained memory. - pub(super) fn totals() -> (usize, usize) { + pub(crate) fn totals() -> (usize, usize) { COUNTERS.with(|counter| { let value = counter.get().unwrap(); (value.allocations, value.allocated_bytes) @@ -155,7 +156,7 @@ mod tests { }); } - pub(super) fn measure<T>(run: impl FnOnce() -> T) -> (T, usize) { + pub(crate) fn measure<T>(run: impl FnOnce() -> T) -> (T, usize) { struct Reset; impl Drop for Reset { fn drop(&mut self) { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
