fix(scheduler): execute null-aware anti joins in one task

ballista/client/tests/null_aware.rs

@@ -0,0 +1,138 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#[cfg(test)]
+#[cfg(feature = "standalone")]
+mod null_aware {
+    use std::fs;
+    use std::path::Path;
+
+    use ballista::prelude::SessionContextExt;
+    use datafusion::arrow::util::pretty::pretty_format_batches;
+    use datafusion::prelude::*;
+
+    fn write_tables(dir: &Path, t2_has_null: bool) {
+        // t1: 4 files -> 4 partitions, values 0..19
+        let t1 = dir.join("t1");
+        fs::create_dir_all(&t1).unwrap();
+        for i in 0..4i32 {
+            let mut s = String::from("a\n");
+            for v in 0..5i32 {
+                s.push_str(&format!("{}\n", i * 5 + v));
+            }
+            fs::write(t1.join(format!("p{i}.csv")), s).unwrap();
+        }
+
+        // t2: 4 files -> 4 partitions. One file optionally holds a NULL key.
+        // The second column keeps the empty field unambiguous, since a bare
+        // blank line is skipped by the CSV reader rather than read as NULL.
+        let t2 = dir.join("t2");
+        fs::create_dir_all(&t2).unwrap();
+        fs::write(t2.join("p0.csv"), "b,tag\n0,x\n1,x\n").unwrap();
+        fs::write(
+            t2.join("p1.csv"),
+            if t2_has_null {
+                "b,tag\n,x\n"
+            } else {
+                "b,tag\n2,x\n"
+            },
+        )
+        .unwrap();
+        fs::write(t2.join("p2.csv"), "b,tag\n3,x\n").unwrap();
+        fs::write(t2.join("p3.csv"), "b,tag\n4,x\n").unwrap();
+    }
+
+    async fn register(ctx: &SessionContext, dir: &Path) {
+        for name in ["t1", "t2"] {
+            ctx.register_csv(
+                name,
+                dir.join(name).to_str().unwrap(),
+                CsvReadOptions::new()
+                    .has_header(true)
+                    .file_extension(".csv"),
+            )
+            .await
+            .unwrap();
+        }
+    }
+
+    const QUERY: &str = "select a from t1 where a not in (select b from t2) order by a";
+
+    async fn run_case(case: &str, t2_has_null: bool) {
+        let dir = std::env::temp_dir().join(format!("ballista_null_aware_{case}"));
+        let _ = fs::remove_dir_all(&dir);
+        fs::create_dir_all(&dir).unwrap();
+        write_tables(&dir, t2_has_null);
+
+        let df_ctx = SessionContext::new_with_config(
+            SessionConfig::new().with_target_partitions(4),
+        );
+        register(&df_ctx, &dir).await;
+        let df_batches = df_ctx.sql(QUERY).await.unwrap().collect().await.unwrap();
+        let expected = pretty_format_batches(&df_batches).unwrap().to_string();
+
+        let mut mismatches = vec![];
+        // The default Ballista setting selects SortMergeJoinExec and loses the
+        // null-aware flag before scheduler lowering, so this test covers the
+        // hash-join planning paths that retain the flag.
+        for variant in ["prefer_hash_join", "aqe"] {
+            let ctx = SessionContext::standalone().await.unwrap();
+            if variant == "aqe" {
+                ctx.sql("SET ballista.planner.adaptive.enabled = true")
+                    .await
+                    .unwrap()
+                    .collect()
+                    .await
+                    .unwrap();
+            }
+            ctx.sql("SET datafusion.optimizer.prefer_hash_join = true")
+                .await
+                .unwrap()
+                .collect()
+                .await
+                .unwrap();
+            register(&ctx, &dir).await;
+
+            match ctx.sql(QUERY).await.unwrap().collect().await {
+                Ok(batches) => {
+                    let actual = pretty_format_batches(&batches).unwrap().to_string();
+                    if actual.trim() != expected.trim() {
+                        mismatches.push(format!(
+                            "[{case}/{variant}] expected:\n{expected}\nactual:\n{actual}"
+                        ));
+                    }
+                }
+                Err(error) => {
+                    mismatches.push(format!("[{case}/{variant}] failed: {error}"));
+                }
+            }
+        }
+
+        let _ = fs::remove_dir_all(&dir);
+        assert!(mismatches.is_empty(), "{}", mismatches.join("\n\n"));
+    }
+
+    #[tokio::test]
+    async fn not_in_with_null_in_subquery() {
+        run_case("with_null", true).await;
+    }
+
+    #[tokio::test]
+    async fn not_in_without_null_in_subquery() {
+        run_case("without_null", false).await;
+    }
+}

ballista/core/src/config.rs

@@ -97,7 +97,9 @@ pub const BALLISTA_SHUFFLE_SORT_BASED_MEMORY_LIMIT_PER_TASK_BYTES: &str =
     "ballista.shuffle.sort_based.memory_limit_per_task_bytes";
 /// Configuration key for the byte-size threshold below which a hash join's
 /// smaller side is promoted to `CollectLeft` and lowered via the broadcast
-/// pattern in the distributed planner. Set to `0` to disable promotion.
+/// pattern in the distributed planner. It also caps null-aware anti joins with
+/// a known build size because they require single-task `CollectLeft` execution.
+/// Set to `0` to disable promotion and reject null-aware anti joins.
 pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES: &str =
     "ballista.optimizer.broadcast_join_threshold_bytes";
 
@@ -255,8 +257,10 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
                          "Byte-size threshold below which a hash join's smaller side is \
                           promoted to CollectLeft and lowered via the broadcast pattern. \
                           Governs broadcast selection under both the static distributed \
-                          planner and adaptive query planning (AQE). Set to 0 to disable \
-                          promotion.".to_string(),
+                          planner and adaptive query planning (AQE). It also caps \
+                          null-aware anti joins with a known build size because they require \
+                          single-task CollectLeft execution. Set to 0 to disable promotion \
+                          and reject null-aware anti joins.".to_string(),
                          DataType::UInt64,
                          Some((10 * 1024 * 1024).to_string())),
         ConfigEntry::new(BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS.to_string(),

ballista/scheduler/src/physical_optimizer/join_selection.rs

@@ -322,10 +322,10 @@ pub(crate) fn partitioned_hash_join(
     {
         hash_join.swap_inputs(PartitionMode::Partitioned)
     } else {
-        // Null-aware anti joins must use CollectLeft mode because they track probe-side state
-        // (probe_side_non_empty, probe_side_has_null) per-partition, but need global knowledge
-        // for correct null handling. With partitioning, a partition might not see probe rows
-        // even if the probe side is globally non-empty, leading to incorrect NULL row handling.
+        // Keep this copied rule aligned with upstream DataFusion. An Auto join
+        // that cannot be selected by size still requires CollectLeft when it is
+        // null-aware; Ballista's distributed planner later coalesces the probe
+        // side so this mode runs in exactly one task.
         let partition_mode = if hash_join.null_aware {
             PartitionMode::CollectLeft
         } else {
@@ -368,7 +368,12 @@ fn statistical_join_selection_subrule(
             PartitionMode::Partitioned => {
                 let left = hash_join.left();
                 let right = hash_join.right();
+                // Keep this branch aligned with upstream DataFusion: a
+                // null-aware join is not swapped, but its partition mode is not
+                // changed here. Ballista's distributed planner is responsible
+                // for lowering it to a single-task CollectLeft join.
                 if hash_join.join_type().supports_swap()
+                    && !hash_join.null_aware
                     && should_swap_join_order(&**left, &**right)?
                 {
                     hash_join
@@ -583,6 +588,7 @@ pub fn hash_join_swap_subrule(
     if let Some(hash_join) = input.downcast_ref::<HashJoinExec>()
         && hash_join.left.boundedness().is_unbounded()
         && !hash_join.right.boundedness().is_unbounded()
+        && !hash_join.null_aware
         && matches!(
             *hash_join.join_type(),
             JoinType::Inner | JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti
@@ -772,6 +778,122 @@ mod test {
         );
     }
 
+    #[test]
+    fn partitioned_null_aware_anti_join_is_not_swapped() {
+        use datafusion::{
+            common::NullEquality,
+            physical_optimizer::PhysicalOptimizerRule,
+            physical_plan::joins::{HashJoinExec, PartitionMode},
+        };
+
+        use crate::physical_optimizer::join_selection::JoinSelection;
+
+        // The large left and small right sides would normally be swapped by
+        // statistical join selection. Keep this rule aligned with upstream
+        // DataFusion: a null-aware anti join retains both its orientation and
+        // its existing partition mode. Distributed lowering happens later.
+        let (big, small) = create_big_and_small();
+        let join = Arc::new(
+            HashJoinExec::try_new(
+                Arc::clone(&big),
+                Arc::clone(&small),
+                vec![(
+                    Arc::new(Column::new("big_col", 0)) as _,
+                    Arc::new(Column::new("small_col", 0)) as _,
+                )],
+                None,
+                &JoinType::LeftAnti,
+                None,
+                PartitionMode::Partitioned,
+                NullEquality::NullEqualsNothing,
+                true,
+            )
+            .unwrap(),
+        ) as Arc<dyn ExecutionPlan>;
+
+        let optimized = JoinSelection::new()
+            .optimize(join, &ConfigOptions::new())
+            .unwrap();
+        let hash_join = optimized
+            .downcast_ref::<HashJoinExec>()
+            .expect("null-aware join should remain a HashJoinExec");
+
+        assert_eq!(*hash_join.join_type(), JoinType::LeftAnti);
+        assert_eq!(*hash_join.partition_mode(), PartitionMode::Partitioned);
+        assert!(hash_join.null_aware);
+    }
+
+    #[test]
+    fn unbounded_input_rule_does_not_swap_null_aware_anti_join() {
+        use datafusion::{
+            arrow::datatypes::SchemaRef,
+            common::NullEquality,
+            execution::{SendableRecordBatchStream, TaskContext},
+            physical_plan::{
+                EmptyRecordBatchStream,
+                joins::{HashJoinExec, PartitionMode},
+                streaming::{PartitionStream, StreamingTableExec},
+            },
+        };
+
+        use crate::physical_optimizer::join_selection::hash_join_swap_subrule;
+
+        #[derive(Debug)]
+        struct EmptyPartitionStream(SchemaRef);
+
+        impl PartitionStream for EmptyPartitionStream {
+            fn schema(&self) -> &SchemaRef {
+                &self.0
+            }
+
+            fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
+                Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.0)))
+            }
+        }
+
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
+        let left = Arc::new(
+            StreamingTableExec::try_new(
+                Arc::clone(&schema),
+                vec![Arc::new(EmptyPartitionStream(Arc::clone(&schema)))],
+                None,
+                vec![],
+                true,
+                None,
+            )
+            .unwrap(),
+        ) as Arc<dyn ExecutionPlan>;
+        let right = Arc::new(StatisticsExec::new(
+            Statistics::new_unknown(&schema),
+            schema.as_ref().clone(),
+        )) as Arc<dyn ExecutionPlan>;
+        let join = Arc::new(
+            HashJoinExec::try_new(
+                Arc::clone(&left),
+                Arc::clone(&right),
+                vec![(
+                    Arc::new(Column::new("key", 0)) as _,
+                    Arc::new(Column::new("key", 0)) as _,
+                )],
+                None,
+                &JoinType::LeftAnti,
+                None,
+                PartitionMode::Partitioned,
+                NullEquality::NullEqualsNothing,
+                true,
+            )
+            .unwrap(),
+        ) as Arc<dyn ExecutionPlan>;
+
+        let optimized = hash_join_swap_subrule(Arc::clone(&join), &ConfigOptions::new())
+            .expect(
+                "the unbounded-input rule must not try to create a null-aware RightAnti",
+            );
+
+        assert!(Arc::ptr_eq(&optimized, &join));
+    }
+
     fn create_big_and_small() -> (Arc<dyn ExecutionPlan>, Arc<dyn ExecutionPlan>) {
         let big = Arc::new(StatisticsExec::new(
             big_statistics(),