step 2

Author: AdamGS

datafusion/expr/src/logical_plan/plan.rs

@@ -3909,13 +3909,13 @@ pub struct Join {
     pub schema: DFSchemaRef,
     /// Defines the null equality for the join.
     pub null_equality: NullEquality,
-    /// Whether this is a null-aware anti join (for NOT IN semantics).
+    /// Whether this join needs null-aware NOT IN semantics.
     ///
-    /// Only applies to LeftAnti joins. When true, implements SQL NOT IN semantics where:
-    /// - If the right side (subquery) contains any NULL in join keys, no rows are output
-    /// - Left side rows with NULL in join keys are not output
+    /// For `LeftAnti`, if the right side contains any NULL in join keys, no rows are output and
+    /// left rows with NULL join keys are also excluded.
     ///
-    /// This is required for correct NOT IN subquery behavior with three-valued logic.
+    /// For `LeftMark`, the generated `mark` column becomes nullable so unmatched rows can produce
+    /// `NULL` rather than `false` when SQL three-valued logic requires it.
     pub null_aware: bool,
 }
 
@@ -3934,7 +3934,7 @@ impl Join {
     /// * `join_type` - Type of join (Inner, Left, Right, etc.)
     /// * `join_constraint` - Join constraint (On, Using)
     /// * `null_equality` - How to handle nulls in join comparisons
-    /// * `null_aware` - Whether this is a null-aware anti join (for NOT IN semantics)
+    /// * `null_aware` - Whether this join needs null-aware NOT IN semantics
     ///
     /// # Returns
     ///

datafusion/optimizer/src/decorrelate_predicate_subquery.rs

@@ -371,6 +371,8 @@ fn build_join(
         .values()
         .for_each(|cols| all_correlated_cols.extend(cols.clone()));
 
+    let has_correlated_join_filter = !pull_up.join_filters.is_empty();
+
     // alias the join filter
     let join_filter_opt = conjunction(pull_up.join_filters)
         .map_or(Ok(None), |filter| {
@@ -440,9 +442,27 @@ fn build_join(
             sub_query_alias.clone()
         };
 
-        // Mark joins don't use null-aware semantics (they use three-valued logic with mark column)
+        // For simple uncorrelated NOT IN disjunctions, propagate null-aware semantics into the
+        // nullable mark column. Correlated mark joins still use the legacy path because the
+        // runtime state is global to the probe side rather than per-left-row.
+        let null_aware = join_type == JoinType::LeftMark
+            && in_predicate_opt.is_some()
+            && !has_correlated_join_filter
+            && join_keys_may_be_null(
+                &join_filter,
+                left.schema(),
+                right_projected.schema(),
+            )?;
+
         let new_plan = LogicalPlanBuilder::from(left.clone())
-            .join_on(right_projected, join_type, Some(join_filter))?
+            .join_detailed_with_options(
+                right_projected,
+                join_type,
+                (Vec::<Column>::new(), Vec::<Column>::new()),
+                Some(join_filter),
+                NullEquality::NullEqualsNothing,
+                null_aware,
+            )?
             .build()?;
 
         debug!(
@@ -461,13 +481,9 @@ fn build_join(
     //
     // Additionally, if the join keys are non-nullable on both sides, we don't need
     // null-aware semantics because NULLs cannot exist in the data.
-    let null_aware = dbg!(matches!(join_type, JoinType::LeftAnti | JoinType::LeftMark))
-        && dbg!(in_predicate_opt.is_some())
-        && dbg!(join_keys_may_be_null(
-            &join_filter,
-            left.schema(),
-            sub_query_alias.schema()
-        )?);
+    let null_aware = matches!(join_type, JoinType::LeftAnti)
+        && in_predicate_opt.is_some()
+        && join_keys_may_be_null(&join_filter, left.schema(), sub_query_alias.schema())?;
 
     // join our sub query into the main plan
     let new_plan = if null_aware {
@@ -1740,8 +1756,8 @@ mod tests {
             plan,
             @r"
         Projection: customer.c_custkey [c_custkey:Int64]
-          Filter: __correlated_sq_1.mark OR customer.c_custkey = Int32(1) [c_custkey:Int64, c_name:Utf8, mark:Boolean]
-            LeftMark Join:  Filter: Boolean(true) [c_custkey:Int64, c_name:Utf8, mark:Boolean]
+          Filter: __correlated_sq_1.mark OR customer.c_custkey = Int32(1) [c_custkey:Int64, c_name:Utf8, mark:Boolean;N]
+            LeftMark Join:  Filter: Boolean(true) [c_custkey:Int64, c_name:Utf8, mark:Boolean;N]
               TableScan: customer [c_custkey:Int64, c_name:Utf8]
               SubqueryAlias: __correlated_sq_1 [o_custkey:Int64]
                 Projection: orders.o_custkey [o_custkey:Int64]

datafusion/physical-plan/src/joins/hash_join/exec.rs

@@ -215,12 +215,8 @@ pub(super) struct JoinLeftData {
     /// Shared atomic flag indicating if any probe partition saw data (for null-aware anti/mark joins)
     /// This is shared across all probe partitions to provide global knowledge
     pub(super) probe_side_non_empty: AtomicBool,
-    /// Shared atomic flag indicating if any probe partition saw NULL in join keys (for null-aware anti joins)
+    /// Shared atomic flag indicating if any probe partition saw NULL in join keys
     pub(super) probe_side_has_null: AtomicBool,
-    /// Shared atomic flag indicating if any build partition saw NULL in join keys (for null-aware mark joins)
-    pub(super) build_side_has_nulls: AtomicBool,
-    /// Not sure how to use this yet
-    pub(super) build_side_is_empty: AtomicBool,
 }
 
 impl JoinLeftData {
@@ -409,15 +405,15 @@ impl HashJoinExecBuilder {
         // Validate null_aware flag
         if exec.null_aware {
             let join_type = exec.join_type();
-            if !matches!(join_type, JoinType::LeftAnti) {
+            if !matches!(join_type, JoinType::LeftAnti | JoinType::LeftMark) {
                 return plan_err!(
-                    "null_aware can only be true for LeftAnti joins, got {join_type}"
+                    "null_aware can only be true for LeftAnti or LeftMark joins, got {join_type}"
                 );
             }
             let on = exec.on();
             if on.len() != 1 {
                 return plan_err!(
-                    "null_aware anti join only supports single column join key, got {} columns",
+                    "null_aware joins only support single column join key, got {} columns",
                     on.len()
                 );
             }
@@ -2079,9 +2075,6 @@ async fn collect_left_input(
         bounds = None;
     }
 
-    let build_side_has_nulls = batch.columns().iter().any(|col| col.null_count() > 0);
-    let build_side_is_empty = batch.num_rows() == 0;
-
     let data = JoinLeftData {
         map,
         batch,
@@ -2093,8 +2086,6 @@ async fn collect_left_input(
         membership,
         probe_side_non_empty: AtomicBool::new(false),
         probe_side_has_null: AtomicBool::new(false),
-        build_side_has_nulls: AtomicBool::new(build_side_has_nulls),
-        build_side_is_empty: AtomicBool::new(build_side_is_empty),
     };
 
     Ok(data)
@@ -6067,7 +6058,7 @@ mod tests {
         Ok(())
     }
 
-    /// Test that null_aware validation rejects non-LeftAnti join types
+    /// Test that null_aware validation rejects unsupported join types
     #[tokio::test]
     async fn test_null_aware_validation_wrong_join_type() {
         let left =
@@ -6098,7 +6089,7 @@ mod tests {
             result
                 .unwrap_err()
                 .to_string()
-                .contains("null_aware can only be true for LeftAnti joins")
+                .contains("null_aware can only be true for LeftAnti or LeftMark joins")
         );
     }
 
@@ -6138,8 +6129,114 @@ mod tests {
             result
                 .unwrap_err()
                 .to_string()
-                .contains("null_aware anti join only supports single column join key")
+                .contains("null_aware joins only support single column join key")
+        );
+    }
+
+    /// Test null-aware left mark join when probe side contains NULL.
+    /// Expected:
+    /// - matched rows => true
+    /// - unmatched non-NULL rows => NULL
+    /// - NULL build keys with non-empty probe side => NULL
+    #[apply(hash_join_exec_configs)]
+    #[tokio::test]
+    async fn test_null_aware_left_mark_probe_null(batch_size: usize) -> Result<()> {
+        let task_ctx = prepare_task_ctx(batch_size, false);
+
+        let left = build_table_two_cols(
+            ("c1", &vec![Some(1), Some(4), None]),
+            ("dummy", &vec![Some(10), Some(40), Some(0)]),
+        );
+
+        let right = build_table_two_cols(
+            ("c2", &vec![Some(1), Some(2), None]),
+            ("dummy", &vec![Some(100), Some(200), Some(300)]),
+        );
+
+        let on = vec![(
+            Arc::new(Column::new_with_schema("c1", &left.schema())?) as _,
+            Arc::new(Column::new_with_schema("c2", &right.schema())?) as _,
+        )];
+
+        let join = HashJoinExec::try_new(
+            left,
+            right,
+            on,
+            None,
+            &JoinType::LeftMark,
+            None,
+            PartitionMode::CollectLeft,
+            NullEquality::NullEqualsNothing,
+            true, // null_aware = true
+        )?;
+
+        let stream = join.execute(0, task_ctx)?;
+        let batches = common::collect(stream).await?;
+
+        allow_duplicates! {
+            assert_snapshot!(batches_to_sort_string(&batches), @r"
+            +----+-------+------+
+            | c1 | dummy | mark |
+            +----+-------+------+
+            |    | 0     |      |
+            | 1  | 10    | true |
+            | 4  | 40    |      |
+            +----+-------+------+
+            ");
+        }
+
+        Ok(())
+    }
+
+    /// Test null-aware left mark join when probe side is empty.
+    /// Expected: all rows are marked false, including NULL build keys.
+    #[apply(hash_join_exec_configs)]
+    #[tokio::test]
+    async fn test_null_aware_left_mark_empty_probe(batch_size: usize) -> Result<()> {
+        let task_ctx = prepare_task_ctx(batch_size, false);
+
+        let left = build_table_two_cols(
+            ("c1", &vec![Some(1), None]),
+            ("dummy", &vec![Some(10), Some(0)]),
         );
+
+        let right = build_table_two_cols(
+            ("c2", &Vec::<Option<i32>>::new()),
+            ("dummy", &Vec::<Option<i32>>::new()),
+        );
+
+        let on = vec![(
+            Arc::new(Column::new_with_schema("c1", &left.schema())?) as _,
+            Arc::new(Column::new_with_schema("c2", &right.schema())?) as _,
+        )];
+
+        let join = HashJoinExec::try_new(
+            left,
+            right,
+            on,
+            None,
+            &JoinType::LeftMark,
+            None,
+            PartitionMode::CollectLeft,
+            NullEquality::NullEqualsNothing,
+            true, // null_aware = true
+        )?;
+
+        let stream = join.execute(0, task_ctx)?;
+        let batches = common::collect(stream).await?;
+
+        allow_duplicates! {
+            assert_snapshot!(batches_to_sort_string(&batches), @r"
+            +----+-------+-------+
+            | c1 | dummy | mark  |
+            +----+-------+-------+
+            |    | 0     | false |
+            | 1  | 10    | false |
+            +----+-------+-------+
+            ");
+        }
+
+        Ok(())
     }
 
     #[test]