PERF : modify SMJ shuffle file reader to skip validation

datafusion/physical-plan/Cargo.toml

@@ -86,3 +86,7 @@ name = "partial_ordering"
 [[bench]]
 harness = false
 name = "spill_io"
+
+[[bench]]
+harness = false
+name = "sort_merge_join"

datafusion/physical-plan/benches/sort_merge_join.rs

@@ -0,0 +1,117 @@
+// 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_schema::SortOptions;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion_common::JoinType::Inner;
+use datafusion_execution::config::SessionConfig;
+use datafusion_execution::disk_manager::DiskManagerConfig;
+use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+use datafusion_execution::TaskContext;
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_plan::common::collect;
+use datafusion_physical_plan::joins::SortMergeJoinExec;
+use datafusion_physical_plan::test::{build_table_i32, TestMemoryExec};
+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();
+    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 _,
+    )];
+    let sort_options = vec![SortOptions::default(); on.len()];
+
+    SortMergeJoinExec::try_new(left, right, 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();
+
+    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)),
+    );
+
+    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()
+            }))
+        })
+    });
+    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_main!(benches);

datafusion/physical-plan/src/joins/sort_merge_join.rs

@@ -2313,7 +2313,7 @@ fn fetch_right_columns_from_batch_by_idxs(
                 Vec::with_capacity(buffered_indices.len());
 
             let file = BufReader::new(File::open(spill_file.path())?);
-            let reader = StreamReader::try_new(file, None)?;
+            let reader = unsafe {StreamReader::try_new(file, None)?.with_skip_validation(true)};
 
             for batch in reader {
                 batch?.columns().iter().for_each(|column| {

datafusion/physical-plan/src/lib.rs

@@ -92,5 +92,4 @@ pub mod udaf {
 }
 
 pub mod coalesce;
-#[cfg(test)]
 pub mod test;