move to slt tests as much as possible

Author: buraksenn

datafusion/core/tests/fuzz_cases/join_fuzz.rs

@@ -20,7 +20,7 @@ use std::time::SystemTime;
 
 use crate::fuzz_cases::join_fuzz::JoinTestType::{HjSmj, NljHj};
 
-use arrow::array::{ArrayRef, BinaryArray, Int32Array, StringArray};
+use arrow::array::{ArrayRef, BinaryArray, Int32Array};
 use arrow::compute::SortOptions;
 use arrow::datatypes::Schema;
 use arrow::record_batch::RecordBatch;
@@ -30,24 +30,22 @@ use datafusion::datasource::memory::MemorySourceConfig;
 use datafusion::datasource::source::DataSourceExec;
 use datafusion::logical_expr::{JoinType, Operator};
 use datafusion::physical_expr::expressions::BinaryExpr;
+use datafusion::physical_plan::collect;
 use datafusion::physical_plan::expressions::Column;
 use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter};
 use datafusion::physical_plan::joins::{
     HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec,
 };
-use datafusion::physical_plan::{ExecutionPlan, collect};
 use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_common::{NullEquality, ScalarValue};
 use datafusion_execution::TaskContext;
 use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};
-use datafusion_execution::memory_pool::FairSpillPool;
 use datafusion_execution::runtime_env::RuntimeEnvBuilder;
 use datafusion_physical_expr::PhysicalExprRef;
 use datafusion_physical_expr::expressions::Literal;
 
 use itertools::Itertools;
-use rand::rngs::SmallRng;
-use rand::{Rng, SeedableRng};
+use rand::Rng;
 use test_utils::stagger_batch_with_seed;
 
 // Determines what Fuzz tests needs to run
@@ -1128,131 +1126,133 @@ impl JoinFuzzTestCase {
     }
 }
 
-/// Compare a guaranteed-spilling SMJ against an unlimited-memory hash join
-/// for filtered materializing joins.
+/// Fuzz test: compare SMJ (with spilling) against HJ (no spill) for filtered
+/// outer joins under memory pressure. This exercises the deferred filtering +
+/// spill read-back path that unit tests can't easily cover with random data.
 #[tokio::test]
 async fn test_filtered_join_spill_fuzz() {
-    let join_types = [
-        JoinType::Inner,
-        JoinType::Left,
-        JoinType::Right,
-        JoinType::Full,
-    ];
-    let input1 = make_spill_join_batches(256, 32, 512, 1);
-    let input2 = make_spill_join_batches(256, 32, 512, 2);
-    let schema1 = input1[0].schema();
-    let schema2 = input2[0].schema();
-    let filter = col_lt_col_filter(Arc::clone(&schema1), Arc::clone(&schema2));
-    let on = vec![
-        (
-            Arc::new(Column::new_with_schema("a", &schema1).unwrap()) as _,
-            Arc::new(Column::new_with_schema("a", &schema2).unwrap()) as _,
-        ),
-        (
-            Arc::new(Column::new_with_schema("b", &schema1).unwrap()) as _,
-            Arc::new(Column::new_with_schema("b", &schema2).unwrap()) as _,
-        ),
-    ];
+    let join_types = [JoinType::Left, JoinType::Right, JoinType::Full];
 
     let runtime_spill = RuntimeEnvBuilder::new()
-        .with_memory_pool(Arc::new(FairSpillPool::new(1024)))
+        .with_memory_limit(4096, 1.0)
         .with_disk_manager_builder(
             DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory),
         )
         .build_arc()
         .unwrap();
 
-    for join_type in join_types {
-        for batch_size in [2, 50] {
-            let session_config = SessionConfig::new().with_batch_size(batch_size);
+    for join_type in &join_types {
+        for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] {
+            let input1 = make_staggered_batches_i32(1000, left_extra);
+            let input2 = make_staggered_batches_i32(1000, right_extra);
 
-            let left_hj = MemorySourceConfig::try_new_exec(
-                std::slice::from_ref(&input1),
-                Arc::clone(&schema1),
-                None,
-            )
-            .unwrap();
-            let right_hj = MemorySourceConfig::try_new_exec(
-                std::slice::from_ref(&input2),
-                Arc::clone(&schema2),
-                None,
-            )
-            .unwrap();
-            let hj = Arc::new(
-                HashJoinExec::try_new(
-                    left_hj,
-                    right_hj,
-                    on.clone(),
-                    Some(filter.clone()),
-                    &join_type,
+            let schema1 = input1[0].schema();
+            let schema2 = input2[0].schema();
+            let filter = col_lt_col_filter(schema1.clone(), schema2.clone());
+
+            let on = vec![
+                (
+                    Arc::new(Column::new_with_schema("a", &schema1).unwrap()) as _,
+                    Arc::new(Column::new_with_schema("a", &schema2).unwrap()) as _,
+                ),
+                (
+                    Arc::new(Column::new_with_schema("b", &schema1).unwrap()) as _,
+                    Arc::new(Column::new_with_schema("b", &schema2).unwrap()) as _,
+                ),
+            ];
+
+            for batch_size in [2, 49, 100] {
+                let session_config = SessionConfig::new().with_batch_size(batch_size);
+
+                // HJ baseline (no memory limit)
+                let left_hj = MemorySourceConfig::try_new_exec(
+                    std::slice::from_ref(&input1),
+                    schema1.clone(),
                     None,
-                    PartitionMode::Partitioned,
-                    NullEquality::NullEqualsNothing,
-                    false,
                 )
-                .unwrap(),
-            );
-            let ctx_hj = SessionContext::new_with_config(session_config.clone());
-            let hj_collected = collect(hj, ctx_hj.task_ctx()).await.unwrap();
+                .unwrap();
+                let right_hj = MemorySourceConfig::try_new_exec(
+                    std::slice::from_ref(&input2),
+                    schema2.clone(),
+                    None,
+                )
+                .unwrap();
+                let hj = Arc::new(
+                    HashJoinExec::try_new(
+                        left_hj,
+                        right_hj,
+                        on.clone(),
+                        Some(filter.clone()),
+                        join_type,
+                        None,
+                        PartitionMode::Partitioned,
+                        NullEquality::NullEqualsNothing,
+                        false,
+                    )
+                    .unwrap(),
+                );
+                let ctx_hj = SessionContext::new_with_config(session_config.clone());
+                let hj_collected = collect(hj, ctx_hj.task_ctx()).await.unwrap();
 
-            let left_smj = MemorySourceConfig::try_new_exec(
-                std::slice::from_ref(&input1),
-                Arc::clone(&schema1),
-                None,
-            )
-            .unwrap();
-            let right_smj = MemorySourceConfig::try_new_exec(
-                std::slice::from_ref(&input2),
-                Arc::clone(&schema2),
-                None,
-            )
-            .unwrap();
-            let smj = Arc::new(
-                SortMergeJoinExec::try_new(
-                    left_smj,
-                    right_smj,
-                    on.clone(),
-                    Some(filter.clone()),
-                    join_type,
-                    vec![SortOptions::default(); on.len()],
-                    NullEquality::NullEqualsNothing,
+                // SMJ with spilling
+                let left_smj = MemorySourceConfig::try_new_exec(
+                    std::slice::from_ref(&input1),
+                    schema1.clone(),
+                    None,
                 )
-                .unwrap(),
-            );
-            let task_ctx_spill = Arc::new(
-                TaskContext::default()
-                    .with_session_config(session_config)
-                    .with_runtime(Arc::clone(&runtime_spill)),
-            );
-            let smj_collected =
-                collect(Arc::clone(&smj) as Arc<dyn ExecutionPlan>, task_ctx_spill)
-                    .await
-                    .unwrap();
-
-            assert!(
-                smj.metrics().unwrap().spill_count().unwrap_or_default() > 0,
-                "expected SMJ to spill for {join_type:?} batch_size={batch_size}",
-            );
+                .unwrap();
+                let right_smj = MemorySourceConfig::try_new_exec(
+                    std::slice::from_ref(&input2),
+                    schema2.clone(),
+                    None,
+                )
+                .unwrap();
+                let smj = Arc::new(
+                    SortMergeJoinExec::try_new(
+                        left_smj,
+                        right_smj,
+                        on.clone(),
+                        Some(filter.clone()),
+                        *join_type,
+                        vec![SortOptions::default(); on.len()],
+                        NullEquality::NullEqualsNothing,
+                    )
+                    .unwrap(),
+                );
+                let task_ctx_spill = Arc::new(
+                    TaskContext::default()
+                        .with_session_config(session_config)
+                        .with_runtime(Arc::clone(&runtime_spill)),
+                );
+                let smj_collected = collect(smj, task_ctx_spill).await.unwrap();
 
-            let hj_rows: usize = hj_collected.iter().map(|b| b.num_rows()).sum();
-            let smj_rows: usize = smj_collected.iter().map(|b| b.num_rows()).sum();
-            assert_eq!(
-                hj_rows, smj_rows,
-                "row count mismatch for {join_type:?} batch_size={batch_size}: \
-                 HJ={hj_rows} SMJ={smj_rows}",
-            );
+                let hj_rows: usize = hj_collected.iter().map(|b| b.num_rows()).sum();
+                let smj_rows: usize = smj_collected.iter().map(|b| b.num_rows()).sum();
 
-            if hj_rows > 0 {
-                let hj_fmt = pretty_format_batches(&hj_collected).unwrap().to_string();
-                let smj_fmt = pretty_format_batches(&smj_collected).unwrap().to_string();
-                let mut hj_sorted: Vec<&str> = hj_fmt.trim().lines().collect();
-                hj_sorted.sort_unstable();
-                let mut smj_sorted: Vec<&str> = smj_fmt.trim().lines().collect();
-                smj_sorted.sort_unstable();
                 assert_eq!(
-                    hj_sorted, smj_sorted,
-                    "content mismatch for {join_type:?} batch_size={batch_size}",
+                    hj_rows, smj_rows,
+                    "Row count mismatch for {join_type:?} batch_size={batch_size} \
+                     left_extra={left_extra} right_extra={right_extra}: \
+                     HJ={hj_rows} SMJ={smj_rows}"
                 );
+
+                if hj_rows > 0 {
+                    let hj_fmt =
+                        pretty_format_batches(&hj_collected).unwrap().to_string();
+                    let smj_fmt =
+                        pretty_format_batches(&smj_collected).unwrap().to_string();
+
+                    let mut hj_sorted: Vec<&str> = hj_fmt.trim().lines().collect();
+                    hj_sorted.sort_unstable();
+                    let mut smj_sorted: Vec<&str> = smj_fmt.trim().lines().collect();
+                    smj_sorted.sort_unstable();
+
+                    assert_eq!(
+                        hj_sorted, smj_sorted,
+                        "Content mismatch for {join_type:?} batch_size={batch_size} \
+                         left_extra={left_extra} right_extra={right_extra}"
+                    );
+                }
             }
         }
     }
@@ -1347,39 +1347,3 @@ fn make_staggered_batches_binary(
     // preserve your existing randomized partitioning
     stagger_batch_with_seed(batch, 42)
 }
-
-/// Sorted, low-cardinality inputs whose wide payloads force SMJ key-group
-/// spilling. `(a, b)` is sorted (`b` is constant) and `x` is nullable to
-/// exercise filter NULL handling.
-fn make_spill_join_batches(
-    len: usize,
-    num_keys: i32,
-    payload_len: usize,
-    seed: u64,
-) -> Vec<RecordBatch> {
-    let mut rng = SmallRng::seed_from_u64(seed);
-
-    let mut keys: Vec<i32> = (0..len).map(|_| rng.random_range(0..num_keys)).collect();
-    keys.sort_unstable();
-    let a = Int32Array::from_iter_values(keys);
-    let b = Int32Array::from_iter_values(std::iter::repeat_n(0, len));
-    let x = Int32Array::from_iter((0..len).map(|_| {
-        if rng.random_range(0..10) == 0 {
-            None
-        } else {
-            Some(rng.random_range(0..1000))
-        }
-    }));
-
-    let payload = "a".repeat(payload_len);
-    let p = StringArray::from_iter_values(std::iter::repeat_n(payload.as_str(), len));
-    let batch = RecordBatch::try_from_iter(vec![
-        ("a", Arc::new(a) as ArrayRef),
-        ("b", Arc::new(b) as ArrayRef),
-        ("x", Arc::new(x) as ArrayRef),
-        ("p", Arc::new(p) as ArrayRef),
-    ])
-    .unwrap();
-
-    stagger_batch_with_seed(batch, 7)
-}