diff --git a/.ai/skills/review-comet-shuffle-pr/SKILL.md b/.ai/skills/review-comet-shuffle-pr/SKILL.md index b768923ae7f..4604f5da6b2 100644 --- a/.ai/skills/review-comet-shuffle-pr/SKILL.md +++ b/.ai/skills/review-comet-shuffle-pr/SKILL.md @@ -88,13 +88,24 @@ 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 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, 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 search over comparable-row-format keys. A change to the comparison or the row encoding must 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..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 @@ -305,21 +300,136 @@ 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 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 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. + +#### `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 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. +`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. +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 +smaller, and `RunIterator` builds each output chunk by slicing and concatenating runs. A run that +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`, 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 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 + `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 @@ -370,21 +480,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/docs/source/user-guide/latest/metrics.md b/docs/source/user-guide/latest/metrics.md index 5404344e987..9e0aeeed462 100644 --- a/docs/source/user-guide/latest/metrics.md +++ b/docs/source/user-guide/latest/metrics.md @@ -139,7 +139,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. | diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 1f94a892613..df0982d9570 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; @@ -3660,15 +3660,33 @@ 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 + // 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, + group_rows: rr_partition.positional_group_rows as usize, + // Kept for the case where the schema rules positional placement out. + max_hash_columns, + } } else { - rr_partition.max_hash_columns as usize + RoundRobinStrategy::HashAll { max_hash_columns } }; Ok(CometPartitioning::RoundRobin( rr_partition.num_partitions as usize, - max_hash_columns, + strategy, )) } } diff --git a/native/proto/src/proto/partitioning.proto b/native/proto/src/proto/partitioning.proto index e70b8264f02..19b809f9a62 100644 --- a/native/proto/src/proto/partitioning.proto +++ b/native/proto/src/proto/partitioning.proto @@ -57,4 +57,16 @@ 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 + // (positional_start_partition + i / positional_group_rows) % num_partitions. + // 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. 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`. + int32 positional_start_partition = 5; } diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index cc94375436e..c099fc0413f 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}; @@ -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; @@ -73,19 +74,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 +122,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 +148,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 +171,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, @@ -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); @@ -270,6 +275,107 @@ 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 `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; + 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: BATCH_SIZE / NUM_PARTITIONS, + max_hash_columns: 0, + }, + ), + ( + "RowGroups(8192)", + RoundRobinStrategy::RowGroups { + start_partition: 0, + group_rows: BATCH_SIZE, + max_hash_columns: 0, + }, + ), + ]; + + // `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. 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", + (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), + ), + ]; + + 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. @@ -372,13 +478,50 @@ 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, +} + +/// `count` separately built batches, each filled from where the previous one left off. +fn nested_batches( + num_rows: usize, + count: usize, + num_cols: usize, + depth: usize, + fill: Fill, +) -> Vec { + (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); + let array = nested_struct_array(first_row, num_rows, depth, fill); fields.push(Field::new( format!("col{col}"), array.data_type().clone(), @@ -390,23 +533,45 @@ fn nested_schema_batch(num_rows: usize) -> RecordBatch { 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) -> 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 { + Fill::Constant => ( + vec![1_i64; num_rows], + vec!["x".to_string(); num_rows], + vec![1.0_f64; num_rows], + ), + 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 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, ), ])); @@ -428,6 +593,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/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..5b15f2e49d6 100644 --- a/native/shuffle/src/comet_partitioning.rs +++ b/native/shuffle/src/comet_partitioning.rs @@ -17,8 +17,81 @@ 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)`. 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 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. 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. + max_hash_columns: 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, + } + } +} + +/// 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. +/// +/// `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, +) -> 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; + 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 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. #[derive(Debug, Clone)] pub enum CometPartitioning { @@ -32,10 +105,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 +141,23 @@ mod tests { let expected = vec![69, 5, 193, 171, 115]; assert_eq!(result, expected); } + + /// 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() { + // Rows 2..4 of the task finish group 0; rows 4..8 are group 1. + assert_eq!( + 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() { + assert_eq!( + positional_runs(0, 100, 3, 8192, 200).collect::>(), + vec![(3, 0..100)] + ); + } } diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 1158a2b1e2e..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; @@ -32,7 +35,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..05b14b53243 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; 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}; @@ -101,7 +104,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 +125,9 @@ 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, + /// 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 +194,42 @@ 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. + 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. // 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_starts: vec![0; num_output_partitions + 1], + 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: if places_rows_individually { + vec![0; num_output_partitions + 1] + } else { + vec![] + }, }; let reservation = MemoryConsumer::new(format!("ShuffleRepartitioner[{partition}]")) @@ -211,7 +238,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 +252,7 @@ impl MultiPartitionShuffleRepartitioner { max_buffer_bytes, tracing_enabled, pinned_buffers: HashSet::new(), + row_seq: 0, }) } @@ -349,20 +381,41 @@ 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, + RoundRobinStrategy::RowGroups { + start_partition, + group_rows, + .. + }, + ) => { + 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, + RoundRobinStrategy::HashAll { max_hash_columns }, + ) => { let mut scratch = std::mem::take(&mut self.scratch); + let num_rows = input.num_rows(); + + // 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(); - 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 @@ -389,8 +442,9 @@ impl MultiPartitionShuffleRepartitioner { 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. + // 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); @@ -426,13 +480,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 +503,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 +513,67 @@ impl MultiPartitionShuffleRepartitioner { mem_growth += after_size.saturating_sub(before_size); } + self.reserve_and_may_spill(mem_growth) + } + + /// 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 { + return Err(DataFusionError::Internal( + "positional placement against a row-indexed repartitioner".to_string(), + )); + }; + { + 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) + } + + /// 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 +610,9 @@ 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 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(); + let indices = std::mem::replace(&mut self.partition_indices, empty); PartitionedBatchesProducer::new(buffered_batches, indices, self.batch_size) } @@ -518,7 +629,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 +698,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(); @@ -678,36 +789,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(), - CometPartitioning::RoundRobin(2, 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 - .iter() - .map(|indices| indices.allocated_size()) - .sum::(); - let reserved_bytes = repartitioner.reservation.size(); - assert_eq!(repartitioner.spill_count(), 0); - assert_eq!(reserved_bytes, input_bytes + index_bytes); + // 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, + }, + ), + 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] @@ -788,7 +914,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 +931,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 +946,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 +960,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 +978,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 +1027,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, @@ -989,36 +1104,363 @@ 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(); + 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!(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(()) + } + } + + /// 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, + ) -> MultiPartitionShuffleRepartitioner { + let runtime = Arc::new( + datafusion::execution::runtime_env::RuntimeEnvBuilder::new() + .build() + .unwrap(), + ); + MultiPartitionShuffleRepartitioner::try_new( + 0, + W::default(), + CometPartitioning::RoundRobin( + num_partitions, + RoundRobinStrategy::RowGroups { + start_partition, + group_rows, + max_hash_columns: 0, + }, + ), + ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + runtime, + batch_size, + false, + None, + ) + .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 batch = int64_batch(next..next + *rows as i64); + next += *rows as i64; + 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" + ); + } + } + + /// 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 groups = 5; + + let stage = |starts: Vec| async move { + let mut covered = std::collections::BTreeSet::new(); + for start in starts { + 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() + }; + + // 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, - FailingPartitionWriter::default(), - CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2), - ShufflePartitionerMetrics::new(&metrics_set, 0), - Arc::clone(&runtime), + 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, - Some(1), + max_buffer_bytes, ) .unwrap(); - for start in (0..num_rows).step_by(input_batch_rows) { + let mut next = 0i64; + for rows in framing { repartitioner - .insert_batch(batch.slice(start, input_batch_rows)) + .insert_batch(int64_batch(next..next + rows as i64)) .await .unwrap(); + next += rows as i64; } - 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!(repartitioner.partition_indices.iter().all(Vec::is_empty)); - spill_bytes.push(repartitioner.metrics.memory_spilled_bytes.value()); + 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" + ); } - assert!(spill_bytes[0] > 0); - assert_eq!(spill_bytes[0], spill_bytes[1]); + } + + /// 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 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" + ); + + 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); + 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..6151c336fe2 100644 --- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs +++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs @@ -16,23 +16,87 @@ // 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 { + /// 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]), + } + } + + 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 { @@ -49,6 +113,9 @@ 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], @@ -62,17 +129,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, + interleave_time, + )), + PartitionIndices::Runs(runs) => PartitionedBatchIterator::Runs(RunIterator::new( + &runs[partition_id], + refs, + self.batch_size, + interleave_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)], @@ -84,7 +179,7 @@ pub(crate) struct PartitionedBatchIterator<'a> { interleave_time: &'a Time, } -impl<'a> PartitionedBatchIterator<'a> { +impl<'a> RowIterator<'a> { fn new( indices: &'a [(u32, u32)], record_batches: &'a [&'a RecordBatch], @@ -113,7 +208,7 @@ impl<'a> PartitionedBatchIterator<'a> { } } -impl Iterator for PartitionedBatchIterator<'_> { +impl Iterator for RowIterator<'_> { type Item = datafusion::common::Result; fn next(&mut self) -> Option { @@ -144,6 +239,107 @@ 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, + interleave_time: &'a Time, +} + +impl<'a> RunIterator<'a> { + fn new( + runs: &'a [BufferedRun], + record_batches: &'a [&'a RecordBatch], + batch_size: usize, + interleave_time: &'a Time, + ) -> Self { + Self { + record_batches, + batch_size, + runs, + chunk_scratch: Vec::with_capacity(runs.len().min(batch_size)), + pos: 0, + consumed: 0, + interleave_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.interleave_time.timer(); + + // 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 + { + 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, 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 { + 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 +382,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 +420,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 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() { + 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..0e9085916de 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, Schema, SchemaRef}, error::Result, execution::context::TaskContext, physical_plan::{ @@ -387,6 +387,63 @@ 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, + } +} + +/// 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( @@ -400,6 +457,7 @@ fn create_repartitioner( max_buffer_bytes: Option, ) -> Result> { let partition_count = partitioning.partition_count(); + let partitioning = partitioning_for_schema(partitioning, &schema); if schema.fields().is_empty() { log::debug!( @@ -440,7 +498,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 +1201,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 +1269,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 +1673,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 +1757,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 +1794,94 @@ 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() + ))); + } + + /// 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 + } + ) + )); + } } diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index d371f1ba44c..4766a30b373 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -475,6 +475,46 @@ 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, 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 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 + .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. 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 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( + 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..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 @@ -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.positionalRoundRobinSpec]] 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'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. + * + * 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 + 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..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 @@ -421,6 +421,22 @@ 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) + // Per task, unlike the two above: which partition this mapper's first group goes to. + // 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(mapPartitionId, effectivePartitionCount)) + } 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..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 @@ -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 a re-executed map task places rows exactly as the attempt it replaces. See + * `CometShuffleExchangeExec.positionalRoundRobinSpec`. + */ + positionalRoundRobin: Option[PositionalRoundRobin] = None) + +/** + * Parameters for positional round-robin placement, resolved on the driver. + * + * @param groupRows + * 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) /** * 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..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 @@ -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 @@ -115,6 +115,20 @@ 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. Only the native writer places positionally. + */ + @transient private[shuffle] lazy val positionalRoundRobin: Option[PositionalRoundRobin] = + 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) @@ -127,7 +141,8 @@ case class CometShuffleExchangeExec( ctx.numPartitions, ctx.shuffleScanIndices, CometMetricNode(metrics, Seq(nativeChildMetricNode)), - ctx.perPartitionByKey) + ctx.perPartitionByKey, + 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. @@ -206,7 +221,11 @@ case class CometShuffleExchangeExec( outputPartitioning, serializer, metrics, - NativeShuffleSpec(nativeChild.nativeOp, nativeChildMetricNode, ctx)) + NativeShuffleSpec( + nativeChild.nativeOp, + nativeChildMetricNode, + ctx, + positionalRoundRobin)) case None => CometShuffleExchangeExec.prepareShuffleDependency( inputRDD.asInstanceOf[RDD[ColumnarBatch]], @@ -298,6 +317,108 @@ object CometShuffleExchangeExec if (shuffleSupported(op).isDefined) Compatible() else Unsupported() } + /** + * 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 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` + * guard mirrors `isRoundRobin` in `prepareJVMShuffleDependency`: with one output partition + * there is no placement to get wrong. + */ + 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() && + !isCometCelebornShuffleManagerEnabled(conf) && + replaysRowsInOrder(child) + if (eligible) { + Some( + PositionalRoundRobin( + 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 + * 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 + + /** + * Whether re-executing this subtree yields the same rows in the same order. + * + * 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. + * + * 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 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 => + p.projectList.forall(_.deterministic) && replaysRowsInOrder(p.child) + case f: CometFilterExec => f.condition.deterministic && 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/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 new file mode 100644 index 00000000000..77f9cedffbc --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativePositionalRoundRobinSuite.scala @@ -0,0 +1,291 @@ +/* + * 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.{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, rand, udf} +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) + + /** 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 && + e.outputPartitioning.isInstanceOf[RoundRobinPartitioning] => + e + } + assert( + exchanges.size == 1, + s"expected one native round-robin exchange in\n${df.queryExecution.executedPlan}") + 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 + 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("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( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_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(!e.usesPositionalRoundRobin)) + } + } + } + + test("positional placement keeps every row exactly once") { + withPositionalRoundRobin() { + withParquetTable(5000) { t => + checkSparkAnswer(spark.table(t).repartition(numPartitions).selectExpr("id", "s", "g")) + } + } + } + + 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. + 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))) + } + + withPositionalRoundRobin( + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_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") + } + } + } + + 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 + // 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") + } + } + } +} 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 =>