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