codex help

Signed-off-by: Adam Gutglick <adamgsal@gmail.com>

Author: AdamGS

datafusion/common/src/join_type.rs

@@ -61,9 +61,10 @@ pub enum JoinType {
     /// [1]. This join type is used to decorrelate EXISTS subqueries used inside disjunctive
     /// predicates.
     ///
-    /// Note: This we currently do not implement the full null semantics for the mark join described
-    /// in [1] which will be needed if we and ANY subqueries. In our version the mark column will
-    /// only be true for had a match and false when no match was found, never null.
+    /// For scalar `NOT IN`, DataFusion can plan a null-aware hash mark join where the
+    /// mark column is nullable: TRUE for a match, NULL for SQL UNKNOWN, and FALSE
+    /// otherwise. Row-valued multi-column `NOT IN` and non-hash residual predicate
+    /// null-aware mark semantics are not implemented.
     ///
     /// [1]: http://btw2017.informatik.uni-stuttgart.de/slidesandpapers/F1-10-37/paper_web.pdf
     LeftMark,

datafusion/core/src/physical_planner.rs

@@ -3384,11 +3384,27 @@ mod tests {
         ctx.sql(query).await?.create_physical_plan().await
     }
 
+    async fn plan_sql_with_config(
+        query: &str,
+        config: SessionConfig,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let ctx = SessionContext::new_with_config(config);
+        ctx.sql(query).await?.create_physical_plan().await
+    }
+
     async fn collect_sql(query: &str) -> Result<Vec<RecordBatch>> {
         let ctx = SessionContext::new();
         ctx.sql(query).await?.collect().await
     }
 
+    async fn collect_sql_with_config(
+        query: &str,
+        config: SessionConfig,
+    ) -> Result<Vec<RecordBatch>> {
+        let ctx = SessionContext::new_with_config(config);
+        ctx.sql(query).await?.collect().await
+    }
+
     #[tokio::test]
     async fn test_all_operators() -> Result<()> {
         let logical_plan = test_csv_scan()
@@ -3523,6 +3539,62 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn correlated_not_in_is_null_uses_null_aware_hash_mark_join() -> Result<()> {
+        let query = "
+            SELECT value
+            FROM (
+              VALUES
+                (1, 1, 'a'),
+                (3, 1, 'b'),
+                (1, 2, 'c'),
+                (NULL, 1, 'd'),
+                (5, 3, 'e'),
+                (2, 1, 'f'),
+                (NULL, 2, 'g')
+            ) AS outer_corr_table(id, grp, value)
+            WHERE (id NOT IN (
+              SELECT id
+              FROM (
+                VALUES
+                  (2, 1),
+                  (NULL, 1),
+                  (1, 2)
+              ) AS inner_corr_table(id, grp)
+              WHERE inner_corr_table.grp = outer_corr_table.grp
+            )) IS NULL
+            ORDER BY value";
+
+        let config = SessionConfig::new()
+            .with_target_partitions(4)
+            .set_bool("datafusion.optimizer.prefer_hash_join", false);
+
+        let plan = plan_sql_with_config(query, config.clone()).await?;
+        let formatted = displayable(plan.as_ref()).indent(true).to_string();
+        assert_contains!(
+            &formatted,
+            "HashJoinExec: mode=CollectLeft, join_type=LeftMark"
+        );
+        assert!(!formatted.contains("SortMergeJoinExec"), "{formatted}");
+
+        let batches = collect_sql_with_config(query, config).await?;
+        assert_batches_eq!(
+            &[
+                "+-------+",
+                "| value |",
+                "+-------+",
+                "| a     |",
+                "| b     |",
+                "| d     |",
+                "| g     |",
+                "+-------+",
+            ],
+            &batches
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn scalar_subquery_in_projection_and_filter_plans() -> Result<()> {
         let plan = plan_sql(

datafusion/optimizer/src/decorrelate_predicate_subquery.rs

@@ -21,6 +21,7 @@ use std::ops::Deref;
 use std::sync::Arc;
 
 use crate::decorrelate::PullUpCorrelatedExpr;
+use crate::extract_equijoin_predicate::split_eq_and_noneq_join_predicate;
 use crate::optimizer::ApplyOrder;
 use crate::utils::replace_qualified_name;
 use crate::{OptimizerConfig, OptimizerRule};
@@ -379,8 +380,6 @@ 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| {
@@ -450,12 +449,25 @@ fn build_join(
             sub_query_alias.clone()
         };
 
-        // 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 mark_filter_is_hashable_only =
+            if join_type == JoinType::LeftMark && in_predicate_opt.is_some() {
+                let (_, residual_filter) = split_eq_and_noneq_join_predicate(
+                    join_filter.clone(),
+                    left.schema(),
+                    right_projected.schema(),
+                )?;
+                residual_filter.is_none()
+            } else {
+                false
+            };
+
+        // For scalar NOT IN mark joins, propagate null-aware semantics into the
+        // nullable mark column when the predicate can be implemented by hash keys.
+        // Non-equality correlated filters stay on the legacy path because hash join
+        // execution cannot mark UNKNOWN candidates for residual predicates.
         let null_aware = join_type == JoinType::LeftMark
             && in_predicate_opt.is_some()
-            && !has_correlated_join_filter
+            && mark_filter_is_hashable_only
             && join_keys_may_be_null(
                 &join_filter,
                 left.schema(),
@@ -591,6 +603,43 @@ mod tests {
         ))
     }
 
+    fn nullable_scalar_mark_scan(name: &str) -> Result<LogicalPlan> {
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int32, true),
+            Field::new("grp", DataType::Int32, true),
+        ]);
+        table_scan(Some(name), &schema, None)?.build()
+    }
+
+    fn has_null_aware_left_mark_join(plan: &LogicalPlan) -> bool {
+        if let LogicalPlan::Join(join) = plan
+            && join.join_type == JoinType::LeftMark
+        {
+            return join.null_aware;
+        }
+
+        plan.inputs().into_iter().any(has_null_aware_left_mark_join)
+    }
+
+    fn has_non_null_aware_left_mark_join(plan: &LogicalPlan) -> bool {
+        if let LogicalPlan::Join(join) = plan
+            && join.join_type == JoinType::LeftMark
+        {
+            return !join.null_aware;
+        }
+
+        plan.inputs()
+            .into_iter()
+            .any(has_non_null_aware_left_mark_join)
+    }
+
+    fn optimize_with_decorrelate(plan: LogicalPlan) -> Result<LogicalPlan> {
+        let optimizer = crate::Optimizer::with_rules(vec![Arc::new(
+            DecorrelatePredicateSubquery::new(),
+        )]);
+        optimizer.optimize(plan, &crate::OptimizerContext::new(), |_, _| {})
+    }
+
     /// Test for several IN subquery expressions
     #[test]
     fn in_subquery_multiple() -> Result<()> {
@@ -1221,6 +1270,62 @@ mod tests {
         )
     }
 
+    #[test]
+    fn correlated_not_in_mark_join_is_null_aware_for_hashable_filter() -> Result<()> {
+        let outer_scan = nullable_scalar_mark_scan("outer_t")?;
+        let inner_scan = nullable_scalar_mark_scan("inner_t")?;
+
+        let subquery = Arc::new(
+            LogicalPlanBuilder::from(inner_scan)
+                .filter(
+                    out_ref_col(DataType::Int32, "outer_t.grp").eq(col("inner_t.grp")),
+                )?
+                .project(vec![col("inner_t.id")])?
+                .build()?,
+        );
+
+        let plan = LogicalPlanBuilder::from(outer_scan)
+            .filter(not_in_subquery(col("outer_t.id"), subquery).is_null())?
+            .build()?;
+
+        let optimized = optimize_with_decorrelate(plan)?;
+        assert!(
+            has_null_aware_left_mark_join(&optimized),
+            "{}",
+            optimized.display_indent_schema()
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn correlated_not_in_mark_join_is_not_null_aware_for_residual_filter() -> Result<()> {
+        let outer_scan = nullable_scalar_mark_scan("outer_t")?;
+        let inner_scan = nullable_scalar_mark_scan("inner_t")?;
+
+        let subquery = Arc::new(
+            LogicalPlanBuilder::from(inner_scan)
+                .filter(
+                    out_ref_col(DataType::Int32, "outer_t.grp").lt(col("inner_t.grp")),
+                )?
+                .project(vec![col("inner_t.id")])?
+                .build()?,
+        );
+
+        let plan = LogicalPlanBuilder::from(outer_scan)
+            .filter(not_in_subquery(col("outer_t.id"), subquery).is_null())?
+            .build()?;
+
+        let optimized = optimize_with_decorrelate(plan)?;
+        assert!(
+            has_non_null_aware_left_mark_join(&optimized),
+            "{}",
+            optimized.display_indent_schema()
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn in_subquery_both_side_expr() -> Result<()> {
         let table_scan = test_table_scan()?;

datafusion/optimizer/src/extract_equijoin_predicate.rs

@@ -174,7 +174,7 @@ impl OptimizerRule for ExtractEquijoinPredicate {
 ///
 /// According to the above rule, `expr1` is the equijoin predicate, while `expr2` and `expr3` are not.
 /// The function returns Ok(\[expr1\], Some(expr2 AND expr3))
-fn split_eq_and_noneq_join_predicate(
+pub(crate) fn split_eq_and_noneq_join_predicate(
     filter: Expr,
     left_schema: &DFSchema,
     right_schema: &DFSchema,