From d3b148870d173935346371361583afb613609bed Mon Sep 17 00:00:00 2001 From: peterxcli Date: Wed, 9 Sep 2026 22:05:25 +0800 Subject: [PATCH 1/4] 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/4] 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/4] 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/4] 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