increase sample data size

Author: getChan

datafusion/core/benches/sort_merge_join.rs

@@ -15,9 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow_schema::SortOptions;
+use arrow::array::{Int32Array, RecordBatch};
+use arrow_schema::{DataType, Field, Schema, SortOptions};
 use criterion::{criterion_group, criterion_main, Criterion};
 use datafusion_common::JoinType::Inner;
+use datafusion_datasource::memory::MemorySourceConfig;
 use datafusion_execution::config::SessionConfig;
 use datafusion_execution::disk_manager::DiskManagerConfig;
 use datafusion_execution::runtime_env::RuntimeEnvBuilder;
@@ -29,88 +31,78 @@ use datafusion_physical_plan::ExecutionPlan;
 use std::sync::Arc;
 use tokio::runtime::Runtime;
 
-fn create_test_data() -> SortMergeJoinExec {
-    let left_batch = build_table_i32(
-        ("a1", &vec![0, 1, 2, 3, 4, 5]),
-        ("b1", &vec![1, 2, 3, 4, 5, 6]),
-        ("c1", &vec![4, 5, 6, 7, 8, 9]),
-    );
-    let left_schema = left_batch.schema();
-    let left =
-        TestMemoryExec::try_new_exec(&[vec![left_batch]], left_schema, None).unwrap();
-    let right_batch = build_table_i32(
-        ("a2", &vec![0, 10, 20, 30, 40]),
-        ("b2", &vec![1, 3, 4, 6, 8]),
-        ("c2", &vec![50, 60, 70, 80, 90]),
-    );
-    let right_schema = right_batch.schema();
-    let right =
-        TestMemoryExec::try_new_exec(&[vec![right_batch]], right_schema, None).unwrap();
+fn create_smj_exec(array_len: usize, batch_size: usize) -> SortMergeJoinExec {
+    // define a schema.
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("a1", DataType::Int32, false),
+        Field::new("b1", DataType::Int32, false),
+        Field::new("c1", DataType::Int32, false),
+    ]));
+    // define data.
+    let batches = (0..array_len / batch_size)
+        .map(|i| {
+            RecordBatch::try_new(
+                Arc::clone(&schema),
+                vec![
+                    Arc::new(Int32Array::from(vec![i as i32; batch_size])),
+                    Arc::new(Int32Array::from(vec![i as i32; batch_size])),
+                    Arc::new(Int32Array::from(vec![i as i32; batch_size])),
+                ],
+            )
+            .unwrap()
+        })
+        .collect::<Vec<_>>();
+    let datasource_exec =
+        MemorySourceConfig::try_new_exec(&vec![batches], Arc::clone(&schema), None)
+            .unwrap();
+
     let on = vec![(
-        Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _,
-        Arc::new(Column::new_with_schema("b2", &right.schema()).unwrap()) as _,
+        Arc::new(Column::new_with_schema("b1", &schema).unwrap()) as _,
+        Arc::new(Column::new_with_schema("b1", &schema).unwrap()) as _,
     )];
     let sort_options = vec![SortOptions::default(); on.len()];
-
-    SortMergeJoinExec::try_new(left, right, on, None, Inner, sort_options, false).unwrap()
+    SortMergeJoinExec::try_new(
+        datasource_exec.clone(),
+        datasource_exec.clone(),
+        on,
+        None,
+        Inner,
+        sort_options,
+        false,
+    )
+    .unwrap()
 }
 
 // `cargo bench --bench sort_merge_join`
 fn bench_spill(c: &mut Criterion) {
-    let sort_merge_join_exec = create_test_data();
-
-    let mut group = c.benchmark_group("sort_merge_join_spill");
     let rt = Runtime::new().unwrap();
+    c.bench_function("SortMergeJoinExec_spill", |b| {
+        let join_exec = create_smj_exec(1_048_576, 4096);
 
-    let runtime = RuntimeEnvBuilder::new()
-        .with_memory_limit(100, 1.0) // Set memory limit to 100 bytes
-        .with_disk_manager(DiskManagerConfig::NewOs) // Enable DiskManager to allow spilling
-        .build_arc()
-        .unwrap();
-    let session_config = SessionConfig::default();
-    let task_ctx = Arc::new(
-        TaskContext::default()
-            .with_session_config(session_config.clone())
-            .with_runtime(Arc::clone(&runtime)),
-    );
+        // create a session context. enable spilling
+        let runtime_env = RuntimeEnvBuilder::new()
+            .with_memory_limit(1024 * 1024, 1.0) // Set memory limit to 1MB
+            .with_disk_manager(DiskManagerConfig::NewOs) // Enable DiskManager to allow spilling
+            .build_arc()
+            .unwrap();
+        let task_ctx = Arc::new(
+            TaskContext::default()
+                .with_session_config(SessionConfig::new())
+                .with_runtime(Arc::clone(&runtime_env)),
+        );
 
-    group.bench_function("SortMergeJoinExec_spill", |b| {
         b.iter(|| {
-            criterion::black_box(rt.block_on(async {
-                let stream = sort_merge_join_exec
-                    .execute(0, Arc::clone(&task_ctx))
-                    .unwrap();
-                collect(stream).await.unwrap()
-            }))
-        })
+            let stream = join_exec.execute(0, Arc::clone(&task_ctx)).unwrap();
+            criterion::black_box(rt.block_on(collect(stream))).unwrap();
+        });
+        // check if spilling happened
+        assert!(join_exec.metrics().unwrap().spilled_rows().unwrap() > 0);
     });
-    group.finish();
-
-    assert!(
-        sort_merge_join_exec
-            .metrics()
-            .unwrap()
-            .spill_count()
-            .unwrap()
-            > 0
-    );
-    assert!(
-        sort_merge_join_exec
-            .metrics()
-            .unwrap()
-            .spilled_bytes()
-            .unwrap()
-            > 0
-    );
-    assert!(
-        sort_merge_join_exec
-            .metrics()
-            .unwrap()
-            .spilled_rows()
-            .unwrap()
-            > 0
-    );
 }
 
-criterion_group!(benches, bench_spill);
+criterion_group!(
+    name = benches;
+    config = Criterion::default().sample_size(10);
+    targets = bench_spill
+);
 criterion_main!(benches);