sunchao commented on code in PR #5449:
URL: https://github.com/apache/datafusion-comet/pull/5449#discussion_r4066933153
##########
native/shuffle/benches/shuffle_writer.rs:
##########
@@ -161,18 +162,117 @@ fn criterion_benchmark(c: &mut Criterion) {
let exec = create_shuffle_writer_exec(
CompressionCodec::None,
CometPartitioning::SinglePartition,
- rows_per_batch,
- num_batches,
+ create_batches(rows_per_batch, num_batches),
);
+ let rt = Runtime::new().unwrap();
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();
});
},
);
}
+
+ // RoundRobin on a wide nested schema: compares the row-level
hash-all-columns strategy
+ // against the whole-batch strategy. The nested schema is where the win is
largest:
+ // `create_murmur3_hashes` recurses into every Struct child per row on the
row-level path,
+ // while the whole-batch path does one mod per batch and never inspects
row contents.
+ // 40 top-level struct columns of shallow depth reflect the real-world
shape where this
+ // optimization matters (event log workload was ~194 nested columns).
+ let num_partitions = 50usize;
+ let wide_batches: Vec<RecordBatch> = (0..8).map(|_|
nested_schema_batch(8192, 40, 2)).collect();
+ let wide_schema = wide_batches[0].schema();
+ let round_robin_strategies = [
+ ("hash_all_columns", RoundRobinStrategy::default()),
+ (
+ "whole_batch",
+ RoundRobinStrategy::WholeBatch { start_partition: 0 },
+ ),
+ ];
+ for (label, strategy) in &round_robin_strategies {
+ group.bench_function(
+ format!("shuffle_writer: RoundRobin nested schema
(strategy={label})"),
+ |b| {
+ let ctx = SessionContext::new();
+ let rt = Runtime::new().unwrap();
+ let exec = create_shuffle_writer_exec(
+ CompressionCodec::None,
+ CometPartitioning::RoundRobin(num_partitions,
strategy.clone()),
+ wide_batches.clone(),
+ );
+ b.iter(|| {
+ let task_ctx = ctx.task_ctx();
+ let stream = exec.execute(0, task_ctx).unwrap();
+ rt.block_on(collect(stream)).unwrap();
Review Comment:
### Performance
[P2] Build a fresh writer for each RoundRobin benchmark iteration
The new end-to-end RoundRobin benchmark constructs `exec` before `b.iter`,
so both strategies reuse one `Arc<PartitionOffsets>` across iterations.
`ShuffleWriterExec::execute` clones that destination, and
`LocalPartitionWriter::finish_all` unconditionally publishes into its
`OnceLock`. After the first successful iteration, the next
`collect(stream).unwrap()` therefore panics with `partition offsets were
already published`. This prevents the documented `cargo bench -p
datafusion-comet-shuffle --bench shuffle_writer -- RoundRobin` command from
producing the new comparison, despite the PR description saying this case now
uses `iter_batched`. Move creation of the RoundRobin exec into a per-iteration
setup closure, then run the focused benchmark through warmup and sampling. The
finding is limited to this new benchmark; the older benchmarks' equivalent
problem can remain a separate fix.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
