From 509a68808fb8b64e24188a51c108122240d86558 Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Mon, 27 Jul 2026 11:28:09 +0900 Subject: [PATCH] fix(scheduler): execute null-aware anti joins in one task Null-aware LeftAnti joins cannot be swapped and DataFusion coordinates their visited-row and NULL state only within one process. Preserve CollectLeft, collect the build side, and coalesce the probe side so static and AQE plans run the join in exactly one task. Reject disabled or known oversized single-task builds through the Ballista broadcast threshold. Add optimizer, stage-lowering, threshold, and standalone distributed NOT IN regression tests. Closes #2187 Signed-off-by: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> --- ballista/client/tests/null_aware.rs | 138 ++++++++ ballista/core/src/config.rs | 10 +- .../src/physical_optimizer/join_selection.rs | 130 ++++++- ballista/scheduler/src/planner.rs | 329 ++++++++++++++---- .../state/aqe/execution_plan/dynamic_join.rs | 56 ++- .../optimizer_rule/distributed_exchange.rs | 73 ++++ .../aqe/optimizer_rule/join_selection.rs | 152 +++++++- docs/source/user-guide/configs.md | 2 +- docs/source/user-guide/tuning-guide.md | 13 +- 9 files changed, 803 insertions(+), 100 deletions(-) create mode 100644 ballista/client/tests/null_aware.rs diff --git a/ballista/client/tests/null_aware.rs b/ballista/client/tests/null_aware.rs new file mode 100644 index 0000000000..507d9e7753 --- /dev/null +++ b/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; + } +} diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index fd3060297d..2154076370 100644 --- a/ballista/core/src/config.rs +++ b/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> = 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(), diff --git a/ballista/scheduler/src/physical_optimizer/join_selection.rs b/ballista/scheduler/src/physical_optimizer/join_selection.rs index 20d4649283..7e6ffb01c5 100644 --- a/ballista/scheduler/src/physical_optimizer/join_selection.rs +++ b/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::() && 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; + + let optimized = JoinSelection::new() + .optimize(join, &ConfigOptions::new()) + .unwrap(); + let hash_join = optimized + .downcast_ref::() + .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) -> 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; + let right = Arc::new(StatisticsExec::new( + Statistics::new_unknown(&schema), + schema.as_ref().clone(), + )) as Arc; + 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; + + 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, Arc) { let big = Arc::new(StatisticsExec::new( big_statistics(), diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 7283558991..5f6b4bb92a 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -266,6 +266,107 @@ impl DefaultDistributedPlanner { self.next_stage_id } + /// Returns `Some(true/false)` when statistics can determine whether an + /// input fits Ballista's broadcast byte limit, or `None` when its size is + /// unknown. Falls back to a conservative row-width estimate when possible. + fn broadcast_size_under_threshold( + plan: &dyn ExecutionPlan, + threshold: usize, + ) -> Option { + let Ok(stats) = plan.partition_statistics(None) else { + debug!( + "broadcast check: partition_statistics returned error for {}", + plan.name() + ); + return None; + }; + debug!( + "broadcast check: {} total_byte_size={:?} num_rows={:?} threshold={}", + plan.name(), + stats.total_byte_size, + stats.num_rows, + threshold, + ); + if let Some(bytes) = stats.total_byte_size.get_value() + && *bytes != 0 + { + Some(*bytes < threshold) + } else if let Some(rows) = stats.num_rows.get_value() + && *rows != 0 + { + let schema = plan.schema(); + let bytes_per_row: usize = schema + .fields() + .iter() + .map(|f| match f.data_type() { + DataType::Boolean => 1, + DataType::Int8 | DataType::UInt8 => 1, + DataType::Int16 | DataType::UInt16 => 2, + DataType::Int32 | DataType::UInt32 | DataType::Float32 => 4, + DataType::Int64 | DataType::UInt64 | DataType::Float64 => 8, + DataType::Date32 => 4, + DataType::Date64 => 8, + DataType::Decimal128(_, _) => 16, + DataType::Decimal256(_, _) => 32, + _ => 32, // conservative estimate for variable-length types + }) + .sum(); + let estimated_bytes = *rows * bytes_per_row.max(8); + debug!( + "broadcast check: estimated {estimated_bytes} bytes ({rows} rows * {bytes_per_row} bytes/row from {} columns)", + schema.fields().len(), + ); + Some(estimated_bytes < threshold) + } else { + None + } + } + + /// Lowers a null-aware anti join to the only shape supported correctly by + /// DataFusion's in-process hash join: collect the build side and coalesce + /// the probe side so one task owns all shared null/visited state. + fn lower_null_aware_join( + hash_join: &HashJoinExec, + threshold_bytes: usize, + ) -> Result> { + // Unknown size is allowed because file scans commonly lose exact byte + // statistics before this point. Known oversized inputs and an explicit + // threshold of zero fail clearly instead of running an unsafe plan. + if threshold_bytes == 0 + || matches!( + Self::broadcast_size_under_threshold( + &**hash_join.left(), + threshold_bytes + ), + Some(false) + ) + { + return Err(BallistaError::General(format!( + "Null-aware anti join requires single-task execution, but its build side does not fit ballista.optimizer.broadcast_join_threshold_bytes ({threshold_bytes} bytes)" + ))); + } + + // Always keep an explicit coalesce. Some scans report one output + // partition during planning but expand to multiple partitions when the + // distributed stage is built. + let right: Arc = if hash_join + .right() + .downcast_ref::() + .is_some() + { + hash_join.right().clone() + } else { + Arc::new(CoalescePartitionsExec::new(hash_join.right().clone())) + }; + + hash_join + .builder() + .with_partition_mode(PartitionMode::CollectLeft) + .with_new_children(vec![hash_join.left().clone(), right])? + .build_exec() + .map_err(Into::into) + } + /// Reconciles a join's partition mode with the Ballista broadcast /// threshold (`ballista.optimizer.broadcast_join_threshold_bytes`). /// @@ -279,8 +380,9 @@ impl DefaultDistributedPlanner { /// under the Ballista threshold (including a threshold of `0`, which /// disables broadcast joins). This makes the Ballista key authoritative /// even when it is overridden at runtime below the DataFusion session - /// value. Null-aware anti joins are never demoted (they require - /// `CollectLeft`). + /// value. Null-aware anti joins are instead lowered to a single-task + /// `CollectLeft` join unless the build side is known to exceed the + /// Ballista threshold. /// /// Otherwise returns the input unchanged. fn maybe_promote_to_broadcast( @@ -303,6 +405,15 @@ impl DefaultDistributedPlanner { if let Some(hash_join) = plan.downcast_ref::() && *hash_join.partition_mode() == PartitionMode::CollectLeft { + // A null-aware anti join cannot run once per probe partition: + // DataFusion's visited/null state is shared only within one process. + // Collect the build side and coalesce the probe side so the join has + // exactly one task. Reject a build side known to exceed the normal + // broadcast threshold; an unsupported-plan error is safer than a + // wrong answer. + if hash_join.null_aware { + return Self::lower_null_aware_join(hash_join, threshold_bytes); + } // Broadcasting is only correct for probe-driven join types. If the // join type is not broadcast-safe, demote it back to a partitioned // (shuffle) join. Correctness guard, independent of the threshold. @@ -313,19 +424,19 @@ impl DefaultDistributedPlanner { ); return Self::demote_collect_left_to_partitioned(hash_join, config); } - // Null-aware anti joins must stay `CollectLeft`: they track - // probe-side state that a partitioned join cannot reconstruct, so - // never demote them regardless of the threshold. - if hash_join.null_aware { - return Ok(plan); - } // Safe join type: honor the Ballista broadcast threshold. DataFusion // decided `CollectLeft` using its own session threshold, which can // exceed a user's runtime `broadcast_join_threshold_bytes` override. // If broadcasts are disabled (0) or the build (left) side is not // under the Ballista threshold, demote so the Ballista key is // authoritative in the static planner path too. - if threshold_bytes == 0 || !under(&**hash_join.left(), threshold_bytes) { + if threshold_bytes == 0 + || !Self::broadcast_size_under_threshold( + &**hash_join.left(), + threshold_bytes, + ) + .unwrap_or(false) + { debug!( "broadcast check: demoting CollectLeft join to Partitioned; build side not under Ballista threshold={threshold_bytes} (or broadcasts disabled)", ); @@ -334,6 +445,16 @@ impl DefaultDistributedPlanner { return Ok(plan); } + // An already-partitioned null-aware join is not changed by the copied + // DataFusion optimizer rule. Lower it here, where Ballista can also + // enforce single-task distributed execution. + if let Some(hash_join) = plan.downcast_ref::() + && *hash_join.partition_mode() == PartitionMode::Partitioned + && hash_join.null_aware + { + return Self::lower_null_aware_join(hash_join, threshold_bytes); + } + if threshold_bytes == 0 { debug!("broadcast check: threshold is 0, broadcast disabled"); return Ok(plan); @@ -352,63 +473,13 @@ impl DefaultDistributedPlanner { if *hash_join.partition_mode() != PartitionMode::Partitioned { return Ok(plan); } - if hash_join.null_aware { - return Ok(plan); - } - let left = hash_join.left(); let right = hash_join.right(); - fn under(plan: &dyn ExecutionPlan, threshold: usize) -> bool { - let Ok(stats) = plan.partition_statistics(None) else { - debug!( - "broadcast check: partition_statistics returned error for {}", - plan.name() - ); - return false; - }; - debug!( - "broadcast check: {} total_byte_size={:?} num_rows={:?} threshold={}", - plan.name(), - stats.total_byte_size, - stats.num_rows, - threshold, - ); - if let Some(bytes) = stats.total_byte_size.get_value() { - *bytes != 0 && *bytes < threshold - } else if let Some(rows) = stats.num_rows.get_value() { - let schema = plan.schema(); - let bytes_per_row: usize = schema - .fields() - .iter() - .map(|f| { - match f.data_type() { - DataType::Boolean => 1, - DataType::Int8 | DataType::UInt8 => 1, - DataType::Int16 | DataType::UInt16 => 2, - DataType::Int32 | DataType::UInt32 | DataType::Float32 => 4, - DataType::Int64 | DataType::UInt64 | DataType::Float64 => 8, - DataType::Date32 => 4, - DataType::Date64 => 8, - DataType::Decimal128(_, _) => 16, - DataType::Decimal256(_, _) => 32, - _ => 32, // conservative estimate for variable-length types - } - }) - .sum(); - let estimated_bytes = *rows * bytes_per_row.max(8); - debug!( - "broadcast check: estimated {estimated_bytes} bytes ({rows} rows * {bytes_per_row} bytes/row from {} columns)", - schema.fields().len(), - ); - estimated_bytes != 0 && estimated_bytes < threshold - } else { - false - } - } - - let left_under = under(&**left, threshold_bytes); - let right_under = under(&**right, threshold_bytes); + let left_under = Self::broadcast_size_under_threshold(&**left, threshold_bytes) + .unwrap_or(false); + let right_under = Self::broadcast_size_under_threshold(&**right, threshold_bytes) + .unwrap_or(false); if !left_under && !right_under { debug!("broadcast check: neither side under threshold, skipping promotion"); return Ok(plan); @@ -1087,6 +1158,138 @@ order by Ok(()) } + #[test] + fn null_aware_collect_left_join_coalesces_probe_to_one_partition() { + use datafusion::{ + arrow::datatypes::{DataType, Field, Schema}, + common::{ + ColumnStatistics, JoinType, NullEquality, Statistics, stats::Precision, + }, + physical_plan::{ + Partitioning, coalesce_partitions::CoalescePartitionsExec, + joins::PartitionMode, repartition::RepartitionExec, + test::exec::StatisticsExec, + }, + }; + + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)])); + let stats = Statistics { + num_rows: Precision::Exact(20), + total_byte_size: Precision::Exact(80), + column_statistics: vec![ColumnStatistics::new_unknown()], + }; + let left = Arc::new(StatisticsExec::new(stats.clone(), schema.as_ref().clone())) + as Arc; + let right = Arc::new( + RepartitionExec::try_new( + Arc::new(StatisticsExec::new(stats, schema.as_ref().clone())), + Partitioning::RoundRobinBatch(4), + ) + .unwrap(), + ) as Arc; + let plan = Arc::new( + HashJoinExec::try_new( + left, + right, + vec![( + Arc::new(Column::new("key", 0)) as _, + Arc::new(Column::new("key", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(), + ) as Arc; + + let planned = DefaultDistributedPlanner::maybe_promote_to_broadcast( + plan, + &datafusion::config::ConfigOptions::new(), + ) + .unwrap(); + let hash_join = planned + .downcast_ref::() + .expect("null-aware join should remain a HashJoinExec"); + + assert_eq!(*hash_join.join_type(), JoinType::LeftAnti); + assert_eq!(*hash_join.partition_mode(), PartitionMode::CollectLeft); + assert!(hash_join.null_aware); + assert!( + hash_join + .right() + .downcast_ref::() + .is_some(), + "probe side must be coalesced so the join runs in one task" + ); + assert_eq!( + hash_join + .right() + .properties() + .output_partitioning() + .partition_count(), + 1 + ); + } + + #[test] + fn null_aware_join_rejects_known_oversized_build_side() { + use datafusion::{ + arrow::datatypes::{DataType, Field, Schema}, + common::{ + ColumnStatistics, JoinType, NullEquality, Statistics, stats::Precision, + }, + physical_plan::{joins::PartitionMode, test::exec::StatisticsExec}, + }; + + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)])); + let left = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Exact(5_000_000), + total_byte_size: Precision::Exact(20 * 1024 * 1024), + column_statistics: vec![ColumnStatistics::new_unknown()], + }, + schema.as_ref().clone(), + )) as Arc; + let right = Arc::new(StatisticsExec::new( + Statistics::new_unknown(&schema), + schema.as_ref().clone(), + )) as Arc; + let plan = Arc::new( + HashJoinExec::try_new( + left, + right, + vec![( + Arc::new(Column::new("key", 0)) as _, + Arc::new(Column::new("key", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(), + ) as Arc; + + let error = DefaultDistributedPlanner::maybe_promote_to_broadcast( + plan, + &datafusion::config::ConfigOptions::new(), + ) + .unwrap_err(); + assert!( + error.to_string().contains( + "build side does not fit ballista.optimizer.broadcast_join_threshold_bytes" + ), + "{error}" + ); + } + #[tokio::test] async fn distributed_broadcast_join_plan() -> Result<(), BallistaError> { use datafusion::physical_plan::joins::PartitionMode; diff --git a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs index 40ab7e7854..bec08bfc69 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -19,7 +19,10 @@ use ballista_core::config::BallistaConfig; use datafusion::{ arrow::compute::SortOptions, arrow::datatypes::{DataType, Schema}, - common::{ColumnStatistics, JoinType, NullEquality, Result, exec_err, internal_err}, + common::{ + ColumnStatistics, JoinType, NullEquality, Result, exec_err, internal_err, + plan_err, + }, config::ConfigOptions, execution::{SendableRecordBatchStream, TaskContext}, physical_expr_common::physical_expr::fmt_sql, @@ -248,10 +251,13 @@ impl DynamicJoinSelectionExec { // actually builds from rather than `self.left` unconditionally. // `supports_swap_join_order` is true when the *left* is the larger side, // so a swap moves the build onto `self.right`. - let swap_inputs = SelectJoinRule::supports_swap_join_order( - self.left.as_ref(), - self.right.as_ref(), - )?; + // Null-aware anti joins are only valid as LeftAnti and therefore cannot + // participate in the size-driven input swap. + let swap_inputs = !self.null_aware + && SelectJoinRule::supports_swap_join_order( + self.left.as_ref(), + self.right.as_ref(), + )?; let build_side = if swap_inputs { &self.right } else { &self.left }; let build_max_partition_bytes = max_per_partition_build_bytes(build_side); @@ -284,15 +290,41 @@ impl DynamicJoinSelectionExec { self.join_type }; - let partition_mode = - if under_threshold && collect_left_broadcast_safe(build_side_join_type) { - PartitionMode::CollectLeft - } else { - PartitionMode::Partitioned - }; - let stats_left = self.left.partition_statistics(None)?; let stats_right = self.right.partition_statistics(None)?; + let build_stats = if swap_inputs { + &stats_right + } else { + &stats_left + }; + let build_size_known = match build_stats.total_byte_size.get_value() { + Some(bytes) => *bytes != 0, + None => build_stats + .num_rows + .get_value() + .is_some_and(|rows| *rows != 0), + }; + + // A null-aware anti join must collect its build side and run its probe + // side in one task. Reject a disabled broadcast or a build side known + // to exceed the threshold. Unknown sizes are allowed because file scans + // commonly lose exact statistics before AQE first resolves the join. + if self.null_aware + && (threshold_collect_left_join_bytes == 0 + || (build_size_known && !under_threshold)) + { + return plan_err!( + "Null-aware anti join requires single-task execution, but its build side does not fit ballista.optimizer.broadcast_join_threshold_bytes ({threshold_collect_left_join_bytes} bytes)" + ); + } + + let partition_mode = if self.null_aware + || (under_threshold && collect_left_broadcast_safe(build_side_join_type)) + { + PartitionMode::CollectLeft + } else { + PartitionMode::Partitioned + }; let action = match (&self.selection_state, partition_mode) { (JoinInputState::Unknown, PartitionMode::CollectLeft) => self diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs index a07312d14e..2b5f36eeea 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs @@ -19,6 +19,7 @@ use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::{ExecutionPlan, execution_plan}; @@ -45,6 +46,40 @@ impl DistributedExchangeRule { &self, execution_plan: Arc, ) -> datafusion::error::Result>> { + // DataFusion's null-aware hash join coordinates visited rows and + // probe-side NULL state with in-process atomics. A CollectLeft join + // running once per probe partition therefore produces duplicate or + // incorrect output. This is the final plan-mutating rule, so enforce an + // explicit probe-side coalesce here after DataFusion's optimizers have + // finished. Add an exchange when the subtree has no existing boundary. + if let Some(hash_join) = execution_plan.downcast_ref::() + && hash_join.null_aware + && *hash_join.partition_mode() == PartitionMode::CollectLeft + && hash_join + .right() + .downcast_ref::() + .is_none() + { + let left = hash_join.left().clone(); + let right = hash_join.right().clone(); + let right = if right.downcast_ref::().is_none() + && !matches!(nearest_exchange_status(&right), ExchangeStatus::Unresolved) + { + Arc::new(ExchangeExec::new( + right, + None, + self.plan_id_generator + .fetch_add(1, std::sync::atomic::Ordering::Relaxed), + )) as Arc + } else { + right + }; + let right = Arc::new(CoalescePartitionsExec::new(right)); + return Ok(Transformed::yes( + execution_plan.with_new_children(vec![left, right])?, + )); + } + if let Some(coalesce) = execution_plan.downcast_ref::() { let input = coalesce.input(); if input.downcast_ref::().is_none() @@ -208,6 +243,44 @@ mod tests { Arc::new(SortPreservingMergeExec::new(ordering, input)) } + #[test] + fn null_aware_join_coalesces_probe_after_other_optimizers() { + use datafusion::common::{JoinType, NullEquality}; + + let left = leaf_exec(); + let right = leaf_exec(); + let join = Arc::new( + HashJoinExec::try_new( + left, + right, + vec![( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("a", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(), + ) as Arc; + + let result = DistributedExchangeRule::default() + .optimize(join, &config()) + .unwrap(); + + assert_plan!(result.as_ref(), @ r" + AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false + HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(a@0, a@0)] + StatisticsExec: col_count=1, row_count=Absent + CoalescePartitionsExec + ExchangeExec: partitioning=None, plan_id=0, stage_id=pending, stage_resolved=false + StatisticsExec: col_count=1, row_count=Absent + "); + } + // --- CoalescePartitionsExec --- #[test] diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs index 53f44dae61..e3807ab711 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs @@ -151,10 +151,12 @@ impl PhysicalOptimizerRule for SelectJoinRule { // at this point we know there are two exchanges // as we added them beforehand JoinSelectionAction::LateCollectLeft(hash_join_exec) => { - if Self::supports_swap_join_order( - hash_join_exec.left.as_ref(), - hash_join_exec.right.as_ref(), - )? { + if !hash_join_exec.null_aware + && Self::supports_swap_join_order( + hash_join_exec.left.as_ref(), + hash_join_exec.right.as_ref(), + )? + { let left = hash_join_exec.left.clone(); let right = hash_join_exec.right.clone(); @@ -199,10 +201,11 @@ impl PhysicalOptimizerRule for SelectJoinRule { } JoinSelectionAction::CollectLeft(hash_join_exec) => { - let plan = if Self::supports_swap_join_order( - hash_join_exec.left.as_ref(), - hash_join_exec.right.as_ref(), - )? { + let plan = if !hash_join_exec.null_aware + && Self::supports_swap_join_order( + hash_join_exec.left.as_ref(), + hash_join_exec.right.as_ref(), + )? { hash_join_exec .swap_inputs(PartitionMode::CollectLeft)? } else { @@ -293,10 +296,11 @@ impl PhysicalOptimizerRule for SelectJoinRule { Ok(Transformed::yes(dynamic_join)) } JoinSelectionAction::Hash(hash_join_exec) => { - let hash_join_exec = if Self::supports_swap_join_order( - hash_join_exec.left.as_ref(), - hash_join_exec.right.as_ref(), - )? { + let hash_join_exec = if !hash_join_exec.null_aware + && Self::supports_swap_join_order( + hash_join_exec.left.as_ref(), + hash_join_exec.right.as_ref(), + )? { hash_join_exec .swap_inputs(*hash_join_exec.partition_mode())? } else { @@ -610,6 +614,130 @@ mod tests { assert_plan!(optimized.as_ref(), @ "DataSourceExec: partitions=1, partition_sizes=[1]"); } + #[test] + fn null_aware_anti_join_is_not_swapped_by_aqe() { + use datafusion::physical_expr::expressions::Column; + use datafusion::{ + arrow::datatypes::{DataType, Field, Schema}, + common::{ + ColumnStatistics, JoinType, NullEquality, Statistics, stats::Precision, + }, + physical_plan::{ + Partitioning, + joins::{HashJoinExec, PartitionMode}, + repartition::RepartitionExec, + test::exec::StatisticsExec, + }, + }; + + fn stats_exec(name: &str, bytes: usize) -> Arc { + Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(bytes / 4), + total_byte_size: Precision::Inexact(bytes), + column_statistics: vec![ColumnStatistics::new_unknown()], + }, + Schema::new(vec![Field::new(name, DataType::Int32, true)]), + )) + } + + // A normal join would swap these inputs to build from the smaller right + // side. Doing that to a null-aware LeftAnti creates an invalid + // null-aware RightAnti join. + let left = stats_exec("big_key", 1024 * 1024); + let right = Arc::new( + RepartitionExec::try_new( + stats_exec("small_key", 1024), + Partitioning::RoundRobinBatch(4), + ) + .unwrap(), + ) as Arc; + let join = HashJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + vec![( + Arc::new(Column::new("big_key", 0)) as _, + Arc::new(Column::new("small_key", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(); + let dynamic = DynamicJoinSelectionExec::from_hash_join(&join, 0).unwrap() + as Arc; + + let resolved = SelectJoinRule::default() + .optimize(dynamic, &ConfigOptions::default()) + .expect("AQE must not swap a null-aware LeftAnti join"); + let hash_join = resolved + .downcast_ref::() + .expect("AQE should resolve to HashJoinExec"); + + assert_eq!(*hash_join.join_type(), JoinType::LeftAnti); + assert_eq!(*hash_join.partition_mode(), PartitionMode::CollectLeft); + assert!(hash_join.null_aware); + assert_eq!(hash_join.left().schema().field(0).name(), "big_key"); + assert_eq!(hash_join.right().schema().field(0).name(), "small_key"); + } + + #[test] + fn null_aware_anti_join_rejects_known_oversized_build_side_in_aqe() { + use datafusion::physical_expr::expressions::Column; + use datafusion::{ + arrow::datatypes::{DataType, Field, Schema}, + common::{ + ColumnStatistics, JoinType, NullEquality, Statistics, stats::Precision, + }, + physical_plan::{ + joins::{HashJoinExec, PartitionMode}, + test::exec::StatisticsExec, + }, + }; + + fn stats_exec(name: &str, bytes: usize) -> Arc { + Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(bytes / 4), + total_byte_size: Precision::Inexact(bytes), + column_statistics: vec![ColumnStatistics::new_unknown()], + }, + Schema::new(vec![Field::new(name, DataType::Int32, true)]), + )) + } + + let join = HashJoinExec::try_new( + stats_exec("big_key", 20 * 1024 * 1024), + stats_exec("small_key", 1024), + vec![( + Arc::new(Column::new("big_key", 0)) as _, + Arc::new(Column::new("small_key", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(); + let dynamic = DynamicJoinSelectionExec::from_hash_join(&join, 0).unwrap() + as Arc; + + let error = SelectJoinRule::default() + .optimize(dynamic, &ConfigOptions::default()) + .unwrap_err(); + assert!( + error.to_string().contains( + "build side does not fit ballista.optimizer.broadcast_join_threshold_bytes" + ), + "{error}" + ); + } + /// When `ballista.planner.adaptive_join.enabled = false` the `DelayJoinSelectionRule` /// must be a no-op: a plan containing a `DynamicJoinSelectionExec` node must /// be returned unchanged. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index c40f1a582c..d174102e7f 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -119,7 +119,7 @@ standard DataFusion settings. | ballista.client.pull | Boolean | false | Should client employ pull or push job tracking. In pull mode client will make a request to server in the loop, until job finishes. Pull mode is kept for legacy clients. | | ballista.client.use_tls | Boolean | false | Should connection between client, scheduler, and executors use TLS. | | ballista.job.name | Utf8 | (none) | Sets the job name that will appear in the web user interface for any submitted jobs | -| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 10485760 | 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. | +| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 10485760 | 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). 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. | | ballista.optimizer.broadcast_join_threshold_rows | UInt64 | 1000000 | Row-count threshold below which a hash join's smaller side is promoted to CollectLeft and lowered via the broadcast pattern, used as a fallback when byte-size statistics are unavailable. Applies to adaptive query planning (AQE). Set to 0 to disable promotion via the row-count path. | | ballista.optimizer.hash_join_max_build_partition_bytes | UInt64 | 67108864 | Maximum per-partition hash-join build-side bytes for a Partitioned hash join under AQE. A build partition larger than this falls back to SortMergeJoin (spillable). Defaults to 64 MiB; 0 disables the check, which makes AQE use a hash join regardless of build size. | | ballista.planner.adaptive.enabled | Boolean | false | Enables Adaptive Query Planning (EXPERIMENTAL) | diff --git a/docs/source/user-guide/tuning-guide.md b/docs/source/user-guide/tuning-guide.md index 6e0ba0c511..bf1768922c 100644 --- a/docs/source/user-guide/tuning-guide.md +++ b/docs/source/user-guide/tuning-guide.md @@ -253,11 +253,11 @@ Adaptive Query Planning is EXPERIMENTAL, should be used for testing purposes onl ### Configuration -| key | type | default | description | -| ------------------------------------------------- | ------- | -------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| ballista.planner.adaptive.enabled | Boolean | false | Enables the adaptive planner. Experimental. | -| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 10485760 | Byte-size threshold below which a hash join's smaller side is broadcast (`CollectLeft`). Governs both the static planner and AQE. Set to 0 to disable broadcast joins. | -| ballista.optimizer.broadcast_join_threshold_rows | UInt64 | 1000000 | Row-count fallback threshold used when byte-size statistics are unavailable. Applies to AQE. Set to 0 to disable promotion via the row-count path. | +| key | type | default | description | +| ------------------------------------------------- | ------- | -------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| ballista.planner.adaptive.enabled | Boolean | false | Enables the adaptive planner. Experimental. | +| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 10485760 | Byte-size threshold below which a hash join's smaller side is broadcast (`CollectLeft`). Also caps null-aware anti joins with known build sizes because they run in one task. Set to 0 to disable broadcasts and reject null-aware anti joins. | +| ballista.optimizer.broadcast_join_threshold_rows | UInt64 | 1000000 | Row-count fallback threshold used when byte-size statistics are unavailable. Applies to AQE. Set to 0 to disable promotion via the row-count path. | ### What AQE does today @@ -271,6 +271,9 @@ implemented: - **Broadcast join selection.** When a join input's runtime size falls under `ballista.optimizer.broadcast_join_threshold_bytes` (or the row-count fallback), the smaller side is broadcast (`CollectLeft`) instead of shuffled. + Null-aware anti joins use `CollectLeft` with a single probe task because their + state cannot be coordinated across executors. A known oversized build side is + rejected instead of producing an incorrect distributed result. - **Empty stage elimination.** When a completed stage produces zero rows, its downstream exchange is replaced with an empty execution node, and emptiness is propagated up the plan so downstream stages are skipped entirely.