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

[phillipleblanc]

Which issue does this PR close?

Closes #2187.

Rationale for this change

DataFusion plans NOT IN predicates as null-aware LeftAnti hash joins. These joins cannot be swapped to RightAnti. They also need PartitionMode::CollectLeft.

CollectLeft alone is not enough in a distributed plan. DataFusion coordinates visited build rows and probe-side NULL state with in-process atomics. Running the join in several probe tasks gives each task independent state and produces duplicated or incorrect rows.

The four-partition regression case returned 56 rows instead of 0 when the subquery contained NULL, and 75 rows instead of 15 without the NULL.

Before
------
broadcast build ──┬──> probe task 0 ──> independent null/visited state
                  ├──> probe task 1 ──> independent null/visited state
                  └──> probe task N ──> duplicated or incorrect output

After
-----
broadcast build ───────────────┐
probe shuffle ──> coalesce(1) ─┴──> one null-aware join task

This is a fresh implementation against Apache main, based on the downstream report in spiceai#58 and the distributed reproduction from review.

What changes are included in this PR?

• Match DataFusion's null-aware guards in statistical and unbounded-input join swapping.
• Lower static null-aware joins to LeftAnti(CollectLeft) with a single probe task.
• Apply the same single-task lowering after AQE's other physical optimizer rules finish.
• Use ballista.optimizer.broadcast_join_threshold_bytes as a safety limit when the build size is known. A value of 0 rejects null-aware anti joins instead of running an unsafe plan.
• Add optimizer, static planner, AQE exchange, threshold, and standalone distributed NOT IN regression tests.
• Document the threshold behavior.

Are there any user-facing changes?

NOT IN queries that retain a null-aware hash join now execute with correct NULL semantics under static and adaptive planning. A known oversized build side now returns a clear unsupported-plan error rather than hanging or returning incorrect data.

Ballista's default sort-merge planning can discard the null-aware flag before these scheduler rules run. #2193 tracks that separate path.

Test plan

• cargo test -p ballista-scheduler --lib — 276 passed, 2 ignored
• cargo test -p ballista --test null_aware — 2 passed
• cargo clippy --all-targets --package ballista-scheduler --all-features -- -D warnings
• cargo clippy --all-targets --package ballista --all-features -- -D warnings
• cargo fmt --all -- --check
• ./ci/scripts/rust_config_docs_check.sh
• CI after the revised push

[martin-g]

LGTM

[andygrove]

Thanks for picking this up. The hang is real, and the hash_join_swap_subrule change matches upstream DataFusion exactly, which is nice.

I am worried about the reordering in planner.rs though. Moving the null_aware early return ahead of the collect_left_broadcast_safe demotion lets a LeftAnti CollectLeft join through into the distributed plan, and that check was the thing keeping it out. CollectLeft replicates the build side to every probe task, and LeftAnti output is entirely build-side driven, so each task emits its own copy of the unmatched build rows. DataFusion's shared visited_indices_bitmap, report_probe_completed() and probe_side_has_null atomics only coordinate within one process, so they cannot help once the probe side is spread across tasks. The same reasoning applies to the AQE change in dynamic_join.rs, where self.null_aware || short-circuits the same safety check.

I built the branch and ran a NOT IN end to end against a standalone cluster to check. t1 is 20 rows over 4 CSV files, t2 is 5 rows over 4 files with one NULL key, and the query is select a from t1 where a not in (select b from t2) order by a.

	NULL in t2 (correct: 0 rows)	no NULL in t2 (correct: 15 rows)
single process DataFusion	0	15
Ballista @ main, prefer_hash_join=true	query never returns	n/a
Ballista + this PR, prefer_hash_join=true	56 rows	75 rows
Ballista + this PR, AQE enabled	56 rows	n/a
Ballista + this PR, default config	16 rows (unchanged by the PR)	15

The 56 and 75 row outputs are the same values repeated 3 or 4 times each, including values that should have been anti-joined away entirely. So the branch does fix the hang, it just turns it into a silently wrong answer, which I think is the worse of the two.

My read is that neither CollectLeft nor Partitioned is correct for this join once it spans more than one task. The shape that would work is forcing the whole join into a single task, so the build side is collected and the probe side is coalesced to one partition. Would you be up for going that route here? If that is a bigger change than you want in this PR, I would be happy with landing the parts that are unambiguously right (the hash_join_swap_subrule guard and the AQE swap guards, which only prevent invalid RightAnti plans) and filing a follow on issue for the distributed lowering. Returning a clear unsupported-plan error instead of the hang would also be an improvement over a wrong answer in the meantime.

Two smaller things while you are in here:

• The AQE path now forces CollectLeft regardless of broadcast_join_threshold_bytes and also skips the hash_join_max_build_partition_bytes fit check, since hash_build_fits only guards the Partitioned arm. For a null-aware anti join the build side is the outer relation, which is usually the big one, so a NOT IN against a large table would broadcast that whole table to every probe task with no ceiling and no way to turn it off.
• The Partitioned branch in join_selection.rs now rewrites to CollectLeft, where upstream DataFusion (datafusion-physical-optimizer 54.1, join_selection.rs around line 331) only skips the swap and leaves it Partitioned. Since that file is a vendored fork of the upstream rule it would help to have a comment saying why we diverge.

One more thing you may want to know, unrelated to your changes: Ballista defaults prefer_hash_join to false in ballista/core/src/extension.rs, and with that setting DataFusion's physical planner picks a SortMergeJoinExec and drops the null_aware flag before any of these rules run. That is why the default row in the table above is 16 instead of 0. Separate bug, but it means this fix only reaches users who opt into hash joins or AQE.

Here is the test I used, in case it is useful. Drop it in as ballista/client/tests/null_aware.rs and run with cargo test -p ballista --test null_aware -- --nocapture.

ballista/client/tests/null_aware.rs

// 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 ballista::prelude::SessionContextExt;
    use datafusion::arrow::util::pretty::pretty_format_batches;
    use datafusion::prelude::*;
    use std::fs;
    use std::path::Path;

    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);

        // single process DataFusion is the ground truth
        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 df_out = pretty_format_batches(&df_batches).unwrap().to_string();
        println!("=== [{case}] DataFusion result ===\n{df_out}");

        let mut mismatches = vec![];
        for variant in ["default", "prefer_hash_join", "aqe"] {
            let ctx = SessionContext::standalone().await.unwrap();
            let mut settings = vec![];
            if variant == "aqe" {
                settings.push("SET ballista.planner.adaptive.enabled = true");
            }
            if variant != "default" {
                settings.push("SET datafusion.optimizer.prefer_hash_join = true");
            }
            for stmt in settings {
                ctx.sql(stmt).await.unwrap().collect().await.unwrap();
            }
            register(&ctx, &dir).await;

            match ctx.sql(QUERY).await.unwrap().collect().await {
                Ok(batches) => {
                    let out = pretty_format_batches(&batches).unwrap().to_string();
                    println!("=== [{case}/{variant}] Ballista result ===\n{out}");
                    if out.trim() != df_out.trim() {
                        mismatches.push(format!("[{case}/{variant}] differs from DataFusion"));
                    }
                }
                Err(e) => {
                    println!("=== [{case}/{variant}] Ballista failed ===\n{e}");
                    mismatches.push(format!("[{case}/{variant}] failed: {e}"));
                }
            }
        }
        assert!(mismatches.is_empty(), "{mismatches:#?}");
    }

    #[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("no_null", false).await;
    }
}

Note that the prefer_hash_join and aqe variants hang rather than fail on main, so if you run this against main for comparison you will want a timeout.

[phillipleblanc]

Thanks for picking this up. The hang is real, and the hash_join_swap_subrule change matches upstream DataFusion exactly, which is nice.

Thanks, @andygrove. You were right. I reproduced the 56/75-row incorrect results with your test and revised the PR.

The updated implementation now:

• keeps null-aware joins as unswapped LeftAnti(CollectLeft);
• coalesces the probe side so exactly one task executes the join;
• applies this lowering in both static planning and AQE;
• performs AQE coalescing in DistributedExchangeRule, after DataFusion's optimizer rules;
• restores the vendored Partitioned branch to match upstream DataFusion;
• rejects a disabled broadcast threshold or a build side known to exceed it. Unknown sizes remain allowed because file scans commonly lose exact statistics before lowering.

I adapted your standalone test to cover prefer_hash_join and AQE with four input partitions, both with and without a subquery NULL. All cases now match single-process DataFusion: 0 and 15 rows respectively.

The default sort-merge path still loses the null-aware flag before scheduler planning. I filed #2193 for that separate issue.

[andygrove]

Thanks @phillipleblanc. I'm taking another look at this now.

[andygrove]

I'm going to need more time to review this. I'll look at this some more over the next couple of days.

[andygrove]

It's unfortunate that we have to collapse down to a single task and also reject any outer table > 10 MB, but it seems we have no choice based on current DF design. This is better than producing incorrect results.

Thanks again for fixing this @phillipleblanc

[andygrove]

I filed a follow on issue #2198 to remove the single task restriction. I'm investigating.