From d3b148870d173935346371361583afb613609bed Mon Sep 17 00:00:00 2001 From: peterxcli Date: Wed, 9 Sep 2026 22:05:25 +0800 Subject: [PATCH 1/6] 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 --- native/shuffle/Cargo.toml | 4 + native/shuffle/benches/shuffle_reader.rs | 130 +++++++++++++++++++++++ 2 files changed, 134 insertions(+) create mode 100644 native/shuffle/benches/shuffle_reader.rs diff --git a/native/shuffle/Cargo.toml b/native/shuffle/Cargo.toml index 71be932422c..9504834ef4a 100644 --- a/native/shuffle/Cargo.toml +++ b/native/shuffle/Cargo.toml @@ -78,3 +78,7 @@ harness = false [[bench]] name = "row_columnar" harness = false + +[[bench]] +name = "shuffle_reader" +harness = false diff --git a/native/shuffle/benches/shuffle_reader.rs b/native/shuffle/benches/shuffle_reader.rs new file mode 100644 index 00000000000..6d7f6ce8aa4 --- /dev/null +++ b/native/shuffle/benches/shuffle_reader.rs @@ -0,0 +1,130 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shuffle read benchmarks. +//! +//! Every shuffle block is a self-contained Arrow IPC stream, so the reader parses the schema +//! flatbuffer once per block. These benchmarks measure what that costs relative to decoding the +//! block, across the shapes that make the per-block share largest: wide schemas and few rows per +//! block, which is what high partition counts and repeated spilling produce. + +use arrow::array::{Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +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 std::hint::black_box; +use std::io::Cursor; +use std::sync::Arc; + +/// Comet prefixes each block with an 8-byte compressed length and an 8-byte field count. +/// `read_ipc_compressed` expects the bytes after that header. +const BLOCK_HEADER_LEN: usize = 16; + +/// Half `Int64`, half `Utf8`, which keeps the schema flatbuffer representative of a real shuffle +/// rather than one repeated field type. +fn schema_of(num_columns: usize) -> SchemaRef { + Arc::new(Schema::new( + (0..num_columns) + .map(|i| { + let data_type = if i % 2 == 0 { + DataType::Int64 + } else { + DataType::Utf8 + }; + Field::new(format!("column_{i}"), data_type, false) + }) + .collect::>(), + )) +} + +fn batch_of(num_columns: usize, num_rows: usize) -> RecordBatch { + let schema = schema_of(num_columns); + let columns = (0..num_columns) + .map(|i| { + if i % 2 == 0 { + Arc::new( + (0..num_rows) + .map(|r| Some(r as i64)) + .collect::(), + ) as arrow::array::ArrayRef + } else { + Arc::new( + (0..num_rows) + .map(|r| Some(format!("value_{r}"))) + .collect::(), + ) as arrow::array::ArrayRef + } + }) + .collect::>(); + RecordBatch::try_new(schema, columns).unwrap() +} + +/// One encoded block, with the 16-byte Comet header stripped. +fn encode_block(batch: &RecordBatch, codec: CompressionCodec) -> Vec { + let writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), codec).unwrap(); + let mut context = IpcWriteContext::default(); + let mut buffer = Vec::new(); + let mut cursor = Cursor::new(&mut buffer); + writer + .write_batch(batch, &mut cursor, &mut context, &Time::default()) + .unwrap(); + buffer[BLOCK_HEADER_LEN..].to_vec() +} + +fn criterion_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("shuffle_reader"); + + // Rows per block shrink as partition count rises, so the narrow cases stand in for wide + // shuffles. Column counts bracket a typical projection and a wide one. + 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"); + + // Full decode of one block: schema parse plus record batch decode. + group.bench_with_input( + BenchmarkId::new("decode_block", &id), + &uncompressed, + |b, block| b.iter(|| black_box(read_ipc_compressed(black_box(block)).unwrap())), + ); + + // Schema parse alone. `StreamReader::try_new` reads and parses the schema message and + // stops before the record batch, so this is the portion a cached schema would remove. + // The 4-byte codec tag that `read_ipc_compressed` strips is skipped here as well. + group.bench_with_input( + BenchmarkId::new("parse_schema_only", &id), + &uncompressed, + |b, block| { + b.iter(|| { + let mut ipc = &black_box(block)[4..]; + black_box(StreamReader::try_new(&mut ipc, None).unwrap().schema()) + }) + }, + ); + } + } + + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); From 266b4d48666c91603d852cff7955ecefe1b574be Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 10 Sep 2026 01:15:47 +0800 Subject: [PATCH 2/6] 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 --- native/Cargo.lock | 1 + native/Cargo.toml | 1 + native/shuffle/Cargo.toml | 1 + native/shuffle/benches/shuffle_reader.rs | 18 +- native/shuffle/src/ipc.rs | 452 ++++++++++++++++++++++- native/shuffle/src/lib.rs | 2 +- 6 files changed, 457 insertions(+), 18 deletions(-) diff --git a/native/Cargo.lock b/native/Cargo.lock index e8369070804..ff88c2a93a1 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2090,6 +2090,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 ac83bec7844..82c2536dec6 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 6d7f6ce8aa4..8c3256b2f17 100644 --- a/native/shuffle/benches/shuffle_reader.rs +++ b/native/shuffle/benches/shuffle_reader.rs @@ -28,7 +28,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; @@ -107,6 +109,20 @@ fn criterion_benchmark(c: &mut Criterion) { |b, block| b.iter(|| black_box(read_ipc_compressed(black_box(block)).unwrap())), ); + // The same decode with the schema cache cleared first, so every iteration re-parses + // the schema. Measured in the same run as `decode_block` so machine drift moves both + // together and the difference between them is the cache's effect. + 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. `StreamReader::try_new` reads and parses the schema message and // stops before the record batch, so this is the portion a cached schema would remove. // The 4-byte codec tag that `read_ipc_compressed` strips is skipped here as well. diff --git a/native/shuffle/src/ipc.rs b/native/shuffle/src/ipc.rs index 97890f50148..8edefd1a4da 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,283 @@ 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. +/// +/// One is enough for a single shuffle, but a reduce task can interleave blocks from more than one +/// shuffle (a join reading both of its sides, say), and a size-one cache would thrash between +/// them. The cache is keyed on the raw schema message rather than a parsed schema, 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 handed to the fast path, which only runs for blocks that carry 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)?; + // Keep the most recently used entry 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. +/// +/// Exists so benchmarks can measure the cached and uncached decode paths against each other in a +/// single run, where machine drift affects both equally. 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 starting at `offset`, or `None` at a clean end of stream (an explicit +/// end-of-stream marker, or running out of bytes exactly on a message boundary). +/// +/// Returns `Ok(None)` only for a well-formed end; anything truncated or inconsistent is an error, +/// so a corrupt block cannot be mistaken for a short one. +fn read_message(block: &[u8], offset: usize) -> Result>> { + fn corrupt(what: &str) -> DataFusionError { + DataFusionError::Execution(format!("Failed to decode batch: {what}")) + } + + // Ending exactly on a message boundary is the legacy stream ending, which 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` reports both an end-of-stream marker and a clean boundary as "no more +/// messages", which on its own would let trailing bytes after the marker pass unnoticed. +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, avoiding a second parse of the schema +/// flatbuffer. +/// +/// Returns `Ok(None)` when the block is not the simple `[schema][record batch][end]` shape the +/// fast path handles - a dictionary message, more than one record batch, or anything unexpected - +/// so the caller can fall back to the general decoder rather than this reimplementing its rules. +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 fast path taken by the general decoder below. + 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 the block's schema message has +/// been seen before on this thread. +fn decode_block(block: Buffer, validate: bool) -> Result { + if let Some(batch) = try_decode_with_cached_schema(&block, validate) { + return Ok(batch); + } + + // General path: unchanged behaviour, and the only path that parses a schema. Its parsed + // schema is cached so later blocks carrying the same schema message take the fast path. + 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. +/// +/// This never reports an error of its own. Anything it does not handle - a cache miss, a +/// dictionary message, more than one record batch, trailing bytes, or a block that fails to +/// decode - yields `None` so the general decoder runs instead. Validation behaviour and every +/// error message therefore stay exactly as they were, and the fast path is always safe to skip. +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 be the message right after the schema, with nothing but an end of + // stream behind it. A dictionary message lands here instead and takes the general path. + 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, - )?, + // The block is materialized before decoding so its messages can be walked in place. The + // decoded arrays borrow this buffer, so it is the same allocation the general decoder would + // have made for the record batch body rather than an extra copy. + 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 +328,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 +358,16 @@ 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, over the decoded block. +/// +/// Also returns the parsed schema and the raw schema message it came from, so the caller can +/// cache them and let later blocks with the same schema skip this parse. +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 +375,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 +393,22 @@ fn read_single_batch(input: R, validate: bool) -> Result { "Failed to decode batch: trailing data after IPC stream".to_owned(), )); } - Ok(batch) + + // Only cache a leading schema message; anything else is not 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 +454,133 @@ mod tests { bytes } + /// Encodes one batch the way a Comet shuffle block carries it, without the outer 16-byte + /// Comet header that `read_ipc_compressed` does not see. + 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() + } + + /// The second decode of a block reuses the cached schema. It has to produce exactly what the + /// first one did, on every codec and on both the trusted and validated 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-carrying block never takes the fast path, but must still decode correctly + /// once its schema is cached by an earlier block. + #[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 once the schema is + /// cached. A fast path that treated "no further message" as "clean end" would accept them. + #[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 record batch body must fail whether or not its schema is + /// already cached. Dropping only the end-of-stream marker is not truncation: a stream ending + /// on a message boundary is valid, and both paths accept it. + #[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 the cache, then 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, as it was before. + 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}; From bf16d59703e30f246703a41db3ef65cb5c783a80 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 10 Sep 2026 10:35:17 +0800 Subject: [PATCH 3/6] review: trim comments to what they need to say Co-Authored-By: Claude Opus 5 --- native/shuffle/benches/shuffle_reader.rs | 23 +++++++---------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/native/shuffle/benches/shuffle_reader.rs b/native/shuffle/benches/shuffle_reader.rs index 6d7f6ce8aa4..81efcda7aa1 100644 --- a/native/shuffle/benches/shuffle_reader.rs +++ b/native/shuffle/benches/shuffle_reader.rs @@ -15,12 +15,8 @@ // specific language governing permissions and limitations // under the License. -//! Shuffle read benchmarks. -//! -//! Every shuffle block is a self-contained Arrow IPC stream, so the reader parses the schema -//! flatbuffer once per block. These benchmarks measure what that costs relative to decoding the -//! block, across the shapes that make the per-block share largest: wide schemas and few rows per -//! block, which is what high partition counts and repeated spilling produce. +//! Shuffle read benchmarks: the per-block schema parse measured against a full block decode, +//! across column counts and rows per block. use arrow::array::{Int64Array, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; @@ -33,12 +29,10 @@ use std::hint::black_box; use std::io::Cursor; use std::sync::Arc; -/// Comet prefixes each block with an 8-byte compressed length and an 8-byte field count. -/// `read_ipc_compressed` expects the bytes after that header. +/// 8-byte compressed length plus 8-byte field count; `read_ipc_compressed` expects what follows. const BLOCK_HEADER_LEN: usize = 16; -/// Half `Int64`, half `Utf8`, which keeps the schema flatbuffer representative of a real shuffle -/// rather than one repeated field type. +/// Alternating `Int64` and `Utf8`. fn schema_of(num_columns: usize) -> SchemaRef { Arc::new(Schema::new( (0..num_columns) @@ -91,8 +85,7 @@ fn encode_block(batch: &RecordBatch, codec: CompressionCodec) -> Vec { fn criterion_benchmark(c: &mut Criterion) { let mut group = c.benchmark_group("shuffle_reader"); - // Rows per block shrink as partition count rises, so the narrow cases stand in for wide - // shuffles. Column counts bracket a typical projection and a wide one. + // 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); @@ -100,16 +93,14 @@ fn criterion_benchmark(c: &mut Criterion) { let id = format!("{num_columns}col_{num_rows}row"); - // Full decode of one block: schema parse plus record batch decode. + // 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())), ); - // Schema parse alone. `StreamReader::try_new` reads and parses the schema message and - // stops before the record batch, so this is the portion a cached schema would remove. - // The 4-byte codec tag that `read_ipc_compressed` strips is skipped here as well. + // schema parse alone: `try_new` stops before the record batch. Skips the codec tag. group.bench_with_input( BenchmarkId::new("parse_schema_only", &id), &uncompressed, From cc6b3d8059a851ebf53f0c7f87d41844aefbc531 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 10 Sep 2026 10:38:27 +0800 Subject: [PATCH 4/6] review: trim comments to what they need to say Co-Authored-By: Claude Opus 5 --- native/shuffle/src/ipc.rs | 96 ++++++++++++++------------------------- 1 file changed, 33 insertions(+), 63 deletions(-) diff --git a/native/shuffle/src/ipc.rs b/native/shuffle/src/ipc.rs index 8edefd1a4da..3389f9d9df5 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -40,19 +40,14 @@ pub fn read_ipc_compressed_validated(bytes: &[u8]) -> Result { /// Arrow IPC continuation marker introducing a message length. const CONTINUATION_MARKER: [u8; 4] = [0xff, 0xff, 0xff, 0xff]; -/// Distinct schemas cached per thread. -/// -/// One is enough for a single shuffle, but a reduce task can interleave blocks from more than one -/// shuffle (a join reading both of its sides, say), and a size-one cache would thrash between -/// them. The cache is keyed on the raw schema message rather than a parsed schema, so a hit costs -/// one memcmp. +/// 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 handed to the fast path, which only runs for blocks that carry no - /// dictionary messages. + /// Empty dictionary map; the fast path only runs for blocks with no dictionary messages. static NO_DICTIONARIES: HashMap = HashMap::new(); } @@ -62,7 +57,7 @@ fn cached_schema(schema_message: &[u8]) -> Option { let hit = cache .iter() .position(|(message, _)| message.as_ref() == schema_message)?; - // Keep the most recently used entry first so an alternating pair stays resident. + // most recently used first, so an alternating pair stays resident if hit != 0 { cache.swap(0, hit); } @@ -86,10 +81,8 @@ fn cache_schema(schema_message: &[u8], schema: SchemaRef) { }); } -/// Empties this thread's schema cache, so the next decode re-parses its schema. -/// -/// Exists so benchmarks can measure the cached and uncached decode paths against each other in a -/// single run, where machine drift affects both equally. Not part of the decode contract. +/// 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()); @@ -105,17 +98,14 @@ struct IpcMessage<'a> { end: usize, } -/// Reads the message starting at `offset`, or `None` at a clean end of stream (an explicit -/// end-of-stream marker, or running out of bytes exactly on a message boundary). -/// -/// Returns `Ok(None)` only for a well-formed end; anything truncated or inconsistent is an error, -/// so a corrupt block cannot be mistaken for a short one. +/// 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 exactly on a message boundary is the legacy stream ending, which is valid. + // ending on a message boundary is the legacy stream ending, and is valid if offset == block.len() { return Ok(None); } @@ -171,10 +161,8 @@ fn read_message(block: &[u8], offset: usize) -> Result>> { })) } -/// Confirms nothing follows the record batch but a well-formed end of stream. -/// -/// `read_message` reports both an end-of-stream marker and a clean boundary as "no more -/// messages", which on its own would let trailing bytes after the marker pass unnoticed. +/// 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( @@ -206,12 +194,8 @@ fn expect_end_of_stream(block: &[u8], offset: usize) -> Result<()> { Ok(()) } -/// Decodes a block whose schema is already known, avoiding a second parse of the schema -/// flatbuffer. -/// -/// Returns `Ok(None)` when the block is not the simple `[schema][record batch][end]` shape the -/// fast path handles - a dictionary message, more than one record batch, or anything unexpected - -/// so the caller can fall back to the general decoder rather than this reimplementing its rules. +/// 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, @@ -239,7 +223,7 @@ fn decode_with_known_schema( let decoder = if validate { decoder } else { - // Matches the trusted-local fast path taken by the general decoder below. + // 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) @@ -250,15 +234,13 @@ fn decode_with_known_schema( Ok(Some(batch)) } -/// Decodes one decompressed block, reusing a cached schema when the block's schema message has -/// been seen before on this thread. +/// 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: unchanged behaviour, and the only path that parses a schema. Its parsed - // schema is cached so later blocks carrying the same schema message take the fast path. + // 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); @@ -268,10 +250,8 @@ fn decode_block(block: Buffer, validate: bool) -> Result { /// Decodes a block against an already-parsed schema, or `None` if it cannot. /// -/// This never reports an error of its own. Anything it does not handle - a cache miss, a -/// dictionary message, more than one record batch, trailing bytes, or a block that fails to -/// decode - yields `None` so the general decoder runs instead. Validation behaviour and every -/// error message therefore stay exactly as they were, and the fast path is always safe to skip. +/// 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(); @@ -285,8 +265,7 @@ fn try_decode_with_cached_schema(block: &Buffer, validate: bool) -> Option Result DataFusionError::Execution("Failed to decode batch: truncated compression codec".to_owned()) })?; let mut encoded = &bytes[4..]; - // The block is materialized before decoding so its messages can be walked in place. The - // decoded arrays borrow this buffer, so it is the same allocation the general decoder would - // have made for the record batch body rather than an extra copy. + // 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( @@ -358,10 +335,8 @@ impl Read for RequireLz4EndMark { } } -/// General decoder: the original `StreamReader` path, over the decoded block. -/// -/// Also returns the parsed schema and the raw schema message it came from, so the caller can -/// cache them and let later blocks with the same schema skip this parse. +/// 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, @@ -394,7 +369,7 @@ fn read_single_batch_cached( )); } - // Only cache a leading schema message; anything else is not a key the fast path can match. + // 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) @@ -454,8 +429,7 @@ mod tests { bytes } - /// Encodes one batch the way a Comet shuffle block carries it, without the outer 16-byte - /// Comet header that `read_ipc_compressed` does not see. + /// 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(); @@ -497,8 +471,7 @@ mod tests { RecordBatch::try_new(schema, vec![Arc::new(dictionary)]).unwrap() } - /// The second decode of a block reuses the cached schema. It has to produce exactly what the - /// first one did, on every codec and on both the trusted and validated entry points. + /// 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() { @@ -521,8 +494,7 @@ mod tests { } } - /// A dictionary-carrying block never takes the fast path, but must still decode correctly - /// once its schema is cached by an earlier block. + /// 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() { @@ -533,8 +505,7 @@ mod tests { } } - /// Trailing bytes after the end-of-stream marker must stay an error once the schema is - /// cached. A fast path that treated "no further message" as "clean end" would accept them. + /// 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() { @@ -544,7 +515,7 @@ mod tests { writer.write(&batch).unwrap(); writer.finish().unwrap(); - // Warm the cache with the well-formed block first. + // warm the cache with the well-formed block first let good = encode(b"NONE", &payload); assert_eq!(read_ipc_compressed(&good).unwrap(), batch); @@ -557,24 +528,23 @@ mod tests { ); } - /// A block truncated inside its record batch body must fail whether or not its schema is - /// already cached. Dropping only the end-of-stream marker is not truncation: a stream ending - /// on a message boundary is valid, and both paths accept it. + /// 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. + // cold, before anything is cached let cut_into_body = &block[..block.len() - 24]; assert!(read_ipc_compressed(cut_into_body).is_err()); - // Warm the cache, then the same truncation must still fail. + // 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, as it was before. + // dropping just the end-of-stream marker stays valid assert_eq!( read_ipc_compressed(&block[..block.len() - 8]).unwrap(), batch From a5290ca76c480820f629be511c151a3f33fbdb4b Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 15 Sep 2026 00:06:57 +0800 Subject: [PATCH 5/6] 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 --- native/shuffle/benches/shuffle_reader.rs | 142 +++- native/shuffle/src/ipc.rs | 925 +++++++++++++++-------- 2 files changed, 703 insertions(+), 364 deletions(-) diff --git a/native/shuffle/benches/shuffle_reader.rs b/native/shuffle/benches/shuffle_reader.rs index 3771d9fcc3d..504cff3c159 100644 --- a/native/shuffle/benches/shuffle_reader.rs +++ b/native/shuffle/benches/shuffle_reader.rs @@ -16,10 +16,11 @@ // 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 arrow::ipc::writer::IpcWriteContext; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; @@ -34,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) }) @@ -50,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 { @@ -59,13 +75,26 @@ fn batch_of(num_columns: usize, num_rows: usize) -> RecordBatch { (0..num_rows) .map(|r| Some(r as i64)) .collect::(), - ) as arrow::array::ArrayRef + ) as ArrayRef } else { - Arc::new( - (0..num_rows) - .map(|r| Some(format!("value_{r}"))) - .collect::(), - ) as arrow::array::ArrayRef + match strings { + Strings::Plain => Arc::new( + (0..num_rows) + .map(|r| Some(format!("value_{r}"))) + .collect::(), + ) as ArrayRef, + // a small dictionary that every row's key points into + Strings::Dictionary => { + let values: Vec = + (0..num_rows).map(|r| format!("value_{}", r % 16)).collect(); + Arc::new( + values + .iter() + .map(String::as_str) + .collect::>(), + ) as ArrayRef + } + } } }) .collect::>(); @@ -87,37 +116,43 @@ fn encode_block(batch: &RecordBatch, codec: CompressionCodec) -> Vec { 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"); - - // 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())), - ); + // 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); + } + } - // 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()) - }) - }, - ); + // 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..]; @@ -131,5 +166,28 @@ 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())) + }); + + // 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 3389f9d9df5..3abc9ed6ec7 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -16,15 +16,17 @@ // under the License. use arrow::array::{ArrayRef, RecordBatch}; -use arrow::buffer::Buffer; +use arrow::buffer::{Buffer, MutableBuffer}; use arrow::datatypes::SchemaRef; -use arrow::ipc::reader::{RecordBatchDecoder, StreamReader}; -use arrow::ipc::{root_as_message, MessageHeader}; +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::{Cursor, Error, ErrorKind, Read}; +use std::io::{Error, ErrorKind, Read}; use std::sync::Arc; /// Decode trusted local Comet output without revalidating every Arrow array value or offset. @@ -38,281 +40,431 @@ pub fn read_ipc_compressed_validated(bytes: &[u8]) -> Result { } /// Arrow IPC continuation marker introducing a message length. -const CONTINUATION_MARKER: [u8; 4] = [0xff, 0xff, 0xff, 0xff]; +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. Keyed on the raw schema message, so a hit costs one memcmp. +/// from several shuffles, and a single entry would thrash. 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(); +/// 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, + #[cfg(test)] + stats: SchemaCacheStats, } -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)) - }) +thread_local! { + static STATE: RefCell = RefCell::new(DecoderState::default()); } -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) +/// 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)] { - return; - } - if cache.len() == SCHEMA_CACHE_CAPACITY { - cache.pop(); + state.stats = SchemaCacheStats::default(); } - 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()); +/// 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) } -/// 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, +#[cfg(test)] +fn scratch_capacity() -> usize { + STATE.with_borrow(|state| state.scratch.capacity()) } -/// 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}")) +fn cached_schema( + schemas: &mut [(Box<[u8]>, SchemaRef)], + schema_message: &[u8], +) -> Option { + 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)) +} - // ending on a message boundary is the legacy stream ending, and is valid - if offset == block.len() { - return Ok(None); +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)); +} - let mut cursor = offset; - let first = block - .get(cursor..cursor + 4) - .ok_or_else(|| corrupt("truncated IPC message length"))?; - cursor += 4; +fn decode_error(what: &str) -> DataFusionError { + DataFusionError::Execution(format!("Failed to decode batch: {what}")) +} - 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 +fn parse_message(metadata: &[u8]) -> Result> { + root_as_message(metadata) + .map_err(|error| decode_error(&format!("unable to get root as message: {error:?}"))) +} + +fn body_length(message: &Message<'_>) -> Result { + 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 { + let codec = bytes + .get(..4) + .ok_or_else(|| decode_error("truncated compression codec"))?; + let mut encoded = &bytes[4..]; + let batch = match codec { + 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" => 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(decode_error(&format!( + "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(decode_error("trailing data after compressed stream")); + } + Ok(batch) +} - 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, - })) +fn decode<'b, S: BlockSource<'b>>(source: S, validate: bool) -> Result { + 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 + }) } -/// 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(), - ) +/// 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 { + 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 + } }; - if offset == block.len() { - return Ok(()); + // dictionaries belong to the block that carries them, never to the cached schema + let mut dictionaries: HashMap = 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 mut cursor = offset; - let first = block.get(cursor..cursor + 4).ok_or_else(trailing)?; - cursor += 4; + 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) -> Result> + where + 'b: 'a; + + /// The next message's body, `len` bytes long. + fn body(&mut self, len: usize) -> Result; + + /// 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> { let length_bytes = if first == CONTINUATION_MARKER { - let bytes = block.get(cursor..cursor + 4).ok_or_else(trailing)?; - cursor += 4; - bytes + next()? } 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()); + 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}"))), } - 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() - })?; +/// A block read through a decompressor. +struct Streamed(R); - Ok(Some(batch)) +impl Streamed { + 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() + } + }) + } } -/// 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); +impl BlockSource<'static> for Streamed { + fn next_metadata<'a>(&mut self, scratch: &'a mut Vec) -> Result> + 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())) } - // 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); + fn body(&mut self, len: usize) -> Result { + let mut body = MutableBuffer::from_len_zeroed(len); + self.read_exact(&mut body, "truncated IPC body")?; + Ok(body.into()) } - 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; + fn expect_exhausted(&mut self) -> Result<()> { + if self.0.read(&mut [0])? != 0 { + return Err(decode_error("trailing data after IPC stream")); + } + Ok(()) } +} - let schema = cached_schema(schema_message.metadata)?; +/// An uncompressed block, walked in place. +struct Sliced<'b> { + block: &'b [u8], + offset: usize, +} - // 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()?; +impl<'b> Sliced<'b> { + fn new(block: &'b [u8]) -> Self { + Self { block, offset: 0 } + } - decode_with_known_schema(block, schema, &batch_message, validate).ok()? + 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) + } } -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..]; - // 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" => decompress(zstd::Decoder::with_buffer(&mut encoded)?)?, - b"NONE" => { - let block = Buffer::from(encoded); - encoded = &[]; - block +impl<'b> BlockSource<'b> for Sliced<'b> { + fn next_metadata<'a>(&mut self, _scratch: &'a mut Vec) -> Result> + where + 'b: 'a, + { + if self.offset == self.block.len() { + return Ok(None); } - 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(), - )); + 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 { + // an exactly sized copy, with no zero fill before it + Ok(Buffer::from(self.take(len, "truncated IPC body")?)) } - 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)) + 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. @@ -335,60 +487,25 @@ impl Read for RequireLz4EndMark { } } -/// 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 - } else { - // 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()) - })?; - - // 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(), - )); - } - - // 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::{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, SchemaCacheStats, SCHEMA_CACHE_CAPACITY, SCRATCH_RETAIN_LIMIT, + }; + 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, 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 { let schema = Arc::new(Schema::new(vec![Field::new("n", DataType::Int32, false)])); let batch = RecordBatch::try_new( @@ -429,13 +546,18 @@ mod tests { bytes } - /// One encoded block, without the 16-byte Comet header. - fn block_for(batch: &RecordBatch, codec: &[u8; 4]) -> Vec { + /// One batch as a complete IPC stream. + fn ipc_bytes(batch: &RecordBatch) -> 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) + payload + } + + /// One encoded block, without the 16-byte Comet header. + fn block_for(batch: &RecordBatch, codec: &[u8; 4]) -> Vec { + encode(codec, &ipc_bytes(batch)) } fn mixed_batch() -> RecordBatch { @@ -455,69 +577,177 @@ mod tests { .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(); + /// 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 = values.iter().copied().collect(); let schema = Arc::new(Schema::new(vec![Field::new( "d", dictionary.data_type().clone(), - false, + true, )])); RecordBatch::try_new(schema, vec![Arc::new(dictionary)]).unwrap() } - /// A warm decode must equal a cold one, on every codec and both entry points. + fn strings(batch: &RecordBatch) -> Vec { + let values = arrow::compute::cast(batch.column(0), &DataType::Utf8).unwrap(); + let values = values.as_any().downcast_ref::().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::>(); + 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 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"] { + 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!(cold, batch, "cold decode differs, codec {codec:?}"); - assert_eq!(warm, batch, "warm decode differs, codec {codec:?}"); - assert_eq!(warm.schema(), batch.schema()); - + 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!( - validated, batch, - "validated decode differs, codec {codec:?}" + schema_cache_stats(), + stats(2, 1), + "codec {codec:?}, validate {validate}" ); } } } - /// A dictionary block never takes the fast path, but must decode with a warm cache. + /// 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> = (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"); + } + + /// 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 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); + fn decoded_arrays_report_the_same_memory_size_as_stream_reader() { + let num_rows = 100_000; + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new((0..num_rows).collect::()), + Arc::new( + (0..num_rows) + .map(|i| Some(format!("value_{i}"))) + .collect::(), + ), + ], + ) + .unwrap(); + 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] - #[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(); + let payload = ipc_bytes(&batch); - // warm the cache with the well-formed block first - let good = encode(b"NONE", &payload); - assert_eq!(read_ipc_compressed(&good).unwrap(), 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]); @@ -526,23 +756,26 @@ mod tests { 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] - #[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"); + reset_schema_cache(); - // cold, before anything is cached + // 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!( @@ -551,6 +784,57 @@ mod tests { ); } + /// 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!"] { @@ -562,7 +846,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() @@ -584,7 +868,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(); @@ -595,7 +879,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"); @@ -623,10 +907,7 @@ mod tests { 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 = [0_i32, 3, 6] .into_iter() @@ -640,7 +921,7 @@ 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"] { + for codec in CODECS { assert!(read_ipc_compressed_validated(&encode(codec, &payload)).is_err()); } } @@ -648,7 +929,7 @@ mod tests { #[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(); From fe93ff220e9c5199d083899a5620e1f3e87227a7 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 15 Sep 2026 16:38:49 +0800 Subject: [PATCH 6/6] 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 --- native/shuffle/benches/shuffle_reader.rs | 10 ++- native/shuffle/src/ipc.rs | 108 ++++++++++++++++++++--- native/shuffle/src/writers/mod.rs | 2 +- native/shuffle/src/writers/rss/mod.rs | 9 +- 4 files changed, 112 insertions(+), 17 deletions(-) diff --git a/native/shuffle/benches/shuffle_reader.rs b/native/shuffle/benches/shuffle_reader.rs index 504cff3c159..3af39e6b182 100644 --- a/native/shuffle/benches/shuffle_reader.rs +++ b/native/shuffle/benches/shuffle_reader.rs @@ -26,7 +26,8 @@ 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, reset_schema_cache, CompressionCodec, ShuffleBlockWriter, + read_ipc_compressed, read_ipc_compressed_validated, reset_schema_cache, CompressionCodec, + ShuffleBlockWriter, }; use std::hint::black_box; use std::io::Cursor; @@ -176,6 +177,13 @@ fn bench_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), diff --git a/native/shuffle/src/ipc.rs b/native/shuffle/src/ipc.rs index 3abc9ed6ec7..7e54367a330 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -491,13 +491,15 @@ impl Read for RequireLz4EndMark { mod tests { use super::{ read_ipc_compressed, read_ipc_compressed_validated, reset_schema_cache, schema_cache_stats, - scratch_capacity, SchemaCacheStats, SCHEMA_CACHE_CAPACITY, SCRATCH_RETAIN_LIMIT, + 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::{Cursor, Write}; + use std::io::{Cursor, Read, Write}; use std::sync::Arc; const CODECS: [&[u8; 4]; 4] = [b"NONE", b"LZ4_", b"ZSTD", b"SNAP"]; @@ -693,17 +695,13 @@ mod tests { assert_eq!(schema_cache_stats(), stats(5, 6), "most recent stays"); } - /// 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 num_rows = 100_000; + /// 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), ])); - let batch = RecordBatch::try_new( + RecordBatch::try_new( schema, vec![ Arc::new((0..num_rows).collect::()), @@ -714,7 +712,82 @@ mod tests { ), ], ) - .unwrap(); + .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(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() @@ -898,9 +971,11 @@ 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), @@ -921,8 +996,19 @@ 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()); + 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:?}"); } } diff --git a/native/shuffle/src/writers/mod.rs b/native/shuffle/src/writers/mod.rs index fb3af2c991a..4586e46c25b 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 f4de8b850be..9e0750ee73d 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; @@ -45,7 +45,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; @@ -126,7 +127,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) @@ -156,7 +157,7 @@ mod tests { }); } - pub(super) fn measure(run: impl FnOnce() -> T) -> (T, usize) { + pub(crate) fn measure(run: impl FnOnce() -> T) -> (T, usize) { struct Reset; impl Drop for Reset { fn drop(&mut self) {