From ca94ed7e4e974f8610a474ae29643ebe1c74165b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 17:14:43 -0600 Subject: [PATCH 01/12] feat: positional round robin shuffle keyed on a row ordinal Comet implements Spark's round-robin shuffle as hash partitioning over every column of every row. On a wide nested schema that dominates the shuffle write: `create_murmur3_hashes` recurses into every struct child per row, and the resulting row-level scatter forces `interleave_record_batch` to walk every column and child again on flush. It is also not round robin. Placement is a pure function of a row's contents, so a column of one repeated value lands entirely on one reducer where Spark's round robin spreads it. Adds `RoundRobinStrategy::RowGroups`, which places rows the way Spark does: the row at task-global ordinal i goes to `(mapPartitionId + i / groupRows) % numPartitions`. The counter is over rows, not batches, and carries across batch boundaries, so placement does not depend on how the reader frames its input. That matters because no Spark contract covers framing: `DETERMINATE` promises the same rows in the same order and says nothing about chunking, so an operator that spills can reframe under different memory pressure while still honouring it. Keying on a row ordinal reduces the residual assumption to exactly the one Spark's own round robin makes. Positional placement is used only where that assumption is established, in two independent places that both have to hold: * `CometShuffleExchangeExec.replaysRowsInOrder` walks the native subtree fused into the writer, which the RDD graph cannot see because the subtree collapses into one `CometNativeShuffleInputRDD`. Short allowlist rather than a denylist: a native scan under nothing but projections and filters. * `CometNativeShuffleInputRDD.getOutputDeterministicLevel` applies Spark's own `isOrderSensitive` rule to everything below that RDD, reporting INDETERMINATE over a non-determinate parent so the DAGScheduler rolls the stage back instead of re-running one task into consumed output. Anything else keeps `HashAll`, which is safe to re-execute whatever its input does. Both are behind `spark.comet.shuffle.native.partitioning.roundrobin.positional.enabled`, default false. The index representation gains a second shape: positional placement records `(batch, start, len)` runs instead of one `(batch, row)` pair per row, which is smaller against the spill reservation and lets the flush copy whole ranges. `RunIterator` builds a chunk by slicing and concatenating runs, and passes a run covering an entire buffered batch straight through without copying. A schema containing `Utf8View` or `BinaryView` falls back to `HashAll`, because those are the one family whose buffers the IPC writer does not truncate for a slice. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../contributor-guide/native_shuffle.md | 114 +++- native/core/src/execution/planner.rs | 23 +- native/proto/src/proto/partitioning.proto | 8 + native/shuffle/src/bin/shuffle_bench.rs | 8 +- native/shuffle/src/comet_partitioning.rs | 269 ++++++++- native/shuffle/src/lib.rs | 2 +- .../src/partitioners/multi_partition.rs | 558 ++++++++++++++---- .../partitioned_batch_iterator.rs | 379 +++++++++++- native/shuffle/src/rss_execution_tests.rs | 8 +- native/shuffle/src/shuffle_writer.rs | 108 +++- .../scala/org/apache/comet/CometConf.scala | 36 ++ .../shuffle/CometNativeShuffleInputRDD.scala | 38 +- .../shuffle/CometNativeShuffleWriter.scala | 6 + .../shuffle/CometShuffleDependency.scala | 19 +- .../shuffle/CometShuffleExchangeExec.scala | 85 ++- ...CometNativePositionalRoundRobinSuite.scala | 201 +++++++ .../CometNativeShuffleInputRDDSuite.scala | 31 + 19 files changed, 1720 insertions(+), 175 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index a9c184a20b0..da9a940d9d0 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -507,6 +507,7 @@ jobs: org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite + org.apache.spark.sql.comet.execution.shuffle.CometNativePositionalRoundRobinSuite org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite org.apache.comet.exec.CometShuffleEncryptionSuite org.apache.comet.exec.CometShuffleManagerSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index b47ed5a46f6..d74491f09a3 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -155,6 +155,7 @@ jobs: org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite + org.apache.spark.sql.comet.execution.shuffle.CometNativePositionalRoundRobinSuite org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite org.apache.comet.exec.CometShuffleEncryptionSuite org.apache.comet.exec.CometShuffleManagerSuite diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index 8d84aa2f6b8..ef71b168d1d 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -305,7 +305,12 @@ batch is written as a single block that may exceed the batch size. ### Round Robin Partitioning -Comet implements round robin partitioning using hash-based assignment for determinism: +`CometPartitioning::RoundRobin` carries a `RoundRobinStrategy` that decides how rows reach output +partitions. The default is `HashAll`; `RowGroups` is opt-in through +`spark.comet.shuffle.native.partitioning.roundrobin.positional.enabled`, and is used only where the +planner can establish that the map task replays its rows in the same order. + +#### `HashAll`: hash-based assignment (default) 1. Computes a Murmur3 hash of columns (using seed 42) 2. Assigns partitions directly using the hash: `partition_id = hash % num_partitions` @@ -314,13 +319,86 @@ This approach guarantees determinism across retries, which is critical for fault However, unlike true round robin which cycles through partitions row-by-row, hash-based assignment only provides even distribution when the data has sufficient variation in the hashed columns. Data with low cardinality or identical values may result in skewed partition -sizes. +sizes: because placement is a pure function of a row's contents, a column of one repeated value +lands entirely on one reducer. Because Spark assigns round robin partitions by sorting rows on their binary `UnsafeRow` form, which Arrow's layout does not reproduce, unsorted output can land in different partitions than Spark's. Sorted output is identical. That difference is why `spark.comet.shuffle.native.partitioning.roundrobin.enabled` defaults to `false`. +`spark.comet.shuffle.native.partitioning.roundrobin.maxHashColumns` caps how many leading columns +are hashed. `0`, the default, hashes all of them. + +#### `RowGroups`: positional assignment + +Hashing every column of every row dominates the shuffle write on wide nested schemas, because +`create_murmur3_hashes` recurses into every struct child per row and the resulting row-level +scatter forces `interleave_record_batch` to walk every column and child again on flush. +`RowGroups` places rows the way Spark's own round robin does: the row at task-global ordinal `i` +goes to `(mapPartitionId + i / groupRows) % numPartitions`. That removes the per-row hash, and it +replaces the per-row gather with a bulk copy per contiguous run, because adjacent rows now stay +together. It also spreads duplicate rows evenly, which `HashAll` cannot. + +The counter is over **rows**, not batches, and it carries across batch boundaries: a group that one +input batch leaves part-way through is finished by the next. That is deliberate. Spark's +determinism contract, `DeterministicLevel`, describes row _order_ and says nothing about how a +downstream operator frames rows into batches, so an operator that spills can reframe under +different memory pressure while still honouring `DETERMINATE`. Keying on a row ordinal means the +strategy depends only on the property Spark actually publishes. + +`start_partition` is the Spark map partition id, filled in by `PhysicalPlanner::create_partitioning` +from the planner's partition because `ShuffleWriterExec::execute` cannot supply it (`jni_api` runs +every native root plan with partition 0, one Comet execution per Spark task). It has to be distinct +across mappers, or every task starts at partition 0 and a task emitting fewer groups than there are +output partitions leaves the tail empty stage-wide; and it has to be a pure function of the map +partition, or a re-executed task does not reproduce its own placement. Spark seeds +`XORShiftRandom(partitionId)` for the same two reasons. + +`groupRows` trades balance against copying. Imbalance between any two output partitions is bounded +by `groupRows` rows however the reader frames its batches, so small groups balance better; large +groups produce fewer, longer runs to copy, and a group as large as the batch size lets a whole +input batch pass through to one partition untouched. `0`, the default, derives it as +`clamp(batch_size / num_partitions, 64, batch_size)`. + +Internally, `MultiPartitionShuffleRepartitioner` records `(batch, start, len)` runs rather than +one `(batch, row)` pair per row, so the index list charged against the spill reservation is +smaller, and `RunIterator` builds each output chunk by slicing and concatenating runs. A run that +covers an entire buffered batch and already fills a chunk is passed through without copying. + +#### Retry safety under `RowGroups` + +Positional assignment is not a function of the rows, so it is reproducible only when the map task +replays the same rows in the same order. Re-executing one map task against differently ordered +input writes a different partitioning of the same rows, and once any consumer has fetched the +output that attempt replaces, the reduce side silently gets some rows twice and others not at all +([SPARK-23207](https://issues.apache.org/jira/browse/SPARK-23207)). Spark faces the same problem +with its own round robin. Comet establishes the condition in two places, both of which must hold: + +- **In the plan.** `CometShuffleExchangeExec.replaysRowsInOrder` walks the native subtree fused + into the writer, which the RDD graph cannot see because the whole subtree collapses into one + `CometNativeShuffleInputRDD`. It is a short allowlist, not a denylist: a native scan under + nothing but projections and filters. Operators that spill are the interesting exclusion, since + an aggregate or sort under memory pressure emits output in an order that depends on how many + times it spilled, which differs between attempts. Anything else keeps `HashAll`, which is safe + to re-execute whatever its input does. + +- **In the RDD graph.** Spark wraps a round-robin repartition in a `MapPartitionsRDD` with + `isOrderSensitive = true`, which reports `INDETERMINATE` whenever its parent is `UNORDERED`; the + `DAGScheduler` then rolls the whole stage back rather than re-running a single task. The native + path has no `MapPartitionsRDD` to carry the flag, so + `CometNativeShuffleInputRDD.getOutputDeterministicLevel` applies the same rule directly. A + determinate parent such as a plain scan stays determinate and keeps cheap per-task retry; + anything below another exchange is unordered, because reduce tasks see shuffle blocks in arrival + order, and goes indeterminate. + +Neither applies to `HashAll`, whose output is a pure function of the rows it sees. + +One schema-level restriction is applied in `create_repartitioner`: positional placement is the only +strategy that hands a sliced array to the IPC writer, and while the writer truncates a slice's +buffers for every other type, for `Utf8View` and `BinaryView` it serializes every shared data +buffer in full. A schema containing a view type anywhere therefore falls back to `HashAll`. + ## Memory Management Native shuffle uses DataFusion's memory management with spilling support: @@ -370,21 +448,23 @@ independently compressed, allowing parallel decompression during reads. ## Configuration -| Config | Default | Description | -| ------------------------------------------------------------------- | ------- | ------------------------------------------------------------- | -| `spark.comet.shuffle.enabled` | `true` | Enable Comet shuffle | -| `spark.comet.shuffle.mode` | `auto` | Shuffle mode: `native`, `jvm`, or `auto` | -| `spark.comet.shuffle.directRead.enabled` | `true` | Decode shuffle blocks in native code, bypassing Arrow FFI | -| `spark.comet.shuffle.compression.codec` | `lz4` | Compression codec | -| `spark.comet.shuffle.compression.zstd.level` | `1` | Zstd compression level | -| `spark.comet.shuffle.native.writeBufferSize` | `1MB` | Write buffer size | -| `spark.comet.shuffle.native.maxBufferBytes` | `0` | Fixed spill threshold. `0` disables it, leaving pool pressure | -| `spark.comet.shuffle.native.partitioning.hash.enabled` | `true` | Allow `HashPartitioning` on the native path | -| `spark.comet.shuffle.native.partitioning.hash.nested.enabled` | `false` | Allow struct and array hash keys, and map keys on Spark 4.0+ | -| `spark.comet.shuffle.native.partitioning.range.enabled` | `true` | Allow `RangePartitioning` on the native path | -| `spark.comet.shuffle.native.partitioning.roundrobin.enabled` | `false` | Allow `RoundRobinPartitioning` on the native path | -| `spark.comet.shuffle.native.partitioning.roundrobin.maxHashColumns` | `0` | Columns to hash for round robin. `0` hashes all of them | -| `spark.comet.shuffle.jvm.batchSize` | `8192` | Target rows per batch | +| Config | Default | Description | +| ------------------------------------------------------------------------- | ------- | ----------------------------------------------------------------------------- | +| `spark.comet.shuffle.enabled` | `true` | Enable Comet shuffle | +| `spark.comet.shuffle.mode` | `auto` | Shuffle mode: `native`, `jvm`, or `auto` | +| `spark.comet.shuffle.directRead.enabled` | `true` | Decode shuffle blocks in native code, bypassing Arrow FFI | +| `spark.comet.shuffle.compression.codec` | `lz4` | Compression codec | +| `spark.comet.shuffle.compression.zstd.level` | `1` | Zstd compression level | +| `spark.comet.shuffle.native.writeBufferSize` | `1MB` | Write buffer size | +| `spark.comet.shuffle.native.maxBufferBytes` | `0` | Fixed spill threshold. `0` disables it, leaving pool pressure | +| `spark.comet.shuffle.native.partitioning.hash.enabled` | `true` | Allow `HashPartitioning` on the native path | +| `spark.comet.shuffle.native.partitioning.hash.nested.enabled` | `false` | Allow struct and array hash keys, and map keys on Spark 4.0+ | +| `spark.comet.shuffle.native.partitioning.range.enabled` | `true` | Allow `RangePartitioning` on the native path | +| `spark.comet.shuffle.native.partitioning.roundrobin.enabled` | `false` | Allow `RoundRobinPartitioning` on the native path | +| `spark.comet.shuffle.native.partitioning.roundrobin.maxHashColumns` | `0` | Columns to hash for round robin. `0` hashes all of them | +| `spark.comet.shuffle.native.partitioning.roundrobin.positional.enabled` | `false` | Place round-robin rows by position where the plan allows it | +| `spark.comet.shuffle.native.partitioning.roundrobin.positional.groupRows` | `0` | Rows per positional group. `0` derives it from batch size and partition count | +| `spark.comet.shuffle.jvm.batchSize` | `8192` | Target rows per batch | ## Comparison with JVM Shuffle diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 8f030da455b..c3092537228 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -98,7 +98,7 @@ use datafusion_comet_spark_expr::{ use iceberg::expr::Bind; use crate::execution::operators::ExecutionError::GeneralError; -use crate::execution::shuffle::{CometPartitioning, CompressionCodec}; +use crate::execution::shuffle::{CometPartitioning, CompressionCodec, RoundRobinStrategy}; use crate::execution::spark_plan::SparkPlan; use crate::parquet::objectstore::s3_blob_fs_support::normalize_object_store_url; use crate::parquet::parquet_support::prepare_object_store_with_configs; @@ -3652,15 +3652,26 @@ 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.positional { + // The Spark map partition id, not the DataFusion one: `jni_api` runs every + // native root plan with partition 0 (one Comet execution per Spark task), so + // `ShuffleWriterExec::execute` cannot supply it. See + // `RoundRobinStrategy::RowGroups` for why it has to be this value. + RoundRobinStrategy::RowGroups { + start_partition: self.partition.max(0) as usize, + // Negative or zero means "derive it from the batch size and partition + // count", which the repartitioner does once it knows both. + group_rows: rr_partition.positional_group_rows.max(0) as usize, + } } else { - rr_partition.max_hash_columns as usize + // Treat negative max_hash_columns as 0 (no limit). + RoundRobinStrategy::HashAll { + max_hash_columns: rr_partition.max_hash_columns.max(0) as usize, + } }; 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..5cb29218c7a 100644 --- a/native/proto/src/proto/partitioning.proto +++ b/native/proto/src/proto/partitioning.proto @@ -57,4 +57,12 @@ 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, place rows by position rather than by hashing their contents: the row at + // task-global ordinal i goes to (mapPartitionId + i / positional_group_rows) % num_partitions. + // Only set when the planner has established that the map task replays rows in the same order, + // which is the condition Spark's own round robin relies on. + bool positional = 3; + // Rows per contiguous group under positional placement. 0 means derive it from the batch size + // and the partition count. + int32 positional_group_rows = 4; } diff --git a/native/shuffle/src/bin/shuffle_bench.rs b/native/shuffle/src/bin/shuffle_bench.rs index 051e18fbdd8..c2c9db30f0b 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}; @@ -583,7 +585,9 @@ fn build_partitioning( ) -> CometPartitioning { match scheme { "single" => CometPartitioning::SinglePartition, - "round-robin" => CometPartitioning::RoundRobin(num_partitions, 0), + "round-robin" => { + CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::default()) + } "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..f6b9ed257a8 100644 --- a/native/shuffle/src/comet_partitioning.rs +++ b/native/shuffle/src/comet_partitioning.rs @@ -19,6 +19,131 @@ use arrow::row::{OwnedRow, RowConverter}; use datafusion::physical_expr::{LexOrdering, PhysicalExpr}; use std::sync::Arc; +/// How [`CometPartitioning::RoundRobin`] decides which output partition a row belongs to. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RoundRobinStrategy { + /// Hash each row over its leading `max_hash_columns` columns (`0` meaning all of them) and + /// place it at `pmod(hash, num_partitions)`. + /// + /// Placement is a pure function of a row's contents, so a re-executed map task reproduces it + /// no matter what its input does. The price is a murmur3 pass per row that recurses into + /// every struct child, plus a per-row gather on flush because adjacent rows scatter across + /// every partition. It is also not really round robin: identical rows always hash to the same + /// partition, so low-cardinality input skews where Spark's round robin spreads evenly. + HashAll { max_hash_columns: usize }, + + /// Place rows positionally, in contiguous groups of `group_rows` rows, counting rows across + /// input batch boundaries: the row at task-global ordinal `i` goes to output partition + /// `(start_partition + i / group_rows) % num_partitions`. + /// + /// This is Spark's own round robin at a coarser granularity — Spark seeds a counter with + /// `XORShiftRandom(partitionId)` and bumps it per row, which is the `group_rows == 1` case — + /// and it inherits Spark's determinism condition exactly: placement is reproducible when the + /// upstream operator replays rows in the same *order*. It deliberately does not depend on how + /// those rows are framed into batches, because no Spark contract covers framing; + /// `DeterministicLevel::DETERMINATE` promises the same rows in the same order and says + /// nothing about how a downstream operator chunks them, so an operator that spills can reframe + /// under different memory pressure while still honouring it. Keying on a row ordinal rather + /// than a batch ordinal is what lets this strategy rely on the level Spark already publishes + /// instead of an assumption nothing checks. + /// + /// `start_partition` must be the Spark map partition id. It has to be distinct across mappers, + /// or every task starts at partition 0 and a task emitting fewer groups than there are output + /// partitions leaves the tail empty stage-wide; and it has to be a pure function of the map + /// partition, or a re-executed task does not reproduce its own placement. Spark seeds + /// `XORShiftRandom(partitionId)` for the same two reasons. + /// + /// `group_rows` trades balance against copying. Imbalance between any two output partitions is + /// bounded by `group_rows` rows regardless of how the reader frames batches, so small groups + /// balance better; large groups produce fewer, longer runs to copy on flush, and a group as + /// large as the batch size lets a whole input batch pass through to one partition untouched. + /// [`Self::AUTO_GROUP_ROWS`] picks a value from the batch size and partition count. + RowGroups { + start_partition: usize, + group_rows: usize, + }, +} + +impl Default for RoundRobinStrategy { + /// Hashing every column, which is what Comet's round robin did before `RowGroups` existed. + fn default() -> Self { + Self::HashAll { + max_hash_columns: 0, + } + } +} + +impl RoundRobinStrategy { + /// `group_rows` sentinel asking for a value derived from the batch size and partition count. + pub const AUTO_GROUP_ROWS: usize = 0; + + /// Smallest automatically chosen group. A multiple of 8 so that a run starts on a byte + /// boundary of a validity bitmap, which keeps the per-run copy a memcpy rather than a + /// bit-shift for every column. + const MIN_AUTO_GROUP_ROWS: usize = 64; + + /// Resolves [`Self::AUTO_GROUP_ROWS`] against the runtime batch size and partition count. + /// + /// One batch spread over `num_partitions` groups is the finest split that still gives every + /// output partition a run, so `batch_size / num_partitions` balances without fragmenting the + /// copy any further than it has to. + pub fn resolve_group_rows( + group_rows: usize, + batch_size: usize, + num_partitions: usize, + ) -> usize { + let batch_size = batch_size.max(1); + if group_rows != Self::AUTO_GROUP_ROWS { + return group_rows.min(batch_size); + } + (batch_size / num_partitions.max(1)) + .clamp(Self::MIN_AUTO_GROUP_ROWS.min(batch_size), batch_size) + } +} + +/// A contiguous span of rows within one input batch, bound for one output partition. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct PositionalRun { + pub partition: usize, + pub start: u32, + pub len: u32, +} + +/// Splits the rows `[row_seq, row_seq + num_rows)` of a task's input into the runs that +/// [`RoundRobinStrategy::RowGroups`] placement produces, appending them to `out` in row order. +/// +/// `row_seq` is the count of rows the task has already placed, which is what makes the split +/// independent of where batch boundaries happen to fall: a group straddling two input batches +/// comes back as a trailing run of the first and a leading run of the second, and the rows land +/// on the same partition either way. +pub(crate) fn positional_runs( + row_seq: u64, + num_rows: usize, + start_partition: usize, + group_rows: usize, + num_partitions: usize, + out: &mut Vec, +) { + out.clear(); + let group_rows = group_rows.max(1) as u64; + let num_partitions = num_partitions.max(1) as u64; + let num_rows = num_rows as u64; + let mut offset = 0u64; + while offset < num_rows { + let global = row_seq + offset; + // Rows left in the group `global` falls into, so the first run of a batch picks up a + // group that a previous batch left part-way through. + let remaining_in_group = group_rows - (global % group_rows); + let len = remaining_in_group.min(num_rows - offset); + out.push(PositionalRun { + partition: ((start_partition as u64 + global / group_rows) % num_partitions) as usize, + start: offset as u32, + len: len as u32, + }); + offset += len; + } +} + /// Partitioning scheme for distributing rows across shuffle output partitions. #[derive(Debug, Clone)] pub enum CometPartitioning { @@ -32,10 +157,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 1) the number of partitions and 2) the strategy that + /// decides where each row goes. See [`RoundRobinStrategy`] for the trade-offs. + RoundRobin(usize, RoundRobinStrategy), } impl CometPartitioning { @@ -69,4 +193,141 @@ mod tests { let expected = vec![69, 5, 193, 171, 115]; assert_eq!(result, expected); } + + /// Collects the partition of every row in `[row_seq, row_seq + num_rows)` by expanding the + /// runs, which is the property the runs are a compressed encoding of. + fn placement( + row_seq: u64, + num_rows: usize, + group_rows: usize, + num_partitions: usize, + ) -> Vec { + let mut runs = vec![]; + positional_runs(row_seq, num_rows, 0, group_rows, num_partitions, &mut runs); + runs.iter() + .flat_map(|run| std::iter::repeat_n(run.partition, run.len as usize)) + .collect() + } + + #[test] + fn positional_runs_cover_every_row_once_in_order() { + let mut runs = vec![]; + positional_runs(0, 10, 0, 4, 3, &mut runs); + assert_eq!( + runs, + vec![ + PositionalRun { + partition: 0, + start: 0, + len: 4 + }, + PositionalRun { + partition: 1, + start: 4, + len: 4 + }, + PositionalRun { + partition: 2, + start: 8, + len: 2 + }, + ] + ); + } + + /// The point of counting rows rather than batches: however the reader frames the same rows, + /// each row lands on the same partition. + #[test] + fn positional_placement_is_independent_of_batch_framing() { + let group_rows = 7; + let num_partitions = 5; + let total = 100; + + let whole = placement(0, total, group_rows, num_partitions); + + for framing in [ + vec![100], + vec![1; 100], + vec![8; 12].into_iter().chain([4]).collect::>(), + vec![64, 36], + vec![7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 2], + ] { + assert_eq!(framing.iter().sum::(), total, "bad framing fixture"); + let mut row_seq = 0u64; + let mut refrained = vec![]; + for rows in framing.iter() { + refrained.extend(placement(row_seq, *rows, group_rows, num_partitions)); + row_seq += *rows as u64; + } + assert_eq!( + refrained, whole, + "framing {framing:?} placed rows differently" + ); + } + } + + /// A group that a previous batch left part-way through is finished by the next batch, rather + /// than restarting at a group boundary. + #[test] + fn positional_runs_resume_a_partial_group() { + let mut runs = vec![]; + positional_runs(2, 6, 0, 4, 3, &mut runs); + assert_eq!( + runs, + vec![ + // rows 2..4 finish group 0 + PositionalRun { + partition: 0, + start: 0, + len: 2 + }, + PositionalRun { + partition: 1, + start: 2, + len: 4 + }, + ] + ); + } + + #[test] + fn positional_runs_wrap_and_offset_by_start_partition() { + let mut runs = vec![]; + positional_runs(0, 6, 2, 2, 3, &mut runs); + assert_eq!( + runs.iter().map(|r| r.partition).collect::>(), + vec![2, 0, 1], + "start_partition offsets the sequence and it wraps at num_partitions" + ); + } + + #[test] + fn positional_runs_group_larger_than_batch_yields_one_run() { + let mut runs = vec![]; + positional_runs(0, 100, 3, 8192, 200, &mut runs); + assert_eq!( + runs, + vec![PositionalRun { + partition: 3, + start: 0, + len: 100 + }] + ); + } + + #[test] + fn resolve_group_rows_auto_splits_a_batch_across_partitions() { + use RoundRobinStrategy as S; + // One batch spread over the output partitions, floored at the 64-row minimum. + assert_eq!(S::resolve_group_rows(S::AUTO_GROUP_ROWS, 8192, 16), 512); + assert_eq!(S::resolve_group_rows(S::AUTO_GROUP_ROWS, 8192, 200), 64); + assert_eq!(S::resolve_group_rows(S::AUTO_GROUP_ROWS, 8192, 10_000), 64); + // A batch smaller than the minimum group still resolves to something usable. + assert_eq!(S::resolve_group_rows(S::AUTO_GROUP_ROWS, 32, 200), 32); + // An explicit request is honoured, but never exceeds the batch size: a group longer than + // a batch can never see a second batch's rows anyway, since `insert_batch` slices input + // down to `batch_size` before placing it. + assert_eq!(S::resolve_group_rows(1, 8192, 200), 1); + assert_eq!(S::resolve_group_rows(100_000, 8192, 200), 8192); + } } diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 1158a2b1e2e..93183051ecb 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -32,7 +32,7 @@ pub mod spark_unsafe; pub(crate) mod writers; pub use codec_context::ShuffleCodecContext; -pub use comet_partitioning::CometPartitioning; +pub use comet_partitioning::{CometPartitioning, RoundRobinStrategy}; 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; diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 37f67eacbd6..054b84fe463 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -15,11 +15,14 @@ // specific language governing permissions and limitations // under the License. +use crate::comet_partitioning::{positional_runs, PositionalRun}; use crate::metrics::ShufflePartitionerMetrics; -use crate::partitioners::partitioned_batch_iterator::PartitionedBatchesProducer; +use crate::partitioners::partitioned_batch_iterator::{ + BufferedRun, PartitionIndices, 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}; @@ -48,6 +51,9 @@ struct ScratchSpace { /// partition_starts[K + 1] are the start and end indices of partition K in partition_row_indices. /// The length of this array is 1 + the number of partitions. partition_starts: Vec, + /// The runs the current batch splits into under positional round robin. Only ever non-empty + /// for [`RoundRobinStrategy::RowGroups`], which uses none of the row-level buffers above. + positional_runs: Vec, } impl ScratchSpace { @@ -101,7 +107,7 @@ impl ScratchSpace { /// A partitioner that uses a hash function to partition data into multiple partitions pub(crate) struct MultiPartitionShuffleRepartitioner { buffered_batches: Vec, - partition_indices: Vec>, + partition_indices: PartitionIndices, partition_writer: T, /// Partitioning scheme to use partitioning: CometPartitioning, @@ -122,6 +128,12 @@ pub(crate) struct MultiPartitionShuffleRepartitioner { /// allocation once rather than once per slice that references it. Cleared whenever the /// buffered batches drain (spill / shuffle_write). See `count_new_buffers`. pinned_buffers: HashSet, + /// `Some((start_partition, group_rows))` under [`RoundRobinStrategy::RowGroups`], with + /// `group_rows` already resolved against the batch size and partition count. + positional: Option<(usize, usize)>, + /// Rows this task has placed so far, which is the ordinal positional round robin keys on. + /// A `u64` because it counts a whole task's input, not one batch. + row_seq: u64, } /// Sum of the capacities of the backing buffers reachable from `batch` whose start address is @@ -188,21 +200,54 @@ impl MultiPartitionShuffleRepartitioner { "Use SinglePartitionShufflePartitioner for 1 output partition." ); + // Positional round robin is the one strategy that never looks at a row's contents. It + // needs none of the row-level scratch (~64 KB a task), and it records contiguous runs + // rather than individual rows. Resolve it here, before `partitioning` moves into the + // struct below. + let positional = match &partitioning { + CometPartitioning::RoundRobin( + _, + RoundRobinStrategy::RowGroups { + start_partition, + group_rows, + }, + ) => Some(( + *start_partition, + RoundRobinStrategy::resolve_group_rows( + *group_rows, + batch_size, + num_output_partitions, + ), + )), + _ => None, + }; + let places_rows_individually = positional.is_none(); + // 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. 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. + // Positional 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 places_rows_individually { + vec![0; batch_size] + } else { + vec![] + }, + partition_row_indices: if places_rows_individually { + vec![0; batch_size] + } else { + vec![] + }, partition_starts: vec![0; num_output_partitions + 1], + positional_runs: vec![], }; let reservation = MemoryConsumer::new(format!("ShuffleRepartitioner[{partition}]")) @@ -211,7 +256,11 @@ impl MultiPartitionShuffleRepartitioner { Ok(Self { buffered_batches: vec![], - partition_indices: vec![vec![]; num_output_partitions], + partition_indices: if places_rows_individually { + PartitionIndices::Rows(vec![vec![]; num_output_partitions]) + } else { + PartitionIndices::Runs(vec![vec![]; num_output_partitions]) + }, partition_writer, partitioning, metrics, @@ -221,6 +270,8 @@ impl MultiPartitionShuffleRepartitioner { max_buffer_bytes, tracing_enabled, pinned_buffers: HashSet::new(), + positional, + row_seq: 0, }) } @@ -349,65 +400,95 @@ impl MultiPartitionShuffleRepartitioner { .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) => { let mut scratch = std::mem::take(&mut self.scratch); - let (partition_starts, partition_row_indices): (&Vec, &Vec) = { - let mut timer = self.metrics.repart_time.timer(); - - let num_rows = input.num_rows(); - - // 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()) - }; - 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, - )) - }?; + let num_rows = input.num_rows(); + + match strategy { + RoundRobinStrategy::RowGroups { .. } => { + let (start_partition, group_rows) = self + .positional + .expect("positional resolved in try_new for RowGroups"); + { + let mut timer = self.metrics.repart_time.timer(); + positional_runs( + self.row_seq, + num_rows, + start_partition, + group_rows, + *num_output_partitions, + &mut scratch.positional_runs, + ); + timer.stop(); + } + // Count rows, not batches: a group left part-way through by this batch is + // finished by the next one, so placement does not depend on where the + // reader put the batch boundary. See `RoundRobinStrategy::RowGroups`. + self.row_seq += num_rows as u64; + self.scratch = scratch; + self.buffer_positional_batch_may_spill(input).await?; + } + RoundRobinStrategy::HashAll { max_hash_columns } => { + // Comet implements this flavour of "round robin" as hash partitioning on + // columns: even distribution without semantic grouping, and deterministic + // for fault tolerance regardless of input order. + // + // Note: This produces different partition assignments than Spark's round + // robin, which sorts by UnsafeRow binary representation before assigning + // partitions. + let (partition_starts, partition_row_indices): (&Vec, &Vec) = { + let mut timer = self.metrics.repart_time.timer(); + + // 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()) + }; + 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; + }); - self.buffer_partitioned_batch_may_spill( - input, - partition_row_indices, - partition_starts, - ) - .await?; - self.scratch = scratch; + // 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, + )) + }?; + + self.buffer_partitioned_batch_may_spill( + input, + partition_row_indices, + partition_starts, + ) + .await?; + self.scratch = scratch; + } + } } other => { // this should be unreachable as long as the validation logic @@ -426,13 +507,13 @@ impl MultiPartitionShuffleRepartitioner { partition_row_indices: &[u32], partition_starts: &[u32], ) -> datafusion::common::Result<()> { - // Charge both the reservation and the data_size metric for the buffers this batch newly - // pins; `count_new_buffers` dedups buffers shared across already-buffered batches. - let new_buffer_bytes = count_new_buffers(&input, &mut self.pinned_buffers); - self.metrics.data_size.add(new_buffer_bytes); - let mut mem_growth: usize = new_buffer_bytes; - let buffered_partition_idx = self.buffered_batches.len() as u32; - self.buffered_batches.push(input); + let (buffered_partition_idx, mut mem_growth) = self.buffer_input(input); + + let PartitionIndices::Rows(partition_indices) = &mut self.partition_indices else { + return Err(DataFusionError::Internal( + "row-level placement against a run-indexed repartitioner".to_string(), + )); + }; // 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 @@ -449,7 +530,7 @@ impl MultiPartitionShuffleRepartitioner { // 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 indices = &mut partition_indices[partition_id]; let before_size = indices.allocated_size(); indices.reserve(row_indices.len()); for row_idx in row_indices { @@ -459,6 +540,59 @@ impl MultiPartitionShuffleRepartitioner { mem_growth += after_size.saturating_sub(before_size); } + self.reserve_and_may_spill(mem_growth) + } + + /// Buffers `input` against the runs positional round robin split it into, which + /// `partitioning_batch` left in `scratch.positional_runs`. + async fn buffer_positional_batch_may_spill( + &mut self, + input: RecordBatch, + ) -> datafusion::common::Result<()> { + let (buffered_partition_idx, mut mem_growth) = self.buffer_input(input); + + // Taken out and put back so that the loop can hold the runs and `partition_indices` at + // once; `positional_runs` clears before refilling, so the capacity survives. + let runs = std::mem::take(&mut self.scratch.positional_runs); + let result = match &mut self.partition_indices { + PartitionIndices::Runs(partition_runs) => { + for run in &runs { + let indices = &mut partition_runs[run.partition]; + let before_size = indices.allocated_size(); + indices.push(BufferedRun { + batch: buffered_partition_idx, + start: run.start, + len: run.len, + }); + let after_size = indices.allocated_size(); + mem_growth += after_size.saturating_sub(before_size); + } + Ok(()) + } + PartitionIndices::Rows(_) => Err(DataFusionError::Internal( + "positional placement against a row-indexed repartitioner".to_string(), + )), + }; + self.scratch.positional_runs = runs; + result?; + + self.reserve_and_may_spill(mem_growth) + } + + /// Takes ownership of `input`, charging the reservation and the `data_size` metric for the + /// buffers it newly pins. Returns its index in `buffered_batches` and those bytes. + fn buffer_input(&mut self, input: RecordBatch) -> (u32, usize) { + // `count_new_buffers` dedups buffers shared across already-buffered batches. + let new_buffer_bytes = count_new_buffers(&input, &mut self.pinned_buffers); + self.metrics.data_size.add(new_buffer_bytes); + let buffered_partition_idx = self.buffered_batches.len() as u32; + self.buffered_batches.push(input); + (buffered_partition_idx, new_buffer_bytes) + } + + /// Grows the spill reservation by `mem_growth`, spilling if that is refused or if the fixed + /// buffer limit has been reached. + fn reserve_and_may_spill(&mut self, mem_growth: usize) -> datafusion::common::Result<()> { // A rejected reservation does not include this batch's memory, even though the batch // and its partition indices have already been buffered and must be counted as spilled. let reservation_failed = self.reservation.try_grow(mem_growth).is_err(); @@ -495,13 +629,10 @@ impl MultiPartitionShuffleRepartitioner { /// ShuffleRepartitioner to a new PartitionedBatches struct. The returned PartitionedBatches struct /// can be used to produce shuffled batches. fn partitioned_batches(&mut self) -> PartitionedBatchesProducer { - let num_output_partitions = self.partition_indices.len(); + let num_output_partitions = self.partition_indices.num_partitions(); let buffered_batches = std::mem::take(&mut self.buffered_batches); - // let indices = std::mem::take(&mut self.partition_indices); - let indices = std::mem::replace( - &mut self.partition_indices, - vec![vec![]; num_output_partitions], - ); + let empty = self.partition_indices.empty_like(num_output_partitions); + let indices = std::mem::replace(&mut self.partition_indices, empty); PartitionedBatchesProducer::new(buffered_batches, indices, self.batch_size) } @@ -518,7 +649,7 @@ impl MultiPartitionShuffleRepartitioner { } with_trace("shuffle_spill", self.tracing_enabled, || { - let num_output_partitions = self.partition_indices.len(); + let num_output_partitions = self.partition_indices.num_partitions(); let write_result = { let partitioned_batches = self.partitioned_batches(); // Build the batch-ref slice once and share it across all partitions. @@ -587,7 +718,7 @@ impl ShufflePartitioner for MultiPartitionShuffleRepartition let partitioned_batches = self.partitioned_batches(); self.pinned_buffers.clear(); - let num_output_partitions = self.partition_indices.len(); + let num_output_partitions = self.partition_indices.num_partitions(); // Build the batch-ref slice once and share it across all partitions. let batch_refs = partitioned_batches.batch_refs(); @@ -683,7 +814,12 @@ mod tests { let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( 0, FailingPartitionWriter::default(), - CometPartitioning::RoundRobin(2, 1), + CometPartitioning::RoundRobin( + 2, + RoundRobinStrategy::HashAll { + max_hash_columns: 1, + }, + ), ShufflePartitionerMetrics::new(&metrics_set, 0), Arc::clone(&runtime), 64, @@ -692,11 +828,7 @@ mod tests { ) .unwrap(); repartitioner.insert_batch(batch).await.unwrap(); - let index_bytes = repartitioner - .partition_indices - .iter() - .map(|indices| indices.allocated_size()) - .sum::(); + let index_bytes = repartitioner.partition_indices.allocated_size(); let reserved_bytes = repartitioner.reservation.size(); assert_eq!(repartitioner.spill_count(), 0); assert_eq!(reserved_bytes, input_bytes + index_bytes); @@ -788,7 +920,7 @@ mod tests { let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( 0, FailingPartitionWriter::default(), - CometPartitioning::RoundRobin(2, 0), + CometPartitioning::RoundRobin(2, RoundRobinStrategy::default()), ShufflePartitionerMetrics::new(&metrics_set, 0), Arc::clone(&runtime), 2, @@ -805,7 +937,7 @@ mod tests { assert_eq!(runtime.memory_pool.reserved(), 0); assert!(repartitioner.pinned_buffers.is_empty()); assert!(repartitioner.buffered_batches.is_empty()); - assert!(repartitioner.partition_indices.iter().all(Vec::is_empty)); + assert_eq!(repartitioner.partition_indices.entry_count(), 0); assert_eq!(repartitioner.partition_writer.write_calls, 2); assert_eq!(backing_buffer.strong_count(), input_owners); let successful_spill_bytes = repartitioner.metrics.memory_spilled_bytes.value(); @@ -820,11 +952,7 @@ mod tests { .await .unwrap(); let reservation_before_failure = repartitioner.reservation.size(); - let index_bytes_before_failure = repartitioner - .partition_indices - .iter() - .map(|indices| indices.allocated_size()) - .sum::(); + let index_bytes_before_failure = repartitioner.partition_indices.allocated_size(); let metrics_before_failure = ( repartitioner.spill_count(), repartitioner.metrics.memory_spilled_bytes.value(), @@ -838,14 +966,7 @@ mod tests { assert_eq!(runtime.memory_pool.reserved(), reservation_before_failure); assert_eq!(repartitioner.pinned_buffers.len(), 1); assert_eq!(repartitioner.buffered_batches.len(), 1); - assert_eq!( - repartitioner - .partition_indices - .iter() - .map(Vec::len) - .sum::(), - 2 - ); + assert_eq!(repartitioner.partition_indices.entry_count(), 2); assert_eq!(repartitioner.partition_writer.write_calls, 2); assert_eq!( metrics_before_failure, @@ -863,7 +984,7 @@ mod tests { assert_eq!(runtime.memory_pool.reserved(), 0); assert!(repartitioner.pinned_buffers.is_empty()); assert!(repartitioner.buffered_batches.is_empty()); - assert!(repartitioner.partition_indices.iter().all(Vec::is_empty)); + assert_eq!(repartitioner.partition_indices.entry_count(), 0); assert_eq!(repartitioner.partition_writer.write_calls, 3); assert_eq!(backing_buffer.strong_count(), input_owners); assert_eq!( @@ -912,7 +1033,7 @@ mod tests { let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( 0, FailingPartitionWriter::default(), - CometPartitioning::RoundRobin(2, 0), + CometPartitioning::RoundRobin(2, RoundRobinStrategy::default()), ShufflePartitionerMetrics::new(&metrics_set, 0), Arc::clone(&runtime), 64, @@ -1015,10 +1136,233 @@ mod tests { assert_eq!(runtime.memory_pool.reserved(), 0); assert!(repartitioner.pinned_buffers.is_empty()); assert!(repartitioner.buffered_batches.is_empty()); - assert!(repartitioner.partition_indices.iter().all(Vec::is_empty)); + assert_eq!(repartitioner.partition_indices.entry_count(), 0); spill_bytes.push(repartitioner.metrics.memory_spilled_bytes.value()); } assert!(spill_bytes[0] > 0); assert_eq!(spill_bytes[0], spill_bytes[1]); } + + /// Collects every batch handed to it, per partition, so a whole write can be compared. + #[derive(Default)] + struct CollectingPartitionWriter { + written: std::collections::BTreeMap>, + } + + impl PartitionWriter for CollectingPartitionWriter { + fn write( + &mut self, + pid: usize, + iter: &mut I, + _metrics: &ShufflePartitionerMetrics, + ) -> datafusion::common::Result<()> + where + I: Iterator>, + { + for batch in iter { + let batch = batch?; + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("int64 test column"); + self.written.entry(pid).or_default().extend(values.values()); + } + Ok(()) + } + + fn finish_partition( + &mut self, + pid: usize, + iter: &mut I, + metrics: &ShufflePartitionerMetrics, + ) -> datafusion::common::Result<()> + where + I: Iterator>, + { + self.write(pid, iter, metrics) + } + + fn finish_all( + &mut self, + _metrics: &ShufflePartitionerMetrics, + ) -> datafusion::common::Result<()> { + Ok(()) + } + } + + /// Runs `total_rows` sequential i64 values through a positional repartitioner, framed into + /// batches of the given sizes, and returns the values each output partition received. + async fn positional_placement( + framing: &[usize], + num_partitions: usize, + group_rows: usize, + start_partition: usize, + batch_size: usize, + ) -> std::collections::BTreeMap> { + let schema = Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("v", arrow::datatypes::DataType::Int64, false), + ])); + let runtime = Arc::new( + datafusion::execution::runtime_env::RuntimeEnvBuilder::new() + .build() + .unwrap(), + ); + let metrics_set = ExecutionPlanMetricsSet::new(); + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + CollectingPartitionWriter::default(), + CometPartitioning::RoundRobin( + num_partitions, + RoundRobinStrategy::RowGroups { + start_partition, + group_rows, + }, + ), + ShufflePartitionerMetrics::new(&metrics_set, 0), + Arc::clone(&runtime), + batch_size, + false, + None, + ) + .unwrap(); + + let mut next = 0i64; + for rows in framing { + let values: Vec = (next..next + *rows as i64).collect(); + next += *rows as i64; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values)) as ArrayRef], + ) + .unwrap(); + repartitioner.insert_batch(batch).await.unwrap(); + } + repartitioner.shuffle_write().unwrap(); + repartitioner.partition_writer().written.clone() + } + + /// The property the whole strategy rests on: the same rows, framed differently, produce the + /// same partitioning. A retried map task whose upstream reframes (a spilling operator under + /// different memory pressure, say) still writes what the attempt it replaces wrote. + #[tokio::test] + async fn positional_placement_survives_reframing() { + let baseline = positional_placement(&[1000], 8, 64, 0, 256).await; + + for framing in [ + vec![1000], + vec![1; 1000], + vec![256, 256, 256, 232], + vec![7; 142].into_iter().chain([6]).collect::>(), + vec![500, 500], + ] { + assert_eq!(framing.iter().sum::(), 1000, "bad framing fixture"); + assert_eq!( + positional_placement(&framing, 8, 64, 0, 256).await, + baseline, + "framing {framing:?} placed rows differently" + ); + } + } + + /// Every input row is written exactly once, and each partition gets its rows in input order. + #[tokio::test] + async fn positional_placement_is_a_partition_of_the_input() { + let written = positional_placement(&[300, 17, 683], 8, 64, 3, 256).await; + + let mut all: Vec = written.values().flatten().copied().collect(); + all.sort_unstable(); + assert_eq!(all, (0..1000).collect::>()); + for (pid, values) in &written { + assert!( + values.windows(2).all(|w| w[0] < w[1]), + "partition {pid} received rows out of input order" + ); + } + } + + /// `start_partition` offsets which partition a task starts on, so mappers do not all pile + /// their first group onto partition 0. + #[tokio::test] + async fn positional_placement_starts_at_the_map_partition() { + for start in 0..4usize { + let written = positional_placement(&[64], 4, 64, start, 256).await; + assert_eq!( + written.keys().copied().collect::>(), + vec![start], + "one group should land on partition {start} alone" + ); + } + } + + /// Imbalance stays within one group regardless of how the input was framed, which is what + /// counting rows rather than batches buys over assigning whole batches. + #[tokio::test] + async fn positional_placement_is_balanced_within_one_group() { + let group_rows = 64; + // Ragged framing: whole-batch assignment would give partitions of wildly different sizes. + let framing = vec![1000, 3, 7, 200, 1, 1, 1, 500, 87]; + let total: usize = framing.iter().sum(); + let written = positional_placement(&framing, 8, group_rows, 0, 256).await; + + let sizes: Vec = written.values().map(Vec::len).collect(); + assert_eq!(sizes.iter().sum::(), total); + let spread = sizes.iter().max().unwrap() - sizes.iter().min().unwrap(); + assert!( + spread <= group_rows, + "partition sizes {sizes:?} spread by {spread}, more than one group of {group_rows}" + ); + } + + /// Positional placement records runs, not rows, and skips the row-level scratch entirely. + #[tokio::test] + async fn positional_placement_records_runs_not_rows() { + let schema = Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("v", arrow::datatypes::DataType::Int64, false), + ])); + let runtime = Arc::new( + datafusion::execution::runtime_env::RuntimeEnvBuilder::new() + .build() + .unwrap(), + ); + let metrics_set = ExecutionPlanMetricsSet::new(); + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + FailingPartitionWriter::default(), + CometPartitioning::RoundRobin( + 8, + RoundRobinStrategy::RowGroups { + start_partition: 0, + group_rows: 64, + }, + ), + ShufflePartitionerMetrics::new(&metrics_set, 0), + Arc::clone(&runtime), + 256, + false, + None, + ) + .unwrap(); + + assert!( + repartitioner.scratch.partition_ids.is_empty() + && repartitioner.scratch.partition_row_indices.is_empty() + && repartitioner.scratch.hashes_buf.is_empty(), + "positional placement should allocate no per-row scratch" + ); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from((0..256i64).collect::>())) as ArrayRef], + ) + .unwrap(); + repartitioner.insert_batch(batch).await.unwrap(); + + // 256 rows in groups of 64 is four runs, not 256 row entries. + assert_eq!(repartitioner.partition_indices.entry_count(), 4); + assert!(matches!( + repartitioner.partition_indices, + PartitionIndices::Runs(_) + )); + } } diff --git a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs index 424fa827d52..1799ce9cf5c 100644 --- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs +++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs @@ -16,23 +16,84 @@ // under the License. use arrow::array::RecordBatch; -use arrow::compute::interleave_record_batch; +use arrow::compute::{concat_batches, interleave_record_batch}; +#[cfg(test)] +use datafusion::common::utils::proxy::VecAllocExt; use datafusion::common::DataFusionError; use datafusion::physical_plan::metrics::Time; +/// A contiguous run of rows within one buffered batch, bound for one output partition. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct BufferedRun { + pub batch: u32, + pub start: u32, + pub len: u32, +} + +/// Per-partition record of which buffered rows belong to that partition. +/// +/// The two shapes are not interchangeable and a repartitioner picks one for its lifetime, from +/// the partitioning it was built with. Anything that places rows individually has to name them +/// individually; positional round robin places whole spans and can say so in a twelfth of the +/// space at a 64-row group, which matters because this list is charged against the spill +/// reservation. The payoff is at flush: a run is copied with one bulk copy per buffer, where a +/// row list has to be gathered a row at a time through `interleave_record_batch`, re-walking +/// every column and every nested child. +pub(crate) enum PartitionIndices { + /// One `(batch, row)` pair per row, in the order the rows should be written. + Rows(Vec>), + /// One `(batch, start, len)` run per contiguous span, in the order they should be written. + Runs(Vec>), +} + +impl PartitionIndices { + pub(crate) fn empty_like(&self, num_partitions: usize) -> Self { + match self { + Self::Rows(_) => Self::Rows(vec![vec![]; num_partitions]), + Self::Runs(_) => Self::Runs(vec![vec![]; num_partitions]), + } + } + + pub(crate) fn num_partitions(&self) -> usize { + match self { + Self::Rows(indices) => indices.len(), + Self::Runs(runs) => runs.len(), + } + } + + /// Bytes the per-partition lists have allocated, which the spill reservation is charged for. + #[cfg(test)] + pub(crate) fn allocated_size(&self) -> usize { + match self { + Self::Rows(indices) => indices.iter().map(|i| i.allocated_size()).sum(), + Self::Runs(runs) => runs.iter().map(|r| r.allocated_size()).sum(), + } + } + + /// Number of entries recorded across all partitions: rows for [`Self::Rows`], runs for + /// [`Self::Runs`]. + #[cfg(test)] + pub(crate) fn entry_count(&self) -> usize { + match self { + Self::Rows(indices) => indices.iter().map(Vec::len).sum(), + Self::Runs(runs) => runs.iter().map(Vec::len).sum(), + } + } +} + /// A helper struct to produce shuffled batches. /// This struct takes ownership of the buffered batches and partition indices from the /// ShuffleRepartitioner, and provides an iterator over the batches in the specified partitions. pub(super) struct PartitionedBatchesProducer { buffered_batches: Vec, - partition_indices: Vec>, + partition_indices: PartitionIndices, batch_size: usize, } impl PartitionedBatchesProducer { pub(super) fn new( buffered_batches: Vec, - indices: Vec>, + indices: PartitionIndices, batch_size: usize, ) -> Self { Self { @@ -53,7 +114,7 @@ impl PartitionedBatchesProducer { &'a self, refs: &'a [&'a RecordBatch], partition_id: usize, - interleave_time: &'a Time, + copy_time: &'a Time, ) -> PartitionedBatchIterator<'a> { // Partition indices index into `buffered_batches`; a refs slice built from a // different producer would silently interleave wrong rows. @@ -62,17 +123,45 @@ impl PartitionedBatchesProducer { self.buffered_batches.len(), "refs slice must cover every buffered batch" ); - PartitionedBatchIterator::new( - &self.partition_indices[partition_id], - refs, - self.batch_size, - interleave_time, - ) + match &self.partition_indices { + PartitionIndices::Rows(indices) => PartitionedBatchIterator::Rows(RowIterator::new( + &indices[partition_id], + refs, + self.batch_size, + copy_time, + )), + PartitionIndices::Runs(runs) => PartitionedBatchIterator::Runs(RunIterator::new( + &runs[partition_id], + refs, + self.batch_size, + copy_time, + )), + } } } /// Iterates over the shuffled record batches belonging to a single output partition. -pub(crate) struct PartitionedBatchIterator<'a> { +/// +/// One concrete type covering both index shapes, because [`crate::writers::PartitionWriter`] is +/// generic over a single iterator type rather than taking a trait object. +pub(crate) enum PartitionedBatchIterator<'a> { + Rows(RowIterator<'a>), + Runs(RunIterator<'a>), +} + +impl Iterator for PartitionedBatchIterator<'_> { + type Item = datafusion::common::Result; + + fn next(&mut self) -> Option { + match self { + Self::Rows(iter) => iter.next(), + Self::Runs(iter) => iter.next(), + } + } +} + +/// Produces a partition's output by gathering individually named rows. +pub(crate) struct RowIterator<'a> { record_batches: &'a [&'a RecordBatch], batch_size: usize, indices: &'a [(u32, u32)], @@ -81,15 +170,15 @@ pub(crate) struct PartitionedBatchIterator<'a> { /// (capacity at most `batch_size`) rather than re-materializing its whole index list. chunk_scratch: Vec<(usize, usize)>, pos: usize, - interleave_time: &'a Time, + copy_time: &'a Time, } -impl<'a> PartitionedBatchIterator<'a> { +impl<'a> RowIterator<'a> { fn new( indices: &'a [(u32, u32)], record_batches: &'a [&'a RecordBatch], batch_size: usize, - interleave_time: &'a Time, + copy_time: &'a Time, ) -> Self { if indices.is_empty() { // Avoid unnecessary allocations when the partition is empty @@ -99,7 +188,7 @@ impl<'a> PartitionedBatchIterator<'a> { indices: &[], chunk_scratch: vec![], pos: 0, - interleave_time, + copy_time, }; } Self { @@ -108,12 +197,12 @@ impl<'a> PartitionedBatchIterator<'a> { indices, chunk_scratch: Vec::with_capacity(batch_size.min(indices.len())), pos: 0, - interleave_time, + copy_time, } } } -impl Iterator for PartitionedBatchIterator<'_> { +impl Iterator for RowIterator<'_> { type Item = datafusion::common::Result; fn next(&mut self) -> Option { @@ -128,7 +217,7 @@ impl Iterator for PartitionedBatchIterator<'_> { .iter() .map(|(i_batch, i_row)| (*i_batch as usize, *i_row as usize)), ); - let mut timer = self.interleave_time.timer(); + let mut timer = self.copy_time.timer(); let result = interleave_record_batch(self.record_batches, &self.chunk_scratch); timer.stop(); match result { @@ -144,6 +233,118 @@ impl Iterator for PartitionedBatchIterator<'_> { } } +/// Produces a partition's output by copying contiguous runs of rows. +pub(crate) struct RunIterator<'a> { + record_batches: &'a [&'a RecordBatch], + batch_size: usize, + runs: &'a [BufferedRun], + /// Scratch for the slices making up the current output chunk. + chunk_scratch: Vec, + /// Index of the next run to consume. + pos: usize, + /// Rows already taken from `runs[pos]`, non-zero only when a run straddled a chunk boundary. + consumed: u32, + copy_time: &'a Time, +} + +impl<'a> RunIterator<'a> { + fn new( + runs: &'a [BufferedRun], + record_batches: &'a [&'a RecordBatch], + batch_size: usize, + copy_time: &'a Time, + ) -> Self { + if runs.is_empty() { + return Self { + record_batches: &[], + batch_size, + runs: &[], + chunk_scratch: vec![], + pos: 0, + consumed: 0, + copy_time, + }; + } + Self { + record_batches, + batch_size, + runs, + chunk_scratch: vec![], + pos: 0, + consumed: 0, + copy_time, + } + } +} + +impl Iterator for RunIterator<'_> { + type Item = datafusion::common::Result; + + fn next(&mut self) -> Option { + if self.pos >= self.runs.len() { + return None; + } + let mut timer = self.copy_time.timer(); + + // Zero-copy path: the next run is an entire buffered batch and already fills a chunk on + // its own, so hand the batch straight through. This is the case a group as large as the + // batch size is chosen to hit. Returning the batch rather than a slice of it also keeps + // `Utf8View`/`BinaryView` columns off the sliced-array path in the IPC writer, which + // truncates the views buffer but serializes every shared data buffer in full. + if self.consumed == 0 { + let run = self.runs[self.pos]; + let source = self.record_batches[run.batch as usize]; + if run.start == 0 + && run.len as usize == source.num_rows() + && run.len as usize >= self.batch_size + { + self.pos += 1; + timer.stop(); + return Some(Ok((*source).clone())); + } + } + + // Otherwise accumulate whole runs until the chunk is full, splitting the run that + // straddles the boundary. Chunks stay `batch_size` rows so that output block sizes do not + // depend on how long the runs happen to be. + self.chunk_scratch.clear(); + let mut rows = 0usize; + while self.pos < self.runs.len() && rows < self.batch_size { + let run = self.runs[self.pos]; + let source = self.record_batches[run.batch as usize]; + let available = (run.len - self.consumed) as usize; + let take = available.min(self.batch_size - rows); + self.chunk_scratch + .push(source.slice((run.start + self.consumed) as usize, take)); + rows += take; + if take == available { + self.pos += 1; + self.consumed = 0; + } else { + self.consumed += take as u32; + } + } + + // `concat_batches` over a single slice returns the slice rather than compacting it + // (arrow's `concat` short-circuits at one input), so skip the call and let the slice + // through: the IPC writer truncates a sliced array's buffers per type, and the one + // family it does not, the view types, cannot reach here (see `create_repartitioner`). + let result = if self.chunk_scratch.len() == 1 { + Ok(self.chunk_scratch.pop().expect("one slice")) + } else { + concat_batches(&self.chunk_scratch[0].schema(), self.chunk_scratch.iter()) + }; + timer.stop(); + match result { + Ok(batch) => Some(Ok(batch)), + Err(e) => Some(Err(DataFusionError::ArrowError( + Box::from(e), + Some(DataFusionError::get_back_trace()), + ))), + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -186,7 +387,7 @@ mod tests { let batch_size = 4; // chunks of 4, 4, and a tail of 2 let producer = PartitionedBatchesProducer::new( buffered.clone(), - vec![indices.clone(), Vec::new()], + PartitionIndices::Rows(vec![indices.clone(), Vec::new()]), batch_size, ); let refs = producer.batch_refs(); @@ -224,10 +425,148 @@ mod tests { #[should_panic(expected = "refs slice must cover every buffered batch")] fn produce_rejects_mismatched_refs() { let buffered = batches(); - let producer = PartitionedBatchesProducer::new(buffered, vec![vec![(0, 0), (2, 1)]], 4); + let producer = PartitionedBatchesProducer::new( + buffered, + PartitionIndices::Rows(vec![vec![(0, 0), (2, 1)]]), + 4, + ); let refs = producer.batch_refs(); let truncated = &refs[..refs.len() - 1]; let time = Time::default(); let _ = producer.produce(truncated, 0, &time); } + + fn run_values(runs: Vec, batch_size: usize) -> (Vec>, Vec) { + let buffered = batches(); + let producer = PartitionedBatchesProducer::new( + buffered, + PartitionIndices::Runs(vec![runs]), + batch_size, + ); + let refs = producer.batch_refs(); + let time = Time::default(); + let produced: Vec = producer + .produce(&refs, 0, &time) + .collect::>() + .unwrap(); + let values = produced + .iter() + .map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + (values, produced) + } + + /// Runs are emitted in order, concatenated up to `batch_size`, with the run that straddles a + /// chunk boundary split across the two chunks. + #[test] + fn runs_concatenate_into_fixed_size_chunks() { + let runs = vec![ + BufferedRun { + batch: 0, + start: 1, + len: 3, + }, // 1, 2, 3 + BufferedRun { + batch: 2, + start: 0, + len: 2, + }, // 200, 201 + BufferedRun { + batch: 1, + start: 3, + len: 2, + }, // 103, 104 + ]; + let (values, _) = run_values(runs, 4); + assert_eq!(values, vec![vec![1, 2, 3, 200], vec![201, 103, 104]]); + } + + /// A run covering a whole buffered batch, long enough to be a chunk on its own, is handed + /// through without copying. Identity rather than equality, because the point is that the + /// output shares the input's buffers. + #[test] + fn whole_batch_run_is_returned_without_copying() { + let buffered = batches(); + let source_ptr = buffered[1].column(0).as_ref() as *const dyn arrow::array::Array; + let producer = PartitionedBatchesProducer::new( + buffered, + PartitionIndices::Runs(vec![vec![BufferedRun { + batch: 1, + start: 0, + len: 5, + }]]), + 5, + ); + let refs = producer.batch_refs(); + let time = Time::default(); + let produced: Vec = producer + .produce(&refs, 0, &time) + .collect::>() + .unwrap(); + assert_eq!(produced.len(), 1); + assert!(std::ptr::addr_eq( + produced[0].column(0).as_ref() as *const dyn arrow::array::Array, + source_ptr + )); + } + + /// A single run that is a strict sub-range is sliced rather than copied, which is the cheap + /// outcome and one the IPC writer truncates correctly for every non-view type. + #[test] + fn single_sub_range_run_stays_a_slice() { + let buffered = batches(); + let source = buffered[2].column(0).to_data().buffers()[0].clone(); + let producer = PartitionedBatchesProducer::new( + buffered, + PartitionIndices::Runs(vec![vec![BufferedRun { + batch: 2, + start: 1, + len: 3, + }]]), + 8, + ); + let refs = producer.batch_refs(); + let time = Time::default(); + let produced: Vec = producer + .produce(&refs, 0, &time) + .collect::>() + .unwrap(); + + assert_eq!( + produced[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(), + vec![201, 202, 203] + ); + // Three i32s starting one element into the source allocation: `data_ptr` (the allocation) + // is unchanged and `as_ptr` (the slice) has advanced, so nothing was copied. + let produced_buffer = produced[0].column(0).to_data().buffers()[0].clone(); + assert_eq!( + produced_buffer.data_ptr().as_ptr() as usize, + source.data_ptr().as_ptr() as usize, + "expected the source allocation, not a copy" + ); + assert_eq!( + produced_buffer.as_ptr() as usize, + source.as_ptr() as usize + size_of::() + ); + } + + #[test] + fn empty_run_list_produces_nothing() { + let (values, _) = run_values(vec![], 4); + assert!(values.is_empty()); + } } diff --git a/native/shuffle/src/rss_execution_tests.rs b/native/shuffle/src/rss_execution_tests.rs index e9269edc9e5..f18409d0baf 100644 --- a/native/shuffle/src/rss_execution_tests.rs +++ b/native/shuffle/src/rss_execution_tests.rs @@ -16,7 +16,7 @@ // under the License. use crate::{ - read_ipc_compressed, CometPartitioning, CompressionCodec, PartitionOffsets, + read_ipc_compressed, CometPartitioning, CompressionCodec, PartitionOffsets, RoundRobinStrategy, ShuffleWriterDestination, ShuffleWriterExec, }; use arrow::array::{Array, Int32Array, RecordBatch, RecordBatchOptions}; @@ -231,7 +231,7 @@ fn rss_multi_partition_supports_hash_range_and_round_robin() { for partitioning in [ CometPartitioning::Hash(vec![expression], 4), CometPartitioning::RangePartitioning(ordering, 4, Arc::new(converter), boundaries), - CometPartitioning::RoundRobin(4, 0), + CometPartitioning::RoundRobin(4, RoundRobinStrategy::default()), ] { let pusher = Arc::new(RecordingPusher::default()); let execution = rss_execution( @@ -271,7 +271,7 @@ fn rss_empty_schema_preserves_row_counts_in_partition_zero() { let execution = rss_execution( vec![batch.clone(), batch], schema, - CometPartitioning::RoundRobin(4, 0), + CometPartitioning::RoundRobin(4, RoundRobinStrategy::default()), pusher.clone(), CompressionCodec::None, 1024 * 1024, @@ -299,7 +299,7 @@ fn rss_empty_schema_without_rows_does_not_push_frames() { let execution = rss_execution( vec![batch], schema, - CometPartitioning::RoundRobin(4, 0), + CometPartitioning::RoundRobin(4, RoundRobinStrategy::default()), pusher.clone(), CompressionCodec::None, 1024 * 1024, diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index ca88509a816..296747421cb 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -23,7 +23,7 @@ use crate::partitioners::{ SinglePartitionShufflePartitioner, }; use crate::writers::{LocalPartitionWriter, PartitionWriter, RssPartitionWriter}; -use crate::{CometPartitioning, CompressionCodec, ShuffleBlockWriter}; +use crate::{CometPartitioning, CompressionCodec, RoundRobinStrategy, ShuffleBlockWriter}; use async_trait::async_trait; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{exec_datafusion_err, DataFusionError}; @@ -31,7 +31,7 @@ use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExp use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::{apply_expression_roots, EmptyRecordBatchStream}; use datafusion::{ - arrow::datatypes::SchemaRef, + arrow::datatypes::{DataType, SchemaRef}, error::Result, execution::context::TaskContext, physical_plan::{ @@ -387,6 +387,38 @@ async fn external_shuffle( Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone(&schema))) as SendableRecordBatchStream) } +/// True when `data_type` is, or contains, one of Arrow's view layouts. +/// +/// Arrow's IPC writer truncates a sliced array's buffers per type — numeric and temporal values +/// through `get_or_truncate_buffer`, byte arrays through `reencode_offsets`, list children through +/// `get_list_array_buffers`, boolean bitmaps through `bit_slice`, and struct children because +/// `ArrayData::slice` pushes the slice down into them. The view types are the exception: it slices +/// the views buffer but serializes every shared data buffer in full, since proving that no +/// surviving view references a buffer is not cheap. +/// +/// [`RoundRobinStrategy::RowGroups`] is the only placement that can hand the writer a sliced +/// array; everything else materializes a fresh batch through `interleave_record_batch`. So a view +/// column anywhere in the schema would let a short run drag a whole batch's data buffers into the +/// shuffle output, which is the opposite of what the strategy is for. +fn contains_view_type(data_type: &DataType) -> bool { + match data_type { + DataType::Utf8View | DataType::BinaryView => true, + DataType::List(field) + | DataType::LargeList(field) + | DataType::ListView(field) + | DataType::LargeListView(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => contains_view_type(field.data_type()), + DataType::Struct(fields) => fields.iter().any(|f| contains_view_type(f.data_type())), + DataType::Union(fields, _) => fields + .iter() + .any(|(_, f)| contains_view_type(f.data_type())), + DataType::Dictionary(_, value) => contains_view_type(value), + DataType::RunEndEncoded(_, values) => contains_view_type(values.data_type()), + _ => false, + } +} + /// Constructs the existing schema-appropriate partitioner for either writer backend. #[allow(clippy::too_many_arguments)] fn create_repartitioner( @@ -401,6 +433,23 @@ fn create_repartitioner( ) -> Result> { let partition_count = partitioning.partition_count(); + // The planner decides whether positional placement is safe to retry; this decides whether it + // is worth doing on this schema. See `contains_view_type`. + let partitioning = match &partitioning { + CometPartitioning::RoundRobin(n, RoundRobinStrategy::RowGroups { .. }) + if schema + .fields() + .iter() + .any(|f| contains_view_type(f.data_type())) => + { + log::debug!( + "schema contains a view type, falling back from positional to hash round robin" + ); + CometPartitioning::RoundRobin(*n, RoundRobinStrategy::default()) + } + _ => partitioning, + }; + if schema.fields().is_empty() { log::debug!( "found empty schema, overriding {partitioning:?} partitioning with EmptySchemaShufflePartitioner" @@ -440,7 +489,7 @@ fn contextualize_shuffle_error(error: DataFusionError, phase: &str) -> DataFusio #[cfg(test)] mod test { use super::*; - use crate::{read_ipc_compressed, ShuffleBlockWriter, ShuffleCodecContext}; + use crate::{read_ipc_compressed, RoundRobinStrategy, ShuffleBlockWriter, ShuffleCodecContext}; use arrow::array::{Array, Int64Array, StringArray, StringBuilder}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; @@ -1143,7 +1192,7 @@ mod test { Arc::new(row_converter), owned_rows, ), - CometPartitioning::RoundRobin(num_partitions, 0), + CometPartitioning::RoundRobin(num_partitions, RoundRobinStrategy::default()), ] { let batches = (0..num_batches).map(|_| batch.clone()).collect::>(); @@ -1211,7 +1260,7 @@ 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::default()), CompressionCodec::Zstd(1), data_file.clone(), false, @@ -1615,7 +1664,7 @@ 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::default()), CompressionCodec::Zstd(1), data_file.to_str().unwrap().to_string(), false, @@ -1699,7 +1748,7 @@ 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::default()), CompressionCodec::Zstd(1), data_file.to_str().unwrap().to_string(), false, @@ -1736,4 +1785,49 @@ mod test { assert_eq!(*offset, 0, "All offsets should be 0 with zero rows"); } } + + /// Positional round robin is the only placement that hands a sliced array to the IPC writer, + /// and the view types are the one family the writer does not truncate. Find them anywhere in + /// the schema, not just at the top level, since the motivating schemas are deeply nested. + #[test] + fn view_types_are_detected_at_any_depth() { + let leaf = |dt: DataType| Field::new("leaf", dt, true); + let nested = DataType::Struct( + vec![ + Field::new("a", DataType::Int64, true), + Field::new( + "b", + DataType::List(Arc::new(leaf(DataType::Utf8View))), + true, + ), + ] + .into(), + ); + assert!(contains_view_type(&nested)); + assert!(contains_view_type(&DataType::Utf8View)); + assert!(contains_view_type(&DataType::BinaryView)); + assert!(contains_view_type(&DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::BinaryView, true), + ] + .into() + ), + false + )), + false + ))); + + assert!(!contains_view_type(&DataType::Utf8)); + assert!(!contains_view_type(&DataType::Struct( + vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::List(Arc::new(leaf(DataType::Utf8))), true), + ] + .into() + ))); + } } diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index d371f1ba44c..a73b81229ef 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -475,6 +475,42 @@ 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_POSITIONAL_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.shuffle.native.partitioning.roundrobin.positional.enabled") + .category(CATEGORY_SHUFFLE) + .doc( + "When true, Comet's native round-robin shuffle places rows by position rather than by " + + "hashing their contents, the way Spark's own round robin does: the row at " + + "task-global ordinal i goes to output partition " + + "(mapPartitionId + i / groupRows) % numPartitions. This skips a murmur3 pass over " + + "every column of every row and replaces the per-row gather on flush with a bulk copy " + + "per run, which is what dominates the shuffle write on wide nested schemas. It also " + + "spreads duplicate rows evenly, where hashing sends them all to one partition. " + + "Positional placement is only reproducible when the map task replays rows in the " + + "same order, so it is used only where Comet can establish that from the plan: a " + + "native scan under nothing but projections and filters. Any other plan silently " + + "keeps content-hash placement. " + + s"Has no effect unless ${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key} " + + "is also true.") + .booleanConf + .createWithDefault(false) + + val COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_GROUP_ROWS: ConfigEntry[Int] = + conf("spark.comet.shuffle.native.partitioning.roundrobin.positional.groupRows") + .category(CATEGORY_SHUFFLE) + .doc( + "Rows per contiguous group under positional round robin. Imbalance between any two " + + "output partitions is bounded by this many rows however the reader frames its " + + "batches, so smaller groups balance better while larger groups produce fewer, longer " + + "runs to copy. When set to 0 (the default) Comet derives it from the batch size and " + + "the partition count. Only applies when " + + s"${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key} is true.") + .intConf + .checkValue( + v => v >= 0, + "The group size for positional round robin partitioning must be non-negative.") + .createWithDefault(0) + 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/CometNativeShuffleInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala index bced55b34df..c3ede338e38 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.comet.execution.shuffle import org.apache.spark._ -import org.apache.spark.rdd.RDD +import org.apache.spark.rdd.{DeterministicLevel, RDD} import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -33,6 +33,10 @@ import org.apache.comet.CometShuffleBlockIterator * [[CometNativeShuffleInputIterator]]. The iterator reports `hasNext = false`; * [[CometNativeShuffleWriter]] downcasts it and reads those slots directly to drive the unified * `ShuffleWriter(child = childNativeOp)` plan. + * + * @param positionalRoundRobin + * whether the writer fed by this RDD places rows by position rather than by content; see + * [[CometShuffleExchangeExec.usesPositionalRoundRobin]] and `getOutputDeterministicLevel`. */ private[shuffle] class CometNativeShuffleInputRDD( sc: SparkContext, @@ -40,7 +44,8 @@ private[shuffle] class CometNativeShuffleInputRDD( numPartitionsParam: Int, shuffleScanIndices: Set[Int], spillMetricNode: CometMetricNode, - @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty) + @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty, + positionalRoundRobin: Boolean = false) extends RDD[Product2[Int, ColumnarBatch]]( sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { @@ -57,7 +62,34 @@ private[shuffle] class CometNativeShuffleInputRDD( numPartitionsParam, shuffleScanIndices, spillMetricNode, - perPartitionByKey) + perPartitionByKey, + positionalRoundRobin) + + /** + * Spark handles the retry hazard of positional round robin declaratively rather than + * per-operator: it wraps the repartition in a `MapPartitionsRDD` with `isOrderSensitive = true` + * (Comet's own JVM path does this in `prepareJVMShuffleDependency`), and that RDD reports + * `INDETERMINATE` whenever its parent is `UNORDERED`, which makes the DAGScheduler roll the + * whole stage back instead of re-running one task into a partially consumed output. The native + * path has no `MapPartitionsRDD` to carry the flag, so apply the same rule here. + * + * This covers everything below the RDD boundary; it cannot see the operators fused into the + * native plan above it, because the whole subtree collapses into this one RDD and `inputRDDs` + * are its leaves. `CometShuffleExchangeExec.replaysRowsInOrder` covers those. Both run, and + * positional placement needs both to agree. + * + * Letting the parent level discriminate is what keeps a plain scan on the cheap per-task retry + * path: a determinate parent stays determinate, while anything below another exchange is + * unordered, because reduce tasks see shuffle blocks in arrival order, and goes indeterminate. + */ + override protected def getOutputDeterministicLevel: DeterministicLevel.Value = { + val inheritedLevel = super.getOutputDeterministicLevel + if (positionalRoundRobin && inheritedLevel != DeterministicLevel.DETERMINATE) { + DeterministicLevel.INDETERMINATE + } else { + inheritedLevel + } + } override protected def getPartitions: Array[Partition] = (0 until numPartitionsParam).map { i => 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 fce4291deb6..461cd842546 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 @@ -421,6 +421,12 @@ class CometNativeShuffleWriter[K, V]( partitioning.setNumPartitions(effectivePartitionCount) partitioning.setMaxHashColumns( CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_MAX_HASH_COLUMNS.get()) + // Decided on the driver, from the shape of the plan fused into this writer; the executor + // cannot re-derive it. See `CometShuffleExchangeExec.positionalRoundRobinSpec`. + spec.positionalRoundRobin.foreach { positional => + partitioning.setPositional(true) + partitioning.setPositionalGroupRows(positional.groupRows) + } val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() shuffleWriterBuilder.setPartitioning( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index 34f1407fc9b..25208fc4089 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala @@ -43,7 +43,24 @@ import org.apache.comet.serde.OperatorOuterClass case class NativeShuffleSpec( childNativeOp: OperatorOuterClass.Operator, childMetricNode: CometMetricNode, - execContext: NativeExecContext) + execContext: NativeExecContext, + /** + * Set when round-robin placement is positional rather than content-hashed. Both the decision + * and the group size are resolved once on the driver: the decision because it depends on the + * shape of the plan fused into `childNativeOp`, which the executor never sees, and the group + * size so that it cannot disagree with the decision. See + * `CometShuffleExchangeExec.usesPositionalRoundRobin`. + */ + positionalRoundRobin: Option[PositionalRoundRobin] = None) + +/** + * Parameters for positional round-robin placement, resolved on the driver. + * + * @param groupRows + * rows per contiguous group, or 0 to let the native side derive it from the batch size and the + * partition count. + */ +case class PositionalRoundRobin(groupRows: Int) /** * A [[ShuffleDependency]] that allows us to identify the shuffle dependency as a Comet shuffle. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 2a539643755..6df8b5ff42e 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Exp import org.apache.spark.sql.catalyst.expressions.codegen.LazilyGeneratedOrdering import org.apache.spark.sql.catalyst.plans.logical.Statistics import org.apache.spark.sql.catalyst.plans.physical._ -import org.apache.spark.sql.comet.{CometMetricNode, CometNativeExec, CometPlan, CometSinkPlaceHolder, NativeExecContext} +import org.apache.spark.sql.comet.{CometFilterExec, CometMetricNode, CometNativeExec, CometNativeScanExec, CometPlan, CometProjectExec, CometSinkPlaceHolder, NativeExecContext} import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec @@ -127,7 +127,8 @@ case class CometShuffleExchangeExec( ctx.numPartitions, ctx.shuffleScanIndices, CometMetricNode(metrics, Seq(nativeChildMetricNode)), - ctx.perPartitionByKey) + ctx.perPartitionByKey, + CometShuffleExchangeExec.usesPositionalRoundRobin(outputPartitioning, child)) case None => // Non-native child (e.g. CometSparkToColumnarExec): no subtree to inline. The dep gets // built via the convenience overload below; we just need a real RDD of batches. @@ -206,7 +207,11 @@ case class CometShuffleExchangeExec( outputPartitioning, serializer, metrics, - NativeShuffleSpec(nativeChild.nativeOp, nativeChildMetricNode, ctx)) + NativeShuffleSpec( + nativeChild.nativeOp, + nativeChildMetricNode, + ctx, + CometShuffleExchangeExec.positionalRoundRobinSpec(outputPartitioning, child))) case None => CometShuffleExchangeExec.prepareShuffleDependency( inputRDD.asInstanceOf[RDD[ColumnarBatch]], @@ -298,6 +303,80 @@ object CometShuffleExchangeExec if (shuffleSupported(op).isDefined) Compatible() else Unsupported() } + /** + * True when this exchange will run the native round-robin writer in its positional mode, where + * the row at task-global ordinal `i` goes to `(mapPartitionId + i / groupRows) % numPartitions` + * rather than to `pmod(hash(row), numPartitions)`. + * + * Positional placement is reproducible exactly when the map task replays its rows in the same + * order, which is the same condition Spark's own round robin depends on. Spark answers it in + * two places and so does Comet: `replaysRowsInOrder` below establishes it for the operators + * fused into this native plan, which the RDD graph cannot see because the whole subtree + * collapses into one `CometNativeShuffleInputRDD`; and + * `CometNativeShuffleInputRDD.getOutputDeterministicLevel` establishes it for everything below + * that RDD, where the leaves are. Both have to hold. + * + * Must stay in step with `PhysicalPlanner::create_partitioning`, which turns the `positional` + * proto field into `RoundRobinStrategy::RowGroups`. + * + * The `numPartitions > 1` guard mirrors `isRoundRobin` in `prepareJVMShuffleDependency`. With a + * single output partition every row lands in the same place, so there is no placement to get + * wrong, and native routes that case to `SinglePartitionShufflePartitioner` regardless. + */ + def usesPositionalRoundRobin(outputPartitioning: Partitioning, child: SparkPlan): Boolean = + positionalRoundRobinSpec(outputPartitioning, child).isDefined + + /** + * [[usesPositionalRoundRobin]] together with the group size to use, both read on the driver so + * that they cannot disagree. `CometConf.get()` resolves against the thread-local `SQLConf`, + * which on an executor is rebuilt from the task's local properties; reading the group size + * there returned the default rather than the session value. + */ + def positionalRoundRobinSpec( + outputPartitioning: Partitioning, + child: SparkPlan): Option[PositionalRoundRobin] = { + val eligible = outputPartitioning.isInstanceOf[RoundRobinPartitioning] && + outputPartitioning.numPartitions > 1 && + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.get() && + replaysRowsInOrder(child) + if (eligible) { + Some( + PositionalRoundRobin( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_GROUP_ROWS.get())) + } else { + None + } + } + + /** + * Whether re-executing this subtree yields the same rows in the same order. + * + * Deliberately a short allowlist rather than a denylist of known-bad operators, because the + * cost of being wrong is silent data loss rather than a failure: a re-executed map task that + * orders rows differently writes a different partitioning of them, and once any consumer has + * fetched the output that attempt replaces, the reduce side gets some rows twice and others not + * at all. Anything not named here keeps content-hash placement, which is safe to re-execute + * whatever its input does. + * + * A native scan replays its partition because the file splits are fixed on the driver when the + * RDD is built, and projections and filters are row-wise. Operators that spill are the + * interesting exclusion: an aggregate or a sort under memory pressure emits its output in an + * order that depends on how many times it spilled, which differs between attempts on different + * executors. Note that this says nothing about how rows are framed into batches: positional + * placement counts rows across batch boundaries precisely so that framing does not have to be + * part of this judgement. + * + * Other leaf scans (Iceberg, DSv2 batch, in-memory) plausibly qualify too, but each needs its + * own argument that a re-executed task reads the same rows in the same order, so they are left + * out until someone makes it. + */ + private def replaysRowsInOrder(plan: SparkPlan): Boolean = plan match { + case _: CometNativeScanExec => true + case p: CometProjectExec => replaysRowsInOrder(p.child) + case f: CometFilterExec => replaysRowsInOrder(f.child) + case _ => false + } + override def createExec( nativeOp: OperatorOuterClass.Operator, op: ShuffleExchangeExec): CometNativeExec = { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala new file mode 100644 index 00000000000..3bd711705c0 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala @@ -0,0 +1,201 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet.execution.shuffle + +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.catalyst.plans.physical.RoundRobinPartitioning +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.functions.{col, lit} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +/** + * Positional round robin places rows by their ordinal within the map task rather than by hashing + * their contents, which is what Spark's own round robin does. Placement is reproducible only when + * the map task replays rows in the same order, so the gating that decides where it is used is as + * much the feature as the placement itself, and most of what is tested here. + * + * Lives in the `execution.shuffle` package so it can reach + * `CometShuffleExchangeExec.usesPositionalRoundRobin` directly rather than inferring the decision + * from output. + */ +class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + private val numPartitions = 8 + + private def withPositionalRoundRobin(extra: (String, String)*)(f: => Unit): Unit = + withSQLConf( + Seq( + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key -> "true", + // Keep AQE from coalescing the round robin away, so the exchange under test survives + // into the executed plan. + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") ++ extra: _*)(f) + + /** Whether the round-robin exchange in `df`'s executed plan chose positional placement. */ + private def isPositional(df: DataFrame): Boolean = { + val exchanges = collect(df.queryExecution.executedPlan) { + case e: CometShuffleExchangeExec + if e.shuffleType == CometNativeShuffle && + e.outputPartitioning.isInstanceOf[RoundRobinPartitioning] => + e + } + assert( + exchanges.size == 1, + s"expected one native round-robin exchange in\n${df.queryExecution.executedPlan}") + CometShuffleExchangeExec.usesPositionalRoundRobin( + exchanges.head.outputPartitioning, + exchanges.head.child) + } + + private def withParquetTable(rows: Int)(f: String => Unit): Unit = { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(rows) + .selectExpr("id", "cast(id % 7 as string) as s", "id % 3 as g") + .write + .parquet(path) + withTempView("t") { + spark.read.parquet(path).createOrReplaceTempView("t") + f("t") + } + } + } + + test("a scan under projections and filters takes positional placement") { + withPositionalRoundRobin() { + withParquetTable(1000) { t => + assert(isPositional(spark.table(t).repartition(numPartitions))) + assert(isPositional(spark.table(t).filter("id > 10").repartition(numPartitions))) + assert(isPositional( + spark.table(t).filter("id > 10").selectExpr("id + 1 as id").repartition(numPartitions))) + } + } + } + + test("a plan whose replay order Comet cannot establish keeps content-hash placement") { + withPositionalRoundRobin() { + withParquetTable(1000) { t => + // An aggregate under memory pressure emits groups in an order that depends on how many + // times it spilled, which differs between attempts. + assert(!isPositional(spark.table(t).groupBy("g").count().repartition(numPartitions))) + // A sort is not excluded because it reorders, but because Comet has not established that + // its tie-breaking survives a differing spill count. + assert(!isPositional(spark.table(t).sort("s").repartition(numPartitions))) + } + } + } + + test("positional placement is off unless its own config is on") { + withParquetTable(100) { t => + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + assert(!isPositional(spark.table(t).repartition(numPartitions))) + } + } + } + + test("a hash repartition is never positional") { + withPositionalRoundRobin() { + withParquetTable(100) { t => + val df = spark.table(t).repartition(numPartitions, spark.table(t)("g")) + val exchanges = collect(df.queryExecution.executedPlan) { + case e: CometShuffleExchangeExec => e + } + assert(exchanges.nonEmpty) + exchanges.foreach { e => + assert( + !CometShuffleExchangeExec.usesPositionalRoundRobin(e.outputPartitioning, e.child)) + } + } + } + } + + test("positional placement keeps every row exactly once") { + withPositionalRoundRobin() { + withParquetTable(5000) { t => + checkSparkAnswer(spark.table(t).repartition(numPartitions).selectExpr("id", "s", "g")) + } + } + } + + test("a group size larger than a batch still keeps every row") { + // Exercises the whole-batch fast path, where a run covers an entire input batch and is handed + // through without a copy. + withPositionalRoundRobin( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_GROUP_ROWS.key -> "8192", + CometConf.COMET_BATCH_SIZE.key -> "128") { + withParquetTable(5000) { t => + checkSparkAnswer(spark.table(t).repartition(numPartitions)) + } + } + } + + /** + * Rows per output partition, in partition order. + * + * Taken off the RDD rather than by grouping on `spark_partition_id()`, because the aggregate + * that grouping introduces brings its own exchange and the planner drops the repartition under + * it, leaving the measurement describing the scan's partitioning instead. + */ + private def partitionSizes(df: DataFrame): Seq[Int] = + df.rdd.mapPartitions(rows => Iterator(rows.size)).collect().toSeq + + test("duplicate rows spread evenly, where hashing sends them all to one partition") { + // The behaviour that makes positional placement round robin rather than hash partitioning: + // `pmod(murmur3(row), n)` is a function of the row's contents, so a column of one repeated + // value collapses onto a single reducer. Spark's round robin spreads it. + withParquetTable(100) { t => + val duplicated = spark.table(t).select(lit(1).as("c")) + val distinct = spark.table(t).select(col("id")) + + // An explicit small group, because the derived default is batchSize / numPartitions and + // would put this whole 100-row fixture in one group. + withPositionalRoundRobin( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_GROUP_ROWS.key -> "8") { + val duplicatedSizes = partitionSizes(duplicated.repartition(numPartitions)) + assert(duplicatedSizes.sum == 100) + assert( + duplicatedSizes.count(_ > 0) > 1, + s"expected the rows to spread across partitions, got $duplicatedSizes") + // Placement ignores row contents entirely, so a column of one repeated value partitions + // exactly like a column of distinct ones. + assert(duplicatedSizes == partitionSizes(distinct.repartition(numPartitions))) + } + + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val sizes = partitionSizes(duplicated.repartition(numPartitions)) + assert( + sizes.count(_ > 0) == 1, + s"content hashing should collapse identical rows onto one partition, got $sizes") + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala index 99d9fe93e0c..e2894d1891a 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala @@ -67,6 +67,37 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { } } + test("positional round robin declares itself indeterminate over a non-determinate parent") { + // Spark's own round robin is positional and gets this from the `isOrderSensitive` flag on + // `MapPartitionsRDD`; the native path has no such RDD, so the rule is applied here. A + // determinate parent keeps the cheap per-task retry, anything else forces the DAGScheduler to + // roll the whole stage back rather than re-run one task into a partially consumed output. + Seq( + DeterministicLevel.DETERMINATE -> DeterministicLevel.DETERMINATE, + DeterministicLevel.UNORDERED -> DeterministicLevel.INDETERMINATE, + DeterministicLevel.INDETERMINATE -> DeterministicLevel.INDETERMINATE).foreach { + case (parentLevel, expected) => + val parent = new RDD[AnyRef](spark.sparkContext, Nil) { + override protected def getOutputDeterministicLevel: DeterministicLevel.Value = + parentLevel + override protected def getPartitions: Array[Partition] = Array.empty + override def compute(split: Partition, context: TaskContext): Iterator[AnyRef] = + Iterator.empty + } + val input = new CometNativeShuffleInputRDD( + spark.sparkContext, + Seq(parent), + 0, + Set.empty, + CometMetricNode(Map.empty), + positionalRoundRobin = true) + assert(input.outputDeterministicLevel == expected, s"parent was $parentLevel") + // The flag has to survive the copy, or a local-shuffle fallback silently drops the + // declaration and the scheduler goes back to re-running single tasks. + assert(input.copyForLocalShuffle().outputDeterministicLevel == expected) + } + } + test("local shuffle input is an independent sibling with the same partition inputs") { val upstream = new RDD[AnyRef](spark.sparkContext, Nil) { override protected def getPartitions: Array[Partition] = Array.tabulate(2) { i => From 7fe578bd8489c91468fca9fafc8c2088c8fda91a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 17:32:20 -0600 Subject: [PATCH 02/12] fix: honour a positional group size larger than the batch size Capping `groupRows` at `batch_size` silently ignored what the user asked for. A longer group is meaningful and works as written: it sends several consecutive input batches to the same output partition, which is a legitimate way to trade balance for fewer, larger shuffle blocks. --- native/shuffle/src/comet_partitioning.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/native/shuffle/src/comet_partitioning.rs b/native/shuffle/src/comet_partitioning.rs index f6b9ed257a8..86d092fcc50 100644 --- a/native/shuffle/src/comet_partitioning.rs +++ b/native/shuffle/src/comet_partitioning.rs @@ -86,7 +86,8 @@ impl RoundRobinStrategy { /// /// One batch spread over `num_partitions` groups is the finest split that still gives every /// output partition a run, so `batch_size / num_partitions` balances without fragmenting the - /// copy any further than it has to. + /// copy any further than it has to. An explicit request is taken as given, including one + /// larger than a batch, which sends several consecutive input batches to the same partition. pub fn resolve_group_rows( group_rows: usize, batch_size: usize, @@ -94,7 +95,7 @@ impl RoundRobinStrategy { ) -> usize { let batch_size = batch_size.max(1); if group_rows != Self::AUTO_GROUP_ROWS { - return group_rows.min(batch_size); + return group_rows; } (batch_size / num_partitions.max(1)) .clamp(Self::MIN_AUTO_GROUP_ROWS.min(batch_size), batch_size) @@ -324,10 +325,9 @@ mod tests { assert_eq!(S::resolve_group_rows(S::AUTO_GROUP_ROWS, 8192, 10_000), 64); // A batch smaller than the minimum group still resolves to something usable. assert_eq!(S::resolve_group_rows(S::AUTO_GROUP_ROWS, 32, 200), 32); - // An explicit request is honoured, but never exceeds the batch size: a group longer than - // a batch can never see a second batch's rows anyway, since `insert_batch` slices input - // down to `batch_size` before placing it. + // An explicit request is taken as given. A group longer than a batch is meaningful: it + // sends several consecutive input batches to the same output partition. assert_eq!(S::resolve_group_rows(1, 8192, 200), 1); - assert_eq!(S::resolve_group_rows(100_000, 8192, 200), 8192); + assert_eq!(S::resolve_group_rows(100_000, 8192, 200), 100_000); } } From 196c19d28181dbb46f9e6b049169a1ae81fc1a41 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 20:37:50 -0600 Subject: [PATCH 03/12] refactor: simplify positional round robin plumbing Resolve the positional group size into the partitioning itself instead of a parallel field, give each round-robin strategy its own match arm, and buffer runs without moving the scratch vector in and out. Skip the unused partition_starts scratch under positional placement, pre-size the run iterator's chunk scratch, and compute the driver-side positional decision once per exchange so the RDD and the writer read the same value. --- .../src/partitioners/multi_partition.rs | 343 +++++++++--------- .../partitioned_batch_iterator.rs | 17 +- .../shuffle/CometShuffleExchangeExec.scala | 11 +- ...CometNativePositionalRoundRobinSuite.scala | 14 +- 4 files changed, 180 insertions(+), 205 deletions(-) diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 054b84fe463..8bc2693f556 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -128,9 +128,6 @@ pub(crate) struct MultiPartitionShuffleRepartitioner { /// allocation once rather than once per slice that references it. Cleared whenever the /// buffered batches drain (spill / shuffle_write). See `count_new_buffers`. pinned_buffers: HashSet, - /// `Some((start_partition, group_rows))` under [`RoundRobinStrategy::RowGroups`], with - /// `group_rows` already resolved against the batch size and partition count. - positional: Option<(usize, usize)>, /// Rows this task has placed so far, which is the ordinal positional round robin keys on. /// A `u64` because it counts a whole task's input, not one batch. row_seq: u64, @@ -202,26 +199,31 @@ impl MultiPartitionShuffleRepartitioner { // Positional round robin is the one strategy that never looks at a row's contents. It // needs none of the row-level scratch (~64 KB a task), and it records contiguous runs - // rather than individual rows. Resolve it here, before `partitioning` moves into the - // struct below. - let positional = match &partitioning { + // rather than individual rows. Its group size is resolved once, here. + let partitioning = match partitioning { CometPartitioning::RoundRobin( - _, + n, RoundRobinStrategy::RowGroups { start_partition, group_rows, }, - ) => Some(( - *start_partition, - RoundRobinStrategy::resolve_group_rows( - *group_rows, - batch_size, - num_output_partitions, - ), - )), - _ => None, + ) => CometPartitioning::RoundRobin( + n, + RoundRobinStrategy::RowGroups { + start_partition, + group_rows: RoundRobinStrategy::resolve_group_rows( + group_rows, + batch_size, + num_output_partitions, + ), + }, + ), + other => other, }; - let places_rows_individually = positional.is_none(); + let places_rows_individually = !matches!( + partitioning, + CometPartitioning::RoundRobin(_, RoundRobinStrategy::RowGroups { .. }) + ); // 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. @@ -246,7 +248,11 @@ impl MultiPartitionShuffleRepartitioner { } else { vec![] }, - partition_starts: vec![0; num_output_partitions + 1], + partition_starts: if places_rows_individually { + vec![0; num_output_partitions + 1] + } else { + vec![] + }, positional_runs: vec![], }; @@ -270,7 +276,6 @@ impl MultiPartitionShuffleRepartitioner { max_buffer_bytes, tracing_enabled, pinned_buffers: HashSet::new(), - positional, row_seq: 0, }) } @@ -400,95 +405,95 @@ impl MultiPartitionShuffleRepartitioner { .await?; self.scratch = scratch; } - CometPartitioning::RoundRobin(num_output_partitions, strategy) => { + CometPartitioning::RoundRobin( + num_output_partitions, + RoundRobinStrategy::RowGroups { + start_partition, + group_rows, + }, + ) => { + let num_rows = input.num_rows(); + { + let mut timer = self.metrics.repart_time.timer(); + positional_runs( + self.row_seq, + num_rows, + *start_partition, + *group_rows, + *num_output_partitions, + &mut self.scratch.positional_runs, + ); + timer.stop(); + } + // Count rows, not batches: a group left part-way through by this batch is + // finished by the next one, so placement does not depend on where the reader + // put the batch boundary. See `RoundRobinStrategy::RowGroups`. + self.row_seq += num_rows as u64; + self.buffer_positional_batch_may_spill(input).await?; + } + CometPartitioning::RoundRobin( + num_output_partitions, + RoundRobinStrategy::HashAll { max_hash_columns }, + ) => { let mut scratch = std::mem::take(&mut self.scratch); let num_rows = input.num_rows(); - match strategy { - RoundRobinStrategy::RowGroups { .. } => { - let (start_partition, group_rows) = self - .positional - .expect("positional resolved in try_new for RowGroups"); - { - let mut timer = self.metrics.repart_time.timer(); - positional_runs( - self.row_seq, - num_rows, - start_partition, - group_rows, - *num_output_partitions, - &mut scratch.positional_runs, - ); - timer.stop(); - } - // Count rows, not batches: a group left part-way through by this batch is - // finished by the next one, so placement does not depend on where the - // reader put the batch boundary. See `RoundRobinStrategy::RowGroups`. - self.row_seq += num_rows as u64; - self.scratch = scratch; - self.buffer_positional_batch_may_spill(input).await?; - } - RoundRobinStrategy::HashAll { max_hash_columns } => { - // Comet implements this flavour of "round robin" as hash partitioning on - // columns: even distribution without semantic grouping, and deterministic - // for fault tolerance regardless of input order. - // - // Note: This produces different partition assignments than Spark's round - // robin, which sorts by UnsafeRow binary representation before assigning - // partitions. - let (partition_starts, partition_row_indices): (&Vec, &Vec) = { - let mut timer = self.metrics.repart_time.timer(); - - // 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()) - }; - 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; - }); + // Comet implements this flavour of "round robin" as hash partitioning on + // columns: even distribution without semantic grouping, and deterministic + // for fault tolerance regardless of input order. + // + // Note: This produces different partition assignments than Spark's round + // robin, which sorts by UnsafeRow binary representation before assigning + // partitions. + let (partition_starts, partition_row_indices): (&Vec, &Vec) = { + let mut timer = self.metrics.repart_time.timer(); - // 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, - )) - }?; - - self.buffer_partitioned_batch_may_spill( - input, - partition_row_indices, - partition_starts, - ) - .await?; - self.scratch = scratch; - } - } + // 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()) + }; + 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, + )) + }?; + + self.buffer_partitioned_batch_may_spill( + input, + partition_row_indices, + partition_starts, + ) + .await?; + self.scratch = scratch; } other => { // this should be unreachable as long as the validation logic @@ -551,30 +556,22 @@ impl MultiPartitionShuffleRepartitioner { ) -> datafusion::common::Result<()> { let (buffered_partition_idx, mut mem_growth) = self.buffer_input(input); - // Taken out and put back so that the loop can hold the runs and `partition_indices` at - // once; `positional_runs` clears before refilling, so the capacity survives. - let runs = std::mem::take(&mut self.scratch.positional_runs); - let result = match &mut self.partition_indices { - PartitionIndices::Runs(partition_runs) => { - for run in &runs { - let indices = &mut partition_runs[run.partition]; - let before_size = indices.allocated_size(); - indices.push(BufferedRun { - batch: buffered_partition_idx, - start: run.start, - len: run.len, - }); - let after_size = indices.allocated_size(); - mem_growth += after_size.saturating_sub(before_size); - } - Ok(()) - } - PartitionIndices::Rows(_) => Err(DataFusionError::Internal( + let PartitionIndices::Runs(partition_runs) = &mut self.partition_indices else { + return Err(DataFusionError::Internal( "positional placement against a row-indexed repartitioner".to_string(), - )), + )); }; - self.scratch.positional_runs = runs; - result?; + for run in &self.scratch.positional_runs { + let indices = &mut partition_runs[run.partition]; + let before_size = indices.allocated_size(); + indices.push(BufferedRun { + batch: buffered_partition_idx, + start: run.start, + len: run.len, + }); + let after_size = indices.allocated_size(); + mem_growth += after_size.saturating_sub(before_size); + } self.reserve_and_may_spill(mem_growth) } @@ -1191,27 +1188,21 @@ mod tests { } } - /// Runs `total_rows` sequential i64 values through a positional repartitioner, framed into - /// batches of the given sizes, and returns the values each output partition received. - async fn positional_placement( - framing: &[usize], + /// A positional repartitioner over a single non-null i64 column. + fn positional_repartitioner( num_partitions: usize, group_rows: usize, start_partition: usize, batch_size: usize, - ) -> std::collections::BTreeMap> { - let schema = Arc::new(arrow::datatypes::Schema::new(vec![ - arrow::datatypes::Field::new("v", arrow::datatypes::DataType::Int64, false), - ])); + ) -> MultiPartitionShuffleRepartitioner { let runtime = Arc::new( datafusion::execution::runtime_env::RuntimeEnvBuilder::new() .build() .unwrap(), ); - let metrics_set = ExecutionPlanMetricsSet::new(); - let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + MultiPartitionShuffleRepartitioner::try_new( 0, - CollectingPartitionWriter::default(), + W::default(), CometPartitioning::RoundRobin( num_partitions, RoundRobinStrategy::RowGroups { @@ -1219,23 +1210,45 @@ mod tests { group_rows, }, ), - ShufflePartitionerMetrics::new(&metrics_set, 0), - Arc::clone(&runtime), + ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + runtime, batch_size, false, None, ) - .unwrap(); + .unwrap() + } + fn int64_batch(values: std::ops::Range) -> RecordBatch { + let schema = Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("v", arrow::datatypes::DataType::Int64, false), + ])); + RecordBatch::try_new( + schema, + vec![Arc::new(Int64Array::from_iter_values(values)) as ArrayRef], + ) + .unwrap() + } + + /// Runs `total_rows` sequential i64 values through a positional repartitioner, framed into + /// batches of the given sizes, and returns the values each output partition received. + async fn positional_placement( + framing: &[usize], + num_partitions: usize, + group_rows: usize, + start_partition: usize, + batch_size: usize, + ) -> std::collections::BTreeMap> { + let mut repartitioner = positional_repartitioner::( + num_partitions, + group_rows, + start_partition, + batch_size, + ); let mut next = 0i64; for rows in framing { - let values: Vec = (next..next + *rows as i64).collect(); + let batch = int64_batch(next..next + *rows as i64); next += *rows as i64; - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values)) as ArrayRef], - ) - .unwrap(); repartitioner.insert_batch(batch).await.unwrap(); } repartitioner.shuffle_write().unwrap(); @@ -1317,46 +1330,20 @@ mod tests { /// Positional placement records runs, not rows, and skips the row-level scratch entirely. #[tokio::test] async fn positional_placement_records_runs_not_rows() { - let schema = Arc::new(arrow::datatypes::Schema::new(vec![ - arrow::datatypes::Field::new("v", arrow::datatypes::DataType::Int64, false), - ])); - let runtime = Arc::new( - datafusion::execution::runtime_env::RuntimeEnvBuilder::new() - .build() - .unwrap(), - ); - let metrics_set = ExecutionPlanMetricsSet::new(); - let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( - 0, - FailingPartitionWriter::default(), - CometPartitioning::RoundRobin( - 8, - RoundRobinStrategy::RowGroups { - start_partition: 0, - group_rows: 64, - }, - ), - ShufflePartitionerMetrics::new(&metrics_set, 0), - Arc::clone(&runtime), - 256, - false, - None, - ) - .unwrap(); + let mut repartitioner = positional_repartitioner::(8, 64, 0, 256); assert!( repartitioner.scratch.partition_ids.is_empty() && repartitioner.scratch.partition_row_indices.is_empty() + && repartitioner.scratch.partition_starts.is_empty() && repartitioner.scratch.hashes_buf.is_empty(), "positional placement should allocate no per-row scratch" ); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from((0..256i64).collect::>())) as ArrayRef], - ) - .unwrap(); - repartitioner.insert_batch(batch).await.unwrap(); + repartitioner + .insert_batch(int64_batch(0..256)) + .await + .unwrap(); // 256 rows in groups of 64 is four runs, not 256 row entries. assert_eq!(repartitioner.partition_indices.entry_count(), 4); diff --git a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs index 1799ce9cf5c..e65d8b2ed31 100644 --- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs +++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs @@ -254,22 +254,11 @@ impl<'a> RunIterator<'a> { batch_size: usize, copy_time: &'a Time, ) -> Self { - if runs.is_empty() { - return Self { - record_batches: &[], - batch_size, - runs: &[], - chunk_scratch: vec![], - pos: 0, - consumed: 0, - copy_time, - }; - } Self { record_batches, batch_size, runs, - chunk_scratch: vec![], + chunk_scratch: Vec::with_capacity(runs.len().min(batch_size)), pos: 0, consumed: 0, copy_time, @@ -288,9 +277,7 @@ impl Iterator for RunIterator<'_> { // Zero-copy path: the next run is an entire buffered batch and already fills a chunk on // its own, so hand the batch straight through. This is the case a group as large as the - // batch size is chosen to hit. Returning the batch rather than a slice of it also keeps - // `Utf8View`/`BinaryView` columns off the sliced-array path in the IPC writer, which - // truncates the views buffer but serializes every shared data buffer in full. + // batch size is chosen to hit. if self.consumed == 0 { let run = self.runs[self.pos]; let source = self.record_batches[run.batch as usize]; diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 6df8b5ff42e..47301d84dc1 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -115,6 +115,13 @@ case class CometShuffleExchangeExec( case _ => None } + /** + * Positional round-robin decision, computed once so that the RDD's determinism level and the + * writer's placement cannot disagree. + */ + @transient private lazy val positionalRoundRobin: Option[PositionalRoundRobin] = + CometShuffleExchangeExec.positionalRoundRobinSpec(outputPartitioning, child) + @transient private lazy val nativeChildMetricNode: CometMetricNode = CometMetricNode.fromCometPlan(child) @@ -128,7 +135,7 @@ case class CometShuffleExchangeExec( ctx.shuffleScanIndices, CometMetricNode(metrics, Seq(nativeChildMetricNode)), ctx.perPartitionByKey, - CometShuffleExchangeExec.usesPositionalRoundRobin(outputPartitioning, child)) + positionalRoundRobin.isDefined) case None => // Non-native child (e.g. CometSparkToColumnarExec): no subtree to inline. The dep gets // built via the convenience overload below; we just need a real RDD of batches. @@ -211,7 +218,7 @@ case class CometShuffleExchangeExec( nativeChild.nativeOp, nativeChildMetricNode, ctx, - CometShuffleExchangeExec.positionalRoundRobinSpec(outputPartitioning, child))) + positionalRoundRobin)) case None => CometShuffleExchangeExec.prepareShuffleDependency( inputRDD.asInstanceOf[RDD[ColumnarBatch]], diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala index 3bd711705c0..e2437c223f4 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala @@ -108,11 +108,8 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp test("positional placement is off unless its own config is on") { withParquetTable(100) { t => - withSQLConf( - CometConf.COMET_SHUFFLE_MODE.key -> "native", - CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "true", - CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key -> "false", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withPositionalRoundRobin( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key -> "false") { assert(!isPositional(spark.table(t).repartition(numPartitions))) } } @@ -186,11 +183,8 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp assert(duplicatedSizes == partitionSizes(distinct.repartition(numPartitions))) } - withSQLConf( - CometConf.COMET_SHUFFLE_MODE.key -> "native", - CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "true", - CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key -> "false", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withPositionalRoundRobin( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key -> "false") { val sizes = partitionSizes(duplicated.repartition(numPartitions)) assert( sizes.count(_ > 0) == 1, From 8aa290b0992a8b6ab1766a81aee3157c95a4cbfe Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 20:24:52 -0600 Subject: [PATCH 04/12] fix: rebuild the shuffle writer per iteration in end-to-end benches A ShuffleWriterExec publishes its partition offsets through a OnceLock, so re-executing one exec across criterion iterations fails on the second run with "partition offsets were already published". Build a fresh exec per iteration with iter_batched, keeping construction out of the timing. --- native/shuffle/benches/shuffle_writer.rs | 105 +++++++++++------------ 1 file changed, 52 insertions(+), 53 deletions(-) diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index cc94375436e..74d760d8e65 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -19,7 +19,7 @@ use arrow::array::builder::{Date32Builder, Decimal128Builder, Int32Builder}; use arrow::array::{builder::StringBuilder, Array, Int32Array, RecordBatch}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::row::{RowConverter, SortField}; -use criterion::{criterion_group, criterion_main, Criterion}; +use criterion::{criterion_group, criterion_main, BatchSize, Bencher, Criterion}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::physical_expr::expressions::{col, Column}; @@ -73,19 +73,14 @@ fn criterion_benchmark(c: &mut Criterion) { group.bench_function( format!("shuffle_writer: end to end (compression = {compression_codec:?})"), |b| { - let ctx = SessionContext::new(); - let exec = create_shuffle_writer_exec( - compression_codec.clone(), - CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], 16), - 8192, - 10, - ); - b.iter(|| { - let task_ctx = ctx.task_ctx(); - let stream = exec.execute(0, task_ctx).unwrap(); - let rt = Runtime::new().unwrap(); - rt.block_on(collect(stream)).unwrap(); - }); + bench_end_to_end(b, || { + create_shuffle_writer_exec( + compression_codec.clone(), + CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], 16), + 8192, + 10, + ) + }) }, ); } @@ -126,19 +121,14 @@ fn criterion_benchmark(c: &mut Criterion) { group.bench_function( format!("shuffle_writer: end to end (partitioning={partitioning:?})"), |b| { - let ctx = SessionContext::new(); - let exec = create_shuffle_writer_exec( - compression_codec.clone(), - partitioning.clone(), - 8192, - 10, - ); - b.iter(|| { - let task_ctx = ctx.task_ctx(); - let stream = exec.execute(0, task_ctx).unwrap(); - let rt = Runtime::new().unwrap(); - rt.block_on(collect(stream)).unwrap(); - }); + bench_end_to_end(b, || { + create_shuffle_writer_exec( + compression_codec.clone(), + partitioning.clone(), + 8192, + 10, + ) + }) }, ); } @@ -157,19 +147,14 @@ fn criterion_benchmark(c: &mut Criterion) { group.bench_function( format!("shuffle_writer: end to end (partitioning=SinglePartition, rows_per_batch={rows_per_batch})"), |b| { - let ctx = SessionContext::new(); - let exec = create_shuffle_writer_exec( - CompressionCodec::None, - CometPartitioning::SinglePartition, - rows_per_batch, - num_batches, - ); - b.iter(|| { - let task_ctx = ctx.task_ctx(); - let stream = exec.execute(0, task_ctx).unwrap(); - let rt = Runtime::new().unwrap(); - rt.block_on(collect(stream)).unwrap(); - }); + bench_end_to_end(b, || { + create_shuffle_writer_exec( + CompressionCodec::None, + CometPartitioning::SinglePartition, + rows_per_batch, + num_batches, + ) + }) }, ); } @@ -185,25 +170,39 @@ fn criterion_benchmark(c: &mut Criterion) { high_partition_group.bench_function( format!("shuffle_writer: end to end (partitions={num_partitions}, compression=None)"), |b| { - let ctx = SessionContext::new(); - let exec = create_shuffle_writer_exec( - CompressionCodec::None, - CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], num_partitions), - 8192, - 10, - ); - b.iter(|| { - let task_ctx = ctx.task_ctx(); - let stream = exec.execute(0, task_ctx).unwrap(); - let rt = Runtime::new().unwrap(); - rt.block_on(collect(stream)).unwrap(); - }); + bench_end_to_end(b, || { + create_shuffle_writer_exec( + CompressionCodec::None, + CometPartitioning::Hash( + vec![Arc::new(Column::new("a", 0))], + num_partitions, + ), + 8192, + 10, + ) + }) }, ); } high_partition_group.finish(); } +/// Times one execution of a freshly built writer per iteration. A `ShuffleWriterExec` +/// publishes its partition offsets once, so it cannot be re-executed; building it is +/// setup and stays outside the measurement. +fn bench_end_to_end(b: &mut Bencher, make_exec: impl Fn() -> ShuffleWriterExec) { + let ctx = SessionContext::new(); + b.iter_batched( + make_exec, + |exec| { + let stream = exec.execute(0, ctx.task_ctx()).unwrap(); + let rt = Runtime::new().unwrap(); + rt.block_on(collect(stream)).unwrap(); + }, + BatchSize::LargeInput, + ); +} + fn create_shuffle_writer_exec( compression_codec: CompressionCodec, partitioning: CometPartitioning, From e99307b9df2620a743239abfc032f539f21e7042 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 22 Sep 2026 08:15:43 -0600 Subject: [PATCH 05/12] bench: add a partitioning-only benchmark for round robin placement The end-to-end shuffle benches are dominated by IPC encoding and the file write, so they cannot separate two placement strategies. Add a group that stops before both: one arm times placement and per-partition index buffering alone, the other adds the flush through a writer that discards its batches, which is where a run-indexed partitioner diverges from a row-indexed one. benches/ compiles as its own crate, so reaching the partitioners needs a pub seam. Rather than export MultiPartitionShuffleRepartitioner and the PartitionWriter trait, add one opaque handle in a doc(hidden) module and leave the rest crate-private. The nested fixture grows a per-row fill. The existing one repeats a single value down every leaf, so every row hashes alike and a hash strategy would place the whole input on one output partition, never performing the scatter that the gather exists to undo. The encoding benches keep the constant fill and their current numbers. --- native/shuffle/benches/shuffle_writer.rs | 149 +++++++++++++++++++++-- native/shuffle/src/bench_support.rs | 121 ++++++++++++++++++ native/shuffle/src/lib.rs | 3 + 3 files changed, 264 insertions(+), 9 deletions(-) create mode 100644 native/shuffle/src/bench_support.rs diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index 74d760d8e65..e8ea588b02d 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -30,7 +30,8 @@ use datafusion::{ prelude::SessionContext, }; use datafusion_comet_shuffle::{ - CometPartitioning, CompressionCodec, ShuffleBlockWriter, ShuffleCodecContext, ShuffleWriterExec, + bench_support::BenchRepartitioner, CometPartitioning, CompressionCodec, RoundRobinStrategy, + ShuffleBlockWriter, ShuffleCodecContext, ShuffleWriterExec, }; use itertools::Itertools; use std::io::Cursor; @@ -269,6 +270,98 @@ fn create_batch(num_rows: usize, allow_nulls: bool) -> RecordBatch { .unwrap() } +/// Round robin placement in isolation, and then the gather it implies on flush. +/// +/// The end-to-end benches above spend most of their time in IPC encoding and the file write, +/// which on a loaded disk swamps the difference between two placement strategies entirely. +/// These stop short of both. `place` times only the strategy and the per-partition index +/// buffering; `place+gather` adds the flush through a writer that discards its batches, which is +/// where a run-indexed partitioner diverges from a row-indexed one — the row-level scatter makes +/// `interleave_record_batch` walk every column and child again, and a run can instead be sliced, +/// or handed through untouched when it covers a whole buffered batch. +/// +/// 8192 rows per batch into 50 output partitions, so [`RoundRobinStrategy::AUTO_GROUP_ROWS`] +/// resolves to 163. `RowGroups(8192)` is the opposite extreme, one whole input batch per group, +/// where every run covers a buffered batch end to end and the gather copies nothing at all. +fn partitioning_benchmark(c: &mut Criterion) { + const BATCH_SIZE: usize = 8192; + const NUM_BATCHES: usize = 8; + const NUM_PARTITIONS: usize = 50; + + let strategies = [ + ( + "HashAll", + RoundRobinStrategy::HashAll { + max_hash_columns: 0, + }, + ), + // Hashing only the leading column: not a candidate strategy, but it separates the cost + // of recursing through every struct child from the cost of hashing at all. + ( + "HashAll{1}", + RoundRobinStrategy::HashAll { + max_hash_columns: 1, + }, + ), + ( + "RowGroups(auto)", + RoundRobinStrategy::RowGroups { + start_partition: 0, + group_rows: RoundRobinStrategy::AUTO_GROUP_ROWS, + }, + ), + ( + "RowGroups(8192)", + RoundRobinStrategy::RowGroups { + start_partition: 0, + group_rows: BATCH_SIZE, + }, + ), + ]; + + // `plain` is the flat schema the end-to-end benches use. `nested` is the shape that motivates + // positional placement: 40 struct columns over a three-field leaf, so 120 leaf arrays for a + // hash to recurse into and for a gather to walk. + let fixtures = [ + ("plain", create_batches(BATCH_SIZE, NUM_BATCHES)), + ( + "nested", + nested_batches(BATCH_SIZE, NUM_BATCHES, 40, 2, Fill::PerRow), + ), + ]; + + let mut group = c.benchmark_group("shuffle_partitioning"); + for (schema, batches) in &fixtures { + for (label, strategy) in &strategies { + let partitioning = CometPartitioning::RoundRobin(NUM_PARTITIONS, strategy.clone()); + group.bench_function(format!("place ({schema}, {label})"), |b| { + b.iter_batched( + || BenchRepartitioner::try_new(partitioning.clone(), BATCH_SIZE).unwrap(), + |mut repartitioner| { + repartitioner.place(batches).unwrap(); + // Returned so that freeing the buffered batches and the partition index + // lands outside the measurement. + repartitioner + }, + BatchSize::LargeInput, + ); + }); + group.bench_function(format!("place+gather ({schema}, {label})"), |b| { + b.iter_batched( + || BenchRepartitioner::try_new(partitioning.clone(), BATCH_SIZE).unwrap(), + |mut repartitioner| { + repartitioner.place(batches).unwrap(); + repartitioner.gather().unwrap(); + repartitioner + }, + BatchSize::LargeInput, + ); + }); + } + } + group.finish(); +} + /// Benchmarks the per-block IPC encoding cost (schema + record batch) in isolation, using the /// `None` codec so that compression does not obscure the schema-encoding cost. Covers a wide flat /// schema and a deeply nested schema, where the schema flatbuffer is largest. @@ -371,13 +464,38 @@ fn flat_schema_batch(num_rows: usize) -> RecordBatch { /// A schema of several deeply nested struct columns. fn nested_schema_batch(num_rows: usize) -> RecordBatch { - let num_cols = 4; - let depth = 6; + nested_batch(num_rows, 4, 6, Fill::Constant) +} + +/// How a nested fixture fills its leaves. +#[derive(Clone, Copy)] +enum Fill { + /// One value repeated down every leaf. Cheap to build, and enough for the encoding benches + /// that only care about how much there is to encode. Useless for partitioning: identical + /// rows hash alike, so a hash strategy would put the whole input on one output partition + /// and never perform the scatter that a gather has to undo. + Constant, + /// A distinct value per row, so hash placement spreads rows across the output partitions the + /// way real data does. + PerRow, +} + +fn nested_batches( + num_rows: usize, + count: usize, + num_cols: usize, + depth: usize, + fill: Fill, +) -> Vec { + let batch = nested_batch(num_rows, num_cols, depth, fill); + vec![batch; count] +} +fn nested_batch(num_rows: usize, num_cols: usize, depth: usize, fill: Fill) -> 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 { - let array = nested_struct_array(num_rows, depth); + let array = nested_struct_array(num_rows, depth, fill); fields.push(Field::new( format!("col{col}"), array.data_type().clone(), @@ -390,22 +508,35 @@ fn nested_schema_batch(num_rows: usize) -> RecordBatch { } /// Builds a struct array with a multi-field leaf, wrapped in `depth` single-field structs. -fn nested_struct_array(num_rows: usize, depth: usize) -> Arc { +fn nested_struct_array(num_rows: usize, depth: usize, fill: Fill) -> Arc { use arrow::array::{Float64Array, Int64Array, StringArray, StructArray}; + let (ints, strings, floats): (Vec, Vec, Vec) = match fill { + Fill::Constant => ( + vec![1_i64; num_rows], + vec!["x".to_string(); num_rows], + vec![1.0_f64; num_rows], + ), + Fill::PerRow => ( + (0..num_rows as i64).collect(), + (0..num_rows).map(|row| format!("value {row}")).collect(), + (0..num_rows).map(|row| row as f64 * 1.5).collect(), + ), + }; + // Leaf: struct let mut array: Arc = Arc::new(StructArray::from(vec![ ( Arc::new(Field::new("a", DataType::Int64, false)), - Arc::new(Int64Array::from(vec![1_i64; num_rows])) as Arc, + Arc::new(Int64Array::from(ints)) as Arc, ), ( Arc::new(Field::new("b", DataType::Utf8, false)), - Arc::new(StringArray::from(vec!["x"; num_rows])) as Arc, + Arc::new(StringArray::from(strings)) as Arc, ), ( Arc::new(Field::new("c", DataType::Float64, false)), - Arc::new(Float64Array::from(vec![1.0_f64; num_rows])) as Arc, + Arc::new(Float64Array::from(floats)) as Arc, ), ])); @@ -427,6 +558,6 @@ fn config() -> Criterion { criterion_group! { name = benches; config = config(); - targets = criterion_benchmark, schema_encoding_benchmark, ipc_context_reuse_benchmark + targets = criterion_benchmark, partitioning_benchmark, schema_encoding_benchmark, ipc_context_reuse_benchmark } criterion_main!(benches); diff --git a/native/shuffle/src/bench_support.rs b/native/shuffle/src/bench_support.rs new file mode 100644 index 00000000000..a04493adc05 --- /dev/null +++ b/native/shuffle/src/bench_support.rs @@ -0,0 +1,121 @@ +// 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. + +//! A seam onto the shuffle partitioners for `benches/`, which compiles as its own crate and so +//! can only reach `pub` items. +//! +//! The partitioners are `pub(crate)` on purpose: `native/core` drives them through +//! [`ShuffleWriterExec`](crate::ShuffleWriterExec) and nothing else should. Rather than widen +//! their visibility so a benchmark can name them, this exposes one opaque handle that does +//! exactly what the partitioning benchmark needs — place rows, then gather them back out — +//! and keeps every type it is built from private. +//! +//! Not an API, and not used outside `benches/`. + +use crate::metrics::ShufflePartitionerMetrics; +use crate::partitioners::{MultiPartitionShuffleRepartitioner, ShufflePartitioner}; +use crate::writers::PartitionWriter; +use crate::CometPartitioning; +use arrow::record_batch::RecordBatch; +use datafusion::common::Result; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use std::sync::Arc; + +/// Drops every batch it is handed, so a flush measures the gather out of the partition index +/// and nothing downstream of it: no IPC encoding, no compression, no file write. +struct DiscardingPartitionWriter; + +impl PartitionWriter for DiscardingPartitionWriter { + fn write( + &mut self, + _pid: usize, + iter: &mut I, + _metrics: &ShufflePartitionerMetrics, + ) -> Result<()> + where + I: Iterator>, + { + for batch in iter { + std::hint::black_box(batch?); + } + Ok(()) + } + + fn finish_partition( + &mut self, + pid: usize, + iter: &mut I, + metrics: &ShufflePartitionerMetrics, + ) -> Result<()> + where + I: Iterator>, + { + self.write(pid, iter, metrics) + } + + fn finish_all(&mut self, _metrics: &ShufflePartitionerMetrics) -> Result<()> { + Ok(()) + } +} + +/// One shuffle map task's worth of partitioning, with the write side stubbed out. +#[doc(hidden)] +pub struct BenchRepartitioner { + inner: MultiPartitionShuffleRepartitioner, +} + +impl BenchRepartitioner { + /// Builds a repartitioner over an unbounded memory pool and no buffer limit, so nothing + /// spills and every strategy is measured over the same work. + pub fn try_new(partitioning: CometPartitioning, batch_size: usize) -> Result { + let runtime = Arc::new(RuntimeEnvBuilder::new().build()?); + Ok(Self { + inner: MultiPartitionShuffleRepartitioner::try_new( + 0, + DiscardingPartitionWriter, + partitioning, + ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + runtime, + batch_size, + false, + None, + )?, + }) + } + + /// Places every row of every batch and buffers its index, which is everything a shuffle + /// write does before it flushes. The batches are cloned, so a caller can reuse its fixture + /// across iterations. + /// + /// Nothing below here awaits, so a bare executor is enough and the measurement carries no + /// tokio runtime. + pub fn place(&mut self, batches: &[RecordBatch]) -> Result<()> { + futures::executor::block_on(async { + for batch in batches { + self.inner.insert_batch(batch.clone()).await?; + } + Ok(()) + }) + } + + /// Gathers the buffered rows back into output batches, one output partition at a time, and + /// discards them. + pub fn gather(&mut self) -> Result<()> { + self.inner.shuffle_write() + } +} diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 93183051ecb..893951cf6f3 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -15,6 +15,9 @@ // specific language governing permissions and limitations // under the License. +/// Internals reached by `benches/` only. See the module docs. +#[doc(hidden)] +pub mod bench_support; mod codec_context; pub(crate) mod comet_partitioning; pub mod ipc; From 39096347aded2e54b66abb4b921b1fda57d55ac6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 22 Sep 2026 14:12:45 -0600 Subject: [PATCH 06/12] fix: scramble the positional round robin start partition Starting each map task at its own partition id made consecutive tasks start on consecutive partitions. A task walks ceil(rows / groupRows) consecutive partitions from its start, so adjacent starts overlap and the partitions past numMapTasks + groupsPerTask get nothing: ten tasks of 5,000 rows into 200 partitions at the auto group of 64 leave 112 reducers empty. That is the correlation SPARK-21782 fixed. Compute the start the way Spark does, XORShiftRandom(partitionId) .nextInt(numPartitions) + 1, in CometNativeShuffleWriter where the TaskContext is in scope, and pass it in the proto. Still a pure function of the map partition, so a retried task reproduces its own placement, and the + 1 matches Spark's pre-increment so groupRows = 1 places rows exactly where Spark's round robin would. The planner no longer reads self.partition, which removes the jni_api partition-0 caveat. --- native/core/src/execution/planner.rs | 9 ++- native/proto/src/proto/partitioning.proto | 9 ++- .../src/partitioners/multi_partition.rs | 40 ++++++++++- .../shuffle/CometNativeShuffleWriter.scala | 6 ++ .../shuffle/CometShuffleExchangeExec.scala | 22 +++++- ...CometNativePositionalRoundRobinSuite.scala | 68 +++++++++++++++++++ 6 files changed, 145 insertions(+), 9 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index c3092537228..c2029e8d7f4 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3653,12 +3653,11 @@ impl PhysicalPlanner { PartitioningStruct::SinglePartition(_) => Ok(CometPartitioning::SinglePartition), PartitioningStruct::RoundRobinPartition(rr_partition) => { let strategy = if rr_partition.positional { - // The Spark map partition id, not the DataFusion one: `jni_api` runs every - // native root plan with partition 0 (one Comet execution per Spark task), so - // `ShuffleWriterExec::execute` cannot supply it. See - // `RoundRobinStrategy::RowGroups` for why it has to be this value. RoundRobinStrategy::RowGroups { - start_partition: self.partition.max(0) as usize, + // Computed per task on the JVM, where the Spark map partition id is in + // scope, and scrambled the way Spark scrambles it. See + // `CometShuffleExchangeExec.positionalStartPartition`. + start_partition: rr_partition.positional_start_partition.max(0) as usize, // Negative or zero means "derive it from the batch size and partition // count", which the repartitioner does once it knows both. group_rows: rr_partition.positional_group_rows.max(0) as usize, diff --git a/native/proto/src/proto/partitioning.proto b/native/proto/src/proto/partitioning.proto index 5cb29218c7a..daec02dac14 100644 --- a/native/proto/src/proto/partitioning.proto +++ b/native/proto/src/proto/partitioning.proto @@ -58,11 +58,18 @@ message RoundRobinPartition { // Maximum number of columns to hash. 0 means no limit (hash all columns). int32 max_hash_columns = 2; // When true, place rows by position rather than by hashing their contents: the row at - // task-global ordinal i goes to (mapPartitionId + i / positional_group_rows) % num_partitions. + // task-global ordinal i goes to + // (positional_start_partition + i / positional_group_rows) % num_partitions. // Only set when the planner has established that the map task replays rows in the same order, // which is the condition Spark's own round robin relies on. bool positional = 3; // Rows per contiguous group under positional placement. 0 means derive it from the batch size // and the partition count. int32 positional_group_rows = 4; + // Output partition this map task's first group goes to. Filled in per task by + // `CometNativeShuffleWriter.buildUnifiedPlan` as + // `XORShiftRandom(mapPartitionId).nextInt(num_partitions) + 1`, the same scrambled start Spark + // uses (SPARK-21782). A pure function of the map partition id, so a retried task reproduces its + // own placement. + int32 positional_start_partition = 5; } diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 8bc2693f556..d9abfef7f2e 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -1294,10 +1294,10 @@ mod tests { } } - /// `start_partition` offsets which partition a task starts on, so mappers do not all pile + /// `start_partition` offsets which partition a task begins on, so mappers do not all pile /// their first group onto partition 0. #[tokio::test] - async fn positional_placement_starts_at_the_map_partition() { + async fn positional_placement_begins_at_the_start_partition() { for start in 0..4usize { let written = positional_placement(&[64], 4, 64, start, 256).await; assert_eq!( @@ -1308,6 +1308,42 @@ mod tests { } } + /// A task's groups walk *consecutive* partitions from its start, which is what makes the + /// stage-wide spread a function of how the starts are chosen rather than of the data. The JVM + /// picks the starts (`CometShuffleExchangeExec.positionalStartPartition`), but the reason it + /// has to scramble them lives here: with adjacent starts every task's run overlaps its + /// neighbours' and the tail of the partition space gets nothing. + #[tokio::test] + async fn positional_placement_walks_consecutive_partitions_from_its_start() { + let num_partitions = 32; + let group_rows = 16; + let rows = 5 * group_rows; + + let stage = |starts: Vec| async move { + let mut covered = std::collections::BTreeSet::new(); + for start in starts { + let written = + positional_placement(&[rows], num_partitions, group_rows, start, 256).await; + assert_eq!( + written.keys().copied().collect::>(), + (0..5) + .map(|g| (start + g) % num_partitions) + .collect::>() + .into_iter() + .collect::>(), + "a task of five groups from {start} should touch five consecutive partitions" + ); + covered.extend(written.keys().copied()); + } + covered.len() + }; + + // Four tasks of five groups each could reach twenty of the thirty-two partitions. + assert_eq!(stage(vec![0, 8, 16, 24]).await, 20); + // Adjacent starts overlap instead, and twenty-four reducers get nothing. + assert_eq!(stage(vec![0, 1, 2, 3]).await, 8); + } + /// Imbalance stays within one group regardless of how the input was framed, which is what /// counting rows rather than batches buys over assigning whole batches. #[tokio::test] 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 461cd842546..ab86ec2b0c1 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 @@ -426,6 +426,12 @@ class CometNativeShuffleWriter[K, V]( spec.positionalRoundRobin.foreach { positional => partitioning.setPositional(true) partitioning.setPositionalGroupRows(positional.groupRows) + // Per task, unlike the two above: which partition this mapper's first group goes to. + // See `CometShuffleExchangeExec.positionalStartPartition` for why it is scrambled. + partitioning.setPositionalStartPartition( + CometShuffleExchangeExec.positionalStartPartition( + Option(context).map(_.partitionId()).getOrElse(0), + effectivePartitionCount)) } val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 47301d84dc1..d2f36b71751 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -312,7 +312,7 @@ object CometShuffleExchangeExec /** * True when this exchange will run the native round-robin writer in its positional mode, where - * the row at task-global ordinal `i` goes to `(mapPartitionId + i / groupRows) % numPartitions` + * the row at task-global ordinal `i` goes to `(startPartition + i / groupRows) % numPartitions` * rather than to `pmod(hash(row), numPartitions)`. * * Positional placement is reproducible exactly when the map task replays its rows in the same @@ -355,6 +355,26 @@ object CometShuffleExchangeExec } } + /** + * Output partition that the first group of map task `mapPartitionId` goes to. + * + * The starts have to be decorrelated, not merely distinct. Each task walks `ceil(rows / + * groupRows)` consecutive partitions from its start, so if consecutive tasks start on + * consecutive partitions their runs all overlap and the partitions past `numMapTasks + + * groupsPerTask` get nothing: ten map tasks of 5,000 rows into 200 partitions at a group of 64 + * would leave 112 reducers empty. That is the correlation + * [[https://issues.apache.org/jira/browse/SPARK-21782 SPARK-21782]] fixed, and scrambling the + * map partition id through `XORShiftRandom` is how Spark fixes it, both in + * `ShuffleExchangeExec.getPartitionKeyExtractor` and in the JVM path below. + * + * Still a pure function of the map partition id, so a re-executed task reproduces its own + * placement. The `+ 1` matches Spark, which increments the counter before its first use, so at + * `groupRows == 1` this places rows exactly where Spark's round robin would for the same row + * order. + */ + def positionalStartPartition(mapPartitionId: Int, numPartitions: Int): Int = + new XORShiftRandom(mapPartitionId).nextInt(math.max(numPartitions, 1)) + 1 + /** * Whether re-executing this subtree yields the same rows in the same order. * diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala index e2437c223f4..8fda1e23332 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala @@ -192,4 +192,72 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp } } } + + /** + * Rows each reducer receives when `mapTasks` map tasks of `rowsPerTask` rows each place + * positionally into `partitions` output partitions, starting where `start` says. Counted a + * group at a time, which is exact because a group's rows all land together. + */ + private def stageSpread( + mapTasks: Int, + rowsPerTask: Int, + partitions: Int, + groupRows: Int, + start: Int => Int): Array[Int] = { + val counts = Array.fill(partitions)(0) + val groups = (rowsPerTask + groupRows - 1) / groupRows + for (mapPartitionId <- 0 until mapTasks; group <- 0 until groups) { + val rows = math.min(groupRows, rowsPerTask - group * groupRows) + counts(((start(mapPartitionId).toLong + group) % partitions).toInt) += rows + } + counts + } + + test("map tasks start on decorrelated partitions, so the stage leaves no reducer empty") { + // Each task walks ceil(rowsPerTask / groupRows) consecutive partitions from its start, so + // distinct starts are not enough: consecutive starts make every task's run overlap its + // neighbours' and the partitions past mapTasks + groupsPerTask never get a row. This is the + // correlation SPARK-21782 fixed, and `positionalStartPartition` fixes it the same way. + val (mapTasks, rowsPerTask, partitions, groupRows) = (10, 5000, 200, 64) + + val adjacent = stageSpread(mapTasks, rowsPerTask, partitions, groupRows, identity) + assert( + adjacent.count(_ == 0) == 112, + "the hazard this scrambling exists to avoid should still be reachable with adjacent starts") + + val scrambled = stageSpread( + mapTasks, + rowsPerTask, + partitions, + groupRows, + CometShuffleExchangeExec.positionalStartPartition(_, partitions)) + assert(scrambled.sum == mapTasks * rowsPerTask) + assert( + scrambled.count(_ == 0) == 0, + s"every reducer should get rows, got ${scrambled.count(_ == 0)} empty of $partitions") + } + + test("stage-wide balance needs many more groups per task than there are partitions") { + // Why the group size defaults to batchSize / numPartitions rather than to the batch size, + // even though a batch-sized group is far cheaper to flush: the per-task bound does not + // compose. A reducer sees the sum over every map task, and that sum only evens out once each + // task has wrapped the partition space several times. + val (mapTasks, rowsPerTask, partitions) = (50, 1000000, 200) + def spread(groupRows: Int): Double = { + val counts = stageSpread( + mapTasks, + rowsPerTask, + partitions, + groupRows, + CometShuffleExchangeExec.positionalStartPartition(_, partitions)) + assert(counts.sum == mapTasks.toLong * rowsPerTask) + counts.max.toDouble / counts.min + } + + // 64 rows per group is 15625 groups per task, so each task wraps 78 times and the sums + // converge. 8192 is 123 groups, fewer than there are partitions, so a task cannot even cover + // the space once and where the gaps fall is down to the starts. + assert(spread(64) < 1.01, s"expected an even stage at a small group, got ${spread(64)}") + assert(spread(8192) > 1.2, s"expected a batch-sized group to skew, got ${spread(8192)}") + } } From 974932e7509aca8099bc800e75f7f081bd9d3ca9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 22 Sep 2026 14:12:57 -0600 Subject: [PATCH 07/12] docs: correct the positional round robin balance and alignment claims Three wording fixes from review, none of them behavioural. The groupRows bound is per map task. A reducer sees the sum over all of them, which only evens out when each task emits many more groups than there are output partitions, so say that in the config doc rather than promising a stage-wide bound. MIN_AUTO_GROUP_ROWS is a cap on how finely a batch is cut, not an alignment guarantee. A run only starts on a byte boundary when the batch starts on a group boundary, and row_seq counts across batches, so after a filter every run in a batch is offset. The start has to be decorrelated across mappers, not merely distinct, which is what XORShiftRandom is for. Also update the round robin item in the shuffle review skill, which still said a positional round robin was a bug by construction. --- .ai/skills/review-comet-shuffle-pr/SKILL.md | 32 +++++++++++---- .../contributor-guide/native_shuffle.md | 39 ++++++++++++------- native/shuffle/src/comet_partitioning.rs | 37 +++++++++++------- .../scala/org/apache/comet/CometConf.scala | 17 +++++--- 4 files changed, 85 insertions(+), 40 deletions(-) diff --git a/.ai/skills/review-comet-shuffle-pr/SKILL.md b/.ai/skills/review-comet-shuffle-pr/SKILL.md index b768923ae7f..86bce3bd0d7 100644 --- a/.ai/skills/review-comet-shuffle-pr/SKILL.md +++ b/.ai/skills/review-comet-shuffle-pr/SKILL.md @@ -88,13 +88,31 @@ Partitioning is where shuffle silently produces wrong answers rather than failin - [ ] **Hash partitioning uses Murmur3 with seed 42** and `partition_id = hash % num_partitions`, matching Spark. Any change to the hash, the seed, or the modulo changes which rows land in which partition, which breaks a join between a Comet-shuffled side and a Spark-shuffled side. -- [ ] **Round robin is hash-based on purpose, and off by default.** Comet assigns partitions from a - Murmur3 hash rather than cycling row by row, because determinism across task retries is - required for correctness under fault tolerance. A PR that implements "true" round robin to fix - skew breaks that. Two costs are accepted: low-cardinality data distributes unevenly, and - unsorted rows land in different partitions than Spark's `UnsafeRow`-sorted assignment would - put them, which is why `spark.comet.shuffle.native.partitioning.roundrobin.enabled` defaults - to `false`. Sorted output is identical either way. +- [ ] **Round robin defaults to a content hash, and is off by default.** Comet's default + `RoundRobinStrategy::HashAll` assigns partitions from a Murmur3 hash rather than cycling row + by row, because placement has to be reproducible across task retries. Two costs are accepted: + low-cardinality data distributes unevenly, and unsorted rows land in different partitions + than Spark's `UnsafeRow`-sorted assignment would put them, which is why + `spark.comet.shuffle.native.partitioning.roundrobin.enabled` defaults to `false`. Sorted + output is identical either way. +- [ ] **Positional round robin is allowed, but only where retry reproducibility is established.** + `RoundRobinStrategy::RowGroups` places the row at task-global ordinal `i` at + `(startPartition + i / groupRows) % numPartitions`, which is Spark's own round robin at a + coarser granularity. It is not a bug, but it is only correct behind two gates, and a PR that + widens either one needs an argument. `CometShuffleExchangeExec.replaysRowsInOrder` is an + allowlist over the native subtree fused into the writer, which the RDD graph cannot see: a + native scan under nothing but projections and filters, with anything that spills staying out + because it reorders between attempts. `CometNativeShuffleInputRDD.getOutputDeterministicLevel` + applies Spark's `isOrderSensitive` rule to everything below that RDD. +- [ ] **Positional placement keys on a row ordinal, and the start is scrambled.** The counter is + over rows and carries across batch boundaries. A PR that keys on a batch ordinal instead is + relying on framing, which no Spark contract covers: `DETERMINATE` promises the same rows in + the same order and says nothing about how an operator chunks them. `startPartition` must be + decorrelated across mappers, not merely distinct — each task walks consecutive partitions + from its start, so adjacent starts overlap and leave the tail of the partition space empty + (SPARK-21782). And `groupRows` bounds imbalance within one map task only; a reducer sees the + sum over all of them, which only evens out when each task emits many more groups than there + are partitions. - [ ] **Range partitioning bounds come from the driver.** Spark's `RangePartitioner` samples and computes boundaries, they are serialized into the native plan, and native does a binary search over comparable-row-format keys. A change to the comparison or the row encoding must diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index ef71b168d1d..3dbaa320f16 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -336,7 +336,7 @@ Hashing every column of every row dominates the shuffle write on wide nested sch `create_murmur3_hashes` recurses into every struct child per row and the resulting row-level scatter forces `interleave_record_batch` to walk every column and child again on flush. `RowGroups` places rows the way Spark's own round robin does: the row at task-global ordinal `i` -goes to `(mapPartitionId + i / groupRows) % numPartitions`. That removes the per-row hash, and it +goes to `(startPartition + i / groupRows) % numPartitions`. That removes the per-row hash, and it replaces the per-row gather with a bulk copy per contiguous run, because adjacent rows now stay together. It also spreads duplicate rows evenly, which `HashAll` cannot. @@ -347,19 +347,30 @@ downstream operator frames rows into batches, so an operator that spills can ref different memory pressure while still honouring `DETERMINATE`. Keying on a row ordinal means the strategy depends only on the property Spark actually publishes. -`start_partition` is the Spark map partition id, filled in by `PhysicalPlanner::create_partitioning` -from the planner's partition because `ShuffleWriterExec::execute` cannot supply it (`jni_api` runs -every native root plan with partition 0, one Comet execution per Spark task). It has to be distinct -across mappers, or every task starts at partition 0 and a task emitting fewer groups than there are -output partitions leaves the tail empty stage-wide; and it has to be a pure function of the map -partition, or a re-executed task does not reproduce its own placement. Spark seeds -`XORShiftRandom(partitionId)` for the same two reasons. - -`groupRows` trades balance against copying. Imbalance between any two output partitions is bounded -by `groupRows` rows however the reader frames its batches, so small groups balance better; large -groups produce fewer, longer runs to copy, and a group as large as the batch size lets a whole -input batch pass through to one partition untouched. `0`, the default, derives it as -`clamp(batch_size / num_partitions, 64, batch_size)`. +`start_partition` is the output partition a map task's first group goes to. It is computed per task +on the JVM, in `CometNativeShuffleWriter.buildUnifiedPlan` where the Spark map partition id is in +scope, and passed down in the proto. It has to be _decorrelated_ across mappers, not merely +distinct: a task walks `ceil(rows / groupRows)` consecutive partitions from its start, so if +consecutive tasks start on consecutive partitions their runs all overlap and the partitions past +`numMapTasks + groupsPerTask` get nothing — ten map tasks of 5,000 rows into 200 partitions at a +group of 64 would leave 112 reducers empty. It also has to be a pure function of the map partition, +or a re-executed task does not reproduce its own placement. Spark satisfies both by scrambling the +map partition id through `XORShiftRandom` +([SPARK-21782](https://issues.apache.org/jira/browse/SPARK-21782)), and +`CometShuffleExchangeExec.positionalStartPartition` does the same. Spark increments its counter +before the first row uses it, so the start is `nextInt(numPartitions) + 1`, which makes +`groupRows = 1` place rows exactly where Spark's round robin would for the same row order. + +`groupRows` trades balance against copying. Within one map task, imbalance between any two output +partitions is bounded by `groupRows` rows however the reader frames its batches, so small groups +balance better; large groups produce fewer, longer runs to copy, and a group as large as the batch +size lets a whole input batch pass through to one partition untouched. That bound does not compose +across map tasks — a reducer sees the sum over all of them, which is only even when each task emits +many more groups than there are output partitions. `0`, the default, derives it as +`clamp(batch_size / num_partitions, 64, batch_size)`, which keeps a task wrapping around the output +partitions roughly once per batch. The 64-row floor caps how finely a batch is cut: with far more +partitions than rows in a batch, `batch_size / num_partitions` rounds down towards one row and the +flush degenerates into the per-row gather positional placement exists to avoid. Internally, `MultiPartitionShuffleRepartitioner` records `(batch, start, len)` runs rather than one `(batch, row)` pair per row, so the index list charged against the spill reservation is diff --git a/native/shuffle/src/comet_partitioning.rs b/native/shuffle/src/comet_partitioning.rs index 86d092fcc50..b27f113a94e 100644 --- a/native/shuffle/src/comet_partitioning.rs +++ b/native/shuffle/src/comet_partitioning.rs @@ -47,17 +47,23 @@ pub enum RoundRobinStrategy { /// than a batch ordinal is what lets this strategy rely on the level Spark already publishes /// instead of an assumption nothing checks. /// - /// `start_partition` must be the Spark map partition id. It has to be distinct across mappers, - /// or every task starts at partition 0 and a task emitting fewer groups than there are output - /// partitions leaves the tail empty stage-wide; and it has to be a pure function of the map - /// partition, or a re-executed task does not reproduce its own placement. Spark seeds - /// `XORShiftRandom(partitionId)` for the same two reasons. + /// `start_partition` is the output partition this map task's first group goes to. It has to be + /// *decorrelated* across mappers, not merely distinct: a task walks `ceil(rows / group_rows)` + /// consecutive partitions from its start, so if consecutive tasks start on consecutive + /// partitions their runs all overlap and the partitions past `num_map_tasks + groups_per_task` + /// get nothing. It also has to be a pure function of the map partition, or a re-executed task + /// does not reproduce its own placement. Spark satisfies both by scrambling the map partition + /// id through `XORShiftRandom` (SPARK-21782), and the JVM computes this field the same way; see + /// `CometShuffleExchangeExec.positionalStartPartition`. /// - /// `group_rows` trades balance against copying. Imbalance between any two output partitions is - /// bounded by `group_rows` rows regardless of how the reader frames batches, so small groups - /// balance better; large groups produce fewer, longer runs to copy on flush, and a group as - /// large as the batch size lets a whole input batch pass through to one partition untouched. - /// [`Self::AUTO_GROUP_ROWS`] picks a value from the batch size and partition count. + /// `group_rows` trades balance against copying. Within one map task, imbalance between any two + /// output partitions is bounded by `group_rows` rows regardless of how the reader frames + /// batches, so small groups balance better; large groups produce fewer, longer runs to copy on + /// flush, and a group as large as the batch size lets a whole input batch pass through to one + /// partition untouched. That bound does not compose across map tasks: a reducer sees the sum + /// over all of them, which is only even when each task emits many more groups than there are + /// output partitions. [`Self::AUTO_GROUP_ROWS`] picks a value from the batch size and partition + /// count that keeps it so. RowGroups { start_partition: usize, group_rows: usize, @@ -77,9 +83,14 @@ impl RoundRobinStrategy { /// `group_rows` sentinel asking for a value derived from the batch size and partition count. pub const AUTO_GROUP_ROWS: usize = 0; - /// Smallest automatically chosen group. A multiple of 8 so that a run starts on a byte - /// boundary of a validity bitmap, which keeps the per-run copy a memcpy rather than a - /// bit-shift for every column. + /// Smallest automatically chosen group, which is a cap on how finely a batch is cut: with + /// `num_partitions` far larger than the batch size, `batch_size / num_partitions` rounds down + /// towards one row and the flush degenerates into the per-row gather that positional placement + /// exists to avoid. + /// + /// Not an alignment guarantee. A run only starts on a byte boundary of a validity bitmap when + /// the batch itself starts on a group boundary, and `row_seq` counts rows across batches, so + /// after a filter a batch starts at an arbitrary ordinal and every run in it is offset. const MIN_AUTO_GROUP_ROWS: usize = 64; /// Resolves [`Self::AUTO_GROUP_ROWS`] against the runtime batch size and partition count. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a73b81229ef..fb6430ca195 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -482,7 +482,8 @@ object CometConf extends ShimCometConf { "When true, Comet's native round-robin shuffle places rows by position rather than by " + "hashing their contents, the way Spark's own round robin does: the row at " + "task-global ordinal i goes to output partition " + - "(mapPartitionId + i / groupRows) % numPartitions. This skips a murmur3 pass over " + + "(start + i / groupRows) % numPartitions, where start is the map partition id " + + "scrambled the way Spark scrambles it. This skips a murmur3 pass over " + "every column of every row and replaces the per-row gather on flush with a bulk copy " + "per run, which is what dominates the shuffle write on wide nested schemas. It also " + "spreads duplicate rows evenly, where hashing sends them all to one partition. " + @@ -499,11 +500,15 @@ object CometConf extends ShimCometConf { conf("spark.comet.shuffle.native.partitioning.roundrobin.positional.groupRows") .category(CATEGORY_SHUFFLE) .doc( - "Rows per contiguous group under positional round robin. Imbalance between any two " + - "output partitions is bounded by this many rows however the reader frames its " + - "batches, so smaller groups balance better while larger groups produce fewer, longer " + - "runs to copy. When set to 0 (the default) Comet derives it from the batch size and " + - "the partition count. Only applies when " + + "Rows per contiguous group under positional round robin. Within one map task, imbalance " + + "between any two output partitions is bounded by this many rows however the reader " + + "frames its batches, so smaller groups balance better while larger groups produce " + + "fewer, longer runs to copy. That bound does not compose across map tasks: a reducer " + + "sees the sum over all of them, and the stage is only evenly balanced when each task " + + "emits many more groups than there are output partitions, so a group size approaching " + + "a task's whole input will leave some reducers empty. When set to 0 (the default) " + + "Comet derives it from the batch size and the partition count, which keeps a task " + + "wrapping around the output partitions roughly once per batch. Only applies when " + s"${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key} is true.") .intConf .checkValue( From 31335677007d0ff385176b72b7c46fb3730c543c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 23 Sep 2026 08:23:47 -0600 Subject: [PATCH 08/12] refactor: place positional runs straight into the run index positional_runs now returns an iterator of (partition, rows) that the repartitioner consumes directly into BufferedRuns, which drops the PositionalRun struct, the scratch vector holding a batch's runs, and the second pass over it. PartitionIndices::empty_like takes no argument, since its one caller passed the index's own partition count, and the gather timer is named interleave_time again, after the metric it feeds and that the Spark UI and the docs show. RunIterator's zero-copy branch compared a whole buffered batch against batch_size with >=, but insert_batch slices every batch to at most batch_size. It is == now, and the two paths' comments agree that every chunk but a partition's last is batch_size rows. The view-type fallback built RoundRobinStrategy::default(), silently dropping a configured maxHashColumns. RowGroups carries max_hash_columns now, and the fallback, factored into partitioning_for_schema so it can be tested, hashes with it. Nothing spilled under the run-indexed shape in any test. The spill metrics and heterogeneous-batching tests run it alongside the row-indexed shape, and positional_placement_survives_spilling requires placement under both spill triggers, the buffer limit and a pool that refuses to grow, to equal the unspilled placement exactly. Three unit tests the end-to-end placement tests subsumed are gone, and the start-partition test is folded into the walk test, which now asserts every group's partition, a wrapping start included. The rustdoc and proto comments shrink to what each item does and point at native_shuffle.md for the argument. --- native/core/src/execution/planner.rs | 12 +- native/proto/src/proto/partitioning.proto | 11 +- native/shuffle/benches/shuffle_writer.rs | 2 + native/shuffle/src/comet_partitioning.rs | 209 ++--------- .../src/partitioners/multi_partition.rs | 343 +++++++++++------- .../partitioned_batch_iterator.rs | 50 +-- native/shuffle/src/shuffle_writer.rs | 90 ++++- 7 files changed, 371 insertions(+), 346 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index c2029e8d7f4..5be836c5c36 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3652,21 +3652,21 @@ impl PhysicalPlanner { } PartitioningStruct::SinglePartition(_) => Ok(CometPartitioning::SinglePartition), PartitioningStruct::RoundRobinPartition(rr_partition) => { + // Treat negative max_hash_columns as 0 (no limit). + let max_hash_columns = rr_partition.max_hash_columns.max(0) as usize; let strategy = if rr_partition.positional { RoundRobinStrategy::RowGroups { // Computed per task on the JVM, where the Spark map partition id is in - // scope, and scrambled the way Spark scrambles it. See - // `CometShuffleExchangeExec.positionalStartPartition`. + // scope. See `CometShuffleExchangeExec.positionalStartPartition`. start_partition: rr_partition.positional_start_partition.max(0) as usize, // Negative or zero means "derive it from the batch size and partition // count", which the repartitioner does once it knows both. group_rows: rr_partition.positional_group_rows.max(0) as usize, + // Kept for the case where the schema rules positional placement out. + max_hash_columns, } } else { - // Treat negative max_hash_columns as 0 (no limit). - RoundRobinStrategy::HashAll { - max_hash_columns: rr_partition.max_hash_columns.max(0) as usize, - } + RoundRobinStrategy::HashAll { max_hash_columns } }; Ok(CometPartitioning::RoundRobin( rr_partition.num_partitions as usize, diff --git a/native/proto/src/proto/partitioning.proto b/native/proto/src/proto/partitioning.proto index daec02dac14..06923fd973c 100644 --- a/native/proto/src/proto/partitioning.proto +++ b/native/proto/src/proto/partitioning.proto @@ -60,16 +60,13 @@ message RoundRobinPartition { // When true, place rows by position rather than by hashing their contents: the row at // task-global ordinal i goes to // (positional_start_partition + i / positional_group_rows) % num_partitions. - // Only set when the planner has established that the map task replays rows in the same order, - // which is the condition Spark's own round robin relies on. + // Only set where the driver has established that the map task replays its rows in the same + // order; see "Round Robin Partitioning" in the contributor guide's native_shuffle.md. bool positional = 3; // Rows per contiguous group under positional placement. 0 means derive it from the batch size // and the partition count. int32 positional_group_rows = 4; - // Output partition this map task's first group goes to. Filled in per task by - // `CometNativeShuffleWriter.buildUnifiedPlan` as - // `XORShiftRandom(mapPartitionId).nextInt(num_partitions) + 1`, the same scrambled start Spark - // uses (SPARK-21782). A pure function of the map partition id, so a retried task reproduces its - // own placement. + // Output partition this map task's first group goes to, set per task from + // `CometShuffleExchangeExec.positionalStartPartition`. int32 positional_start_partition = 5; } diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index e8ea588b02d..bcc5918df2a 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -308,6 +308,7 @@ fn partitioning_benchmark(c: &mut Criterion) { RoundRobinStrategy::RowGroups { start_partition: 0, group_rows: RoundRobinStrategy::AUTO_GROUP_ROWS, + max_hash_columns: 0, }, ), ( @@ -315,6 +316,7 @@ fn partitioning_benchmark(c: &mut Criterion) { RoundRobinStrategy::RowGroups { start_partition: 0, group_rows: BATCH_SIZE, + max_hash_columns: 0, }, ), ]; diff --git a/native/shuffle/src/comet_partitioning.rs b/native/shuffle/src/comet_partitioning.rs index b27f113a94e..1c5378857b7 100644 --- a/native/shuffle/src/comet_partitioning.rs +++ b/native/shuffle/src/comet_partitioning.rs @@ -17,56 +17,34 @@ use arrow::row::{OwnedRow, RowConverter}; use datafusion::physical_expr::{LexOrdering, PhysicalExpr}; +use std::ops::Range; use std::sync::Arc; /// How [`CometPartitioning::RoundRobin`] decides which output partition a row belongs to. +/// +/// What each strategy trades, and why positional placement is only used where it is, is written +/// up once in the contributor guide's `native_shuffle.md`, under "Round Robin Partitioning". #[derive(Debug, Clone, PartialEq, Eq)] pub enum RoundRobinStrategy { /// Hash each row over its leading `max_hash_columns` columns (`0` meaning all of them) and - /// place it at `pmod(hash, num_partitions)`. - /// - /// Placement is a pure function of a row's contents, so a re-executed map task reproduces it - /// no matter what its input does. The price is a murmur3 pass per row that recurses into - /// every struct child, plus a per-row gather on flush because adjacent rows scatter across - /// every partition. It is also not really round robin: identical rows always hash to the same - /// partition, so low-cardinality input skews where Spark's round robin spreads evenly. + /// place it at `pmod(hash, num_partitions)`. A pure function of the row, so it reproduces + /// whatever order a re-executed map task sees its input in. HashAll { max_hash_columns: usize }, - /// Place rows positionally, in contiguous groups of `group_rows` rows, counting rows across - /// input batch boundaries: the row at task-global ordinal `i` goes to output partition - /// `(start_partition + i / group_rows) % num_partitions`. - /// - /// This is Spark's own round robin at a coarser granularity — Spark seeds a counter with - /// `XORShiftRandom(partitionId)` and bumps it per row, which is the `group_rows == 1` case — - /// and it inherits Spark's determinism condition exactly: placement is reproducible when the - /// upstream operator replays rows in the same *order*. It deliberately does not depend on how - /// those rows are framed into batches, because no Spark contract covers framing; - /// `DeterministicLevel::DETERMINATE` promises the same rows in the same order and says - /// nothing about how a downstream operator chunks them, so an operator that spills can reframe - /// under different memory pressure while still honouring it. Keying on a row ordinal rather - /// than a batch ordinal is what lets this strategy rely on the level Spark already publishes - /// instead of an assumption nothing checks. - /// - /// `start_partition` is the output partition this map task's first group goes to. It has to be - /// *decorrelated* across mappers, not merely distinct: a task walks `ceil(rows / group_rows)` - /// consecutive partitions from its start, so if consecutive tasks start on consecutive - /// partitions their runs all overlap and the partitions past `num_map_tasks + groups_per_task` - /// get nothing. It also has to be a pure function of the map partition, or a re-executed task - /// does not reproduce its own placement. Spark satisfies both by scrambling the map partition - /// id through `XORShiftRandom` (SPARK-21782), and the JVM computes this field the same way; see - /// `CometShuffleExchangeExec.positionalStartPartition`. - /// - /// `group_rows` trades balance against copying. Within one map task, imbalance between any two - /// output partitions is bounded by `group_rows` rows regardless of how the reader frames - /// batches, so small groups balance better; large groups produce fewer, longer runs to copy on - /// flush, and a group as large as the batch size lets a whole input batch pass through to one - /// partition untouched. That bound does not compose across map tasks: a reducer sees the sum - /// over all of them, which is only even when each task emits many more groups than there are - /// output partitions. [`Self::AUTO_GROUP_ROWS`] picks a value from the batch size and partition - /// count that keeps it so. + /// Place the row at task-global ordinal `i` at + /// `(start_partition + i / group_rows) % num_partitions`. The ordinal counts rows across + /// input batch boundaries, so placement does not depend on how the input was framed, but it + /// does depend on row order: only reproducible when the map task replays its rows in the same + /// order, which the planner establishes before choosing this. RowGroups { + /// Output partition the task's first group goes to, chosen per map task by + /// `CometShuffleExchangeExec.positionalStartPartition`. start_partition: usize, + /// Rows per group, or [`Self::AUTO_GROUP_ROWS`]. group_rows: usize, + /// What [`Self::HashAll`] hashes if `create_repartitioner` rules positional placement out + /// for the schema, so that the fallback honours the configured column cap. + max_hash_columns: usize, }, } @@ -95,10 +73,10 @@ impl RoundRobinStrategy { /// Resolves [`Self::AUTO_GROUP_ROWS`] against the runtime batch size and partition count. /// - /// One batch spread over `num_partitions` groups is the finest split that still gives every - /// output partition a run, so `batch_size / num_partitions` balances without fragmenting the - /// copy any further than it has to. An explicit request is taken as given, including one - /// larger than a batch, which sends several consecutive input batches to the same partition. + /// At `batch_size / num_partitions` a task wraps around the output partitions once per + /// batch, which is what keeps a whole stage balanced once each task has several batches. An + /// explicit request is taken as given, including one larger than a batch, which sends several + /// consecutive input batches to the same partition. pub fn resolve_group_rows( group_rows: usize, batch_size: usize, @@ -113,16 +91,9 @@ impl RoundRobinStrategy { } } -/// A contiguous span of rows within one input batch, bound for one output partition. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) struct PositionalRun { - pub partition: usize, - pub start: u32, - pub len: u32, -} - /// Splits the rows `[row_seq, row_seq + num_rows)` of a task's input into the runs that -/// [`RoundRobinStrategy::RowGroups`] placement produces, appending them to `out` in row order. +/// [`RoundRobinStrategy::RowGroups`] placement produces, in row order, each as an output +/// partition and the batch-relative rows bound for it. /// /// `row_seq` is the count of rows the task has already placed, which is what makes the split /// independent of where batch boundaries happen to fall: a group straddling two input batches @@ -134,26 +105,24 @@ pub(crate) fn positional_runs( start_partition: usize, group_rows: usize, num_partitions: usize, - out: &mut Vec, -) { - out.clear(); +) -> impl Iterator)> { let group_rows = group_rows.max(1) as u64; let num_partitions = num_partitions.max(1) as u64; let num_rows = num_rows as u64; let mut offset = 0u64; - while offset < num_rows { + std::iter::from_fn(move || { + if offset >= num_rows { + return None; + } let global = row_seq + offset; // Rows left in the group `global` falls into, so the first run of a batch picks up a // group that a previous batch left part-way through. - let remaining_in_group = group_rows - (global % group_rows); - let len = remaining_in_group.min(num_rows - offset); - out.push(PositionalRun { - partition: ((start_partition as u64 + global / group_rows) % num_partitions) as usize, - start: offset as u32, - len: len as u32, - }); + let len = (group_rows - global % group_rows).min(num_rows - offset); + let partition = (start_partition as u64 + global / group_rows) % num_partitions; + let rows = offset as u32..(offset + len) as u32; offset += len; - } + Some((partition as usize, rows)) + }) } /// Partitioning scheme for distributing rows across shuffle output partitions. @@ -206,124 +175,22 @@ mod tests { assert_eq!(result, expected); } - /// Collects the partition of every row in `[row_seq, row_seq + num_rows)` by expanding the - /// runs, which is the property the runs are a compressed encoding of. - fn placement( - row_seq: u64, - num_rows: usize, - group_rows: usize, - num_partitions: usize, - ) -> Vec { - let mut runs = vec![]; - positional_runs(row_seq, num_rows, 0, group_rows, num_partitions, &mut runs); - runs.iter() - .flat_map(|run| std::iter::repeat_n(run.partition, run.len as usize)) - .collect() - } - - #[test] - fn positional_runs_cover_every_row_once_in_order() { - let mut runs = vec![]; - positional_runs(0, 10, 0, 4, 3, &mut runs); - assert_eq!( - runs, - vec![ - PositionalRun { - partition: 0, - start: 0, - len: 4 - }, - PositionalRun { - partition: 1, - start: 4, - len: 4 - }, - PositionalRun { - partition: 2, - start: 8, - len: 2 - }, - ] - ); - } - - /// The point of counting rows rather than batches: however the reader frames the same rows, - /// each row lands on the same partition. - #[test] - fn positional_placement_is_independent_of_batch_framing() { - let group_rows = 7; - let num_partitions = 5; - let total = 100; - - let whole = placement(0, total, group_rows, num_partitions); - - for framing in [ - vec![100], - vec![1; 100], - vec![8; 12].into_iter().chain([4]).collect::>(), - vec![64, 36], - vec![7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 2], - ] { - assert_eq!(framing.iter().sum::(), total, "bad framing fixture"); - let mut row_seq = 0u64; - let mut refrained = vec![]; - for rows in framing.iter() { - refrained.extend(placement(row_seq, *rows, group_rows, num_partitions)); - row_seq += *rows as u64; - } - assert_eq!( - refrained, whole, - "framing {framing:?} placed rows differently" - ); - } - } - /// A group that a previous batch left part-way through is finished by the next batch, rather /// than restarting at a group boundary. #[test] fn positional_runs_resume_a_partial_group() { - let mut runs = vec![]; - positional_runs(2, 6, 0, 4, 3, &mut runs); - assert_eq!( - runs, - vec![ - // rows 2..4 finish group 0 - PositionalRun { - partition: 0, - start: 0, - len: 2 - }, - PositionalRun { - partition: 1, - start: 2, - len: 4 - }, - ] - ); - } - - #[test] - fn positional_runs_wrap_and_offset_by_start_partition() { - let mut runs = vec![]; - positional_runs(0, 6, 2, 2, 3, &mut runs); + // Rows 2..4 of the task finish group 0; rows 4..8 are group 1. assert_eq!( - runs.iter().map(|r| r.partition).collect::>(), - vec![2, 0, 1], - "start_partition offsets the sequence and it wraps at num_partitions" + positional_runs(2, 6, 0, 4, 3).collect::>(), + vec![(0, 0..2), (1, 2..6)] ); } #[test] fn positional_runs_group_larger_than_batch_yields_one_run() { - let mut runs = vec![]; - positional_runs(0, 100, 3, 8192, 200, &mut runs); assert_eq!( - runs, - vec![PositionalRun { - partition: 3, - start: 0, - len: 100 - }] + positional_runs(0, 100, 3, 8192, 200).collect::>(), + vec![(3, 0..100)] ); } diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index d9abfef7f2e..93c776dbb3f 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::comet_partitioning::{positional_runs, PositionalRun}; +use crate::comet_partitioning::positional_runs; use crate::metrics::ShufflePartitionerMetrics; use crate::partitioners::partitioned_batch_iterator::{ BufferedRun, PartitionIndices, PartitionedBatchesProducer, @@ -51,9 +51,6 @@ struct ScratchSpace { /// partition_starts[K + 1] are the start and end indices of partition K in partition_row_indices. /// The length of this array is 1 + the number of partitions. partition_starts: Vec, - /// The runs the current batch splits into under positional round robin. Only ever non-empty - /// for [`RoundRobinStrategy::RowGroups`], which uses none of the row-level buffers above. - positional_runs: Vec, } impl ScratchSpace { @@ -206,6 +203,7 @@ impl MultiPartitionShuffleRepartitioner { RoundRobinStrategy::RowGroups { start_partition, group_rows, + max_hash_columns, }, ) => CometPartitioning::RoundRobin( n, @@ -216,6 +214,7 @@ impl MultiPartitionShuffleRepartitioner { batch_size, num_output_partitions, ), + max_hash_columns, }, ), other => other, @@ -253,7 +252,6 @@ impl MultiPartitionShuffleRepartitioner { } else { vec![] }, - positional_runs: vec![], }; let reservation = MemoryConsumer::new(format!("ShuffleRepartitioner[{partition}]")) @@ -410,26 +408,18 @@ impl MultiPartitionShuffleRepartitioner { RoundRobinStrategy::RowGroups { start_partition, group_rows, + .. }, ) => { - let num_rows = input.num_rows(); - { - let mut timer = self.metrics.repart_time.timer(); - positional_runs( - self.row_seq, - num_rows, - *start_partition, - *group_rows, - *num_output_partitions, - &mut self.scratch.positional_runs, - ); - timer.stop(); - } - // Count rows, not batches: a group left part-way through by this batch is - // finished by the next one, so placement does not depend on where the reader - // put the batch boundary. See `RoundRobinStrategy::RowGroups`. - self.row_seq += num_rows as u64; - self.buffer_positional_batch_may_spill(input).await?; + let (num_partitions, start_partition, group_rows) = + (*num_output_partitions, *start_partition, *group_rows); + self.buffer_positional_batch_may_spill( + input, + start_partition, + group_rows, + num_partitions, + ) + .await?; } CometPartitioning::RoundRobin( num_output_partitions, @@ -548,12 +538,16 @@ impl MultiPartitionShuffleRepartitioner { self.reserve_and_may_spill(mem_growth) } - /// Buffers `input` against the runs positional round robin split it into, which - /// `partitioning_batch` left in `scratch.positional_runs`. + /// Buffers `input` as the runs positional round robin splits it into, then advances the row + /// ordinal past it. See `RoundRobinStrategy::RowGroups`. async fn buffer_positional_batch_may_spill( &mut self, input: RecordBatch, + start_partition: usize, + group_rows: usize, + num_partitions: usize, ) -> datafusion::common::Result<()> { + let num_rows = input.num_rows(); let (buffered_partition_idx, mut mem_growth) = self.buffer_input(input); let PartitionIndices::Runs(partition_runs) = &mut self.partition_indices else { @@ -561,17 +555,29 @@ impl MultiPartitionShuffleRepartitioner { "positional placement against a row-indexed repartitioner".to_string(), )); }; - for run in &self.scratch.positional_runs { - let indices = &mut partition_runs[run.partition]; - let before_size = indices.allocated_size(); - indices.push(BufferedRun { - batch: buffered_partition_idx, - start: run.start, - len: run.len, - }); - let after_size = indices.allocated_size(); - mem_growth += after_size.saturating_sub(before_size); + { + let mut timer = self.metrics.repart_time.timer(); + for (partition, rows) in positional_runs( + self.row_seq, + num_rows, + start_partition, + group_rows, + num_partitions, + ) { + let runs = &mut partition_runs[partition]; + let before_size = runs.allocated_size(); + runs.push(BufferedRun { + batch: buffered_partition_idx, + start: rows.start, + len: rows.end - rows.start, + }); + mem_growth += runs.allocated_size().saturating_sub(before_size); + } + timer.stop(); } + // Count rows, not batches: a group this batch leaves part-way through is finished by the + // next one, so placement does not depend on where the reader put the batch boundary. + self.row_seq += num_rows as u64; self.reserve_and_may_spill(mem_growth) } @@ -626,9 +632,8 @@ impl MultiPartitionShuffleRepartitioner { /// ShuffleRepartitioner to a new PartitionedBatches struct. The returned PartitionedBatches struct /// can be used to produce shuffled batches. fn partitioned_batches(&mut self) -> PartitionedBatchesProducer { - let num_output_partitions = self.partition_indices.num_partitions(); let buffered_batches = std::mem::take(&mut self.buffered_batches); - let empty = self.partition_indices.empty_like(num_output_partitions); + let empty = self.partition_indices.empty_like(); let indices = std::mem::replace(&mut self.partition_indices, empty); PartitionedBatchesProducer::new(buffered_batches, indices, self.batch_size) } @@ -806,37 +811,51 @@ mod tests { } async fn check_spill_metrics_count_input_buffers(batch: RecordBatch, input_bytes: usize) { - let runtime = Arc::new(RuntimeEnv::default()); - let metrics_set = ExecutionPlanMetricsSet::new(); - let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( - 0, - FailingPartitionWriter::default(), + // Row-indexed and run-indexed placement charge the reservation through the same + // accounting, so hold both to it. + for partitioning in [ CometPartitioning::RoundRobin( 2, RoundRobinStrategy::HashAll { max_hash_columns: 1, }, ), - ShufflePartitionerMetrics::new(&metrics_set, 0), - Arc::clone(&runtime), - 64, - false, - None, - ) - .unwrap(); - repartitioner.insert_batch(batch).await.unwrap(); - let index_bytes = repartitioner.partition_indices.allocated_size(); - let reserved_bytes = repartitioner.reservation.size(); - assert_eq!(repartitioner.spill_count(), 0); - assert_eq!(reserved_bytes, input_bytes + index_bytes); + CometPartitioning::RoundRobin( + 2, + RoundRobinStrategy::RowGroups { + start_partition: 0, + group_rows: 16, + max_hash_columns: 0, + }, + ), + ] { + let runtime = Arc::new(RuntimeEnv::default()); + let metrics_set = ExecutionPlanMetricsSet::new(); + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + FailingPartitionWriter::default(), + partitioning, + ShufflePartitionerMetrics::new(&metrics_set, 0), + Arc::clone(&runtime), + 64, + false, + None, + ) + .unwrap(); + repartitioner.insert_batch(batch.clone()).await.unwrap(); + let index_bytes = repartitioner.partition_indices.allocated_size(); + let reserved_bytes = repartitioner.reservation.size(); + assert_eq!(repartitioner.spill_count(), 0); + assert_eq!(reserved_bytes, input_bytes + index_bytes); - repartitioner.spill(0).unwrap(); + repartitioner.spill(0).unwrap(); - assert_eq!( - repartitioner.metrics.memory_spilled_bytes.value(), - reserved_bytes - ); - assert_eq!(runtime.memory_pool.reserved(), 0); + assert_eq!( + repartitioner.metrics.memory_spilled_bytes.value(), + reserved_bytes + ); + assert_eq!(runtime.memory_pool.reserved(), 0); + } } #[tokio::test] @@ -1107,37 +1126,52 @@ mod tests { // Shared child allocations contribute once per spill, independently of whether the // caller or insert_batch slices the input. They are not counted per output batch. - let mut spill_bytes = Vec::new(); - for input_batch_rows in [num_rows, batch_size] { - let runtime = Arc::new(RuntimeEnv::default()); - let metrics_set = ExecutionPlanMetricsSet::new(); - let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( - 0, - FailingPartitionWriter::default(), - CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2), - ShufflePartitionerMetrics::new(&metrics_set, 0), - Arc::clone(&runtime), - batch_size, - false, - Some(1), - ) - .unwrap(); - for start in (0..num_rows).step_by(input_batch_rows) { - repartitioner - .insert_batch(batch.slice(start, input_batch_rows)) - .await - .unwrap(); + for partitioning in [ + CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2), + // One-row groups alternate between the two partitions, so every spill hands each + // partition two runs to slice out of the batch and concatenate, across every type in + // the schema. + CometPartitioning::RoundRobin( + 2, + RoundRobinStrategy::RowGroups { + start_partition: 0, + group_rows: 1, + max_hash_columns: 0, + }, + ), + ] { + let mut spill_bytes = Vec::new(); + for input_batch_rows in [num_rows, batch_size] { + let runtime = Arc::new(RuntimeEnv::default()); + let metrics_set = ExecutionPlanMetricsSet::new(); + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + FailingPartitionWriter::default(), + partitioning.clone(), + ShufflePartitionerMetrics::new(&metrics_set, 0), + Arc::clone(&runtime), + batch_size, + false, + Some(1), + ) + .unwrap(); + for start in (0..num_rows).step_by(input_batch_rows) { + repartitioner + .insert_batch(batch.slice(start, input_batch_rows)) + .await + .unwrap(); + } + assert_eq!(repartitioner.spill_count(), num_rows / batch_size); + assert_eq!(repartitioner.reservation.size(), 0); + assert_eq!(runtime.memory_pool.reserved(), 0); + assert!(repartitioner.pinned_buffers.is_empty()); + assert!(repartitioner.buffered_batches.is_empty()); + assert_eq!(repartitioner.partition_indices.entry_count(), 0); + spill_bytes.push(repartitioner.metrics.memory_spilled_bytes.value()); } - assert_eq!(repartitioner.spill_count(), num_rows / batch_size); - assert_eq!(repartitioner.reservation.size(), 0); - assert_eq!(runtime.memory_pool.reserved(), 0); - assert!(repartitioner.pinned_buffers.is_empty()); - assert!(repartitioner.buffered_batches.is_empty()); - assert_eq!(repartitioner.partition_indices.entry_count(), 0); - spill_bytes.push(repartitioner.metrics.memory_spilled_bytes.value()); + assert!(spill_bytes[0] > 0); + assert_eq!(spill_bytes[0], spill_bytes[1]); } - assert!(spill_bytes[0] > 0); - assert_eq!(spill_bytes[0], spill_bytes[1]); } /// Collects every batch handed to it, per partition, so a whole write can be compared. @@ -1208,6 +1242,7 @@ mod tests { RoundRobinStrategy::RowGroups { start_partition, group_rows, + max_hash_columns: 0, }, ), ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0), @@ -1294,56 +1329,118 @@ mod tests { } } - /// `start_partition` offsets which partition a task begins on, so mappers do not all pile - /// their first group onto partition 0. - #[tokio::test] - async fn positional_placement_begins_at_the_start_partition() { - for start in 0..4usize { - let written = positional_placement(&[64], 4, 64, start, 256).await; - assert_eq!( - written.keys().copied().collect::>(), - vec![start], - "one group should land on partition {start} alone" - ); - } - } - - /// A task's groups walk *consecutive* partitions from its start, which is what makes the - /// stage-wide spread a function of how the starts are chosen rather than of the data. The JVM - /// picks the starts (`CometShuffleExchangeExec.positionalStartPartition`), but the reason it - /// has to scramble them lives here: with adjacent starts every task's run overlaps its - /// neighbours' and the tail of the partition space gets nothing. + /// A task's groups walk *consecutive* partitions from its start, wrapping at the partition + /// count, which is what makes the stage-wide spread a function of how the starts are chosen + /// rather than of the data. The JVM chooses them + /// (`CometShuffleExchangeExec.positionalStartPartition`); this pins why it has to scramble + /// them, since adjacent starts overlap and leave the tail of the partition space empty. #[tokio::test] async fn positional_placement_walks_consecutive_partitions_from_its_start() { let num_partitions = 32; let group_rows = 16; - let rows = 5 * group_rows; + let groups = 5; let stage = |starts: Vec| async move { let mut covered = std::collections::BTreeSet::new(); for start in starts { - let written = - positional_placement(&[rows], num_partitions, group_rows, start, 256).await; - assert_eq!( - written.keys().copied().collect::>(), - (0..5) - .map(|g| (start + g) % num_partitions) - .collect::>() - .into_iter() - .collect::>(), - "a task of five groups from {start} should touch five consecutive partitions" - ); - covered.extend(written.keys().copied()); + let written = positional_placement( + &[groups * group_rows], + num_partitions, + group_rows, + start, + 256, + ) + .await; + assert_eq!(written.len(), groups, "a task starting at {start}"); + for group in 0..groups { + let rows = (group * group_rows) as i64..((group + 1) * group_rows) as i64; + assert_eq!( + written[&((start + group) % num_partitions)], + rows.collect::>(), + "group {group} of a task starting at {start}" + ); + } + covered.extend(written.into_keys()); } covered.len() }; - // Four tasks of five groups each could reach twenty of the thirty-two partitions. - assert_eq!(stage(vec![0, 8, 16, 24]).await, 20); + // Spread-out starts reach twenty of the thirty-two partitions, the last task wrapping + // from 31 round to 0. + assert_eq!(stage(vec![4, 12, 20, 28]).await, 20); // Adjacent starts overlap instead, and twenty-four reducers get nothing. assert_eq!(stage(vec![0, 1, 2, 3]).await, 8); } + /// Spilling is where the run index is most likely to go wrong: the buffered batches drain, + /// the batches buffered after them are numbered from zero again, and the row ordinal has to + /// carry on regardless. Whether the spill is forced by the buffer limit or by the pool + /// refusing to grow, every row has to land exactly where it would have without spilling, + /// each partition still in input order. + #[tokio::test] + async fn positional_placement_survives_spilling() { + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + + // Ragged, so that groups and slices straddle batch and spill boundaries. + let framing = [300, 17, 683, 64, 1, 191]; + let (num_partitions, group_rows, start_partition, batch_size) = (8, 48, 3, 256); + let baseline = positional_placement( + &framing, + num_partitions, + group_rows, + start_partition, + batch_size, + ) + .await; + + for (trigger, memory_limit, max_buffer_bytes) in [ + ("buffer limit", None, Some(1)), + ("pool refusal", Some(4096), None), + ] { + let mut builder = RuntimeEnvBuilder::new(); + if let Some(limit) = memory_limit { + builder = builder.with_memory_limit(limit, 1.0); + } + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + CollectingPartitionWriter::default(), + CometPartitioning::RoundRobin( + num_partitions, + RoundRobinStrategy::RowGroups { + start_partition, + group_rows, + max_hash_columns: 0, + }, + ), + ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + Arc::new(builder.build().unwrap()), + batch_size, + false, + max_buffer_bytes, + ) + .unwrap(); + let mut next = 0i64; + for rows in framing { + repartitioner + .insert_batch(int64_batch(next..next + rows as i64)) + .await + .unwrap(); + next += rows as i64; + } + assert!( + repartitioner.spill_count() > 1, + "{trigger}: expected several spills, got {}", + repartitioner.spill_count() + ); + repartitioner.shuffle_write().unwrap(); + assert_eq!( + repartitioner.partition_writer().written, + baseline, + "{trigger}: spilling changed where rows landed" + ); + } + } + /// Imbalance stays within one group regardless of how the input was framed, which is what /// counting rows rather than batches buys over assigning whole batches. #[tokio::test] diff --git a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs index e65d8b2ed31..6151c336fe2 100644 --- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs +++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs @@ -47,7 +47,10 @@ pub(crate) enum PartitionIndices { } impl PartitionIndices { - pub(crate) fn empty_like(&self, num_partitions: usize) -> Self { + /// An empty index of the same shape and partition count, for after the buffered batches it + /// pointed into have drained. + pub(crate) fn empty_like(&self) -> Self { + let num_partitions = self.num_partitions(); match self { Self::Rows(_) => Self::Rows(vec![vec![]; num_partitions]), Self::Runs(_) => Self::Runs(vec![vec![]; num_partitions]), @@ -110,11 +113,14 @@ impl PartitionedBatchesProducer { self.buffered_batches.iter().collect() } + /// `interleave_time` is the shuffle writer's `interleave_time` metric, which times the gather + /// out of the partition index for either shape: `interleave_record_batch` over rows, or the + /// slicing and concatenation of runs. pub(super) fn produce<'a>( &'a self, refs: &'a [&'a RecordBatch], partition_id: usize, - copy_time: &'a Time, + interleave_time: &'a Time, ) -> PartitionedBatchIterator<'a> { // Partition indices index into `buffered_batches`; a refs slice built from a // different producer would silently interleave wrong rows. @@ -128,13 +134,13 @@ impl PartitionedBatchesProducer { &indices[partition_id], refs, self.batch_size, - copy_time, + interleave_time, )), PartitionIndices::Runs(runs) => PartitionedBatchIterator::Runs(RunIterator::new( &runs[partition_id], refs, self.batch_size, - copy_time, + interleave_time, )), } } @@ -170,7 +176,7 @@ pub(crate) struct RowIterator<'a> { /// (capacity at most `batch_size`) rather than re-materializing its whole index list. chunk_scratch: Vec<(usize, usize)>, pos: usize, - copy_time: &'a Time, + interleave_time: &'a Time, } impl<'a> RowIterator<'a> { @@ -178,7 +184,7 @@ impl<'a> RowIterator<'a> { indices: &'a [(u32, u32)], record_batches: &'a [&'a RecordBatch], batch_size: usize, - copy_time: &'a Time, + interleave_time: &'a Time, ) -> Self { if indices.is_empty() { // Avoid unnecessary allocations when the partition is empty @@ -188,7 +194,7 @@ impl<'a> RowIterator<'a> { indices: &[], chunk_scratch: vec![], pos: 0, - copy_time, + interleave_time, }; } Self { @@ -197,7 +203,7 @@ impl<'a> RowIterator<'a> { indices, chunk_scratch: Vec::with_capacity(batch_size.min(indices.len())), pos: 0, - copy_time, + interleave_time, } } } @@ -217,7 +223,7 @@ impl Iterator for RowIterator<'_> { .iter() .map(|(i_batch, i_row)| (*i_batch as usize, *i_row as usize)), ); - let mut timer = self.copy_time.timer(); + let mut timer = self.interleave_time.timer(); let result = interleave_record_batch(self.record_batches, &self.chunk_scratch); timer.stop(); match result { @@ -244,7 +250,7 @@ pub(crate) struct RunIterator<'a> { pos: usize, /// Rows already taken from `runs[pos]`, non-zero only when a run straddled a chunk boundary. consumed: u32, - copy_time: &'a Time, + interleave_time: &'a Time, } impl<'a> RunIterator<'a> { @@ -252,7 +258,7 @@ impl<'a> RunIterator<'a> { runs: &'a [BufferedRun], record_batches: &'a [&'a RecordBatch], batch_size: usize, - copy_time: &'a Time, + interleave_time: &'a Time, ) -> Self { Self { record_batches, @@ -261,7 +267,7 @@ impl<'a> RunIterator<'a> { chunk_scratch: Vec::with_capacity(runs.len().min(batch_size)), pos: 0, consumed: 0, - copy_time, + interleave_time, } } } @@ -273,17 +279,19 @@ impl Iterator for RunIterator<'_> { if self.pos >= self.runs.len() { return None; } - let mut timer = self.copy_time.timer(); + let mut timer = self.interleave_time.timer(); - // Zero-copy path: the next run is an entire buffered batch and already fills a chunk on - // its own, so hand the batch straight through. This is the case a group as large as the - // batch size is chosen to hit. + // Zero-copy path: the next run is an entire buffered batch of exactly `batch_size` rows, + // so it is already the chunk the copying path below would build, and can be handed + // straight through. This is the case a group as large as the batch size is chosen to + // hit. A buffered batch is never longer than `batch_size`, since `insert_batch` slices + // its input to that, so both paths emit the same chunk sizes. if self.consumed == 0 { let run = self.runs[self.pos]; let source = self.record_batches[run.batch as usize]; if run.start == 0 && run.len as usize == source.num_rows() - && run.len as usize >= self.batch_size + && run.len as usize == self.batch_size { self.pos += 1; timer.stop(); @@ -292,8 +300,8 @@ impl Iterator for RunIterator<'_> { } // Otherwise accumulate whole runs until the chunk is full, splitting the run that - // straddles the boundary. Chunks stay `batch_size` rows so that output block sizes do not - // depend on how long the runs happen to be. + // straddles the boundary, so that every chunk but a partition's last is `batch_size` rows + // however long the runs happen to be. self.chunk_scratch.clear(); let mut rows = 0usize; while self.pos < self.runs.len() && rows < self.batch_size { @@ -476,8 +484,8 @@ mod tests { assert_eq!(values, vec![vec![1, 2, 3, 200], vec![201, 103, 104]]); } - /// A run covering a whole buffered batch, long enough to be a chunk on its own, is handed - /// through without copying. Identity rather than equality, because the point is that the + /// A run covering a whole buffered batch of `batch_size` rows is already a full chunk, so it + /// is handed through without copying. Identity rather than equality, because the point is that the /// output shares the input's buffers. #[test] fn whole_batch_run_is_returned_without_copying() { diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 296747421cb..0e9085916de 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -31,7 +31,7 @@ use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExp use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::{apply_expression_roots, EmptyRecordBatchStream}; use datafusion::{ - arrow::datatypes::{DataType, SchemaRef}, + arrow::datatypes::{DataType, Schema, SchemaRef}, error::Result, execution::context::TaskContext, physical_plan::{ @@ -419,6 +419,31 @@ fn contains_view_type(data_type: &DataType) -> bool { } } +/// The partitioning to actually use on `schema`. The planner decides whether positional placement +/// is safe to retry; this decides whether it is worth doing on this schema, and falls back to +/// hashing, with the same column cap a hash round robin would have used, where it is not. See +/// [`contains_view_type`]. +fn partitioning_for_schema(partitioning: CometPartitioning, schema: &Schema) -> CometPartitioning { + match partitioning { + CometPartitioning::RoundRobin( + n, + RoundRobinStrategy::RowGroups { + max_hash_columns, .. + }, + ) if schema + .fields() + .iter() + .any(|f| contains_view_type(f.data_type())) => + { + log::debug!( + "schema contains a view type, falling back from positional to hash round robin" + ); + CometPartitioning::RoundRobin(n, RoundRobinStrategy::HashAll { max_hash_columns }) + } + other => other, + } +} + /// Constructs the existing schema-appropriate partitioner for either writer backend. #[allow(clippy::too_many_arguments)] fn create_repartitioner( @@ -432,23 +457,7 @@ fn create_repartitioner( max_buffer_bytes: Option, ) -> Result> { let partition_count = partitioning.partition_count(); - - // The planner decides whether positional placement is safe to retry; this decides whether it - // is worth doing on this schema. See `contains_view_type`. - let partitioning = match &partitioning { - CometPartitioning::RoundRobin(n, RoundRobinStrategy::RowGroups { .. }) - if schema - .fields() - .iter() - .any(|f| contains_view_type(f.data_type())) => - { - log::debug!( - "schema contains a view type, falling back from positional to hash round robin" - ); - CometPartitioning::RoundRobin(*n, RoundRobinStrategy::default()) - } - _ => partitioning, - }; + let partitioning = partitioning_for_schema(partitioning, &schema); if schema.fields().is_empty() { log::debug!( @@ -1830,4 +1839,49 @@ mod test { .into() ))); } + + /// Falling back from positional placement on a view-typed schema keeps the column cap that + /// `maxHashColumns` asked for, rather than silently hashing every column. + #[test] + fn view_type_fallback_keeps_the_hash_column_cap() { + let positional = || { + CometPartitioning::RoundRobin( + 8, + RoundRobinStrategy::RowGroups { + start_partition: 3, + group_rows: 64, + max_hash_columns: 4, + }, + ) + }; + let view_schema = Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("s", DataType::Utf8View, true), + ]); + assert!(matches!( + partitioning_for_schema(positional(), &view_schema), + CometPartitioning::RoundRobin( + 8, + RoundRobinStrategy::HashAll { + max_hash_columns: 4 + } + ) + )); + + let plain_schema = Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("s", DataType::Utf8, true), + ]); + assert!(matches!( + partitioning_for_schema(positional(), &plain_schema), + CometPartitioning::RoundRobin( + 8, + RoundRobinStrategy::RowGroups { + start_partition: 3, + group_rows: 64, + max_hash_columns: 4 + } + ) + )); + } } From aa690840d6230ce6005c97c8835c3a304730365a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 23 Sep 2026 08:24:04 -0600 Subject: [PATCH 09/12] fix: keep positional round robin off under Celeborn Positional placement would be the first path to hand the Celeborn push writer sliced batches, and an indeterminate stage's rollback has not been worked through for push shuffle, so positionalRoundRobinSpec now declines under the Celeborn shuffle manager. The planning suite checks it over a bare native scan, so the manager is the only thing ruling it out. A missing TaskContext now throws instead of starting the task at XORShiftRandom(0), which would put every task at the same partition: the correlation the scrambled start exists to prevent. usesPositionalRoundRobin moves from a public companion predicate to a package-private method on the exec, with shuffleType folded into the decision, since nothing consulted it for a columnar exchange. The two tests that exercised only the placement arithmetic are replaced by one that runs ten real map tasks of 5,000 rows into 200 reducers through the shuffle and requires none of them empty. With the start reverted to the bare map partition id it reports exactly the 112 the simulation predicted. The exchange, RDD and config docs now say what Spark's round robin requires. By default it sorts each map partition first, so a retry only has to produce the same rows; positional placement never sorts and needs them in the same order, which makes replaysRowsInOrder the whole safety argument. The RDD-level isOrderSensitive check is described as the defence in depth it is: a native scan leaf contributes no input RDD, so under today's allowlist it cannot fire. --- .../scala/org/apache/comet/CometConf.scala | 40 ++++---- .../shuffle/CometNativeShuffleInputRDD.scala | 30 +++--- .../shuffle/CometNativeShuffleWriter.scala | 12 ++- .../shuffle/CometShuffleDependency.scala | 2 +- .../shuffle/CometShuffleExchangeExec.scala | 99 ++++++++----------- .../CometCelebornShufflePlanningSuite.scala | 24 ++++- ...CometNativePositionalRoundRobinSuite.scala | 98 ++++++------------ 7 files changed, 132 insertions(+), 173 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index fb6430ca195..3df10978c68 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -480,17 +480,15 @@ object CometConf extends ShimCometConf { .category(CATEGORY_SHUFFLE) .doc( "When true, Comet's native round-robin shuffle places rows by position rather than by " + - "hashing their contents, the way Spark's own round robin does: the row at " + - "task-global ordinal i goes to output partition " + - "(start + i / groupRows) % numPartitions, where start is the map partition id " + - "scrambled the way Spark scrambles it. This skips a murmur3 pass over " + - "every column of every row and replaces the per-row gather on flush with a bulk copy " + - "per run, which is what dominates the shuffle write on wide nested schemas. It also " + - "spreads duplicate rows evenly, where hashing sends them all to one partition. " + - "Positional placement is only reproducible when the map task replays rows in the " + - "same order, so it is used only where Comet can establish that from the plan: a " + - "native scan under nothing but projections and filters. Any other plan silently " + - "keeps content-hash placement. " + + "hashing their contents, sending each map task's rows to the output partitions in " + + "turn, in contiguous groups. This skips a murmur3 pass over every column of every " + + "row and replaces the per-row gather on flush with a bulk copy per group, which is " + + "what dominates the shuffle write on wide nested schemas, and it spreads duplicate " + + "rows evenly where hashing sends them all to one partition. Placement then depends " + + "on the order a map task reads its rows in, so it is only used where Comet can " + + "establish from the plan that a retried task reads them in the same order: a native " + + "scan under nothing but projections and filters, and not with the Celeborn shuffle " + + "manager. Any other plan keeps content-hash placement. " + s"Has no effect unless ${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key} " + "is also true.") .booleanConf @@ -499,17 +497,15 @@ object CometConf extends ShimCometConf { val COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_GROUP_ROWS: ConfigEntry[Int] = conf("spark.comet.shuffle.native.partitioning.roundrobin.positional.groupRows") .category(CATEGORY_SHUFFLE) - .doc( - "Rows per contiguous group under positional round robin. Within one map task, imbalance " + - "between any two output partitions is bounded by this many rows however the reader " + - "frames its batches, so smaller groups balance better while larger groups produce " + - "fewer, longer runs to copy. That bound does not compose across map tasks: a reducer " + - "sees the sum over all of them, and the stage is only evenly balanced when each task " + - "emits many more groups than there are output partitions, so a group size approaching " + - "a task's whole input will leave some reducers empty. When set to 0 (the default) " + - "Comet derives it from the batch size and the partition count, which keeps a task " + - "wrapping around the output partitions roughly once per batch. Only applies when " + - s"${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key} is true.") + .doc("Rows per contiguous group under positional round robin. Smaller groups balance " + + "better and larger ones are cheaper to copy. Within one map task, output partitions " + + "differ by at most this many rows, but a reducer receives groups from every map task, " + + "so the stage is only evenly balanced when each task emits many more groups than " + + "there are output partitions; a group approaching a task's whole input skews the " + + "stage and can leave reducers empty. When set to 0 (the default) Comet derives it " + + "from the batch size and the partition count, which keeps each task wrapping around " + + "the output partitions once per batch. Only applies when " + + s"${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key} is true.") .intConf .checkValue( v => v >= 0, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala index c3ede338e38..8770496bff2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala @@ -36,7 +36,7 @@ import org.apache.comet.CometShuffleBlockIterator * * @param positionalRoundRobin * whether the writer fed by this RDD places rows by position rather than by content; see - * [[CometShuffleExchangeExec.usesPositionalRoundRobin]] and `getOutputDeterministicLevel`. + * [[CometShuffleExchangeExec.positionalRoundRobinSpec]] and `getOutputDeterministicLevel`. */ private[shuffle] class CometNativeShuffleInputRDD( sc: SparkContext, @@ -66,21 +66,21 @@ private[shuffle] class CometNativeShuffleInputRDD( positionalRoundRobin) /** - * Spark handles the retry hazard of positional round robin declaratively rather than - * per-operator: it wraps the repartition in a `MapPartitionsRDD` with `isOrderSensitive = true` - * (Comet's own JVM path does this in `prepareJVMShuffleDependency`), and that RDD reports - * `INDETERMINATE` whenever its parent is `UNORDERED`, which makes the DAGScheduler roll the - * whole stage back instead of re-running one task into a partially consumed output. The native - * path has no `MapPartitionsRDD` to carry the flag, so apply the same rule here. + * Spark's `isOrderSensitive` rule, applied to the RDD graph below the native plan. Spark only + * needs it for its own round robin with `spark.sql.execution.sortBeforeRepartition` off: it + * then wraps the repartition in a `MapPartitionsRDD` with `isOrderSensitive = true`, which + * reports `INDETERMINATE` over an `UNORDERED` parent, so the DAGScheduler rolls the whole stage + * back instead of re-running one task into a partially consumed output. In the default + * configuration Spark sorts each map partition first and the flag is `false`, on its path and + * on Comet's JVM path alike. Positional placement never sorts, so it takes the rule + * unconditionally, and with no `MapPartitionsRDD` on the native path to carry the flag it is + * applied here. * - * This covers everything below the RDD boundary; it cannot see the operators fused into the - * native plan above it, because the whole subtree collapses into this one RDD and `inputRDDs` - * are its leaves. `CometShuffleExchangeExec.replaysRowsInOrder` covers those. Both run, and - * positional placement needs both to agree. - * - * Letting the parent level discriminate is what keeps a plain scan on the cheap per-task retry - * path: a determinate parent stays determinate, while anything below another exchange is - * unordered, because reduce tasks see shuffle blocks in arrival order, and goes indeterminate. + * Defence in depth rather than a live gate. `CometShuffleExchangeExec.replaysRowsInOrder` + * admits only a native scan leaf, which contributes no input RDD, so under that allowlist the + * inherited level is always `DETERMINATE`. This starts to matter once the allowlist admits an + * input that crosses the RDD boundary: anything below another exchange is `UNORDERED`, because + * reduce tasks see shuffle blocks in arrival order, and goes indeterminate. */ override protected def getOutputDeterministicLevel: DeterministicLevel.Value = { val inheritedLevel = super.getOutputDeterministicLevel 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 ab86ec2b0c1..c97dc62744b 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 @@ -427,11 +427,15 @@ class CometNativeShuffleWriter[K, V]( partitioning.setPositional(true) partitioning.setPositionalGroupRows(positional.groupRows) // Per task, unlike the two above: which partition this mapper's first group goes to. - // See `CometShuffleExchangeExec.positionalStartPartition` for why it is scrambled. + // A real task always has a context. Guessing a partition id without one would start + // every task in the same place, the correlation the scrambled start exists to prevent. + val mapPartitionId = Option(context) + .map(_.partitionId()) + .getOrElse(throw new IllegalStateException( + "Positional round robin needs the map task's TaskContext")) partitioning.setPositionalStartPartition( - CometShuffleExchangeExec.positionalStartPartition( - Option(context).map(_.partitionId()).getOrElse(0), - effectivePartitionCount)) + CometShuffleExchangeExec + .positionalStartPartition(mapPartitionId, effectivePartitionCount)) } val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index 25208fc4089..7710cec6f6c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala @@ -49,7 +49,7 @@ case class NativeShuffleSpec( * and the group size are resolved once on the driver: the decision because it depends on the * shape of the plan fused into `childNativeOp`, which the executor never sees, and the group * size so that it cannot disagree with the decision. See - * `CometShuffleExchangeExec.usesPositionalRoundRobin`. + * `CometShuffleExchangeExec.positionalRoundRobinSpec`. */ positionalRoundRobin: Option[PositionalRoundRobin] = None) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index d2f36b71751..e12a747bccb 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -117,10 +117,17 @@ case class CometShuffleExchangeExec( /** * Positional round-robin decision, computed once so that the RDD's determinism level and the - * writer's placement cannot disagree. + * writer's placement cannot disagree. Only the native writer places positionally. */ @transient private lazy val positionalRoundRobin: Option[PositionalRoundRobin] = - CometShuffleExchangeExec.positionalRoundRobinSpec(outputPartitioning, child) + if (shuffleType == CometNativeShuffle) { + CometShuffleExchangeExec.positionalRoundRobinSpec(outputPartitioning, child) + } else { + None + } + + /** Whether this exchange's writer places rows positionally rather than by content. */ + private[shuffle] def usesPositionalRoundRobin: Boolean = positionalRoundRobin.isDefined @transient private lazy val nativeChildMetricNode: CometMetricNode = CometMetricNode.fromCometPlan(child) @@ -311,33 +318,17 @@ object CometShuffleExchangeExec } /** - * True when this exchange will run the native round-robin writer in its positional mode, where - * the row at task-global ordinal `i` goes to `(startPartition + i / groupRows) % numPartitions` - * rather than to `pmod(hash(row), numPartitions)`. - * - * Positional placement is reproducible exactly when the map task replays its rows in the same - * order, which is the same condition Spark's own round robin depends on. Spark answers it in - * two places and so does Comet: `replaysRowsInOrder` below establishes it for the operators - * fused into this native plan, which the RDD graph cannot see because the whole subtree - * collapses into one `CometNativeShuffleInputRDD`; and - * `CometNativeShuffleInputRDD.getOutputDeterministicLevel` establishes it for everything below - * that RDD, where the leaves are. Both have to hold. - * - * Must stay in step with `PhysicalPlanner::create_partitioning`, which turns the `positional` - * proto field into `RoundRobinStrategy::RowGroups`. + * Whether a round-robin exchange over `child` places rows positionally + * (`RoundRobinStrategy::RowGroups` in `PhysicalPlanner::create_partitioning`), and with what + * group size. Read once on the driver, so that the RDD's determinism level, the writer's + * placement and the group size cannot disagree: on an executor `CometConf.get()` resolves + * against a `SQLConf` rebuilt from the task's local properties, which returned the default + * group size rather than the session's. * - * The `numPartitions > 1` guard mirrors `isRoundRobin` in `prepareJVMShuffleDependency`. With a - * single output partition every row lands in the same place, so there is no placement to get - * wrong, and native routes that case to `SinglePartitionShufflePartitioner` regardless. - */ - def usesPositionalRoundRobin(outputPartitioning: Partitioning, child: SparkPlan): Boolean = - positionalRoundRobinSpec(outputPartitioning, child).isDefined - - /** - * [[usesPositionalRoundRobin]] together with the group size to use, both read on the driver so - * that they cannot disagree. `CometConf.get()` resolves against the thread-local `SQLConf`, - * which on an executor is rebuilt from the task's local properties; reading the group size - * there returned the default rather than the session value. + * Only where [[replaysRowsInOrder]] holds, and not under Celeborn, whose push path has not been + * shown to handle sliced batches or an indeterminate stage's rollback. The `numPartitions > 1` + * guard mirrors `isRoundRobin` in `prepareJVMShuffleDependency`: with one output partition + * there is no placement to get wrong. */ def positionalRoundRobinSpec( outputPartitioning: Partitioning, @@ -345,6 +336,7 @@ object CometShuffleExchangeExec val eligible = outputPartitioning.isInstanceOf[RoundRobinPartitioning] && outputPartitioning.numPartitions > 1 && CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.get() && + !isCometCelebornShuffleManagerEnabled(conf) && replaysRowsInOrder(child) if (eligible) { Some( @@ -356,21 +348,11 @@ object CometShuffleExchangeExec } /** - * Output partition that the first group of map task `mapPartitionId` goes to. - * - * The starts have to be decorrelated, not merely distinct. Each task walks `ceil(rows / - * groupRows)` consecutive partitions from its start, so if consecutive tasks start on - * consecutive partitions their runs all overlap and the partitions past `numMapTasks + - * groupsPerTask` get nothing: ten map tasks of 5,000 rows into 200 partitions at a group of 64 - * would leave 112 reducers empty. That is the correlation - * [[https://issues.apache.org/jira/browse/SPARK-21782 SPARK-21782]] fixed, and scrambling the - * map partition id through `XORShiftRandom` is how Spark fixes it, both in - * `ShuffleExchangeExec.getPartitionKeyExtractor` and in the JVM path below. - * - * Still a pure function of the map partition id, so a re-executed task reproduces its own - * placement. The `+ 1` matches Spark, which increments the counter before its first use, so at - * `groupRows == 1` this places rows exactly where Spark's round robin would for the same row - * order. + * Output partition that map task `mapPartitionId` places its first group in. This is Spark's + * own round-robin start, scrambled through `XORShiftRandom` because adjacent starts leave the + * tail of the partition space empty (SPARK-21782), and a pure function of the map partition so + * that a re-executed task reproduces its placement. The `+ 1` is Spark's pre-increment. See + * `native_shuffle.md` for why the starts must be decorrelated rather than merely distinct. */ def positionalStartPartition(mapPartitionId: Int, numPartitions: Int): Int = new XORShiftRandom(mapPartitionId).nextInt(math.max(numPartitions, 1)) + 1 @@ -378,24 +360,21 @@ object CometShuffleExchangeExec /** * Whether re-executing this subtree yields the same rows in the same order. * - * Deliberately a short allowlist rather than a denylist of known-bad operators, because the - * cost of being wrong is silent data loss rather than a failure: a re-executed map task that - * orders rows differently writes a different partitioning of them, and once any consumer has - * fetched the output that attempt replaces, the reduce side gets some rows twice and others not - * at all. Anything not named here keeps content-hash placement, which is safe to re-execute - * whatever its input does. - * - * A native scan replays its partition because the file splits are fixed on the driver when the - * RDD is built, and projections and filters are row-wise. Operators that spill are the - * interesting exclusion: an aggregate or a sort under memory pressure emits its output in an - * order that depends on how many times it spilled, which differs between attempts on different - * executors. Note that this says nothing about how rows are framed into batches: positional - * placement counts rows across batch boundaries precisely so that framing does not have to be - * part of this judgement. + * Positional placement is a function of row order, so this is the only thing standing between + * it and SPARK-23207, and it asks more than Spark's own round robin does: by default Spark + * sorts each map partition before assigning positions + * (`spark.sql.execution.sortBeforeRepartition`), so a retry only has to produce the same rows. + * `CometNativeShuffleInputRDD` mirrors Spark's `isOrderSensitive` rule for the RDD graph, but + * under this allowlist the only leaf is a native scan, which contributes no RDD input, so that + * check cannot fire. Widening this is what would make it live. * - * Other leaf scans (Iceberg, DSv2 batch, in-memory) plausibly qualify too, but each needs its - * own argument that a re-executed task reads the same rows in the same order, so they are left - * out until someone makes it. + * Deliberately a short allowlist rather than a denylist, because being wrong costs silent data + * loss rather than a failure: a re-executed task that orders its rows differently writes a + * different partitioning of them, and once any reducer has fetched from the attempt it + * replaces, some rows arrive twice and others not at all. A native scan replays its partition + * because its file splits are fixed on the driver, and projections and filters are row-wise. + * Anything that spills is out, since it emits rows in an order that depends on how often it + * spilled. Other leaf scans plausibly qualify, but each needs that argument made for it. */ private def replaysRowsInOrder(plan: SparkPlan): Boolean = plan match { case _: CometNativeScanExec => true diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala index b2c6c0f8a3d..9617ec5644d 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Attribut import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.plans.logical.LocalRelation import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning, RoundRobinPartitioning, SinglePartition} -import org.apache.spark.sql.comet.{CometCollectLimitExec, CometHashAggregateExec, CometLocalTableScanExec, CometNativeExec, CometScanWrapper, CometSortExec, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec} +import org.apache.spark.sql.comet.{CometCollectLimitExec, CometHashAggregateExec, CometLocalTableScanExec, CometNativeExec, CometNativeScanExec, CometScanWrapper, CometSortExec, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec} import org.apache.spark.sql.execution.{CollectLimitExec, ColumnarToRowTransition, LocalTableScanExec, SortExec, SparkPlan, TakeOrderedAndProjectExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec import org.apache.spark.sql.execution.aggregate.BaseAggregateExec @@ -191,6 +191,28 @@ class CometCelebornShufflePlanningSuite extends CometTestBase { } } + test("positional round robin stays off under Celeborn") { + // Positional placement would be the first path to hand the push writer sliced batches, and + // an indeterminate stage's rollback has not been shown to hold for push shuffle, so the + // planner keeps content-hash placement even where the plan would otherwise qualify. + withTempPath { dir => + spark.range(100).write.parquet(dir.getAbsolutePath) + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val plan = + spark.read.parquet(dir.getAbsolutePath).repartition(8).queryExecution.executedPlan + val exchanges = cometExchanges(plan) + assert(exchanges.size == 1, s"expected one Comet exchange:\n$plan") + // A bare native scan, so the shuffle manager is the only thing ruling positional out. + assert(exchanges.head.child.isInstanceOf[CometNativeScanExec], s"$plan") + assert(!exchanges.head.usesPositionalRoundRobin) + } + } + } + test("unsupported native partitioning falls back without trying Comet columnar shuffle") { withSQLConf( CometConf.COMET_SHUFFLE_MODE.key -> "native", diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala index 8fda1e23332..3350863f7e4 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala @@ -62,9 +62,7 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp assert( exchanges.size == 1, s"expected one native round-robin exchange in\n${df.queryExecution.executedPlan}") - CometShuffleExchangeExec.usesPositionalRoundRobin( - exchanges.head.outputPartitioning, - exchanges.head.child) + exchanges.head.usesPositionalRoundRobin } private def withParquetTable(rows: Int)(f: String => Unit): Unit = { @@ -123,10 +121,7 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp case e: CometShuffleExchangeExec => e } assert(exchanges.nonEmpty) - exchanges.foreach { e => - assert( - !CometShuffleExchangeExec.usesPositionalRoundRobin(e.outputPartitioning, e.child)) - } + exchanges.foreach(e => assert(!e.usesPositionalRoundRobin)) } } } @@ -193,71 +188,34 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp } } - /** - * Rows each reducer receives when `mapTasks` map tasks of `rowsPerTask` rows each place - * positionally into `partitions` output partitions, starting where `start` says. Counted a - * group at a time, which is exact because a group's rows all land together. - */ - private def stageSpread( - mapTasks: Int, - rowsPerTask: Int, - partitions: Int, - groupRows: Int, - start: Int => Int): Array[Int] = { - val counts = Array.fill(partitions)(0) - val groups = (rowsPerTask + groupRows - 1) / groupRows - for (mapPartitionId <- 0 until mapTasks; group <- 0 until groups) { - val rows = math.min(groupRows, rowsPerTask - group * groupRows) - counts(((start(mapPartitionId).toLong + group) % partitions).toInt) += rows - } - counts - } - test("map tasks start on decorrelated partitions, so the stage leaves no reducer empty") { // Each task walks ceil(rowsPerTask / groupRows) consecutive partitions from its start, so - // distinct starts are not enough: consecutive starts make every task's run overlap its - // neighbours' and the partitions past mapTasks + groupsPerTask never get a row. This is the - // correlation SPARK-21782 fixed, and `positionalStartPartition` fixes it the same way. - val (mapTasks, rowsPerTask, partitions, groupRows) = (10, 5000, 200, 64) - - val adjacent = stageSpread(mapTasks, rowsPerTask, partitions, groupRows, identity) - assert( - adjacent.count(_ == 0) == 112, - "the hazard this scrambling exists to avoid should still be reachable with adjacent starts") - - val scrambled = stageSpread( - mapTasks, - rowsPerTask, - partitions, - groupRows, - CometShuffleExchangeExec.positionalStartPartition(_, partitions)) - assert(scrambled.sum == mapTasks * rowsPerTask) - assert( - scrambled.count(_ == 0) == 0, - s"every reducer should get rows, got ${scrambled.count(_ == 0)} empty of $partitions") - } - - test("stage-wide balance needs many more groups per task than there are partitions") { - // Why the group size defaults to batchSize / numPartitions rather than to the batch size, - // even though a batch-sized group is far cheaper to flush: the per-task bound does not - // compose. A reducer sees the sum over every map task, and that sum only evens out once each - // task has wrapped the partition space several times. - val (mapTasks, rowsPerTask, partitions) = (50, 1000000, 200) - def spread(groupRows: Int): Double = { - val counts = stageSpread( - mapTasks, - rowsPerTask, - partitions, - groupRows, - CometShuffleExchangeExec.positionalStartPartition(_, partitions)) - assert(counts.sum == mapTasks.toLong * rowsPerTask) - counts.max.toDouble / counts.min + // starts that are merely distinct are not enough: consecutive starts would make every task's + // run overlap its neighbours', and at ten tasks of 5000 rows in groups of 64 the 112 + // partitions past the last task's run would get nothing. `positionalStartPartition` + // scrambles the start the way Spark does (SPARK-21782), which leaves none empty. + val (mapTasks, rowsPerTask, reducers) = (10, 5000, 200) + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, mapTasks.toLong * rowsPerTask, 1, mapTasks).write.parquet(path) + // An open cost as large as a split keeps each file in a map task of its own, and the files + // are far smaller than a split, so none of them is divided. + val splitBytes = (128L * 1024 * 1024).toString + withPositionalRoundRobin( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_GROUP_ROWS.key -> "64", + SQLConf.FILES_MAX_PARTITION_BYTES.key -> splitBytes, + SQLConf.FILES_OPEN_COST_IN_BYTES.key -> splitBytes) { + val input = spark.read.parquet(path) + assert(input.rdd.getNumPartitions == mapTasks) + val df = input.repartition(reducers) + assert(isPositional(df)) + + val sizes = partitionSizes(df) + assert(sizes.sum == mapTasks * rowsPerTask) + assert( + sizes.count(_ == 0) == 0, + s"every reducer should get rows, got ${sizes.count(_ == 0)} empty of $reducers") + } } - - // 64 rows per group is 15625 groups per task, so each task wraps 78 times and the sums - // converge. 8192 is 123 groups, fewer than there are partitions, so a task cannot even cover - // the space once and where the gaps fall is down to the starts. - assert(spread(64) < 1.01, s"expected an even stage at a small group, got ${spread(64)}") - assert(spread(8192) > 1.2, s"expected a batch-sized group to skew, got ${spread(8192)}") } } From 2669d2f00e6cbdf322bf05c1fefeb0df729deaa4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 23 Sep 2026 08:24:19 -0600 Subject: [PATCH 10/12] docs: correct the Spark round robin comparison and keep the argument in one place Spark's round robin sorts each map partition before assigning positions unless spark.sql.execution.sortBeforeRepartition is off, so by default a retry only has to produce the same rows. Positional placement never sorts and needs them in the same order. The guide said the two made the same assumption, which undersold how much rests on replaysRowsInOrder, and it said groupRows = 1 matches Spark, which holds only with the sort off. The rationale was written out in about seven places and had already needed a three-way fix. The round robin section of native_shuffle.md now holds it, opening with what Spark does in both modes, and the review skill's item is a pointer plus the checks a reviewer should make. The interleave_time metric description covers the run copy as well. --- .ai/skills/review-comet-shuffle-pr/SKILL.md | 28 ++- .../contributor-guide/native_shuffle.md | 165 ++++++++++-------- docs/source/user-guide/latest/metrics.md | 2 +- 3 files changed, 100 insertions(+), 95 deletions(-) diff --git a/.ai/skills/review-comet-shuffle-pr/SKILL.md b/.ai/skills/review-comet-shuffle-pr/SKILL.md index 86bce3bd0d7..9a0239ff186 100644 --- a/.ai/skills/review-comet-shuffle-pr/SKILL.md +++ b/.ai/skills/review-comet-shuffle-pr/SKILL.md @@ -95,24 +95,16 @@ Partitioning is where shuffle silently produces wrong answers rather than failin than Spark's `UnsafeRow`-sorted assignment would put them, which is why `spark.comet.shuffle.native.partitioning.roundrobin.enabled` defaults to `false`. Sorted output is identical either way. -- [ ] **Positional round robin is allowed, but only where retry reproducibility is established.** - `RoundRobinStrategy::RowGroups` places the row at task-global ordinal `i` at - `(startPartition + i / groupRows) % numPartitions`, which is Spark's own round robin at a - coarser granularity. It is not a bug, but it is only correct behind two gates, and a PR that - widens either one needs an argument. `CometShuffleExchangeExec.replaysRowsInOrder` is an - allowlist over the native subtree fused into the writer, which the RDD graph cannot see: a - native scan under nothing but projections and filters, with anything that spills staying out - because it reorders between attempts. `CometNativeShuffleInputRDD.getOutputDeterministicLevel` - applies Spark's `isOrderSensitive` rule to everything below that RDD. -- [ ] **Positional placement keys on a row ordinal, and the start is scrambled.** The counter is - over rows and carries across batch boundaries. A PR that keys on a batch ordinal instead is - relying on framing, which no Spark contract covers: `DETERMINATE` promises the same rows in - the same order and says nothing about how an operator chunks them. `startPartition` must be - decorrelated across mappers, not merely distinct — each task walks consecutive partitions - from its start, so adjacent starts overlap and leave the tail of the partition space empty - (SPARK-21782). And `groupRows` bounds imbalance within one map task only; a reducer sees the - sum over all of them, which only evens out when each task emits many more groups than there - are partitions. +- [ ] **Positional round robin is allowed, but its allowlist is the whole safety argument.** + `RoundRobinStrategy::RowGroups` is not a bug by construction. It is a function of row order + and never sorts, so it needs a retry to replay the same rows in the same order, which is + more than Spark's own round robin needs under its default `sortBeforeRepartition=true`. + `CometShuffleExchangeExec.replaysRowsInOrder` is what establishes that, and a PR that widens + it needs an argument that each operator it admits replays its rows in order; the RDD-level + determinism check is defence in depth and cannot fire under today's allowlist. Also check + that placement still counts rows rather than batches, and that the start still comes from + `positionalStartPartition`. The reasoning behind all three is in `native_shuffle.md` under + "Round Robin Partitioning". - [ ] **Range partitioning bounds come from the driver.** Spark's `RangePartitioner` samples and computes boundaries, they are serialized into the native plan, and native does a binary search over comparable-row-format keys. A change to the comparison or the row encoding must diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index 3dbaa320f16..acf91eaccd1 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -308,24 +308,34 @@ batch is written as a single block that may exceed the batch size. `CometPartitioning::RoundRobin` carries a `RoundRobinStrategy` that decides how rows reach output partitions. The default is `HashAll`; `RowGroups` is opt-in through `spark.comet.shuffle.native.partitioning.roundrobin.positional.enabled`, and is used only where the -planner can establish that the map task replays its rows in the same order. +planner can establish that a retried map task replays its rows in the same order. + +Both are judged against what Spark does. Spark's round robin keeps a per-task counter that starts +at `XORShiftRandom(partitionId).nextInt(numPartitions)` and is incremented before each row, so the +`k`th row a task sees goes to `(start + 1 + k) % numPartitions`. That is a function of row order, +which a retry does not have to reproduce +([SPARK-23207](https://issues.apache.org/jira/browse/SPARK-23207)), so by default +(`spark.sql.execution.sortBeforeRepartition=true`) Spark first sorts each map partition on the +binary `UnsafeRow` form. Placement then depends only on which rows the partition holds, and a +retry only has to produce the same rows, in any order. With the sort turned off, Spark instead +marks the repartition `isOrderSensitive`, which reports the stage `INDETERMINATE` over an +`UNORDERED` parent, and the `DAGScheduler` rolls the whole stage back rather than re-running one +task into a partially consumed output. + +Neither Comet strategy sorts, and neither places rows where Spark's default would: Arrow's layout +does not reproduce the `UnsafeRow` sort order, so unsorted output can land in different partitions +than Spark's. Sorted output is identical. That difference is why +`spark.comet.shuffle.native.partitioning.roundrobin.enabled` defaults to `false`. #### `HashAll`: hash-based assignment (default) 1. Computes a Murmur3 hash of columns (using seed 42) 2. Assigns partitions directly using the hash: `partition_id = hash % num_partitions` -This approach guarantees determinism across retries, which is critical for fault tolerance. -However, unlike true round robin which cycles through partitions row-by-row, hash-based -assignment only provides even distribution when the data has sufficient variation in the -hashed columns. Data with low cardinality or identical values may result in skewed partition -sizes: because placement is a pure function of a row's contents, a column of one repeated value -lands entirely on one reducer. - -Because Spark assigns round robin partitions by sorting rows on their binary `UnsafeRow` form, -which Arrow's layout does not reproduce, unsorted output can land in different partitions than -Spark's. Sorted output is identical. That difference is why -`spark.comet.shuffle.native.partitioning.roundrobin.enabled` defaults to `false`. +Placement is a pure function of each row, so like Spark's default it only needs a retry to produce +the same rows. However, unlike true round robin, hash-based assignment only provides even +distribution when the data has sufficient variation in the hashed columns: a column of one repeated +value lands entirely on one reducer. `spark.comet.shuffle.native.partitioning.roundrobin.maxHashColumns` caps how many leading columns are hashed. `0`, the default, hashes all of them. @@ -335,80 +345,83 @@ are hashed. `0`, the default, hashes all of them. Hashing every column of every row dominates the shuffle write on wide nested schemas, because `create_murmur3_hashes` recurses into every struct child per row and the resulting row-level scatter forces `interleave_record_batch` to walk every column and child again on flush. -`RowGroups` places rows the way Spark's own round robin does: the row at task-global ordinal `i` -goes to `(startPartition + i / groupRows) % numPartitions`. That removes the per-row hash, and it -replaces the per-row gather with a bulk copy per contiguous run, because adjacent rows now stay -together. It also spreads duplicate rows evenly, which `HashAll` cannot. +`RowGroups` places rows positionally instead: the row at task-global ordinal `i` goes to +`(startPartition + i / groupRows) % numPartitions`. That removes the per-row hash, and it replaces +the per-row gather with a bulk copy per contiguous run, because adjacent rows now stay together. +It also spreads duplicate rows evenly, which `HashAll` cannot. With `groupRows = 1` it is exactly +Spark's round robin with `sortBeforeRepartition=false`. The counter is over **rows**, not batches, and it carries across batch boundaries: a group that one -input batch leaves part-way through is finished by the next. That is deliberate. Spark's -determinism contract, `DeterministicLevel`, describes row _order_ and says nothing about how a -downstream operator frames rows into batches, so an operator that spills can reframe under -different memory pressure while still honouring `DETERMINATE`. Keying on a row ordinal means the -strategy depends only on the property Spark actually publishes. - -`start_partition` is the output partition a map task's first group goes to. It is computed per task -on the JVM, in `CometNativeShuffleWriter.buildUnifiedPlan` where the Spark map partition id is in -scope, and passed down in the proto. It has to be _decorrelated_ across mappers, not merely -distinct: a task walks `ceil(rows / groupRows)` consecutive partitions from its start, so if -consecutive tasks start on consecutive partitions their runs all overlap and the partitions past -`numMapTasks + groupsPerTask` get nothing — ten map tasks of 5,000 rows into 200 partitions at a -group of 64 would leave 112 reducers empty. It also has to be a pure function of the map partition, -or a re-executed task does not reproduce its own placement. Spark satisfies both by scrambling the -map partition id through `XORShiftRandom` -([SPARK-21782](https://issues.apache.org/jira/browse/SPARK-21782)), and -`CometShuffleExchangeExec.positionalStartPartition` does the same. Spark increments its counter -before the first row uses it, so the start is `nextInt(numPartitions) + 1`, which makes -`groupRows = 1` place rows exactly where Spark's round robin would for the same row order. - -`groupRows` trades balance against copying. Within one map task, imbalance between any two output -partitions is bounded by `groupRows` rows however the reader frames its batches, so small groups -balance better; large groups produce fewer, longer runs to copy, and a group as large as the batch -size lets a whole input batch pass through to one partition untouched. That bound does not compose -across map tasks — a reducer sees the sum over all of them, which is only even when each task emits -many more groups than there are output partitions. `0`, the default, derives it as -`clamp(batch_size / num_partitions, 64, batch_size)`, which keeps a task wrapping around the output -partitions roughly once per batch. The 64-row floor caps how finely a batch is cut: with far more -partitions than rows in a batch, `batch_size / num_partitions` rounds down towards one row and the -flush degenerates into the per-row gather positional placement exists to avoid. +input batch leaves part-way through is finished by the next. That is deliberate. +`DeterministicLevel` describes which rows a task produces and in what order, and says nothing about +how an operator frames them into batches, so an operator that spills can reframe under different +memory pressure while still honouring `DETERMINATE`. Keying on a row ordinal means placement +depends on order alone. + +`startPartition` is the output partition a map task's first group goes to, computed per task by +`CometShuffleExchangeExec.positionalStartPartition` in `CometNativeShuffleWriter.buildUnifiedPlan`, +where the map partition id is in scope, and passed down in the proto. It has to be _decorrelated_ +across mappers, not merely distinct: a task walks `ceil(rows / groupRows)` consecutive partitions +from its start, so if consecutive tasks started on consecutive partitions their runs would all +overlap and the partitions past `numMapTasks + groupsPerTask` would get nothing — ten map tasks of +5,000 rows into 200 partitions at a group of 64 would leave 112 reducers empty. It also has to be a +pure function of the map partition, or a re-executed task does not reproduce its own placement. +Spark's own start satisfies both, which is why it is scrambled +([SPARK-21782](https://issues.apache.org/jira/browse/SPARK-21782)), so Comet uses it unchanged, +including the `+ 1` for Spark's pre-increment. + +`groupRows` trades balance against copying. Within one map task, output partitions differ by at +most `groupRows` rows however the reader frames its batches, so small groups balance better; large +groups produce fewer, longer runs to copy, and a group as large as the batch size lets a whole +input batch pass through to one partition untouched. That bound does not compose across map +tasks: a reducer sees the sum over all of them, which is only even when each task emits many more +groups than there are output partitions. Over 50 map tasks of a million rows into 200 partitions, +a batch-sized group of 8192 leaves the largest reducer 1.57 times the smallest, where a group of +64 is within 0.3%. `0`, the default, derives the group as +`clamp(batch_size / num_partitions, 64, batch_size)`, which keeps each task wrapping around the +output partitions once per batch. The 64-row floor caps how finely a batch is cut, so that a +partition count far larger than the batch size cannot turn the flush back into a per-row gather. Internally, `MultiPartitionShuffleRepartitioner` records `(batch, start, len)` runs rather than one `(batch, row)` pair per row, so the index list charged against the spill reservation is smaller, and `RunIterator` builds each output chunk by slicing and concatenating runs. A run that -covers an entire buffered batch and already fills a chunk is passed through without copying. - -#### Retry safety under `RowGroups` - -Positional assignment is not a function of the rows, so it is reproducible only when the map task -replays the same rows in the same order. Re-executing one map task against differently ordered -input writes a different partitioning of the same rows, and once any consumer has fetched the -output that attempt replaces, the reduce side silently gets some rows twice and others not at all -([SPARK-23207](https://issues.apache.org/jira/browse/SPARK-23207)). Spark faces the same problem -with its own round robin. Comet establishes the condition in two places, both of which must hold: - -- **In the plan.** `CometShuffleExchangeExec.replaysRowsInOrder` walks the native subtree fused - into the writer, which the RDD graph cannot see because the whole subtree collapses into one - `CometNativeShuffleInputRDD`. It is a short allowlist, not a denylist: a native scan under - nothing but projections and filters. Operators that spill are the interesting exclusion, since - an aggregate or sort under memory pressure emits output in an order that depends on how many - times it spilled, which differs between attempts. Anything else keeps `HashAll`, which is safe - to re-execute whatever its input does. - -- **In the RDD graph.** Spark wraps a round-robin repartition in a `MapPartitionsRDD` with - `isOrderSensitive = true`, which reports `INDETERMINATE` whenever its parent is `UNORDERED`; the - `DAGScheduler` then rolls the whole stage back rather than re-running a single task. The native - path has no `MapPartitionsRDD` to carry the flag, so - `CometNativeShuffleInputRDD.getOutputDeterministicLevel` applies the same rule directly. A - determinate parent such as a plain scan stays determinate and keeps cheap per-task retry; - anything below another exchange is unordered, because reduce tasks see shuffle blocks in arrival - order, and goes indeterminate. - -Neither applies to `HashAll`, whose output is a pure function of the rows it sees. +covers an entire `batch_size` buffered batch is passed through without copying. One schema-level restriction is applied in `create_repartitioner`: positional placement is the only strategy that hands a sliced array to the IPC writer, and while the writer truncates a slice's buffers for every other type, for `Utf8View` and `BinaryView` it serializes every shared data -buffer in full. A schema containing a view type anywhere therefore falls back to `HashAll`. +buffer in full. A schema containing a view type anywhere therefore falls back to `HashAll`, with +the same `maxHashColumns`. + +#### Retry safety under `RowGroups` + +Positional placement asks more of a retry than Spark's default does. Spark sorts first, so it +needs the same rows; `RowGroups` does not sort, so it needs the same rows in the same order. +Re-executing one map task against differently ordered input writes a different partitioning of the +same rows, and once any consumer has fetched the output that attempt replaces, the reduce side +silently gets some rows twice and others not at all. Comet establishes the condition in two places: + +- **In the plan, which is the gate that matters.** `CometShuffleExchangeExec.replaysRowsInOrder` + walks the native subtree fused into the writer, which the RDD graph cannot see because the whole + subtree collapses into one `CometNativeShuffleInputRDD`. It is a short allowlist, not a + denylist: a native scan under nothing but projections and filters. A native scan replays its + partition because its file splits are fixed on the driver. Operators that spill are the + interesting exclusion, since an aggregate or sort under memory pressure emits output in an order + that depends on how many times it spilled, which differs between attempts. Anything else keeps + `HashAll`. Because nothing sorts behind it, this allowlist is the only thing standing between + `RowGroups` and SPARK-23207, and a change that widens it needs its own argument that the new + operator replays its rows in order. + +- **In the RDD graph, as defence in depth.** + `CometNativeShuffleInputRDD.getOutputDeterministicLevel` applies Spark's `isOrderSensitive` rule + to everything below that RDD: a determinate parent stays determinate, and anything below another + exchange is unordered, because reduce tasks see shuffle blocks in arrival order, and goes + indeterminate. Under today's allowlist it cannot fire, since the only leaf is a native scan, + which contributes no input RDD. It starts to matter once the allowlist admits an input that + crosses the RDD boundary. + +`RowGroups` is also not used with the Celeborn shuffle manager, whose push path has not been shown +to handle sliced batches or an indeterminate stage's rollback. ## Memory Management diff --git a/docs/source/user-guide/latest/metrics.md b/docs/source/user-guide/latest/metrics.md index 2c9f012f321..32973f0cfa3 100644 --- a/docs/source/user-guide/latest/metrics.md +++ b/docs/source/user-guide/latest/metrics.md @@ -99,7 +99,7 @@ Here is a guide to some of the native metrics. | ---------------------- | --------------------------------------------------------------------- | | `elapsed_compute` | Total time excluding any child operators. | | `repart_time` | Time to repartition batches. | -| `interleave_time` | Time to interleave partitioned batches before writing them. | +| `interleave_time` | Time to gather partitioned rows into output batches before writing. | | `ipc_time` | Time to encode batches in IPC format and compress using ZSTD. | | `mempool_time` | Time interacting with memory pool. | | `write_time` | Time spent writing bytes to disk. | From dfc5d2b02430d362bb3c4e3b02dc23fd4de4a839 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 23 Sep 2026 08:24:19 -0600 Subject: [PATCH 11/12] bench: build each partitioning fixture batch separately Both partitioning fixtures were one batch cloned eight times. A clone shares its buffers, so the reservation charged seven of the eight nothing, and the gather kept rereading one cache-resident batch. Each batch is now built from where the previous one left off. On the same machine, the cloned fixture reproduces the numbers first posted for this branch, while distinct batches roughly triple the nested gather wherever rows are copied: HashAll's place+gather goes from 37.7 ms to 106 ms and RowGroups(auto)'s from 8.5 ms to 20.4 ms. The zero-copy RowGroups(8192) arm, which copies nothing, is unchanged. --- native/shuffle/benches/shuffle_writer.rs | 59 ++++++++++++++++++------ 1 file changed, 46 insertions(+), 13 deletions(-) diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index bcc5918df2a..d3dd1ae8f64 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -237,6 +237,11 @@ fn create_batches(size: usize, count: usize) -> Vec { } fn create_batch(num_rows: usize, allow_nulls: bool) -> RecordBatch { + create_batch_from(0, num_rows, allow_nulls) +} + +/// [`create_batch`] with every value offset by `first_row`, so that successive batches differ. +fn create_batch_from(first_row: usize, num_rows: usize, allow_nulls: bool) -> RecordBatch { let schema = Arc::new(Schema::new(vec![ Field::new("c0", DataType::Int32, true), Field::new("c1", DataType::Utf8, true), @@ -249,7 +254,7 @@ fn create_batch(num_rows: usize, allow_nulls: bool) -> RecordBatch { let mut d = Decimal128Builder::new() .with_precision_and_scale(11, 2) .unwrap(); - for i in 0..num_rows { + for i in first_row..first_row + num_rows { a.append_value(i as i32); c.append_value(i as i32); d.append_value((i * 1000000) as i128); @@ -323,9 +328,16 @@ fn partitioning_benchmark(c: &mut Criterion) { // `plain` is the flat schema the end-to-end benches use. `nested` is the shape that motivates // positional placement: 40 struct columns over a three-field leaf, so 120 leaf arrays for a - // hash to recurse into and for a gather to walk. + // hash to recurse into and for a gather to walk. Every batch is built separately, rather than + // cloned, so that none of them share buffers: a clone would be charged nothing by the + // reservation, and the gather would keep rereading one cache-resident batch. let fixtures = [ - ("plain", create_batches(BATCH_SIZE, NUM_BATCHES)), + ( + "plain", + (0..NUM_BATCHES) + .map(|b| create_batch_from(b * BATCH_SIZE, BATCH_SIZE, true)) + .collect::>(), + ), ( "nested", nested_batches(BATCH_SIZE, NUM_BATCHES, 40, 2, Fill::PerRow), @@ -482,6 +494,7 @@ enum Fill { PerRow, } +/// `count` separately built batches, each filled from where the previous one left off. fn nested_batches( num_rows: usize, count: usize, @@ -489,15 +502,26 @@ fn nested_batches( depth: usize, fill: Fill, ) -> Vec { - let batch = nested_batch(num_rows, num_cols, depth, fill); - vec![batch; count] + (0..count) + .map(|b| nested_batch_from(b * num_rows, num_rows, num_cols, depth, fill)) + .collect() } fn nested_batch(num_rows: usize, num_cols: usize, depth: usize, fill: Fill) -> RecordBatch { + nested_batch_from(0, num_rows, num_cols, depth, fill) +} + +fn nested_batch_from( + first_row: usize, + num_rows: usize, + num_cols: usize, + depth: usize, + fill: Fill, +) -> 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 { - let array = nested_struct_array(num_rows, depth, fill); + let array = nested_struct_array(first_row, num_rows, depth, fill); fields.push(Field::new( format!("col{col}"), array.data_type().clone(), @@ -509,8 +533,14 @@ fn nested_batch(num_rows: usize, num_cols: usize, depth: usize, fill: Fill) -> R RecordBatch::try_new(schema, columns).unwrap() } -/// Builds a struct array with a multi-field leaf, wrapped in `depth` single-field structs. -fn nested_struct_array(num_rows: usize, depth: usize, fill: Fill) -> Arc { +/// Builds a struct array with a multi-field leaf, wrapped in `depth` single-field structs. Under +/// [`Fill::PerRow`] the leaf values run from `first_row`. +fn nested_struct_array( + first_row: usize, + num_rows: usize, + depth: usize, + fill: Fill, +) -> Arc { use arrow::array::{Float64Array, Int64Array, StringArray, StructArray}; let (ints, strings, floats): (Vec, Vec, Vec) = match fill { @@ -519,11 +549,14 @@ fn nested_struct_array(num_rows: usize, depth: usize, fill: Fill) -> Arc ( - (0..num_rows as i64).collect(), - (0..num_rows).map(|row| format!("value {row}")).collect(), - (0..num_rows).map(|row| row as f64 * 1.5).collect(), - ), + Fill::PerRow => { + let rows = first_row..first_row + num_rows; + ( + rows.clone().map(|row| row as i64).collect(), + rows.clone().map(|row| format!("value {row}")).collect(), + rows.map(|row| row as f64 * 1.5).collect(), + ) + } }; // Leaf: struct From 8578d6f10b46eb4428cf0e2e5af43a7a478a7487 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 23 Sep 2026 16:34:19 -0600 Subject: [PATCH 12/12] fix: gate positional round robin on deterministic expressions and freeze its group size Admit a native project or filter under positional placement only when its expressions are deterministic, since a nondeterministic UDF can reorder or re-filter rows on a retry while the RDD still reports DETERMINATE. Resolve the default group size on the driver and carry it in the shuffle dependency, so a map stage re-run after spark.comet.batchSize changes places rows exactly as the attempt it replaces. The native planner now rejects a non-positive group size instead of deriving one from the executor's batch size. Qualify the "wraps once per batch" claim for partition counts past the 64-row floor. --- .ai/skills/review-comet-shuffle-pr/SKILL.md | 9 ++- .../contributor-guide/native_shuffle.md | 22 ++++-- native/core/src/execution/planner.rs | 14 +++- native/proto/src/proto/partitioning.proto | 4 +- native/shuffle/benches/shuffle_writer.rs | 6 +- native/shuffle/src/comet_partitioning.rs | 52 +------------ .../src/partitioners/multi_partition.rs | 24 +----- .../scala/org/apache/comet/CometConf.scala | 11 ++- .../shuffle/CometShuffleDependency.scala | 6 +- .../shuffle/CometShuffleExchangeExec.scala | 58 +++++++++++--- ...CometNativePositionalRoundRobinSuite.scala | 78 ++++++++++++++++++- 11 files changed, 170 insertions(+), 114 deletions(-) diff --git a/.ai/skills/review-comet-shuffle-pr/SKILL.md b/.ai/skills/review-comet-shuffle-pr/SKILL.md index 9a0239ff186..4604f5da6b2 100644 --- a/.ai/skills/review-comet-shuffle-pr/SKILL.md +++ b/.ai/skills/review-comet-shuffle-pr/SKILL.md @@ -100,10 +100,11 @@ Partitioning is where shuffle silently produces wrong answers rather than failin and never sorts, so it needs a retry to replay the same rows in the same order, which is more than Spark's own round robin needs under its default `sortBeforeRepartition=true`. `CometShuffleExchangeExec.replaysRowsInOrder` is what establishes that, and a PR that widens - it needs an argument that each operator it admits replays its rows in order; the RDD-level - determinism check is defence in depth and cannot fire under today's allowlist. Also check - that placement still counts rows rather than batches, and that the start still comes from - `positionalStartPartition`. The reasoning behind all three is in `native_shuffle.md` under + it needs an argument that each operator it admits replays its rows in order, including + that the admitted expressions are deterministic; the RDD-level determinism check is defence + in depth and cannot fire under today's allowlist. Also check that placement still counts + rows rather than batches, that the start still comes from `positionalStartPartition`, and + that the group size is still resolved on the driver rather than from executor state. The reasoning behind all three is in `native_shuffle.md` under "Round Robin Partitioning". - [ ] **Range partitioning bounds come from the driver.** Spark's `RangePartitioner` samples and computes boundaries, they are serialized into the native plan, and native does a binary diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index acf91eaccd1..e1cabfd6c61 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -49,7 +49,6 @@ Native shuffle (`CometExchange`) is selected when all of the following condition columnar output. Row-based Spark operators require JVM shuffle. 3. **Supported partitioning type**: Native shuffle supports: - - `HashPartitioning` - `RangePartitioning` - `SinglePartition` @@ -152,7 +151,6 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet- 1. **Plan construction**: `CometNativeShuffleWriter` builds a protobuf operator tree with a `ShuffleWriter` operator at the root and `childNativeOp` as its child. `childNativeOp` takes one of two shapes: - - The child plan's `nativeOp` directly, when `CometShuffleExchangeExec`'s child is a `CometNativeExec` subtree. The upstream operators run inside the same `CometExecIterator` as the writer, with no JVM-to-native batch boundary between them. @@ -165,7 +163,6 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet- 2. **Native execution**: A single `CometExecIterator` per partition runs the unified plan. 3. **Partitioning**: `ShuffleWriterExec` receives batches and routes to the appropriate partitioner: - - `MultiPartitionShuffleRepartitioner`: For hash/range/round-robin partitioning - `SinglePartitionShufflePartitioner`: For single partition (simpler path) @@ -173,7 +170,6 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet- exceeds the threshold, partitions spill to temporary files. 5. **Encoding**: `ShuffleBlockWriter` encodes each partition's data as compressed Arrow IPC: - - Writes compression type header - Writes field count header - Writes compressed IPC stream @@ -203,7 +199,6 @@ read time. See [Direct Read](#direct-read-shufflescan) below for how the choice 1. `CometBlockStoreShuffleReader` fetches shuffle blocks via `ShuffleBlockFetcherIterator`. 2. For each block, `NativeBatchDecoderIterator`: - - Reads the 8-byte compressed length header - Reads the 8-byte field count header - Reads the compressed IPC data @@ -381,6 +376,15 @@ a batch-sized group of 8192 leaves the largest reducer 1.57 times the smallest, `clamp(batch_size / num_partitions, 64, batch_size)`, which keeps each task wrapping around the output partitions once per batch. The 64-row floor caps how finely a batch is cut, so that a partition count far larger than the batch size cannot turn the flush back into a per-row gather. +The floor also bounds that promise: past `batch_size / 64` output partitions a task covers only +`batch_size / 64` of them per batch, so small map tasks can still leave reducers empty. Ten map +tasks of one 8192-row batch each into 1,000 partitions emit 1,280 groups between them, and leave +311 reducers empty. + +The group size is resolved on the driver, in `CometShuffleExchangeExec.resolvePositionalGroupRows`, +and frozen with the shuffle dependency. Deriving it on the executor from the task's batch size +would let a map task re-executed after `spark.comet.batchSize` changed use a different group +size, and so a different placement, from the attempt it replaces. Internally, `MultiPartitionShuffleRepartitioner` records `(batch, start, len)` runs rather than one `(batch, row)` pair per row, so the index list charged against the spill reservation is @@ -404,8 +408,12 @@ silently gets some rows twice and others not at all. Comet establishes the condi - **In the plan, which is the gate that matters.** `CometShuffleExchangeExec.replaysRowsInOrder` walks the native subtree fused into the writer, which the RDD graph cannot see because the whole subtree collapses into one `CometNativeShuffleInputRDD`. It is a short allowlist, not a - denylist: a native scan under nothing but projections and filters. A native scan replays its - partition because its file splits are fixed on the driver. Operators that spill are the + denylist: a native scan under nothing but projections and filters whose expressions are all + deterministic. A native scan replays its partition because its file splits are fixed on the + driver. A nondeterministic expression, such as a UDF marked `asNondeterministic`, is out even + though it is evaluated row by row, because nothing bounds what it does between attempts: one + that reorders or re-filters rows on a retry moves them to different reducers while the RDD still + reports `DETERMINATE`. Operators that spill are the interesting exclusion, since an aggregate or sort under memory pressure emits output in an order that depends on how many times it spilled, which differs between attempts. Anything else keeps `HashAll`. Because nothing sorts behind it, this allowlist is the only thing standing between diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index c87c879acdd..df0982d9570 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3663,13 +3663,21 @@ impl PhysicalPlanner { // Treat negative max_hash_columns as 0 (no limit). let max_hash_columns = rr_partition.max_hash_columns.max(0) as usize; let strategy = if rr_partition.positional { + // Resolved on the driver and frozen with the shuffle dependency. Deriving it + // here from the executor's batch size would let a retried task use a + // different group size, and so a different placement, than the attempt it + // replaces. + if rr_partition.positional_group_rows <= 0 { + return Err(GeneralError(format!( + "Positional round robin needs a positive group size, got {}", + rr_partition.positional_group_rows + ))); + } RoundRobinStrategy::RowGroups { // Computed per task on the JVM, where the Spark map partition id is in // scope. See `CometShuffleExchangeExec.positionalStartPartition`. start_partition: rr_partition.positional_start_partition.max(0) as usize, - // Negative or zero means "derive it from the batch size and partition - // count", which the repartitioner does once it knows both. - group_rows: rr_partition.positional_group_rows.max(0) as usize, + group_rows: rr_partition.positional_group_rows as usize, // Kept for the case where the schema rules positional placement out. max_hash_columns, } diff --git a/native/proto/src/proto/partitioning.proto b/native/proto/src/proto/partitioning.proto index 06923fd973c..19b809f9a62 100644 --- a/native/proto/src/proto/partitioning.proto +++ b/native/proto/src/proto/partitioning.proto @@ -63,8 +63,8 @@ message RoundRobinPartition { // Only set where the driver has established that the map task replays its rows in the same // order; see "Round Robin Partitioning" in the contributor guide's native_shuffle.md. bool positional = 3; - // Rows per contiguous group under positional placement. 0 means derive it from the batch size - // and the partition count. + // Rows per contiguous group under positional placement. Always positive: the driver resolves + // the default once, so that a retried task cannot derive a different one. int32 positional_group_rows = 4; // Output partition this map task's first group goes to, set per task from // `CometShuffleExchangeExec.positionalStartPartition`. diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index d3dd1ae8f64..c099fc0413f 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -285,8 +285,8 @@ fn create_batch_from(first_row: usize, num_rows: usize, allow_nulls: bool) -> Re /// `interleave_record_batch` walk every column and child again, and a run can instead be sliced, /// or handed through untouched when it covers a whole buffered batch. /// -/// 8192 rows per batch into 50 output partitions, so [`RoundRobinStrategy::AUTO_GROUP_ROWS`] -/// resolves to 163. `RowGroups(8192)` is the opposite extreme, one whole input batch per group, +/// 8192 rows per batch into 50 output partitions, so `RowGroups(auto)` uses the 163-row group the +/// driver derives by default (`CometShuffleExchangeExec.resolvePositionalGroupRows`). `RowGroups(8192)` is the opposite extreme, one whole input batch per group, /// where every run covers a buffered batch end to end and the gather copies nothing at all. fn partitioning_benchmark(c: &mut Criterion) { const BATCH_SIZE: usize = 8192; @@ -312,7 +312,7 @@ fn partitioning_benchmark(c: &mut Criterion) { "RowGroups(auto)", RoundRobinStrategy::RowGroups { start_partition: 0, - group_rows: RoundRobinStrategy::AUTO_GROUP_ROWS, + group_rows: BATCH_SIZE / NUM_PARTITIONS, max_hash_columns: 0, }, ), diff --git a/native/shuffle/src/comet_partitioning.rs b/native/shuffle/src/comet_partitioning.rs index 1c5378857b7..5b15f2e49d6 100644 --- a/native/shuffle/src/comet_partitioning.rs +++ b/native/shuffle/src/comet_partitioning.rs @@ -40,7 +40,8 @@ pub enum RoundRobinStrategy { /// Output partition the task's first group goes to, chosen per map task by /// `CometShuffleExchangeExec.positionalStartPartition`. start_partition: usize, - /// Rows per group, or [`Self::AUTO_GROUP_ROWS`]. + /// Rows per group. Resolved on the driver and frozen with the shuffle dependency, so that a + /// re-executed task uses the same group size whatever batch size its executor runs with. group_rows: usize, /// What [`Self::HashAll`] hashes if `create_repartitioner` rules positional placement out /// for the schema, so that the fallback honours the configured column cap. @@ -57,40 +58,6 @@ impl Default for RoundRobinStrategy { } } -impl RoundRobinStrategy { - /// `group_rows` sentinel asking for a value derived from the batch size and partition count. - pub const AUTO_GROUP_ROWS: usize = 0; - - /// Smallest automatically chosen group, which is a cap on how finely a batch is cut: with - /// `num_partitions` far larger than the batch size, `batch_size / num_partitions` rounds down - /// towards one row and the flush degenerates into the per-row gather that positional placement - /// exists to avoid. - /// - /// Not an alignment guarantee. A run only starts on a byte boundary of a validity bitmap when - /// the batch itself starts on a group boundary, and `row_seq` counts rows across batches, so - /// after a filter a batch starts at an arbitrary ordinal and every run in it is offset. - const MIN_AUTO_GROUP_ROWS: usize = 64; - - /// Resolves [`Self::AUTO_GROUP_ROWS`] against the runtime batch size and partition count. - /// - /// At `batch_size / num_partitions` a task wraps around the output partitions once per - /// batch, which is what keeps a whole stage balanced once each task has several batches. An - /// explicit request is taken as given, including one larger than a batch, which sends several - /// consecutive input batches to the same partition. - pub fn resolve_group_rows( - group_rows: usize, - batch_size: usize, - num_partitions: usize, - ) -> usize { - let batch_size = batch_size.max(1); - if group_rows != Self::AUTO_GROUP_ROWS { - return group_rows; - } - (batch_size / num_partitions.max(1)) - .clamp(Self::MIN_AUTO_GROUP_ROWS.min(batch_size), batch_size) - } -} - /// Splits the rows `[row_seq, row_seq + num_rows)` of a task's input into the runs that /// [`RoundRobinStrategy::RowGroups`] placement produces, in row order, each as an output /// partition and the batch-relative rows bound for it. @@ -193,19 +160,4 @@ mod tests { vec![(3, 0..100)] ); } - - #[test] - fn resolve_group_rows_auto_splits_a_batch_across_partitions() { - use RoundRobinStrategy as S; - // One batch spread over the output partitions, floored at the 64-row minimum. - assert_eq!(S::resolve_group_rows(S::AUTO_GROUP_ROWS, 8192, 16), 512); - assert_eq!(S::resolve_group_rows(S::AUTO_GROUP_ROWS, 8192, 200), 64); - assert_eq!(S::resolve_group_rows(S::AUTO_GROUP_ROWS, 8192, 10_000), 64); - // A batch smaller than the minimum group still resolves to something usable. - assert_eq!(S::resolve_group_rows(S::AUTO_GROUP_ROWS, 32, 200), 32); - // An explicit request is taken as given. A group longer than a batch is meaningful: it - // sends several consecutive input batches to the same output partition. - assert_eq!(S::resolve_group_rows(1, 8192, 200), 1); - assert_eq!(S::resolve_group_rows(100_000, 8192, 200), 100_000); - } } diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 93c776dbb3f..05b14b53243 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -196,29 +196,7 @@ impl MultiPartitionShuffleRepartitioner { // Positional round robin is the one strategy that never looks at a row's contents. It // needs none of the row-level scratch (~64 KB a task), and it records contiguous runs - // rather than individual rows. Its group size is resolved once, here. - let partitioning = match partitioning { - CometPartitioning::RoundRobin( - n, - RoundRobinStrategy::RowGroups { - start_partition, - group_rows, - max_hash_columns, - }, - ) => CometPartitioning::RoundRobin( - n, - RoundRobinStrategy::RowGroups { - start_partition, - group_rows: RoundRobinStrategy::resolve_group_rows( - group_rows, - batch_size, - num_output_partitions, - ), - max_hash_columns, - }, - ), - other => other, - }; + // rather than individual rows. let places_rows_individually = !matches!( partitioning, CometPartitioning::RoundRobin(_, RoundRobinStrategy::RowGroups { .. }) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 3df10978c68..4766a30b373 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -487,8 +487,8 @@ object CometConf extends ShimCometConf { "rows evenly where hashing sends them all to one partition. Placement then depends " + "on the order a map task reads its rows in, so it is only used where Comet can " + "establish from the plan that a retried task reads them in the same order: a native " + - "scan under nothing but projections and filters, and not with the Celeborn shuffle " + - "manager. Any other plan keeps content-hash placement. " + + "scan under nothing but deterministic projections and filters, and not with the " + + "Celeborn shuffle manager. Any other plan keeps content-hash placement. " + s"Has no effect unless ${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key} " + "is also true.") .booleanConf @@ -503,8 +503,11 @@ object CometConf extends ShimCometConf { "so the stage is only evenly balanced when each task emits many more groups than " + "there are output partitions; a group approaching a task's whole input skews the " + "stage and can leave reducers empty. When set to 0 (the default) Comet derives it " + - "from the batch size and the partition count, which keeps each task wrapping around " + - "the output partitions once per batch. Only applies when " + + "from the batch size and the partition count when the query is planned, which keeps " + + "each task wrapping around the output partitions about once per batch. The derived " + + "group is at least 64 rows, so with more than a sixty-fourth of the batch size in " + + "output partitions a task needs several batches to wrap once, and map tasks of only a " + + "few batches can still leave reducers empty. Only applies when " + s"${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED.key} is true.") .intConf .checkValue( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index 7710cec6f6c..be91caa0ec9 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala @@ -48,7 +48,7 @@ case class NativeShuffleSpec( * Set when round-robin placement is positional rather than content-hashed. Both the decision * and the group size are resolved once on the driver: the decision because it depends on the * shape of the plan fused into `childNativeOp`, which the executor never sees, and the group - * size so that it cannot disagree with the decision. See + * size so that a re-executed map task places rows exactly as the attempt it replaces. See * `CometShuffleExchangeExec.positionalRoundRobinSpec`. */ positionalRoundRobin: Option[PositionalRoundRobin] = None) @@ -57,8 +57,8 @@ case class NativeShuffleSpec( * Parameters for positional round-robin placement, resolved on the driver. * * @param groupRows - * rows per contiguous group, or 0 to let the native side derive it from the batch size and the - * partition count. + * rows per contiguous group, always positive. A configured `0` is resolved here rather than on + * the executor, whose batch size need not match the one the stage was first run with. */ case class PositionalRoundRobin(groupRows: Int) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index e12a747bccb..c4fbf094cc2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -119,7 +119,7 @@ case class CometShuffleExchangeExec( * Positional round-robin decision, computed once so that the RDD's determinism level and the * writer's placement cannot disagree. Only the native writer places positionally. */ - @transient private lazy val positionalRoundRobin: Option[PositionalRoundRobin] = + @transient private[shuffle] lazy val positionalRoundRobin: Option[PositionalRoundRobin] = if (shuffleType == CometNativeShuffle) { CometShuffleExchangeExec.positionalRoundRobinSpec(outputPartitioning, child) } else { @@ -320,10 +320,12 @@ object CometShuffleExchangeExec /** * Whether a round-robin exchange over `child` places rows positionally * (`RoundRobinStrategy::RowGroups` in `PhysicalPlanner::create_partitioning`), and with what - * group size. Read once on the driver, so that the RDD's determinism level, the writer's - * placement and the group size cannot disagree: on an executor `CometConf.get()` resolves - * against a `SQLConf` rebuilt from the task's local properties, which returned the default - * group size rather than the session's. + * group size. Read once on the driver and frozen with the shuffle dependency, so that the RDD's + * determinism level, the writer's placement and the group size cannot disagree, and so that a + * map task re-executed after the session's batch size changed still uses the group size, and so + * the placement, of the attempt it replaces. On an executor `CometConf.get()` resolves against + * a `SQLConf` rebuilt from the task's local properties, which returned the default group size + * rather than the session's. * * Only where [[replaysRowsInOrder]] holds, and not under Celeborn, whose push path has not been * shown to handle sliced batches or an indeterminate stage's rollback. The `numPartitions > 1` @@ -341,12 +343,42 @@ object CometShuffleExchangeExec if (eligible) { Some( PositionalRoundRobin( - CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_GROUP_ROWS.get())) + resolvePositionalGroupRows( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_GROUP_ROWS.get(), + CometConf.COMET_BATCH_SIZE.get(), + outputPartitioning.numPartitions))) } else { None } } + /** + * Smallest derived group, which caps how finely a batch is cut: with far more output partitions + * than `batchSize / 64`, `batchSize / numPartitions` would round down towards one row and turn + * the flush back into the per-row gather positional placement exists to avoid. Not an alignment + * guarantee: after a filter a batch starts at an arbitrary row ordinal, so its runs start off a + * byte boundary whatever the group size. + */ + private val MinDerivedGroupRows = 64 + + /** + * The group size positional placement uses. An explicit `configured` value is taken as given; + * `0` derives `batchSize / numPartitions`, floored at [[MinDerivedGroupRows]] and capped at a + * batch, so that each task wraps around the output partitions about once per batch. + */ + private[shuffle] def resolvePositionalGroupRows( + configured: Int, + batchSize: Int, + numPartitions: Int): Int = { + if (configured > 0) { + configured + } else { + val batch = math.max(batchSize, 1) + val derived = batch / math.max(numPartitions, 1) + math.min(math.max(derived, math.min(MinDerivedGroupRows, batch)), batch) + } + } + /** * Output partition that map task `mapPartitionId` places its first group in. This is Spark's * own round-robin start, scrambled through `XORShiftRandom` because adjacent starts leave the @@ -372,14 +404,18 @@ object CometShuffleExchangeExec * loss rather than a failure: a re-executed task that orders its rows differently writes a * different partitioning of them, and once any reducer has fetched from the attempt it * replaces, some rows arrive twice and others not at all. A native scan replays its partition - * because its file splits are fixed on the driver, and projections and filters are row-wise. - * Anything that spills is out, since it emits rows in an order that depends on how often it - * spilled. Other leaf scans plausibly qualify, but each needs that argument made for it. + * because its file splits are fixed on the driver, and deterministic projections and filters + * are row-wise. A nondeterministic expression is out even though it is evaluated per row, since + * nothing bounds what it does between attempts: a nondeterministic UDF can drop, keep or + * reorder rows differently on a retry. Anything that spills is out, since it emits rows in an + * order that depends on how often it spilled. Other leaf scans plausibly qualify, but each + * needs that argument made for it. */ private def replaysRowsInOrder(plan: SparkPlan): Boolean = plan match { case _: CometNativeScanExec => true - case p: CometProjectExec => replaysRowsInOrder(p.child) - case f: CometFilterExec => replaysRowsInOrder(f.child) + case p: CometProjectExec => + p.projectList.forall(_.deterministic) && replaysRowsInOrder(p.child) + case f: CometFilterExec => f.condition.deterministic && replaysRowsInOrder(f.child) case _ => false } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala index 3350863f7e4..77f9cedffbc 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala @@ -19,10 +19,12 @@ package org.apache.spark.sql.comet.execution.shuffle +import org.apache.spark.{MapOutputTrackerMaster, SparkEnv} import org.apache.spark.sql.{CometTestBase, DataFrame} import org.apache.spark.sql.catalyst.plans.physical.RoundRobinPartitioning +import org.apache.spark.sql.execution.SQLExecution import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.functions.{col, lit} +import org.apache.spark.sql.functions.{col, lit, rand, udf} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -51,8 +53,8 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp // into the executed plan. SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") ++ extra: _*)(f) - /** Whether the round-robin exchange in `df`'s executed plan chose positional placement. */ - private def isPositional(df: DataFrame): Boolean = { + /** The one native round-robin exchange in `df`'s executed plan. */ + private def roundRobinExchange(df: DataFrame): CometShuffleExchangeExec = { val exchanges = collect(df.queryExecution.executedPlan) { case e: CometShuffleExchangeExec if e.shuffleType == CometNativeShuffle && @@ -62,9 +64,13 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp assert( exchanges.size == 1, s"expected one native round-robin exchange in\n${df.queryExecution.executedPlan}") - exchanges.head.usesPositionalRoundRobin + exchanges.head } + /** Whether the round-robin exchange in `df`'s executed plan chose positional placement. */ + private def isPositional(df: DataFrame): Boolean = + roundRobinExchange(df).usesPositionalRoundRobin + private def withParquetTable(rows: Int)(f: String => Unit): Unit = { withTempPath { dir => val path = dir.getAbsolutePath @@ -104,6 +110,22 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp } } + test("a nondeterministic projection or filter keeps content-hash placement") { + // Evaluated row by row, but nothing bounds what a nondeterministic expression does between + // attempts: one that reorders or re-filters rows on a retry sends them to different reducers. + val sameId = udf((id: Long) => id) + withPositionalRoundRobin() { + withParquetTable(1000) { t => + val df = spark.table(t) + assert(!isPositional( + df.select(sameId.asNondeterministic()(col("id")).as("id")).repartition(numPartitions))) + assert(!isPositional(df.filter(rand(42) < 0.5).repartition(numPartitions))) + // The same UDF marked deterministic is admitted, so it is the flag that excludes it. + assert(isPositional(df.select(sameId(col("id")).as("id")).repartition(numPartitions))) + } + } + } + test("positional placement is off unless its own config is on") { withParquetTable(100) { t => withPositionalRoundRobin( @@ -134,6 +156,54 @@ class CometNativePositionalRoundRobinSuite extends CometTestBase with AdaptiveSp } } + test("the derived group size follows the batch size and partition count") { + import CometShuffleExchangeExec.resolvePositionalGroupRows + // One batch spread over the output partitions, floored at 64 rows and capped at a batch. + assert(resolvePositionalGroupRows(0, 8192, 16) == 512) + assert(resolvePositionalGroupRows(0, 8192, 200) == 64) + assert(resolvePositionalGroupRows(0, 8192, 10000) == 64) + assert(resolvePositionalGroupRows(0, 32, 200) == 32) + // An explicit size is taken as given, including one larger than a batch. + assert(resolvePositionalGroupRows(1, 8192, 200) == 1) + assert(resolvePositionalGroupRows(100000, 8192, 200) == 100000) + } + + test("a re-executed map stage keeps its group size when the batch size changes") { + // The derived group size depends on the batch size, which can change between a stage's first + // run and a retry. Resolving it on the executor would give the retry a different placement + // from the attempt it replaces, so it is frozen with the shuffle dependency instead. + assert( + CometShuffleExchangeExec.resolvePositionalGroupRows(0, 4096, numPartitions) != + CometShuffleExchangeExec.resolvePositionalGroupRows(0, 8192, numPartitions)) + withPositionalRoundRobin(CometConf.COMET_BATCH_SIZE.key -> "8192") { + withParquetTable(5000) { t => + val df = spark.table(t).select("id").repartition(numPartitions) + val exchange = roundRobinExchange(df) + assert(exchange.positionalRoundRobin.map(_.groupRows).contains(8192 / numPartitions)) + + // `queryExecution.toRdd` rather than `df.rdd`, which plans a fresh query and so a + // different exchange from the one whose map output is dropped below. Propagating the SQL + // conf is what lets the tasks see the changed batch size, as they would under an action. + def placement(): Seq[Seq[Long]] = + SQLExecution.withSQLConfPropagated(spark) { + df.queryExecution.toRdd + .mapPartitions(rows => Iterator(rows.map(_.getLong(0)).toVector)) + .collect() + .toSeq + } + val first = placement() + + withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "4096") { + // Drop the map output so the next job re-runs the map stage under the new batch size. + SparkEnv.get.mapOutputTracker + .asInstanceOf[MapOutputTrackerMaster] + .unregisterAllMapAndMergeOutput(exchange.shuffleId) + assert(placement() == first) + } + } + } + } + test("a group size larger than a batch still keeps every row") { // Exercises the whole-batch fast path, where a run covers an entire input batch and is handed // through without a copy.