feat(pwmj): support RightSemi/RightAnti existence joins

datafusion/core/src/physical_planner.rs

@@ -1606,16 +1606,13 @@ impl DefaultPhysicalPlanner {
                         Arc::new(CrossJoinExec::new(physical_left, physical_right))
                     } else if num_range_filters == 1
                         && total_filters == 1
-                        // PWMJ supports classic joins and Left Semi/Anti existence joins.
-                        // Right Semi/Anti and Mark joins are not implemented yet (they
-                        // would require swapping the inputs so the marked side is buffered),
-                        // so exclude them here and let them fall back to NestedLoopJoin.
+                        // PWMJ supports classic joins and Semi/Anti existence joins. Mark
+                        // joins are not implemented yet (they need an extra boolean column
+                        // rather than a subset of one side's rows), so exclude them here
+                        // and let them fall back to NestedLoopJoin.
                         && !matches!(
                             join_type,
-                            JoinType::RightSemi
-                                | JoinType::RightAnti
-                                | JoinType::LeftMark
-                                | JoinType::RightMark
+                            JoinType::LeftMark | JoinType::RightMark
                         )
                         && session_state
                             .config_options()

datafusion/core/tests/fuzz_cases/join_fuzz.rs

@@ -20,9 +20,9 @@ use std::time::SystemTime;
 
 use crate::fuzz_cases::join_fuzz::JoinTestType::{HjSmj, NljHj};
 
-use arrow::array::{Array, ArrayRef, BinaryArray, Int32Array};
+use arrow::array::{Array, ArrayRef, BinaryArray, Float64Array, Int32Array};
 use arrow::compute::SortOptions;
-use arrow::datatypes::Schema;
+use arrow::datatypes::{DataType, Schema};
 use arrow::record_batch::RecordBatch;
 use arrow::util::pretty::pretty_format_batches;
 use datafusion::common::JoinSide;
@@ -1374,40 +1374,98 @@ fn make_staggered_batches_binary(
 // streamed side below is spread round-robin over several partitions as one-row batches, so
 // batches arrive in an order no static test pins down, and the counter that gates the final
 // pass is seeded from a partition count that deliberately disagrees with the `num_partitions`
-// argument.
+// argument. `RightSemi`/`RightAnti` take the mirror path -- every partition emits its own
+// rows, decided against a single buffered key -- and are covered here too, over the same
+// inputs, so the two halves are held to the same oracle. Their buffered side is fanned out as
+// well, since it carries no single-partition requirement: each partition is folded to its own
+// extreme on a separate task and those are then combined, and only randomization varies which
+// partition holds the deciding key, or holds none at all.
+
+/// A key type the differential fuzz test can generate and array-ify. `i32` covers the plain
+/// ordering/duplicate/NULL dimensions; `f64` adds `-0.0` and `NaN`, whose comparison semantics
+/// arrow's `cmp` kernels and its `min`/`max` accumulators must agree on for
+/// `PiecewiseMergeJoinExec`'s buffered-side reduction to be sound (see `right_existence_join`'s
+/// module doc).
+trait PwmjFuzzKey: Copy + 'static {
+    fn data_type() -> DataType;
+    fn to_array(values: &[Option<Self>]) -> ArrayRef;
+    /// Draws a value for the narrow `[0, key_range)` range the rest of the test relies on to
+    /// force duplicates and equal-boundary cases, plus (for floats) the values a uniform draw
+    /// over that range would never produce on its own.
+    fn gen_value(key_range: i32, rng: &mut StdRng) -> Self;
+}
+
+impl PwmjFuzzKey for i32 {
+    fn data_type() -> DataType {
+        DataType::Int32
+    }
+
+    fn to_array(values: &[Option<Self>]) -> ArrayRef {
+        Arc::new(Int32Array::from(values.to_vec()))
+    }
+
+    fn gen_value(key_range: i32, rng: &mut StdRng) -> Self {
+        rng.random_range(0..key_range)
+    }
+}
+
+impl PwmjFuzzKey for f64 {
+    fn data_type() -> DataType {
+        DataType::Float64
+    }
 
-fn pwmj_kv_schema() -> Arc<Schema> {
+    fn to_array(values: &[Option<Self>]) -> ArrayRef {
+        Arc::new(Float64Array::from(values.to_vec()))
+    }
+
+    fn gen_value(key_range: i32, rng: &mut StdRng) -> Self {
+        // One in 8 draws hits `-0.0` or `NaN` instead of the uniform range: both must compare
+        // consistently between the `cmp` kernel used to filter the streamed side and the
+        // `min`/`max` accumulators used to reduce the buffered side, and a plain uniform draw
+        // over `[0, key_range)` would essentially never produce either on its own.
+        match rng.random_range(0..8u32) {
+            0 => -0.0,
+            1 => f64::NAN,
+            _ => f64::from(rng.random_range(0..key_range)),
+        }
+    }
+}
+
+fn pwmj_kv_schema<K: PwmjFuzzKey>() -> Arc<Schema> {
     Arc::new(Schema::new(vec![
-        arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int32, false),
-        arrow::datatypes::Field::new("k", arrow::datatypes::DataType::Int32, true),
+        arrow::datatypes::Field::new("id", DataType::Int32, false),
+        arrow::datatypes::Field::new("k", K::data_type(), true),
     ]))
 }
 
-fn pwmj_kv_batch(ids: &[i32], keys: &[Option<i32>]) -> RecordBatch {
+fn pwmj_kv_batch<K: PwmjFuzzKey>(ids: &[i32], keys: &[Option<K>]) -> RecordBatch {
     RecordBatch::try_new(
-        pwmj_kv_schema(),
-        vec![
-            Arc::new(Int32Array::from(ids.to_vec())),
-            Arc::new(Int32Array::from(keys.to_vec())),
-        ],
+        pwmj_kv_schema::<K>(),
+        vec![Arc::new(Int32Array::from(ids.to_vec())), K::to_array(keys)],
     )
     .unwrap()
 }
 
 /// Single-partition, single-batch input: the buffered side, and the oracle's probe side.
-fn pwmj_single_exec(ids: &[i32], keys: &[Option<i32>]) -> Arc<dyn ExecutionPlan> {
+fn pwmj_single_exec<K: PwmjFuzzKey>(
+    ids: &[i32],
+    keys: &[Option<K>],
+) -> Arc<dyn ExecutionPlan> {
     MemorySourceConfig::try_new_exec(
         &[vec![pwmj_kv_batch(ids, keys)]],
-        pwmj_kv_schema(),
+        pwmj_kv_schema::<K>(),
         None,
     )
     .unwrap()
 }
 
-/// Streamed side spread round-robin across `nparts` partitions, one row per batch.
-fn pwmj_parts_exec(
+/// Rows spread round-robin across `nparts` partitions, one row per batch, with any partition
+/// that draws no row left holding a single empty batch. Used for the streamed side throughout,
+/// and for the buffered side of the right existence joins, which place no single-partition
+/// requirement on it.
+fn pwmj_parts_exec<K: PwmjFuzzKey>(
     ids: &[i32],
-    keys: &[Option<i32>],
+    keys: &[Option<K>],
     nparts: usize,
 ) -> Arc<dyn ExecutionPlan> {
     let nparts = nparts.max(1);
@@ -1417,10 +1475,10 @@ fn pwmj_parts_exec(
     }
     for p in partitions.iter_mut() {
         if p.is_empty() {
-            p.push(pwmj_kv_batch(&[], &[]));
+            p.push(pwmj_kv_batch::<K>(&[], &[]));
         }
     }
-    MemorySourceConfig::try_new_exec(&partitions, pwmj_kv_schema(), None).unwrap()
+    MemorySourceConfig::try_new_exec(&partitions, pwmj_kv_schema::<K>(), None).unwrap()
 }
 
 fn pwmj_plan(
@@ -1430,37 +1488,43 @@ fn pwmj_plan(
     join_type: JoinType,
 ) -> Arc<dyn ExecutionPlan> {
     // Matches `PiecewiseMergeJoinExec::required_input_ordering`: descending for `<`/`<=`,
-    // ascending for `>`/`>=`, NULLs first either way.
-    let sort_options = match op {
-        Operator::Lt | Operator::LtEq => SortOptions::new(true, true),
-        Operator::Gt | Operator::GtEq => SortOptions::new(false, true),
-        other => panic!("not a range operator: {other:?}"),
+    // ascending for `>`/`>=`, NULLs first either way. Right existence joins require no
+    // ordering at all -- they only read the buffered side's min/max -- so they are fed the
+    // left side unsorted, which is the input shape they will see in a real plan.
+    let buffered = match join_type {
+        JoinType::RightSemi | JoinType::RightAnti => left,
+        _ => {
+            let sort_options = match op {
+                Operator::Lt | Operator::LtEq => SortOptions::new(true, true),
+                Operator::Gt | Operator::GtEq => SortOptions::new(false, true),
+                other => panic!("not a range operator: {other:?}"),
+            };
+            let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+                Arc::new(Column::new("k", 1)),
+                sort_options,
+            )])
+            .unwrap();
+            Arc::new(SortExec::new(ordering, left))
+        }
     };
-    let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
-        Arc::new(Column::new("k", 1)),
-        sort_options,
-    )])
-    .unwrap();
-    let sorted_left = Arc::new(SortExec::new(ordering, left));
     let on: (PhysicalExprRef, PhysicalExprRef) =
         (Arc::new(Column::new("k", 1)), Arc::new(Column::new("k", 1)));
     // `num_partitions` is 1 while the streamed side has up to 3: the final-pass counter must
     // come from the streamed side's partition count, not from this argument.
     Arc::new(
-        PiecewiseMergeJoinExec::try_new(sorted_left, right, on, op, join_type, 1)
-            .unwrap(),
+        PiecewiseMergeJoinExec::try_new(buffered, right, on, op, join_type, 1).unwrap(),
     )
 }
 
-fn pwmj_nlj_oracle_plan(
+fn pwmj_nlj_oracle_plan<K: PwmjFuzzKey>(
     left: Arc<dyn ExecutionPlan>,
     right: Arc<dyn ExecutionPlan>,
     op: Operator,
     join_type: JoinType,
 ) -> Arc<dyn ExecutionPlan> {
     let intermediate_schema = Schema::new(vec![
-        arrow::datatypes::Field::new("k", arrow::datatypes::DataType::Int32, true),
-        arrow::datatypes::Field::new("k", arrow::datatypes::DataType::Int32, true),
+        arrow::datatypes::Field::new("k", K::data_type(), true),
+        arrow::datatypes::Field::new("k", K::data_type(), true),
     ]);
     let expr = Arc::new(BinaryExpr::new(
         Arc::new(Column::new("k", 0)),
@@ -1486,7 +1550,10 @@ fn pwmj_nlj_oracle_plan(
 /// Executes every output partition concurrently and returns the join's rows as
 /// `(left id, right id)` pairs, sorted so partition interleaving does not affect the
 /// comparison. `None` means the join filled that side with NULLs, or -- for the existence
-/// joins, whose output carries the left side only -- that the side is absent entirely.
+/// joins, whose output carries only the surviving side -- that the other side is absent
+/// entirely. Both halves of an existence join output that surviving `id` as their first
+/// column: the left side's for `LeftSemi`/`LeftAnti`, the right side's for
+/// `RightSemi`/`RightAnti`.
 ///
 /// Concurrent rather than one partition at a time: the partitions share the watermark and race
 /// to be the one that runs the final pass, which is the part a sequential drain cannot reach.
@@ -1534,15 +1601,17 @@ async fn pwmj_collect_id_pairs(
 }
 
 /// Differential test for every join type `PiecewiseMergeJoin` supports, against a
-/// `NestedLoopJoin` oracle.
+/// `NestedLoopJoin` oracle, run once over `i32` keys and once over `f64` keys via
+/// [`run_pwmj_fuzz`].
 ///
 /// `Left`/`Full` are the ones with teeth: their unmatched buffered rows are derived from the
 /// shared `min_marked` watermark rather than materialized per row, and that encoding is only
 /// valid because every match marks a *suffix* of the buffered side. The dimensions the
 /// cheaper tests do not reach are the ones that matter here -- `pwmj.slt` and the unit tests
 /// both run the streamed side at a single partition and the default batch size, so neither
 /// covers several partitions racing to run the final pass, nor the mid-scan resume path a
-/// small batch size forces.
+/// small batch size forces. The `f64` pass adds `-0.0` and `NaN`, which `pwmj.slt` only
+/// exercises via a handful of fixed rows.
 #[tokio::test(flavor = "multi_thread")]
 async fn fuzz_pwmj_matches_nested_loop() {
     // A small batch size splits output across several coalesced batches even for these tiny
@@ -1551,6 +1620,13 @@ async fn fuzz_pwmj_matches_nested_loop() {
         TaskContext::default()
             .with_session_config(SessionConfig::new().with_batch_size(3)),
     );
+    // `f64` keys run every dimension below a second time with `-0.0` and `NaN` mixed in, on
+    // top of what `i32` already covers with plain duplicates and NULLs.
+    run_pwmj_fuzz::<i32>(&task_ctx).await;
+    run_pwmj_fuzz::<f64>(&task_ctx).await;
+}
+
+async fn run_pwmj_fuzz<K: PwmjFuzzKey + std::fmt::Debug>(task_ctx: &Arc<TaskContext>) {
     let ops = [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq];
     let join_types = [
         JoinType::Inner,
@@ -1559,6 +1635,8 @@ async fn fuzz_pwmj_matches_nested_loop() {
         JoinType::Full,
         JoinType::LeftSemi,
         JoinType::LeftAnti,
+        JoinType::RightSemi,
+        JoinType::RightAnti,
     ];
 
     for seed in 0..60u64 {
@@ -1568,10 +1646,14 @@ async fn fuzz_pwmj_matches_nested_loop() {
         // A narrow key range forces duplicates and equal-boundary cases.
         let key_range = rng.random_range(1..6i32);
         let nparts = rng.random_range(1..4usize);
+        // Independent of the streamed side's, so the two counts disagree across seeds and the
+        // 1-partition buffered case is still drawn. Only the right existence joins can use it:
+        // the others require the buffered side coalesced and globally sorted.
+        let buffered_nparts = rng.random_range(1..4usize);
 
-        let gen_keys = |n: usize, rng: &mut StdRng| -> Vec<Option<i32>> {
+        let gen_keys = |n: usize, rng: &mut StdRng| -> Vec<Option<K>> {
             (0..n)
-                .map(|_| (!rng.random_bool(0.2)).then(|| rng.random_range(0..key_range)))
+                .map(|_| (!rng.random_bool(0.2)).then(|| K::gen_value(key_range, rng)))
                 .collect()
         };
 
@@ -1582,31 +1664,43 @@ async fn fuzz_pwmj_matches_nested_loop() {
 
         for op in ops {
             for join_type in join_types {
+                // Fanned out only for the right existence joins, whose buffered side is
+                // folded one partition per task and then combined -- a reduction the
+                // single-partition shape below never reaches.
+                let buffered = match join_type {
+                    JoinType::RightSemi | JoinType::RightAnti => {
+                        pwmj_parts_exec(&left_ids, &left_keys, buffered_nparts)
+                    }
+                    _ => pwmj_single_exec(&left_ids, &left_keys),
+                };
                 let got = pwmj_collect_id_pairs(
                     pwmj_plan(
-                        pwmj_single_exec(&left_ids, &left_keys),
+                        buffered,
                         pwmj_parts_exec(&right_ids, &right_keys, nparts),
                         op,
                         join_type,
                     ),
-                    Arc::clone(&task_ctx),
+                    Arc::clone(task_ctx),
                 )
                 .await;
                 let want = pwmj_collect_id_pairs(
-                    pwmj_nlj_oracle_plan(
+                    pwmj_nlj_oracle_plan::<K>(
                         pwmj_single_exec(&left_ids, &left_keys),
                         pwmj_single_exec(&right_ids, &right_keys),
                         op,
                         join_type,
                     ),
-                    Arc::clone(&task_ctx),
+                    Arc::clone(task_ctx),
                 )
                 .await;
 
                 assert_eq!(
-                    got, want,
-                    "mismatch seed={seed} op={op:?} join_type={join_type:?} \
-                     nparts={nparts} left_keys={left_keys:?} right_keys={right_keys:?}"
+                    got,
+                    want,
+                    "mismatch key_type={} seed={seed} op={op:?} join_type={join_type:?} \
+                     nparts={nparts} buffered_nparts={buffered_nparts} \
+                      left_keys={left_keys:?} right_keys={right_keys:?}",
+                    std::any::type_name::<K>(),
                 );
             }
         }

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

@@ -32,8 +32,8 @@ use crate::execution_plan::{EmissionType, boundedness_from_children};
 use crate::joins::SharedBitmapBuilder;
 use crate::joins::utils::{
     BuildProbeJoinMetrics, ColumnIndex, JoinFilter, OnceAsync, OnceFut,
-    build_join_schema, check_join_is_valid, estimate_join_statistics,
-    need_produce_right_in_final,
+    boolean_mask_from_filter, build_join_schema, check_join_is_valid,
+    estimate_join_statistics, need_produce_right_in_final,
 };
 use crate::metrics::{
     Count, ExecutionPlanMetricsSet, MetricBuilder, MetricType, MetricsSet, RatioMetrics,
@@ -2988,20 +2988,6 @@ fn apply_filter_to_row_join_batch(
     Ok(bitmap_combined)
 }
 
-/// Convert a boolean filter array into a unified mask bitmap.
-///
-/// Caution: The filter result is NOT a bitmap; it contains true/false/null values.
-/// For example, `1 < NULL` evaluates to NULL. Therefore, we must combine (AND)
-/// the boolean array with its null bitmap to construct a unified bitmap.
-#[inline]
-fn boolean_mask_from_filter(filter_arr: &BooleanArray) -> BooleanArray {
-    let (values, nulls) = filter_arr.clone().into_parts();
-    match nulls {
-        Some(nulls) => BooleanArray::new(nulls.inner() & &values, None),
-        None => BooleanArray::new(values, None),
-    }
-}
-
 /// This function performs the following steps:
 /// 1. Apply filter to probe-side batch
 /// 2. Broadcast the left row (build_side_batch\[build_side_index\]) to the