less code

Author: AdamGS

datafusion/core/src/physical_planner.rs

@@ -3460,27 +3460,11 @@ 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()
@@ -3644,16 +3628,17 @@ mod tests {
         let config = SessionConfig::new()
             .with_target_partitions(4)
             .set_bool("datafusion.optimizer.prefer_hash_join", false);
+        let ctx = SessionContext::new_with_config(config);
 
-        let plan = plan_sql_with_config(query, config.clone()).await?;
+        let plan = ctx.sql(query).await?.create_physical_plan().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?;
+        let batches = ctx.sql(query).await?.collect().await?;
         assert_batches_eq!(
             &[
                 "+-------+",

datafusion/expr/src/logical_plan/builder.rs

@@ -1154,7 +1154,7 @@ impl LogicalPlanBuilder {
             .zip(right_keys)
             .map(|(l, r)| (Expr::Column(l), Expr::Column(r)))
             .collect();
-        let join_schema = build_join_schema_with_null_aware(
+        let join_schema = build_join_schema(
             self.plan.schema(),
             right.schema(),
             &join_type,
@@ -1677,25 +1677,13 @@ fn mark_field(schema: &DFSchema, nullable: bool) -> (Option<TableReference>, Arc
 /// Creates a schema for a join operation.
 /// The fields from the left side are first.
 ///
-/// The `mark` column of a mark join is non-nullable. Use
-/// [`build_join_schema_with_null_aware`] for null-aware mark joins, whose mark
-/// column can be `NULL` (SQL UNKNOWN).
+/// When `null_aware` is set, the `LeftMark`/`RightMark` `mark` column is made
+/// nullable so it can represent SQL UNKNOWN for null-aware `NOT IN` semantics.
+/// `null_aware` has no effect on non-mark join types.
 pub fn build_join_schema(
     left: &DFSchema,
     right: &DFSchema,
     join_type: &JoinType,
-) -> Result<DFSchema> {
-    build_join_schema_with_null_aware(left, right, join_type, false)
-}
-
-/// Like [`build_join_schema`], but makes the `LeftMark`/`RightMark` `mark`
-/// column nullable when `null_aware` is set, so it can represent SQL UNKNOWN
-/// for null-aware `NOT IN` semantics. `null_aware` has no effect on non-mark
-/// join types.
-pub fn build_join_schema_with_null_aware(
-    left: &DFSchema,
-    right: &DFSchema,
-    join_type: &JoinType,
     null_aware: bool,
 ) -> Result<DFSchema> {
     fn nullify_fields<'a>(
@@ -2933,13 +2921,13 @@ mod tests {
         )?;
 
         let join_schema =
-            build_join_schema(&left_schema, &right_schema, &JoinType::Left)?;
+            build_join_schema(&left_schema, &right_schema, &JoinType::Left, false)?;
         assert_eq!(
             join_schema.metadata(),
             &HashMap::from([("key".to_string(), "left".to_string())])
         );
         let join_schema =
-            build_join_schema(&left_schema, &right_schema, &JoinType::Right)?;
+            build_join_schema(&left_schema, &right_schema, &JoinType::Right, false)?;
         assert_eq!(
             join_schema.metadata(),
             &HashMap::from([("key".to_string(), "right".to_string())])

datafusion/expr/src/logical_plan/mod.rs

@@ -28,7 +28,7 @@ pub mod tree_node;
 
 pub use builder::{
     LogicalPlanBuilder, LogicalPlanBuilderOptions, LogicalTableSource, UNNAMED_TABLE,
-    build_join_schema, build_join_schema_with_null_aware, requalify_sides_if_needed,
+    build_join_schema, requalify_sides_if_needed,
     table_scan, union, wrap_projection_for_join_if_necessary,
 };
 pub use ddl::{

datafusion/expr/src/logical_plan/plan.rs

@@ -47,7 +47,7 @@ use crate::utils::{
 use crate::{
     BinaryExpr, CreateMemoryTable, CreateView, Execute, Expr, ExprSchemable, GroupingSet,
     LogicalPlanBuilder, Operator, Prepare, TableProviderFilterPushDown, TableSource,
-    WindowFunctionDefinition, build_join_schema_with_null_aware, expr_vec_fmt,
+    WindowFunctionDefinition, build_join_schema, expr_vec_fmt,
     requalify_sides_if_needed,
 };
 
@@ -668,7 +668,7 @@ impl LogicalPlan {
                 null_equality,
                 null_aware,
             }) => {
-                let schema = build_join_schema_with_null_aware(
+                let schema = build_join_schema(
                     left.schema(),
                     right.schema(),
                     &join_type,
@@ -949,7 +949,7 @@ impl LogicalPlan {
                 ..
             }) => {
                 let (left, right) = self.only_two_inputs(inputs)?;
-                let schema = build_join_schema_with_null_aware(
+                let schema = build_join_schema(
                     left.schema(),
                     right.schema(),
                     join_type,
@@ -4279,7 +4279,7 @@ impl Join {
         null_equality: NullEquality,
         null_aware: bool,
     ) -> Result<Self> {
-        let join_schema = build_join_schema_with_null_aware(
+        let join_schema = build_join_schema(
             left.schema(),
             right.schema(),
             &join_type,
@@ -4335,7 +4335,7 @@ impl Join {
             .map(|(l, r)| (Expr::Column(l), Expr::Column(r)))
             .collect();
 
-        let join_schema = build_join_schema_with_null_aware(
+        let join_schema = build_join_schema(
             left_sch.schema(),
             right_sch.schema(),
             &original_join.join_type,

datafusion/optimizer/src/decorrelate_predicate_subquery.rs

@@ -627,14 +627,6 @@ 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_left_mark_join_with_null_aware(
         plan: &LogicalPlan,
         expected_null_aware: bool,
@@ -1287,81 +1279,33 @@ 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_left_mark_join_with_null_aware(&optimized, true),
-            "{}",
-            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_left_mark_join_with_null_aware(&optimized, false),
-            "{}",
-            optimized.display_indent_schema()
-        );
-
-        Ok(())
-    }
-
-    #[test]
-    fn correlated_not_in_mark_join_not_null_aware_for_non_nullable_value_key()
-    -> Result<()> {
-        // The value key `id` is non-nullable on both sides; only the
-        // correlation key `grp` is nullable. A NULL correlation key just makes
-        // the correlated subquery empty, so NOT IN can never be UNKNOWN and
-        // the mark join does not need null-aware semantics.
+    /// Builds `... WHERE (id NOT IN (SELECT id FROM inner WHERE <grp filter>))
+    /// IS NULL` over scans whose value key `id` has the given nullability (the
+    /// correlation key `grp` is always nullable), optimizes it, and asserts the
+    /// resulting `LeftMark` join's `null_aware` flag.
+    fn assert_correlated_not_in_mark_join_null_aware(
+        value_key_nullable: bool,
+        hashable_filter: bool,
+        expected_null_aware: bool,
+    ) -> Result<()> {
         let schema = Schema::new(vec![
-            Field::new("id", DataType::Int32, false),
+            Field::new("id", DataType::Int32, value_key_nullable),
             Field::new("grp", DataType::Int32, true),
         ]);
         let outer_scan = table_scan(Some("outer_t"), &schema, None)?.build()?;
         let inner_scan = table_scan(Some("inner_t"), &schema, None)?.build()?;
 
+        let outer_grp = out_ref_col(DataType::Int32, "outer_t.grp");
+        let inner_grp = col("inner_t.grp");
+        let grp_filter = if hashable_filter {
+            outer_grp.eq(inner_grp)
+        } else {
+            outer_grp.lt(inner_grp)
+        };
+
         let subquery = Arc::new(
             LogicalPlanBuilder::from(inner_scan)
-                .filter(
-                    out_ref_col(DataType::Int32, "outer_t.grp").eq(col("inner_t.grp")),
-                )?
+                .filter(grp_filter)?
                 .project(vec![col("inner_t.id")])?
                 .build()?,
         );
@@ -1372,14 +1316,29 @@ mod tests {
 
         let optimized = optimize_with_decorrelate(plan)?;
         assert!(
-            has_left_mark_join_with_null_aware(&optimized, false),
+            has_left_mark_join_with_null_aware(&optimized, expected_null_aware),
             "{}",
             optimized.display_indent_schema()
         );
 
         Ok(())
     }
 
+    #[test]
+    fn correlated_not_in_mark_join_null_awareness() -> Result<()> {
+        // Hashable (`=`) correlation filter on a nullable value key: NOT IN can
+        // be UNKNOWN, so the mark join is null-aware.
+        assert_correlated_not_in_mark_join_null_aware(true, true, true)?;
+        // Residual (`<`) correlation filter: the equijoin predicate can't be
+        // extracted, so the mark join is not null-aware.
+        assert_correlated_not_in_mark_join_null_aware(true, false, false)?;
+        // Non-nullable value key `id`: a NULL correlation key just makes the
+        // subquery empty, so NOT IN can never be UNKNOWN and the join is not
+        // null-aware.
+        assert_correlated_not_in_mark_join_null_aware(false, true, false)?;
+        Ok(())
+    }
+
     #[test]
     fn in_subquery_both_side_expr() -> Result<()> {
         let table_scan = test_table_scan()?;

datafusion/optimizer/src/eliminate_cross_join.rs

@@ -374,6 +374,7 @@ fn find_inner_join(
                 left_input.schema(),
                 right_input.schema(),
                 &JoinType::Inner,
+                false,
             )?);
 
             return Ok(LogicalPlan::Join(Join {
@@ -397,6 +398,7 @@ fn find_inner_join(
         left_input.schema(),
         right.schema(),
         &JoinType::Inner,
+        false,
     )?);
 
     Ok(LogicalPlan::Join(Join {
@@ -1402,6 +1404,7 @@ mod tests {
             t1.schema(),
             t2.schema(),
             &JoinType::Inner,
+            false,
         )?);
 
         let inner_join = LogicalPlan::Join(Join {

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -1024,7 +1024,8 @@ mod tests {
             let left_schema = left_child.schema();
             let right_schema = right_child.schema();
             let schema = Arc::new(
-                build_join_schema(left_schema, right_schema, &JoinType::Inner).unwrap(),
+                build_join_schema(left_schema, right_schema, &JoinType::Inner, false)
+                    .unwrap(),
             );
             Self {
                 exprs: vec![],