diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index d109627e825..b3110f8a047 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3358,15 +3358,16 @@ impl PhysicalPlanner { } PartitioningStruct::SinglePartition(_) => Ok(CometPartitioning::SinglePartition), PartitioningStruct::RoundRobinPartition(rr_partition) => { - // Treat negative max_hash_columns as 0 (no limit) - let max_hash_columns = if rr_partition.max_hash_columns <= 0 { - 0 + let strategy = if rr_partition.batch_granular { + crate::execution::shuffle::RoundRobinStrategy::WholeBatch } else { - rr_partition.max_hash_columns as usize + // Treat negative max_hash_columns as 0 (no limit). + let max_hash_columns = rr_partition.max_hash_columns.max(0) as usize; + crate::execution::shuffle::RoundRobinStrategy::HashAll { max_hash_columns } }; Ok(CometPartitioning::RoundRobin( rr_partition.num_partitions as usize, - max_hash_columns, + strategy, )) } } diff --git a/native/proto/src/proto/partitioning.proto b/native/proto/src/proto/partitioning.proto index e70b8264f02..9aa6acc3224 100644 --- a/native/proto/src/proto/partitioning.proto +++ b/native/proto/src/proto/partitioning.proto @@ -57,4 +57,8 @@ message RoundRobinPartition { int32 num_partitions = 1; // Maximum number of columns to hash. 0 means no limit (hash all columns). int32 max_hash_columns = 2; + // When true, assign each incoming RecordBatch as a whole to one output partition, + // skipping per-row hashing entirely. Retry-safe when the upstream operator + // emits deterministic batches under retry. + bool batch_granular = 3; } diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index 2e82f6fc532..704701ad5b1 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -23,15 +23,18 @@ use arrow::row::{RowConverter, SortField}; use criterion::{criterion_group, criterion_main, Criterion}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_expr::expressions::{col, Column}; use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; -use datafusion::physical_plan::metrics::Time; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, Time}; use datafusion::{ physical_plan::{common::collect, ExecutionPlan}, prelude::SessionContext, }; use datafusion_comet_shuffle::{ - CometPartitioning, CompressionCodec, ShuffleBlockWriter, ShuffleWriterExec, + CometPartitioning, CompressionCodec, LocalPartitionWriter, MultiPartitionShuffleRepartitioner, + RoundRobinStrategy, ShuffleBlockWriter, ShufflePartitioner, ShufflePartitionerMetrics, + ShuffleWriterExec, }; use itertools::Itertools; use std::io::Cursor; @@ -78,8 +81,7 @@ fn criterion_benchmark(c: &mut Criterion) { let exec = create_shuffle_writer_exec( compression_codec.clone(), CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], 16), - 8192, - 10, + create_batches(8192, 10), ); b.iter(|| { let task_ctx = ctx.task_ctx(); @@ -131,8 +133,7 @@ fn criterion_benchmark(c: &mut Criterion) { let exec = create_shuffle_writer_exec( compression_codec.clone(), partitioning.clone(), - 8192, - 10, + create_batches(8192, 10), ); b.iter(|| { let task_ctx = ctx.task_ctx(); @@ -162,8 +163,7 @@ fn criterion_benchmark(c: &mut Criterion) { let exec = create_shuffle_writer_exec( CompressionCodec::None, CometPartitioning::SinglePartition, - rows_per_batch, - num_batches, + create_batches(rows_per_batch, num_batches), ); b.iter(|| { let task_ctx = ctx.task_ctx(); @@ -174,15 +174,200 @@ fn criterion_benchmark(c: &mut Criterion) { }, ); } + + // RoundRobin on a wide nested schema: compares the row-level hash-all-columns strategy + // against the whole-batch strategy. The nested schema is where the win is largest: + // `create_murmur3_hashes` recurses into every Struct child per row on the row-level path, + // while the whole-batch path does one mod per batch and never inspects row contents. + // 40 top-level struct columns of shallow depth reflect the real-world shape where this + // optimization matters (event log workload was ~194 nested columns). + let num_partitions = 50usize; + let wide_batches: Vec = (0..8).map(|_| nested_schema_batch(8192, 40, 2)).collect(); + let wide_schema = wide_batches[0].schema(); + let round_robin_strategies = [ + ( + "hash_all_columns", + RoundRobinStrategy::HashAll { + max_hash_columns: 0, + }, + ), + ("whole_batch", RoundRobinStrategy::WholeBatch), + ]; + for (label, strategy) in &round_robin_strategies { + group.bench_function( + format!("shuffle_writer: RoundRobin nested schema (strategy={label})"), + |b| { + let ctx = SessionContext::new(); + let rt = Runtime::new().unwrap(); + let exec = create_shuffle_writer_exec( + CompressionCodec::None, + CometPartitioning::RoundRobin(num_partitions, strategy.clone()), + wide_batches.clone(), + ); + b.iter(|| { + let task_ctx = ctx.task_ctx(); + let stream = exec.execute(0, task_ctx).unwrap(); + rt.block_on(collect(stream)).unwrap(); + }); + }, + ); + } + + // Partitioning-only microbench. `insert_batch` runs partition assignment + index + // buffering but never flushes to disk, so the two RoundRobin strategies can be compared + // without the IPC encode + disk write cost that dominates the end-to-end bench above. + for (label, strategy) in &round_robin_strategies { + group.bench_function( + format!("partitioning_only: RoundRobin nested schema (strategy={label})"), + |b| { + let rt = Runtime::new().unwrap(); + let batches = wide_batches.clone(); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_limit(1024 * 1024 * 1024, 1.0) + .build() + .unwrap(), + ); + let dir = tempfile::tempdir().unwrap(); + let data_path = dir.path().join("data.out").to_str().unwrap().to_string(); + let index_path = dir.path().join("index.out").to_str().unwrap().to_string(); + b.iter(|| { + let block_writer = + ShuffleBlockWriter::try_new(wide_schema.as_ref(), CompressionCodec::None) + .unwrap(); + let writer = LocalPartitionWriter::try_new( + data_path.clone(), + index_path.clone(), + block_writer, + num_partitions, + 8192, + 1024 * 1024, + Arc::clone(&runtime_env), + ) + .unwrap(); + let metrics = + ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + writer, + CometPartitioning::RoundRobin(num_partitions, strategy.clone()), + metrics, + Arc::clone(&runtime_env), + 8192, + false, + None, + ) + .unwrap(); + rt.block_on(async { + for batch in &batches { + repartitioner.insert_batch(batch.clone()).await.unwrap(); + } + }); + }); + }, + ); + } + + // End-to-end bench on a real on-disk parquet dataset at + // /tmp/clickstream_data_smoke (Spark-written, snappy-compressed, wide nested + // clickstream schema — the shape this optimization targets). Loaded once per + // process; each iteration streams the batches through `ShuffleWriterExec` + // for both RoundRobin strategies. Skipped if the dataset is absent. + let clickstream_batches = load_parquet_dir_batches("/tmp/clickstream_data_smoke"); + if let Some(batches) = clickstream_batches { + let clickstream_num_partitions = 50usize; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + let num_cols = batches[0].num_columns(); + eprintln!( + "clickstream bench: {} batches, {} rows total, {} columns", + batches.len(), + total_rows, + num_cols, + ); + for (label, strategy) in &round_robin_strategies { + group.bench_function( + format!("shuffle_writer: RoundRobin clickstream parquet (strategy={label})"), + |b| { + let ctx = SessionContext::new(); + let rt = Runtime::new().unwrap(); + let exec = create_shuffle_writer_exec( + CompressionCodec::Lz4Frame, + CometPartitioning::RoundRobin(clickstream_num_partitions, strategy.clone()), + batches.clone(), + ); + b.iter(|| { + let task_ctx = ctx.task_ctx(); + let stream = exec.execute(0, task_ctx).unwrap(); + rt.block_on(collect(stream)).unwrap(); + }); + }, + ); + } + } else { + eprintln!("clickstream bench: /tmp/clickstream_data_smoke not found; skipping"); + } +} + +/// Loads every `*.parquet` file under `dir` (recursively, one level deep for +/// partitioned datasets like `date=…/part-*.parquet`) and returns their record +/// batches at 8192 rows per batch. Returns `None` if the directory is missing +/// or contains no parquet files. +fn load_parquet_dir_batches(dir: &str) -> Option> { + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + + let root = std::path::Path::new(dir); + if !root.exists() { + return None; + } + + let mut files: Vec = Vec::new(); + let mut stack: Vec = vec![root.to_path_buf()]; + while let Some(p) = stack.pop() { + let entries = match std::fs::read_dir(&p) { + Ok(e) => e, + Err(_) => continue, + }; + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + stack.push(path); + } else if path + .extension() + .and_then(|s| s.to_str()) + .map(|s| s.eq_ignore_ascii_case("parquet")) + .unwrap_or(false) + { + files.push(path); + } + } + } + files.sort(); + if files.is_empty() { + return None; + } + + let mut all = Vec::new(); + for path in files { + let file = std::fs::File::open(&path).unwrap_or_else(|e| { + panic!("failed to open {}: {e}", path.display()); + }); + let reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap_or_else(|e| panic!("parquet builder for {}: {e}", path.display())) + .with_batch_size(8192) + .build() + .unwrap_or_else(|e| panic!("parquet build for {}: {e}", path.display())); + for batch in reader { + all.push(batch.unwrap()); + } + } + Some(all) } fn create_shuffle_writer_exec( compression_codec: CompressionCodec, partitioning: CometPartitioning, - rows_per_batch: usize, - num_batches: usize, + batches: Vec, ) -> ShuffleWriterExec { - let batches = create_batches(rows_per_batch, num_batches); let schema = batches[0].schema(); let partitions = &[batches]; ShuffleWriterExec::try_new( @@ -251,7 +436,7 @@ fn schema_encoding_benchmark(c: &mut Criterion) { for (name, batch) in [ ("flat", flat_schema_batch(8192)), - ("nested", nested_schema_batch(8192)), + ("nested", nested_schema_batch(8192, 4, 6)), ] { let writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap(); @@ -288,11 +473,8 @@ fn flat_schema_batch(num_rows: usize) -> RecordBatch { RecordBatch::try_new(schema, columns).unwrap() } -/// A schema of several deeply nested struct columns. -fn nested_schema_batch(num_rows: usize) -> RecordBatch { - let num_cols = 4; - let depth = 6; - +/// A schema of `num_cols` struct columns, each nested `depth` levels deep. +fn nested_schema_batch(num_rows: usize, num_cols: usize, depth: usize) -> RecordBatch { let mut fields: Vec = Vec::with_capacity(num_cols); let mut columns: Vec> = Vec::with_capacity(num_cols); for col in 0..num_cols { diff --git a/native/shuffle/src/bin/shuffle_bench.rs b/native/shuffle/src/bin/shuffle_bench.rs index cd43d41dbc1..8e613b881ab 100644 --- a/native/shuffle/src/bin/shuffle_bench.rs +++ b/native/shuffle/src/bin/shuffle_bench.rs @@ -45,7 +45,9 @@ use datafusion::physical_plan::common::collect; use datafusion::physical_plan::metrics::{MetricValue, MetricsSet}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::{ParquetReadOptions, SessionContext}; -use datafusion_comet_shuffle::{CometPartitioning, CompressionCodec, ShuffleWriterExec}; +use datafusion_comet_shuffle::{ + CometPartitioning, CompressionCodec, RoundRobinStrategy, ShuffleWriterExec, +}; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use std::fs; use std::path::{Path, PathBuf}; @@ -592,7 +594,12 @@ fn build_partitioning( ) -> CometPartitioning { match scheme { "single" => CometPartitioning::SinglePartition, - "round-robin" => CometPartitioning::RoundRobin(num_partitions, 0), + "round-robin" => CometPartitioning::RoundRobin( + num_partitions, + RoundRobinStrategy::HashAll { + max_hash_columns: 0, + }, + ), "hash" => { let exprs: Vec> = hash_col_indices .iter() diff --git a/native/shuffle/src/comet_partitioning.rs b/native/shuffle/src/comet_partitioning.rs index 15912e6481d..7b331ed3eb9 100644 --- a/native/shuffle/src/comet_partitioning.rs +++ b/native/shuffle/src/comet_partitioning.rs @@ -19,6 +19,20 @@ use arrow::row::{OwnedRow, RowConverter}; use datafusion::physical_expr::{LexOrdering, PhysicalExpr}; use std::sync::Arc; +/// How Comet's RoundRobin partitioning assigns rows to output partitions. +#[derive(Debug, Clone)] +pub enum RoundRobinStrategy { + /// Hash every row across up to `max_hash_columns` columns (0 means no limit) + /// and assign rows individually. Deterministic under retry regardless of upstream + /// order preservation, but pays the per-row hash cost. + HashAll { max_hash_columns: usize }, + /// Assign each incoming RecordBatch as a whole to one output partition, based on + /// the input partition id and a per-task batch sequence number. Skips per-row + /// hashing entirely. Retry-safe only when the upstream operator emits the same + /// batches in the same order under retry. + WholeBatch, +} + /// Partitioning scheme for distributing rows across shuffle output partitions. #[derive(Debug, Clone)] pub enum CometPartitioning { @@ -32,10 +46,9 @@ pub enum CometPartitioning { /// Rows for comparing to 4) OwnedRows that represent the boundaries of each partition, used with /// LexOrdering to bin each value in the RecordBatch to a partition. RangePartitioning(LexOrdering, usize, Arc, Vec), - /// Round robin partitioning. Distributes rows across partitions by sorting them by hash - /// (computed from columns) and then assigning partitions sequentially. Args are: - /// 1) number of partitions, 2) max columns to hash (0 means no limit). - RoundRobin(usize, usize), + /// Round robin partitioning. Args are the number of partitions and the strategy that + /// decides how rows are assigned. See [`RoundRobinStrategy`] for the trade-offs. + RoundRobin(usize, RoundRobinStrategy), } impl CometPartitioning { diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 2263ae0dac0..b8e2369b5e9 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -25,8 +25,17 @@ mod spark_crc32c_hasher; pub mod spark_unsafe; pub(crate) mod writers; -pub use comet_partitioning::CometPartitioning; +pub use comet_partitioning::{CometPartitioning, RoundRobinStrategy}; pub use ipc::read_ipc_compressed; pub use schema_align::SchemaAlignExec; pub use shuffle_writer::ShuffleWriterExec; pub use writers::{CompressionCodec, ShuffleBlockWriter}; + +// Bench-only re-exports. `#[doc(hidden)]` marks these as not part of the stable API contract. +// Consumed by the criterion benches under `native/shuffle/benches/`. +#[doc(hidden)] +pub use metrics::ShufflePartitionerMetrics; +#[doc(hidden)] +pub use partitioners::{MultiPartitionShuffleRepartitioner, ShufflePartitioner}; +#[doc(hidden)] +pub use writers::LocalPartitionWriter; diff --git a/native/shuffle/src/metrics.rs b/native/shuffle/src/metrics.rs index 23dbfe96a15..4bad3242f4e 100644 --- a/native/shuffle/src/metrics.rs +++ b/native/shuffle/src/metrics.rs @@ -20,7 +20,7 @@ use datafusion::physical_plan::metrics::{ }; /// Execution metrics for a shuffle partition operation. -pub(crate) struct ShufflePartitionerMetrics { +pub struct ShufflePartitionerMetrics { /// metrics pub(crate) baseline: BaselineMetrics, @@ -53,7 +53,7 @@ pub(crate) struct ShufflePartitionerMetrics { } impl ShufflePartitionerMetrics { - pub(crate) fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { Self { baseline: BaselineMetrics::new(metrics, partition), repart_time: MetricBuilder::new(metrics).subset_time("repart_time", partition), diff --git a/native/shuffle/src/partitioners/mod.rs b/native/shuffle/src/partitioners/mod.rs index 8f4239a820e..46d51603ad8 100644 --- a/native/shuffle/src/partitioners/mod.rs +++ b/native/shuffle/src/partitioners/mod.rs @@ -22,6 +22,6 @@ mod single_partition; mod traits; pub(crate) use empty_schema::EmptySchemaShufflePartitioner; -pub(crate) use multi_partition::MultiPartitionShuffleRepartitioner; +pub use multi_partition::MultiPartitionShuffleRepartitioner; pub(crate) use single_partition::SinglePartitionShufflePartitioner; -pub(crate) use traits::ShufflePartitioner; +pub use traits::ShufflePartitioner; diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index ccfe14790c2..bd99cff5863 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -19,7 +19,7 @@ use crate::metrics::ShufflePartitionerMetrics; use crate::partitioners::partitioned_batch_iterator::PartitionedBatchesProducer; use crate::partitioners::ShufflePartitioner; use crate::writers::PartitionWriter; -use crate::{comet_partitioning, CometPartitioning}; +use crate::{comet_partitioning, CometPartitioning, RoundRobinStrategy}; use arrow::array::{Array, ArrayData, ArrayRef, RecordBatch}; use datafusion::common::utils::proxy::VecAllocExt; use datafusion::common::{DataFusionError, HashSet}; @@ -99,7 +99,7 @@ impl ScratchSpace { } /// A partitioner that uses a hash function to partition data into multiple partitions -pub(crate) struct MultiPartitionShuffleRepartitioner { +pub struct MultiPartitionShuffleRepartitioner { buffered_batches: Vec, partition_indices: Vec>, partition_writer: T, @@ -129,6 +129,11 @@ pub(crate) struct MultiPartitionShuffleRepartitioner { /// Bytes in the currently buffered batches that were already counted by a previous spill of /// the same outer input batch. Partition-index allocations are never included here. repeated_spill_buffer_bytes: usize, + /// Batch counter for the batch-granular RoundRobin path. Seeded with the input partition + /// id so mappers with adjacent partition ids do not concentrate their first batches on the + /// same output partition, then incremented once per input batch slice that reaches + /// `partitioning_batch`. Unused by other strategies. + round_robin_batch_seq: usize, } /// Sum of the capacities of the backing buffers reachable from `batch` whose start address is @@ -204,7 +209,7 @@ fn count_new_buffers( impl MultiPartitionShuffleRepartitioner { #[allow(clippy::too_many_arguments)] - pub(crate) fn try_new( + pub fn try_new( partition: usize, partition_writer: T, partitioning: CometPartitioning, @@ -223,17 +228,36 @@ impl MultiPartitionShuffleRepartitioner { // Vectors in the scratch space will be filled with valid values before being used, this // initialization code is simply initializing the vectors to the desired size. // The initial values are not used. + // + // `partition_ids` and `partition_row_indices` are only touched by the row-level + // strategies (`Hash`, `RangePartitioning`, and hash-all-columns RoundRobin). The + // whole-batch RoundRobin path never inspects rows, so allocating them there would be + // ~64 KB of dead scratch on every task. + let needs_row_scratch = !matches!( + &partitioning, + CometPartitioning::SinglePartition + | CometPartitioning::RoundRobin(_, RoundRobinStrategy::WholeBatch), + ); let scratch = ScratchSpace { - hashes_buf: match partitioning { - // Allocate hashes_buf for hash and round robin partitioning. - // Round robin hashes all columns to achieve even, deterministic distribution. - CometPartitioning::Hash(_, _) | CometPartitioning::RoundRobin(_, _) => { + hashes_buf: match &partitioning { + // Allocate hashes_buf for hash and hash-all-columns round robin partitioning. + // Whole-batch round robin does no per-row hashing. + CometPartitioning::Hash(_, _) + | CometPartitioning::RoundRobin(_, RoundRobinStrategy::HashAll { .. }) => { vec![0; batch_size] } _ => vec![], }, - partition_ids: vec![0; batch_size], - partition_row_indices: vec![0; batch_size], + partition_ids: if needs_row_scratch { + vec![0; batch_size] + } else { + vec![] + }, + partition_row_indices: if needs_row_scratch { + vec![0; batch_size] + } else { + vec![] + }, partition_starts: vec![0; num_output_partitions + 1], }; @@ -255,6 +279,10 @@ impl MultiPartitionShuffleRepartitioner { pinned_buffers: HashSet::new(), spill_accounted_input_buffers: HashSet::new(), repeated_spill_buffer_bytes: 0, + // Seed with the input partition id so batches from mapper i land on + // partition (i + k) mod N on the k-th batch, spreading concurrent mappers' first + // batches across distinct output partitions. + round_robin_batch_seq: partition, }) } @@ -324,7 +352,7 @@ impl MultiPartitionShuffleRepartitioner { self.buffer_partitioned_batch_may_spill( input, - partition_row_indices, + Some(partition_row_indices), partition_starts, ) .await?; @@ -377,68 +405,85 @@ impl MultiPartitionShuffleRepartitioner { self.buffer_partitioned_batch_may_spill( input, - partition_row_indices, + Some(partition_row_indices), partition_starts, ) .await?; self.scratch = scratch; } - CometPartitioning::RoundRobin(num_output_partitions, max_hash_columns) => { - // Comet implements "round robin" as hash partitioning on columns. - // This achieves the same goal as Spark's round robin (even distribution - // without semantic grouping) while being deterministic for fault tolerance. - // - // Note: This produces different partition assignments than Spark's round robin, - // which sorts by UnsafeRow binary representation before assigning partitions. - // However, both approaches provide even distribution and determinism. + CometPartitioning::RoundRobin(num_output_partitions, strategy) => { + // Two strategies share the scratch-take / timer / spill / scratch-restore + // scaffold. Only the middle "fill partition_row_indices + partition_starts" + // step differs. WholeBatch: assign the whole input batch to one partition and + // pay no per-row cost. HashAll: hash every row (over up to max_hash_columns + // columns) and route rows individually. WholeBatch is retry-safe only when the + // upstream operator emits the same batches in the same order under retry + // (Comet's `CometNativeScan` and other order-preserving operators do so). let mut scratch = std::mem::take(&mut self.scratch); - let (partition_starts, partition_row_indices): (&Vec, &Vec) = { + let num_rows = input.num_rows(); + let partition_row_indices: Option<&[u32]> = { let mut timer = self.metrics.repart_time.timer(); - let num_rows = input.num_rows(); + let indices: Option<&[u32]> = match strategy { + RoundRobinStrategy::WholeBatch => { + let target_idx = self.round_robin_batch_seq % *num_output_partitions; + self.round_robin_batch_seq = self.round_robin_batch_seq.wrapping_add(1); + + // `partition_starts[k]..partition_starts[k+1]` is partition k's + // slice. Slots 0..=target_idx are 0; slots after `target_idx` are + // `num_rows`, so `target_idx` owns the whole [0..num_rows) range. + // `partition_row_indices` is left unmaterialized — the spill path + // treats the target partition's slice as an identity mapping. + let partition_starts = &mut scratch.partition_starts; + partition_starts.clear(); + partition_starts.resize(target_idx + 1, 0); + partition_starts.resize(*num_output_partitions + 1, num_rows as u32); + None + } + RoundRobinStrategy::HashAll { max_hash_columns } => { + // Hash-partition rows into pmod(hash, N). This produces different + // partition assignments than Spark's round robin (which sorts by + // UnsafeRow binary representation before assigning partitions), but + // both approaches provide even distribution and determinism. + // + // max_hash_columns of 0 means no limit (hash all columns). Negative + // values are normalized to 0 in the planner. + let num_columns_to_hash = if *max_hash_columns == 0 { + input.num_columns() + } else { + (*max_hash_columns).min(input.num_columns()) + }; + let columns_to_hash: Vec = (0..num_columns_to_hash) + .map(|i| Arc::clone(input.column(i))) + .collect(); + + // Use identical seed as Spark hash partitioning. + let hashes_buf = &mut scratch.hashes_buf[..num_rows]; + hashes_buf.fill(42_u32); + create_murmur3_hashes(&columns_to_hash, hashes_buf)?; + + let partition_ids = &mut scratch.partition_ids[..num_rows]; + hashes_buf.iter().enumerate().for_each(|(idx, hash)| { + partition_ids[idx] = + comet_partitioning::pmod(*hash, *num_output_partitions) as u32; + }); - // Collect columns for hashing, respecting max_hash_columns limit - // max_hash_columns of 0 means no limit (hash all columns) - // Negative values are normalized to 0 in the planner - let num_columns_to_hash = if *max_hash_columns == 0 { - input.num_columns() - } else { - (*max_hash_columns).min(input.num_columns()) + scratch.map_partition_ids_to_starts_and_indices( + *num_output_partitions, + num_rows, + ); + Some(scratch.partition_row_indices.as_slice()) + } }; - let columns_to_hash: Vec = (0..num_columns_to_hash) - .map(|i| Arc::clone(input.column(i))) - .collect(); - - // Use identical seed as Spark hash partitioning. - let hashes_buf = &mut scratch.hashes_buf[..num_rows]; - hashes_buf.fill(42_u32); - - // Compute hash for selected columns - create_murmur3_hashes(&columns_to_hash, hashes_buf)?; - - // Assign partition IDs based on hash (same as hash partitioning) - let partition_ids = &mut scratch.partition_ids[..num_rows]; - hashes_buf.iter().enumerate().for_each(|(idx, hash)| { - partition_ids[idx] = - comet_partitioning::pmod(*hash, *num_output_partitions) as u32; - }); - - // We now have partition ids for every input row, map that to partition starts - // and partition indices to eventually write these rows to partition buffers. - scratch - .map_partition_ids_to_starts_and_indices(*num_output_partitions, num_rows); timer.stop(); - Ok::<(&Vec, &Vec), DataFusionError>(( - &scratch.partition_starts, - &scratch.partition_row_indices, - )) - }?; + indices + }; self.buffer_partitioned_batch_may_spill( input, partition_row_indices, - partition_starts, + &scratch.partition_starts, ) .await?; self.scratch = scratch; @@ -457,7 +502,7 @@ impl MultiPartitionShuffleRepartitioner { async fn buffer_partitioned_batch_may_spill( &mut self, input: RecordBatch, - partition_row_indices: &[u32], + partition_row_indices: Option<&[u32]>, partition_starts: &[u32], ) -> datafusion::common::Result<()> { // Charge both the reservation and the data_size metric for the buffers this batch newly @@ -473,26 +518,29 @@ impl MultiPartitionShuffleRepartitioner { let buffered_partition_idx = self.buffered_batches.len() as u32; self.buffered_batches.push(input); - // partition_starts conceptually slices partition_row_indices into smaller slices, - // each slice contains the indices of rows in input that will go into the corresponding - // partition. The following loop iterates over the slices and put the row indices into - // the indices array of the corresponding partition. + // `partition_starts` slices the input's rows into per-partition ranges: partition K + // owns rows `partition_starts[K]..partition_starts[K + 1]`. When `partition_row_indices` + // is `Some(indices)`, `indices[start..end]` are the source-row ids in the input for + // partition K (see the hash arms). When it is `None`, the source rows are the identity + // range `start..end` itself (whole-batch RoundRobin — one target partition owns all + // input rows, and the mapping is trivial so we do not materialize it). for (partition_id, (&start, &end)) in partition_starts .iter() .tuple_windows() .enumerate() .filter(|(_, (start, end))| start < end) { - let row_indices = &partition_row_indices[start as usize..end as usize]; - - // Put row indices for the current partition into the indices array of that partition. - // This indices array will be used for calling interleave_record_batch to produce - // shuffled batches. let indices = &mut self.partition_indices[partition_id]; let before_size = indices.allocated_size(); - indices.reserve(row_indices.len()); - for row_idx in row_indices { - indices.push((buffered_partition_idx, *row_idx)); + match partition_row_indices { + Some(row_indices) => indices.extend( + row_indices[start as usize..end as usize] + .iter() + .map(|&row_idx| (buffered_partition_idx, row_idx)), + ), + None => { + indices.extend((start..end).map(|row_idx| (buffered_partition_idx, row_idx))) + } } let after_size = indices.allocated_size(); mem_growth += after_size.saturating_sub(before_size); diff --git a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs index f124d98ff24..92b7a720f74 100644 --- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs +++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs @@ -107,6 +107,36 @@ impl Iterator for PartitionedBatchIterator<'_> { let indices_end = std::cmp::min(self.pos + self.batch_size, self.indices.len()); let indices = &self.indices[self.pos..indices_end]; + + // Whole-source-batch fast path: when this chunk covers every row of one source + // batch in natural order, cloning the source batch is equivalent to `interleave` + // and skips walking every column (and every nested child on wide/nested schemas). + // This is the vectorized path used by RoundRobinStrategy::WholeBatch, but the check + // is generic and fires whenever indices happen to line up this way. + if let (Some(&(first_batch, first_row)), Some(&(last_batch, last_row))) = + (indices.first(), indices.last()) + { + let source_rows = self.record_batches[first_batch].num_rows(); + if first_batch == last_batch + && first_row == 0 + && last_row + 1 == source_rows + && indices.len() == source_rows + { + // Verify the full contiguous invariant. Invariant holds by construction in + // `buffer_partitioned_batch_may_spill`, but the check is O(indices.len()) + // with two integer compares per step and vectorizes easily. Two orders of + // magnitude cheaper than a nested-schema interleave. + let contiguous = indices + .iter() + .enumerate() + .all(|(i, &(b, r))| b == first_batch && r == i); + if contiguous { + self.pos = indices_end; + return Some(Ok(self.record_batches[first_batch].clone())); + } + } + } + let mut timer = self.interleave_time.timer(); let result = interleave_record_batch(&self.record_batches, indices); timer.stop(); diff --git a/native/shuffle/src/partitioners/traits.rs b/native/shuffle/src/partitioners/traits.rs index 9572b70db5b..dc3dfdfb871 100644 --- a/native/shuffle/src/partitioners/traits.rs +++ b/native/shuffle/src/partitioners/traits.rs @@ -19,7 +19,7 @@ use arrow::record_batch::RecordBatch; use datafusion::common::Result; #[async_trait::async_trait] -pub(crate) trait ShufflePartitioner: Send + Sync { +pub trait ShufflePartitioner: Send + Sync { /// Insert a batch into the partitioner async fn insert_batch(&mut self, batch: RecordBatch) -> Result<()>; /// Write shuffle data and shuffle index file to disk diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 46d3b592bc9..561e11b218e 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -269,7 +269,7 @@ async fn external_shuffle( #[cfg(test)] mod test { use super::*; - use crate::{read_ipc_compressed, ShuffleBlockWriter}; + use crate::{read_ipc_compressed, RoundRobinStrategy, ShuffleBlockWriter}; use arrow::array::{Array, Int64Array, StringArray, StringBuilder}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::ipc::writer::CompressionContext; @@ -868,7 +868,12 @@ mod test { Arc::new(row_converter), owned_rows, ), - CometPartitioning::RoundRobin(num_partitions, 0), + CometPartitioning::RoundRobin( + num_partitions, + RoundRobinStrategy::HashAll { + max_hash_columns: 0, + }, + ), ] { let batches = (0..num_batches).map(|_| batch.clone()).collect::>(); @@ -937,7 +942,12 @@ mod test { Arc::new(DataSourceExec::new(Arc::new( MemorySourceConfig::try_new(partitions, batch.schema(), None).unwrap(), ))), - CometPartitioning::RoundRobin(num_partitions, 0), + CometPartitioning::RoundRobin( + num_partitions, + RoundRobinStrategy::HashAll { + max_hash_columns: 0, + }, + ), CompressionCodec::Zstd(1), data_file.clone(), index_file.clone(), @@ -1005,8 +1015,212 @@ mod test { let _ = fs::remove_file("/tmp/rr_index_1.out"); } - /// Test that batch coalescing in BufBatchWriter reduces output size by - /// writing fewer, larger IPC blocks instead of many small ones. + #[test] + #[cfg_attr(miri, ignore)] + fn test_round_robin_batch_granular_retry_deterministic() { + // A retry that sees the same batches in the same order must produce byte-identical + // shuffle output. This is the contract the batch-granular path relies on in place + // of per-row content hashing. + use arrow::array::{Float64Array, Int64Array, StringArray, StructArray}; + use std::fs; + use std::io::Read; + + let num_rows = 1024; + let num_batches = 6; + let num_partitions = 8; + + // A schema with a nested struct to exercise the code path this optimization targets. + let leaf_dt = DataType::Struct( + vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Utf8, false), + Field::new("c", DataType::Float64, false), + ] + .into(), + ); + let wrapped_dt = DataType::Struct(vec![Field::new("leaf", leaf_dt.clone(), false)].into()); + let schema = Arc::new(Schema::new(vec![ + Field::new("scalar", DataType::Int64, false), + Field::new("nested", wrapped_dt.clone(), false), + ])); + + let make_batch = |seed: i64| { + let scalar = Int64Array::from_iter_values((0..num_rows as i64).map(|i| i + seed)); + let leaf = StructArray::from(vec![ + ( + Arc::new(Field::new("a", DataType::Int64, false)), + Arc::new(Int64Array::from_iter_values( + (0..num_rows as i64).map(|i| i * 3 + seed), + )) as Arc, + ), + ( + Arc::new(Field::new("b", DataType::Utf8, false)), + Arc::new(StringArray::from_iter_values( + (0..num_rows).map(|i| format!("v{i}")), + )) as Arc, + ), + ( + Arc::new(Field::new("c", DataType::Float64, false)), + Arc::new(Float64Array::from_iter_values( + (0..num_rows).map(|i| (i as f64) * 0.5), + )) as Arc, + ), + ]); + let nested = StructArray::from(vec![( + Arc::new(Field::new("leaf", leaf_dt.clone(), false)), + Arc::new(leaf) as Arc, + )]); + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(scalar), Arc::new(nested)], + ) + .unwrap() + }; + let batches: Vec = (0..num_batches as i64).map(make_batch).collect(); + + let dir = tempfile::tempdir().unwrap(); + let run = |tag: &str| -> (Vec, Vec) { + let data_file = dir.path().join(format!("{tag}_data.out")); + let index_file = dir.path().join(format!("{tag}_index.out")); + let partitions = std::slice::from_ref(&batches); + let exec = ShuffleWriterExec::try_new( + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(), + ))), + CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::WholeBatch), + CompressionCodec::Zstd(1), + data_file.to_str().unwrap().to_string(), + index_file.to_str().unwrap().to_string(), + false, + 1024 * 1024, + None, + ) + .unwrap(); + + let ctx = SessionContext::new_with_config_rt( + SessionConfig::new(), + Arc::new( + RuntimeEnvBuilder::new() + .with_memory_limit(10 * 1024 * 1024, 1.0) + .build() + .unwrap(), + ), + ); + let stream = exec.execute(0, ctx.task_ctx()).unwrap(); + Runtime::new().unwrap().block_on(collect(stream)).unwrap(); + + let mut data = Vec::new(); + fs::File::open(&data_file) + .unwrap() + .read_to_end(&mut data) + .unwrap(); + let mut index = Vec::new(); + fs::File::open(&index_file) + .unwrap() + .read_to_end(&mut index) + .unwrap(); + (data, index) + }; + + let (data_a, index_a) = run("attempt0"); + let (data_b, index_b) = run("attempt1"); + assert_eq!( + data_a, data_b, + "batch-granular round robin must produce byte-identical shuffle data on retry" + ); + assert_eq!( + index_a, index_b, + "batch-granular round robin must produce byte-identical index on retry" + ); + + // Sanity: the assignment must have distributed the 6 batches across the partitions + // (specifically, the index must show non-trivial partition sizes, not all in one bucket). + let offsets: Vec = index_a + .chunks_exact(8) + .map(|c| i64::from_le_bytes(c.try_into().unwrap())) + .collect(); + assert_eq!(offsets.len(), num_partitions + 1); + let sizes: Vec = offsets.windows(2).map(|w| w[1] - w[0]).collect(); + let nonempty = sizes.iter().filter(|&&s| s > 0).count(); + assert!( + nonempty >= 2, + "batch-granular round robin should populate at least 2 partitions across 6 batches; got sizes {sizes:?}" + ); + + // And every batch's rows must show up somewhere in the output. + let total_rows: usize = read_all_ipc_batches(&data_a) + .iter() + .map(|b| b.num_rows()) + .sum(); + assert_eq!(total_rows, num_rows * num_batches, "no rows may be dropped"); + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_round_robin_batch_granular_whole_batch_per_partition() { + // Each input batch must land entirely on a single output partition. Reading back + // the IPC blocks from any one partition should give whole batches (num_rows each), + // not row-level slices. + use std::fs; + + let num_rows = 500; + let num_batches = 4; + let num_partitions = 8; + + let batch = create_batch(num_rows); + let batches: Vec = (0..num_batches).map(|_| batch.clone()).collect(); + + let dir = tempfile::tempdir().unwrap(); + let data_file = dir.path().join("data.out"); + let index_file = dir.path().join("index.out"); + + let partitions = std::slice::from_ref(&batches); + let exec = ShuffleWriterExec::try_new( + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(partitions, batch.schema(), None).unwrap(), + ))), + CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::WholeBatch), + CompressionCodec::Zstd(1), + data_file.to_str().unwrap().to_string(), + index_file.to_str().unwrap().to_string(), + false, + 1024 * 1024, + None, + ) + .unwrap(); + + let ctx = SessionContext::new_with_config_rt( + SessionConfig::new(), + Arc::new(RuntimeEnvBuilder::new().build().unwrap()), + ); + let stream = exec.execute(0, ctx.task_ctx()).unwrap(); + Runtime::new().unwrap().block_on(collect(stream)).unwrap(); + + let data = fs::read(&data_file).unwrap(); + let index = fs::read(&index_file).unwrap(); + let offsets: Vec = index + .chunks_exact(8) + .map(|c| i64::from_le_bytes(c.try_into().unwrap())) + .collect(); + assert_eq!(offsets.len(), num_partitions + 1); + + for p in 0..num_partitions { + let start = offsets[p] as usize; + let end = offsets[p + 1] as usize; + if start == end { + continue; + } + let batches = read_all_ipc_batches(&data[start..end]); + for b in &batches { + assert_eq!( + b.num_rows(), + num_rows, + "each shuffle block in partition {p} must contain a whole input batch" + ); + } + } + } + #[test] #[cfg_attr(miri, ignore)] fn test_batch_coalescing_reduces_size() { @@ -1312,7 +1526,12 @@ mod test { Arc::new(DataSourceExec::new(Arc::new( MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(), ))), - CometPartitioning::RoundRobin(num_partitions, 0), + CometPartitioning::RoundRobin( + num_partitions, + RoundRobinStrategy::HashAll { + max_hash_columns: 0, + }, + ), CompressionCodec::Zstd(1), data_file.to_str().unwrap().to_string(), index_file.to_str().unwrap().to_string(), @@ -1401,7 +1620,12 @@ mod test { Arc::new(DataSourceExec::new(Arc::new( MemorySourceConfig::try_new(partitions, Arc::clone(&schema), None).unwrap(), ))), - CometPartitioning::RoundRobin(num_partitions, 0), + CometPartitioning::RoundRobin( + num_partitions, + RoundRobinStrategy::HashAll { + max_hash_columns: 0, + }, + ), CompressionCodec::Zstd(1), data_file.to_str().unwrap().to_string(), index_file.to_str().unwrap().to_string(), diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index 3a9a6484dba..ed555b33d98 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -61,7 +61,7 @@ enum DataOutput { /// Writes shuffle output to a single data file plus an index file recording the /// byte offset where each partition begins. See [`DataOutput`] for how the /// single- and multi-partition modes differ. -pub(crate) struct LocalPartitionWriter { +pub struct LocalPartitionWriter { output_index_file: String, data_output: DataOutput, /// Start offset of each partition in the data file, plus a trailing entry @@ -77,7 +77,7 @@ pub(crate) struct LocalPartitionWriter { } impl LocalPartitionWriter { - pub(crate) fn try_new( + pub fn try_new( output_data_file: String, output_index_file: String, shuffle_block_writer: ShuffleBlockWriter, diff --git a/native/shuffle/src/writers/mod.rs b/native/shuffle/src/writers/mod.rs index 6d330fd12aa..488db4a90d4 100644 --- a/native/shuffle/src/writers/mod.rs +++ b/native/shuffle/src/writers/mod.rs @@ -23,6 +23,6 @@ mod shuffle_block_writer; pub(crate) use buf_batch_writer::BufBatchWriter; pub(crate) use checksum::Checksum; -pub(crate) use local::local_partition_writer::LocalPartitionWriter; +pub use local::local_partition_writer::LocalPartitionWriter; pub(crate) use partition_writer::PartitionWriter; pub use shuffle_block_writer::{CompressionCodec, ShuffleBlockWriter}; diff --git a/native/shuffle/src/writers/partition_writer.rs b/native/shuffle/src/writers/partition_writer.rs index 25b0e598df8..374a077b6b3 100644 --- a/native/shuffle/src/writers/partition_writer.rs +++ b/native/shuffle/src/writers/partition_writer.rs @@ -32,7 +32,7 @@ use arrow::record_batch::RecordBatch; /// ascending id order, then a single [`finish_all`](PartitionWriter::finish_all). /// /// [`LocalPartitionWriter`]: crate::writers::local::local_partition_writer::LocalPartitionWriter -pub(crate) trait PartitionWriter: Send + Sync { +pub trait PartitionWriter: Send + Sync { /// Stages the batches from `iter` for partition `pid` without finalizing it. /// /// Used to stream single-partition output and to stage multi-partition diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 475b73ed692..729fa4b2174 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -421,6 +421,18 @@ object CometConf extends ShimCometConf { "The maximum number of columns to hash for round robin partitioning must be non-negative.") .createWithDefault(0) + val COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_BATCH_GRANULAR: ConfigEntry[Boolean] = + conf("spark.comet.shuffle.native.partitioning.roundrobin.batchGranular") + .category(CATEGORY_SHUFFLE) + .doc( + "When true, Comet's native round-robin shuffle assigns each RecordBatch as a whole " + + "to one output partition instead of hashing every column of every row. This is much " + + "cheaper on wide/nested schemas. Distribution is quasi-even at batch granularity. " + + "Retry-safe only when the upstream operator emits the same batches in the same order " + + "under retry (Comet's Parquet scan and other order-preserving operators satisfy this).") + .booleanConf + .createWithDefault(false) + val COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.convertFromSparkPlan.enabled") .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.convertFromSparkPlan.enabled") diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index 3245cc3429b..773d32a98b2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -330,6 +330,8 @@ class CometNativeShuffleWriter[K, V]( partitioning.setNumPartitions(outputPartitioning.numPartitions) partitioning.setMaxHashColumns( CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_MAX_HASH_COLUMNS.get()) + partitioning.setBatchGranular( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_BATCH_GRANULAR.get()) val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() shuffleWriterBuilder.setPartitioning(