chore: Add cargo bench for shuffle writer (#438)

* Add cargo bench for shuffle writer

* add profiling guide

* add sample output

core/Cargo.toml

@@ -125,3 +125,7 @@ harness = false
 [[bench]]
 name = "cast_numeric"
 harness = false
+
+[[bench]]
+name = "shuffle_writer"
+harness = false

core/benches/shuffle_writer.rs

@@ -0,0 +1,81 @@
+// 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.
+
+use arrow_array::{builder::StringBuilder, RecordBatch};
+use arrow_schema::{DataType, Field, Schema};
+use comet::execution::datafusion::shuffle_writer::ShuffleWriterExec;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::{
+    physical_plan::{common::collect, memory::MemoryExec, ExecutionPlan},
+    prelude::SessionContext,
+};
+use datafusion_physical_expr::{expressions::Column, Partitioning};
+use std::sync::Arc;
+use tokio::runtime::Runtime;
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let batch = create_batch();
+    let mut batches = Vec::new();
+    for _ in 0..10 {
+        batches.push(batch.clone());
+    }
+    let partitions = &[batches];
+    let exec = ShuffleWriterExec::try_new(
+        Arc::new(MemoryExec::try_new(partitions, batch.schema(), None).unwrap()),
+        Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 16),
+        "/tmp/data.out".to_string(),
+        "/tmp/index.out".to_string(),
+    )
+    .unwrap();
+
+    let mut group = c.benchmark_group("shuffle_writer");
+    group.bench_function("shuffle_writer", |b| {
+        let ctx = SessionContext::new();
+        b.iter(|| {
+            let task_ctx = ctx.task_ctx();
+            let stream = exec.execute(0, task_ctx).unwrap();
+            let rt = Runtime::new().unwrap();
+            criterion::black_box(rt.block_on(collect(stream)).unwrap());
+        });
+    });
+}
+
+fn create_batch() -> RecordBatch {
+    let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)]));
+    let mut b = StringBuilder::new();
+    for i in 0..8192 {
+        if i % 10 == 0 {
+            b.append_null();
+        } else {
+            b.append_value(format!("{i}"));
+        }
+    }
+    let array = b.finish();
+    let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+    batch
+}
+
+fn config() -> Criterion {
+    Criterion::default()
+}
+
+criterion_group! {
+    name = benches;
+    config = config();
+    targets = criterion_benchmark
+}
+criterion_main!(benches);

core/src/execution/datafusion/mod.rs

@@ -20,6 +20,6 @@
 pub mod expressions;
 mod operators;
 pub mod planner;
-pub(crate) mod shuffle_writer;
+pub mod shuffle_writer;
 mod spark_hash;
 mod util;

docs/source/contributor-guide/profiling_native_code.md

@@ -0,0 +1,94 @@
+<!--
+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.
+-->
+
+# Profiling Native Code
+
+We use `cargo bench` to run benchmarks to measure the performance of individual operators and expressions
+and [cargo-flamegraph](https://github.com/flamegraph-rs/flamegraph) for profiling.
+
+## Running micro benchmarks with cargo bench
+
+When implementing a new operator or expression, it is good practice to add a new microbenchmark under `core/benches`.
+
+It is often easiest to copy an existing benchmark and modify it for the new operator or expression. It is also
+necessary to add a new section to the `Cargo.toml` file, such as:
+
+```toml
+[[bench]]
+name = "shuffle_writer"
+harness = false
+```
+
+These benchmarks are useful when for comparing performance between releases or between feature branches and the
+main branch to help prevent regressions in performance when adding new features or fixing bugs.
+
+Individual benchmarks can be run by name with the following command.
+
+```shell
+cargo bench shuffle_writer
+```
+
+Here is some sample output from running this command.
+
+```
+     Running benches/shuffle_writer.rs (target/release/deps/shuffle_writer-e37b59e37879cce7)
+Gnuplot not found, using plotters backend
+shuffle_writer/shuffle_writer
+                        time:   [2.0880 ms 2.0989 ms 2.1118 ms]
+Found 9 outliers among 100 measurements (9.00%)
+  3 (3.00%) high mild
+  6 (6.00%) high severe
+```
+
+## Profiling with cargo-flamegraph
+
+Install cargo-flamegraph:
+
+```shell
+cargo install flamegraph
+```
+
+Follow the instructions in [cargo-flamegraph](https://github.com/flamegraph-rs/flamegraph) for your platform for
+running flamegraph.
+
+Here is a sample command for running `cargo-flamegraph` on MacOS.
+
+```shell
+cargo flamegraph --root --bench shuffle_writer
+```
+
+This will produce output similar to the following.
+
+```
+dtrace: system integrity protection is on, some features will not be available
+
+dtrace: description 'profile-997 ' matched 1 probe
+Gnuplot not found, using plotters backend
+Testing shuffle_writer/shuffle_writer
+Success
+
+dtrace: pid 66402 has exited
+writing flamegraph to "flamegraph.svg"
+```
+
+The generated flamegraph can now be opened in a browser that supports svg format.
+
+Here is the flamegraph for this example:
+
+![flamegraph](../_static/images/flamegraph.png)

docs/source/index.rst

@@ -57,6 +57,7 @@ as a native runtime to achieve improvement in terms of query efficiency and quer
    Comet Plugin Overview <contributor-guide/plugin_overview>
    Development Guide <contributor-guide/development>
    Debugging Guide <contributor-guide/debugging>
+   Profiling Native Code <contributor-guide/profiling_native_code>
    Github and Issue Tracker <https://github.com/apache/datafusion-comet>
 
 .. _toc.asf-links: