feat: add standalone shuffle benchmark tool (#3752)

Author: andygrove

native/Cargo.lock

@@ -32,6 +32,7 @@ publish = false
 arrow = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }
+clap = { version = "4", features = ["derive"], optional = true }
 crc32c = "0.6.8"
 crc32fast = "1.3.2"
 datafusion = { workspace = true }
@@ -43,6 +44,8 @@ itertools = "0.14.0"
 jni = "0.21"
 log = "0.4"
 lz4_flex = { version = "0.13.0", default-features = false, features = ["frame"] }
+# parquet is only used by the shuffle_bench binary (shuffle-bench feature)
+parquet = { workspace = true, optional = true }
 simd-adler32 = "0.3.9"
 snap = "1.1"
 tokio = { version = "1", features = ["rt-multi-thread"] }
@@ -54,10 +57,18 @@ datafusion = { workspace = true, features = ["parquet_encryption", "sql"] }
 itertools = "0.14.0"
 tempfile = "3.26.0"
 
+[features]
+shuffle-bench = ["clap", "parquet"]
+
 [lib]
 name = "datafusion_comet_shuffle"
 path = "src/lib.rs"
 
+[[bin]]
+name = "shuffle_bench"
+path = "src/bin/shuffle_bench.rs"
+required-features = ["shuffle-bench"]
+
 [[bench]]
 name = "shuffle_writer"
 harness = false

native/shuffle/Cargo.toml

@@ -23,3 +23,44 @@ This crate provides the shuffle writer and reader implementation for Apache Data
 of the [Apache DataFusion Comet] subproject.
 
 [Apache DataFusion Comet]: https://github.com/apache/datafusion-comet/
+
+## Shuffle Benchmark Tool
+
+A standalone benchmark binary (`shuffle_bench`) is included for profiling shuffle write
+performance outside of Spark. It streams input data directly from Parquet files.
+
+### Basic usage
+
+```sh
+cargo run --release --features shuffle-bench --bin shuffle_bench -- \
+  --input /data/tpch-sf100/lineitem/ \
+  --partitions 200 \
+  --codec lz4 \
+  --hash-columns 0,3
+```
+
+### Options
+
+| Option                | Default                    | Description                                            |
+| --------------------- | -------------------------- | ------------------------------------------------------ |
+| `--input`             | _(required)_               | Path to a Parquet file or directory of Parquet files   |
+| `--partitions`        | `200`                      | Number of output shuffle partitions                    |
+| `--partitioning`      | `hash`                     | Partitioning scheme: `hash`, `single`, `round-robin`   |
+| `--hash-columns`      | `0`                        | Comma-separated column indices to hash on (e.g. `0,3`) |
+| `--codec`             | `lz4`                      | Compression codec: `none`, `lz4`, `zstd`, `snappy`     |
+| `--zstd-level`        | `1`                        | Zstd compression level (1–22)                          |
+| `--batch-size`        | `8192`                     | Batch size for reading Parquet data                    |
+| `--memory-limit`      | _(none)_                   | Memory limit in bytes; triggers spilling when exceeded |
+| `--write-buffer-size` | `1048576`                  | Write buffer size in bytes                             |
+| `--limit`             | `0`                        | Limit rows processed per iteration (0 = no limit)      |
+| `--iterations`        | `1`                        | Number of timed iterations                             |
+| `--warmup`            | `0`                        | Number of warmup iterations before timing              |
+| `--output-dir`        | `/tmp/comet_shuffle_bench` | Directory for temporary shuffle output files           |
+
+### Profiling with flamegraph
+
+```sh
+cargo flamegraph --release --features shuffle-bench --bin shuffle_bench -- \
+  --input /data/tpch-sf100/lineitem/ \
+  --partitions 200 --codec lz4
+```