andygrove commented on PR #2188: URL: https://github.com/apache/datafusion-ballista/pull/2188#issuecomment-5091987189
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`. <details> <summary>ballista/client/tests/null_aware.rs</summary> ```rust // 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; } } ``` </details> 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
