diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index 10f64aaaa..63196a95b 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -110,8 +110,10 @@ pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS: &str = /// Configuration key for the maximum per-partition hash-join build-side bytes /// permitted for a Partitioned hash join under AQE. When a build partition -/// exceeds this, the join falls back to SortMergeJoin (spillable). `0` disables -/// the check (hash join is used regardless of build size). +/// exceeds this, the join falls back to SortMergeJoin (spillable). Defaults to +/// 64 MiB. `0` disables the check (hash join is used regardless of build size); +/// because AQE does not consult `datafusion.optimizer.prefer_hash_join`, that +/// also makes SortMergeJoin unreachable for a Partitioned join. pub const BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES: &str = "ballista.optimizer.hash_join_max_build_partition_bytes"; @@ -267,9 +269,10 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| ConfigEntry::new(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES.to_string(), "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). 0 (the default) disables the check.".to_string(), + SortMergeJoin (spillable). Defaults to 64 MiB; 0 disables the check, \ + which makes AQE use a hash join regardless of build size.".to_string(), DataType::UInt64, - Some("0".to_string())), + Some((64 * 1024 * 1024).to_string())), ConfigEntry::new(BALLISTA_CLIENT_PULL.to_string(), "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.".to_string(), DataType::Boolean, @@ -888,11 +891,15 @@ mod tests { Ok(()) } + // The default must stay non-zero: `0` disables the fit check, and since AQE + // no longer consults `prefer_hash_join`, a disabled check would leave the + // Partitioned arm unconditionally on hash join and make SortMergeJoin — the + // spillable strategy — unreachable. #[test] - fn hash_join_max_build_partition_bytes_defaults_to_zero() { + fn hash_join_max_build_partition_bytes_defaults_to_64_mib() { assert_eq!( BallistaConfig::default().hash_join_max_build_partition_bytes(), - 0 + 64 * 1024 * 1024 ); } } diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index 31a9e12af..112ef6db5 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -20,9 +20,10 @@ use crate::config::{ BALLISTA_CLIENT_GRPC_MAX_MESSAGE_SIZE, BALLISTA_CLIENT_USE_TLS, BALLISTA_COALESCE_ENABLED, BALLISTA_COALESCE_MERGED_PARTITION_FACTOR, BALLISTA_COALESCE_SMALL_PARTITION_FACTOR, BALLISTA_COALESCE_TARGET_PARTITION_BYTES, - BALLISTA_JOB_NAME, BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ, - BALLISTA_SHUFFLE_READER_MAX_REQUESTS, BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT, - BALLISTA_STANDALONE_PARALLELISM, BallistaConfig, + BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES, BALLISTA_JOB_NAME, + BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ, BALLISTA_SHUFFLE_READER_MAX_REQUESTS, + BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT, BALLISTA_STANDALONE_PARALLELISM, + BallistaConfig, }; use crate::planner::BallistaQueryPlanner; use crate::serde::protobuf::KeyValuePair; @@ -208,6 +209,11 @@ pub trait SessionConfigExt { /// falling back to SortMergeJoin under AQE. `0` disables the check. fn ballista_hash_join_max_build_partition_bytes(&self) -> usize; + /// Sets the maximum per-partition hash-join build-side bytes before falling + /// back to SortMergeJoin under AQE. Setting `0` disables the check, which + /// leaves AQE on a hash join whatever the build size. + fn with_ballista_hash_join_max_build_partition_bytes(self, max_bytes: usize) -> Self; + /// retrieves grpc client max message size fn ballista_grpc_client_max_message_size(&self) -> usize; @@ -524,6 +530,15 @@ impl SessionConfigExt for SessionConfig { }) } + fn with_ballista_hash_join_max_build_partition_bytes(self, max_bytes: usize) -> Self { + if self.options().extensions.get::().is_some() { + self.set_usize(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES, max_bytes) + } else { + self.with_option_extension(BallistaConfig::default()) + .set_usize(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES, max_bytes) + } + } + fn ballista_shuffle_reader_maximum_concurrent_requests(&self) -> usize { self.options() .extensions 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 fbc09c451..40ab7e785 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -215,7 +215,17 @@ impl DynamicJoinSelectionExec { &self, config: &ConfigOptions, ) -> Result { - let prefer_hash_join = config.optimizer.prefer_hash_join; + // `datafusion.optimizer.prefer_hash_join` is deliberately NOT consulted + // here. It is a *planner* flag: DataFusion applies it when it builds the + // physical plan, and `DelayJoinSelectionRule` then maps both + // `HashJoinExec` and `SortMergeJoinExec` onto this node, so the flag has + // already had its effect by the time AQE runs. Re-reading it would apply + // the same preference twice, and — since the two `CollectLeft` arms below + // produce a hash join unconditionally — it would only ever apply to one + // of the three paths that can emit one. AQE picks its strategy from + // runtime evidence instead: broadcast thresholds and the build-side fit + // check. + // // Both broadcast thresholds come from Ballista's own config so that a // single set of keys governs broadcast selection under both the static // planner (`maybe_promote_to_broadcast`) and AQE. A byte threshold of 0 @@ -230,34 +240,45 @@ impl DynamicJoinSelectionExec { let threshold_collect_left_join_bytes = bc.broadcast_join_threshold_bytes(); let threshold_collect_left_join_rows = bc.broadcast_join_threshold_rows(); let max_build_bytes = bc.hash_join_max_build_partition_bytes(); - let build_max_partition_bytes = max_per_partition_build_bytes(&self.left); + // The resolver collects the smaller side onto the build (left) input, + // swapping the inputs when the right side is the smaller one (the + // `swap_inputs` calls in `SelectJoinRule`). Derive that decision once, + // up front, so every size check below measures the input the executor + // 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(), + )?; + let build_side = if swap_inputs { &self.right } else { &self.left }; + + let build_max_partition_bytes = max_per_partition_build_bytes(build_side); + + // Broadcast eligibility is a property of the build side alone, since that + // is the side replicated to every probe task. Testing both sides and + // OR-ing the results would promote the join whenever *either* side fit, + // including when the swap then picks the other, over-threshold side to + // broadcast. + // // The `!= 0` guard sits here rather than inside // `supports_collect_by_thresholds`: that helper falls back to the row // check when the byte threshold is 0/absent, so gating at this level is // what makes a 0 threshold disable both the byte and row broadcast paths. let under_threshold = threshold_collect_left_join_bytes != 0 - && (Self::supports_collect_by_thresholds( - self.left.as_ref(), - threshold_collect_left_join_bytes, - threshold_collect_left_join_rows, - ) || Self::supports_collect_by_thresholds( - self.right.as_ref(), + && Self::supports_collect_by_thresholds( + build_side.as_ref(), threshold_collect_left_join_bytes, threshold_collect_left_join_rows, - )); + ); - // The resolver collects the smaller side onto the build (left) input, - // swapping the join type when the right side is the smaller one (the - // `swap_inputs` calls in `SelectJoinRule`). A `CollectLeft` join then - // broadcasts that build side to every probe task, which is only correct - // for join types that never emit rows on behalf of the build side. Use - // the same swap decision as the resolver to determine the resulting join - // type and skip the broadcast when it would be unsafe. - let build_side_join_type = if SelectJoinRule::supports_swap_join_order( - self.left.as_ref(), - self.right.as_ref(), - )? { + // A `CollectLeft` join broadcasts the build side to every probe task, + // which is only correct for join types that never emit rows on behalf of + // the build side. Evaluate that against the join type the resolver ends + // up with, which the swap rewrites, and skip the broadcast when it would + // be unsafe. + let build_side_join_type = if swap_inputs { self.join_type.swap() } else { self.join_type @@ -278,8 +299,7 @@ impl DynamicJoinSelectionExec { .to_hash_join(PartitionMode::CollectLeft) .map(JoinSelectionAction::CollectLeft), (JoinInputState::Repartitioned, PartitionMode::Partitioned) - if prefer_hash_join - && hash_build_fits(max_build_bytes, build_max_partition_bytes) => + if hash_build_fits(max_build_bytes, build_max_partition_bytes) => { self.to_hash_join(PartitionMode::Partitioned) .map(JoinSelectionAction::Hash) @@ -320,13 +340,15 @@ impl DynamicJoinSelectionExec { }; debug!( - "AQE join decision plan_id={} action={} partition_mode={:?} under_threshold={} \ + "AQE join decision plan_id={} action={} partition_mode={:?} build_side={} \ + under_threshold={} \ build_max_partition_bytes={:?} hash_join_max_build_partition_bytes={} \ left=(rows={:?}, bytes={:?}) right=(rows={:?}, bytes={:?}) \ byte_threshold={} row_threshold={}", self.plan_id, action_label, partition_mode, + if swap_inputs { "right" } else { "left" }, under_threshold, build_max_partition_bytes, max_build_bytes, @@ -520,17 +542,32 @@ fn hash_build_fits( /// the shape of the Q18 OOM, so an average would have hidden the failure /// mode this check exists to catch. /// +/// `build` must be the *post-swap* build input — `self.right` when +/// `SelectJoinRule::supports_swap_join_order` holds, `self.left` otherwise — +/// since that is the side the build task materializes. Measuring `self.left` +/// unconditionally validates a partition the build may never consume. +/// /// Reachability: at the one call site this feeds — the /// `(JoinInputState::Repartitioned, PartitionMode::Partitioned)` arm of -/// `to_actual_join` — `self.left` is always the `ExchangeExec` that -/// `SelectJoinRule` inserted while resolving the prior `Repartition` action -/// (see `JoinSelectionAction::Repartition` in `join_selection.rs`), and -/// `upstream_resolved()` guarantees that exchange's shuffle has already -/// finished. So the actual, materialized per-partition byte sizes that +/// `to_actual_join` — both children are `ExchangeExec`s that `SelectJoinRule` +/// inserted while resolving the prior `Repartition` action (see +/// `JoinSelectionAction::Repartition` in `join_selection.rs`), and +/// `upstream_resolved()` guarantees those shuffles have already finished. So +/// the actual, materialized per-partition byte sizes that /// `CoalescePartitionsRule` reads (`ExchangeExec::shuffle_partitions()` -> /// `PartitionLocation::partition_stats::num_bytes()`) are reachable here too, /// and are used directly rather than falling back to an average. /// +/// Known gap: the sizes are pre-coalesce. `CoalescePartitionsRule` groups +/// neighboring upstream partitions, so a coalesced build partition can exceed +/// the budget this check passed. The grouping is not knowable here — the rule +/// fires later, per stage, in `AdaptivePlanner::actionable_stages()`, after +/// `SelectJoinRule` has already run inside `replan_stages()` — so a +/// coalesce-aware check would have to re-derive the bin-pack over a leaf set +/// that differs from the rule's own for chained joins. Left as-is: +/// `ballista.planner.coalesce.enabled` is `false` by default, so the gap is +/// inert unless coalescing is explicitly turned on. +/// /// Returns `None` (callers treat this as "fits", not as a forced fallback) /// when `build` isn't an `ExchangeExec`, its shuffle hasn't resolved yet, or /// the resolved shuffle has no partitions at all. A resolved shuffle whose @@ -998,6 +1035,43 @@ mod tests { ); } + // Broadcast eligibility must be judged on the side that actually becomes the + // build input, not OR-ed across both. Byte-size statistics are absent on both + // sides here, so the swap decision falls back to row count: the left has + // fewer rows and therefore stays the build side. It is also by far the wider + // side — roughly 40 KB estimated against the right's 4 KB — so broadcasting + // it replicates more bytes to every probe task than the threshold intended, + // even though the right side on its own passes the check. + #[test] + fn broadcast_eligibility_follows_the_build_side() { + let mut left_fields = vec![Field::new("k", DataType::Int32, false)]; + left_fields.extend( + (0..20).map(|i| Field::new(format!("wide{i}"), DataType::Utf8, false)), + ); + // 100 rows * (4 B key + 20 columns * 20 B default width) ≈ 40 KB + let left = sizeless_stats_exec(100, left_fields); + // 1000 rows * 4 B ≈ 4 KB — under the threshold on its own + let right = + sizeless_stats_exec(1000, vec![Field::new("k", DataType::Int32, false)]); + + let action = run_to_actual_join( + left, + right, + JoinType::Inner, + true, + // Between the right side's ~4 KB and the left side's ~40 KB. + 10_000, + // Row ceiling well clear of both sides, so bytes decide. + 1_000_000, + ); + + assert!( + !is_collected(&action), + "the left side is the build (fewer rows) and is over the byte threshold, \ + so the join must repartition rather than broadcast it" + ); + } + // A threshold of 0 disables broadcast promotion entirely, matching the static // planner, even for a side small enough to collect under any positive cutoff. #[test] @@ -1196,49 +1270,29 @@ mod tests { /// string-keyed `BallistaConfig::set` path (the same path /// `run_to_actual_join` above uses for other Ballista settings) since /// there is no fluent setter for this key yet. + /// + /// Broadcast promotion is switched off (`broadcast_join_threshold_bytes = 0`) + /// so the build-fit tests below exercise the `Partitioned` path without a + /// `CollectLeft` decision shadowing it. + /// + /// `prefer_hash_join` is set to `false` deliberately — the hostile value. AQE + /// does not consult it, so every test below that asserts a `HashJoinExec` + /// would fail if that ever regressed. fn config_with_max_build(max_build_bytes: usize) -> ConfigOptions { let mut config = ConfigOptions::new(); - config.optimizer.prefer_hash_join = true; + config.optimizer.prefer_hash_join = false; let mut bc = BallistaConfig::default(); bc.set( "optimizer.hash_join_max_build_partition_bytes", &max_build_bytes.to_string(), ) .unwrap(); + bc.set("optimizer.broadcast_join_threshold_bytes", "0") + .unwrap(); config.extensions.insert(bc); config } - /// Builds a `DynamicJoinSelectionExec` already in `Repartitioned` state - /// whose build side (`left`) is a resolved `ExchangeExec` reporting the - /// given per-partition byte sizes, so `to_actual_join`'s - /// `(Repartitioned, Partitioned)` arm is reached and - /// `max_per_partition_build_bytes` sees known sizes. - fn repartitioned_join_with_build_partition_bytes( - per_partition_bytes: &[usize], - ) -> DynamicJoinSelectionExec { - // test_exchange_with_partition_bytes wraps test_schema(), whose sole - // field is named "x" (not "k" like stats_exec's schema), so the join - // key below must match that name. - let left = test_exchange_with_partition_bytes(per_partition_bytes); - let right = stats_exec(1_000_000); // large enough to never be broadcast - let on: JoinOn = - vec![(Arc::new(Column::new("x", 0)), Arc::new(Column::new("k", 0)))]; - DynamicJoinSelectionExec { - properties: Arc::clone(left.properties()), - left, - right, - on, - filter: None, - join_type: JoinType::Inner, - projection: None, - null_equality: NullEquality::NullEqualsNothing, - selection_state: JoinInputState::Repartitioned, - null_aware: false, - plan_id: 0, - } - } - /// Renders the `ExecutionPlan` a `JoinSelectionAction` resolves to, so /// tests can assert on `displayable(..).indent(false).to_string()` rather /// than matching on the action's structure. @@ -1258,12 +1312,20 @@ mod tests { } } + /// The probe side used by the build-fit tests below: bigger in total than + /// any build side they set up, so `supports_swap_join_order` is false, no + /// swap happens, and `left` stays the build. Its own per-partition size is + /// irrelevant — only the build side is measured. + const PROBE_PARTITION_BYTES: &[usize] = &[600 * MB]; + // A build partition over the configured max falls back to SortMergeJoin, // even though `prefer_hash_join` is true — the fit check overrides the // preference rather than the other way around. #[test] fn build_over_threshold_falls_back_to_sort_merge_join() { - let exec = repartitioned_join_with_build_partition_bytes(&[500 * MB]); // > 200 MB + // 500 MB build > the 200 MB budget + let exec = + repartitioned_join_between_exchanges(&[500 * MB], PROBE_PARTITION_BYTES); let action = exec .to_actual_join(&config_with_max_build(200 * MB)) .unwrap(); @@ -1277,7 +1339,9 @@ mod tests { // join `prefer_hash_join` selected. #[test] fn build_within_threshold_uses_partitioned_hash_join() { - let exec = repartitioned_join_with_build_partition_bytes(&[50 * MB]); // < 200 MB + // 50 MB build < the 200 MB budget + let exec = + repartitioned_join_between_exchanges(&[50 * MB], PROBE_PARTITION_BYTES); let action = exec .to_actual_join(&config_with_max_build(200 * MB)) .unwrap(); @@ -1288,11 +1352,109 @@ mod tests { assert!(rendered.contains("mode=Partitioned"), "{rendered}"); } + /// Builds a `DynamicJoinSelectionExec` in `Repartitioned` state whose *both* + /// children are resolved `ExchangeExec`s with the given per-partition byte + /// sizes. This is the shape `SelectJoinRule` produces after resolving a + /// `Repartition` action, and the only shape in which the swap can move the + /// build side onto `right`. + fn repartitioned_join_between_exchanges( + left_partition_bytes: &[usize], + right_partition_bytes: &[usize], + ) -> DynamicJoinSelectionExec { + let left = test_exchange_with_partition_bytes(left_partition_bytes); + let right = test_exchange_with_partition_bytes(right_partition_bytes); + // Both exchanges wrap test_schema(), whose sole field is named "x". + let on: JoinOn = + vec![(Arc::new(Column::new("x", 0)), Arc::new(Column::new("x", 0)))]; + DynamicJoinSelectionExec { + properties: Arc::clone(left.properties()), + left, + right, + on, + filter: None, + join_type: JoinType::Inner, + projection: None, + null_equality: NullEquality::NullEqualsNothing, + selection_state: JoinInputState::Repartitioned, + null_aware: false, + plan_id: 0, + } + } + + // The fit check must measure the input the build task actually runs. The + // left side is the larger one in total (400 MB against 303 MB), so + // `SelectJoinRule` swaps the inputs and the right becomes the build — and + // the right carries a 300 MB partition, over the 200 MB budget. Measuring + // the pre-swap left would see only its 100 MB max and keep the hash join, + // leaving the build task to OOM on a partition the check never looked at. + #[test] + fn build_fit_check_measures_the_post_swap_build_side() { + let exec = repartitioned_join_between_exchanges( + &[100 * MB, 100 * MB, 100 * MB, 100 * MB], // total 400 MB, max 100 MB + &[MB, MB, MB, 300 * MB], // total 303 MB, max 300 MB + ); + + let action = exec + .to_actual_join(&config_with_max_build(200 * MB)) + .unwrap(); + let rendered = displayable(resolved_plan(&action).as_ref()) + .indent(false) + .to_string(); + assert!(rendered.contains("SortMergeJoinExec"), "{rendered}"); + } + + // The mirror case: same partition shapes, sides exchanged, so the left is + // now the smaller total and no swap happens. The build stays on the left, + // and its 300 MB partition is what must be measured. Pins that the fix + // selects by the swap decision rather than blanket-switching to the right. + #[test] + fn build_fit_check_stays_on_the_left_when_no_swap_happens() { + let exec = repartitioned_join_between_exchanges( + &[MB, MB, MB, 300 * MB], // total 303 MB, max 300 MB + &[100 * MB, 100 * MB, 100 * MB, 100 * MB], // total 400 MB, max 100 MB + ); + + let action = exec + .to_actual_join(&config_with_max_build(200 * MB)) + .unwrap(); + let rendered = displayable(resolved_plan(&action).as_ref()) + .indent(false) + .to_string(); + assert!(rendered.contains("SortMergeJoinExec"), "{rendered}"); + } + + // `datafusion.optimizer.prefer_hash_join` is a planner flag that DataFusion + // has already applied by the time `DelayJoinSelectionRule` folds both join + // operators into this node, so AQE must not consult it again. A build side + // that fits stays a Partitioned hash join under either setting; the strategy + // follows runtime evidence, not the flag. + #[test] + fn prefer_hash_join_does_not_affect_the_aqe_decision() { + for prefer_hash_join in [true, false] { + let exec = + repartitioned_join_between_exchanges(&[50 * MB], PROBE_PARTITION_BYTES); + let mut config = config_with_max_build(200 * MB); + config.optimizer.prefer_hash_join = prefer_hash_join; + + let action = exec.to_actual_join(&config).unwrap(); + let rendered = displayable(resolved_plan(&action).as_ref()) + .indent(false) + .to_string(); + + assert!( + rendered.contains("HashJoinExec") + && rendered.contains("mode=Partitioned"), + "prefer_hash_join={prefer_hash_join} must not change the decision: {rendered}" + ); + } + } + // `hash_join_max_build_partition_bytes = 0` disables the check entirely, // matching the design contract: 0 always fits regardless of build size. #[test] fn zero_threshold_disables_the_check() { - let exec = repartitioned_join_with_build_partition_bytes(&[500 * MB]); + let exec = + repartitioned_join_between_exchanges(&[500 * MB], PROBE_PARTITION_BYTES); let action = exec.to_actual_join(&config_with_max_build(0)).unwrap(); let rendered = displayable(resolved_plan(&action).as_ref()) .indent(false) diff --git a/ballista/scheduler/src/state/aqe/test/join_selection.rs b/ballista/scheduler/src/state/aqe/test/join_selection.rs index aa4aac860..428cb7c6f 100644 --- a/ballista/scheduler/src/state/aqe/test/join_selection.rs +++ b/ballista/scheduler/src/state/aqe/test/join_selection.rs @@ -75,9 +75,30 @@ fn make_ctx(prefer_hash_join: bool) -> SessionContext { SessionContext::new_with_state(state) } -fn make_ctx_without_collect_left(prefer_hash_join: bool) -> SessionContext { +/// A build budget comfortably above the mock shuffle partitions' 10 bytes, so +/// the build-fit check passes and AQE keeps the `Partitioned` hash join. This is +/// the shipped default. +const FITS_BUILD_BUDGET: usize = 64 * 1024 * 1024; + +/// A build budget below the mock shuffle partitions' 10 bytes, so the build-fit +/// check rejects the hash join and AQE falls back to `SortMergeJoinExec`. Not +/// `0` — `0` disables the check entirely, which would keep the hash join. +const REJECTS_BUILD_BUDGET: usize = 1; + +/// `max_build_partition_bytes` is what selects AQE's join strategy once the +/// inputs are repartitioned: `mock_partitions_with_statistics` reports 10 bytes +/// per shuffle partition, so a budget below that fails the build-fit check and +/// falls back to `SortMergeJoinExec`, while a budget above it keeps the +/// `Partitioned` `HashJoinExec`. +/// +/// `datafusion.optimizer.prefer_hash_join` is deliberately not set here — AQE +/// does not consult it (DataFusion has already applied it by the time +/// `DelayJoinSelectionRule` folds both join operators into +/// `DynamicJoinSelectionExec`), so it cannot express the distinction these +/// tests draw. +fn make_ctx_without_collect_left(max_build_partition_bytes: usize) -> SessionContext { let config = SessionConfig::new() - .set_bool("datafusion.optimizer.prefer_hash_join", prefer_hash_join) + .with_ballista_hash_join_max_build_partition_bytes(max_build_partition_bytes) .set_u64( "datafusion.optimizer.hash_join_single_partition_threshold", 0, @@ -180,7 +201,7 @@ async fn test_left_join_not_collected_left() -> datafusion::common::Result<()> { #[tokio::test] async fn test_hash_join_two_tables_repartition() -> datafusion::common::Result<()> { - let ctx = make_ctx_without_collect_left(true); + let ctx = make_ctx_without_collect_left(FITS_BUILD_BUDGET); register_2tables(&ctx); let lp = ctx @@ -241,7 +262,7 @@ async fn test_hash_join_two_tables_repartition() -> datafusion::common::Result<( #[tokio::test] async fn test_sort_merge_join_two_tables_repartition() -> datafusion::common::Result<()> { - let ctx = make_ctx_without_collect_left(false); + let ctx = make_ctx_without_collect_left(REJECTS_BUILD_BUDGET); register_2tables(&ctx); let lp = ctx @@ -379,7 +400,7 @@ async fn test_hash_join_three_tables_collect_left() -> datafusion::common::Resul #[tokio::test] async fn test_hash_join_three_tables_repartition() -> datafusion::common::Result<()> { - let ctx = make_ctx_without_collect_left(true); + let ctx = make_ctx_without_collect_left(FITS_BUILD_BUDGET); register_3tables(&ctx); let lp = ctx @@ -410,7 +431,7 @@ async fn test_hash_join_three_tables_repartition() -> datafusion::common::Result #[tokio::test] async fn test_sort_merge_join_three_tables_repartition() -> datafusion::common::Result<()> { - let ctx = make_ctx_without_collect_left(false); + let ctx = make_ctx_without_collect_left(REJECTS_BUILD_BUDGET); register_3tables(&ctx); let lp = ctx diff --git a/docs/source/contributors-guide/benchmarking.md b/docs/source/contributors-guide/benchmarking.md index 88352e70d..94b608623 100644 --- a/docs/source/contributors-guide/benchmarking.md +++ b/docs/source/contributors-guide/benchmarking.md @@ -149,20 +149,24 @@ tpch benchmark ballista \ --query 18 \ --path /mnt/bigdata/tpch/sf1000 --format parquet \ --partitions 64 --iterations 1 \ - -c datafusion.optimizer.prefer_hash_join=false \ - -c datafusion.optimizer.enable_dynamic_filter_pushdown=false \ -c ballista.planner.adaptive.enabled=true ``` -Omit `--query` to run all 22. `ballista.planner.adaptive.enabled=true` is the AQE-on -configuration these results use. - -`prefer_hash_join=false` makes DataFusion plan joins as `SortMergeJoin`, which is the -join strategy this page compares across engines (Comet is configured to match, below). - -Note `datafusion.optimizer.enable_dynamic_filter_pushdown=false`: DataFusion's -dynamic filter pushdown assumes single-process execution and can deadlock -distributed execution, so it is pinned off for benchmark runs. +Omit `--query` to run all 22, and add `--skip ` (repeatable) to leave individual +queries out of that run — `--skip 18` is what produced the Ballista column below. +`ballista.planner.adaptive.enabled=true` is the AQE-on configuration these results +use. + +Note that `datafusion.optimizer.prefer_hash_join` is deliberately left at its default. +Under AQE it does not select the join strategy: `DelayJoinSelectionRule` folds both +`HashJoinExec` and `SortMergeJoinExec` into a single `DynamicJoinSelectionExec`, and +the strategy is then chosen from runtime statistics — the broadcast thresholds, and +`ballista.optimizer.hash_join_max_build_partition_bytes` (64 MiB by default), which +lowers a Partitioned hash join to the spillable `SortMergeJoin` when the build side's +largest partition would not fit. Setting `prefer_hash_join=false` does not force +sort-merge; it only changes which conversion path the plan takes, and the +`SortMergeJoinExec` path discards the join's projection, so it measures a slower plan +rather than a different join strategy. ### Spark @@ -222,9 +226,13 @@ spark-submit \ ``` `spark.comet.exec.replaceSortMergeJoin=false` keeps Comet on Spark's `SortMergeJoin` -instead of converting it to a shuffled hash join. This matches Ballista, which plans -`SortMergeJoin` (`prefer_hash_join=false`), so the two engines are compared on the -same join strategy rather than one being handed a different one. +instead of converting it to a shuffled hash join. + +This no longer pins the two engines to the same join strategy. Ballista under AQE +chooses per join at runtime — broadcast, Partitioned hash, or sort-merge — and there +is no supported way to force it to sort-merge throughout (see the note in the Ballista +section above). So a Ballista-vs-Comet gap may partly reflect different join +strategies, not only different execution. Read the comparison with that in mind. [comet]: https://datafusion.apache.org/comet/ @@ -264,61 +272,69 @@ planner. ## Results TPC-H **SF1000**, reference cluster above, **AQE on**, 1 iteration, -`target_partitions=64`, `prefer_hash_join=false`, -`enable_dynamic_filter_pushdown=false`, and the sort-shuffle spill cap overridden to -uncapped (`ballista.shuffle.sort_based.memory_limit_per_task_bytes=0`; the shipped -default is 256 MiB). All three engines plan `SortMergeJoin`. Times in seconds; lower -is better. +`target_partitions=64`, all Ballista join settings at their defaults. Ballista picks +each join's strategy at runtime; Spark and Comet are held to `SortMergeJoin`. Times in +seconds; lower is better. Versions under test: | Engine | Version | | -------- | ----------------------------------------------- | -| Ballista | `#2084` @ `fac1fa22` | +| Ballista | `#2124` @ `5290436d` | | Spark | 3.5.3 (vanilla, Comet disabled) | | Comet | `main` @ `b0165552` (1.0.0-SNAPSHOT, Spark 3.5) | Pin the **exact commit** the numbers came from, not "main": `main` moves, and a row that mixes numbers from different commits silently misattributes a regression. -Ballista Q1–Q17 come from a **full 22-query suite run** (one query after another on -freshly started executors); Q18 exhausts the memory pool and OOM-kills the executors -(see below), which currently wedges the rest of the suite, so **Q19–Q22 were run as -individual single-query jobs** against a fresh cluster of the same shape. Comet and -Spark each completed the full suite in one run. +The Ballista column comes from a **single suite run of 21 queries** on freshly +started executors, using `--skip 18` to leave out the one query that exhausts the +memory pool and OOM-kills the executors (see below), which would otherwise wedge +the rest of the suite. Comet and Spark each completed the full suite in one run. + +Keep the whole column to one continuous run wherever possible. Driving queries as +individual single-query jobs measures something different: each job starts against +a cold page cache, which on this cluster cost the two join-free queries (Q1, Q6) +roughly 20 s and 20 s respectively when the two methods were compared directly. Use +`--skip` rather than a loop of `--query` runs. + +Single iteration means **run-to-run variance is not quantified here**. Comparing the +same commit across the two methods above showed spreads of up to ~20% on individual +queries, so treat a per-query difference under about 20% as noise and read the total +row instead. | Query | Spark | Comet | Ballista (AQE on) | Rows | | --------------------: | ---------: | ---------: | ----------------: | -----: | -| 1 | 427.1 | 46.6 | 59.4 | 4 | -| 2 | 75.9 | 39.2 | 57.0 | 100 | -| 3 | 154.1 | 195.4 | 214.8 | 10 | -| 4 | 82.0 | 42.6 | 65.2 | 5 | -| 5 | 336.4 | 493.4 | 456.6 | 5 | -| 6 | 28.9 | 14.6 | 29.5 | 1 | -| 7 | 180.6 | 149.5 | 369.0 | 4 | -| 8 | 391.8 | 642.1 | 540.3 | 2 | -| 9 | 509.8 | 843.5 | 764.9 | 175 | -| 10 | 151.1 | 104.8 | 168.3 | 20 | -| 11 | 44.7 | 44.1 | 91.4 | 0 [1] | -| 12 | 74.1 | 49.9 | 74.9 | 2 | -| 13 | 98.7 | 58.3 | 91.9 | 30 | -| 14 | 43.0 | 27.4 | 43.0 | 1 | -| 15 | 121.5 | 65.6 | 116.5 | 1 [2] | -| 16 | 24.5 | 17.5 | 29.0 | 27840 | -| 17 | 406.7 | 285.8 | 392.8 | 1 | -| 18 | 428.8 | 370.5 | OOM | 100 | -| 19 | 58.9 | 35.5 | 183.9 | 1 | -| 20 | 105.9 | 67.6 | 112.2 | 110759 | -| 21 | 562.2 | 460.1 | 765.2 | 100 | -| 22 | 36.8 | 21.5 | 35.2 | 7 | -| **Total (excl. Q18)** | **3914.7** | **3705.0** | **4661.0** | | - -The **Total** row sums the 21 queries **excluding Q18**, because Q18 does not complete -on Ballista at this sizing (below), so a 22-query total would not be comparable across -engines. Q18's own times are in its row. - -Row counts agree across all three engines on every query, including Q18 (Spark and -Comet return 100; Ballista OOMs before producing a result). +| 1 | 427.1 | 46.6 | 54.2 | 4 | +| 2 | 75.9 | 39.2 | 60.3 | 100 | +| 3 | 154.1 | 195.4 | 140.0 | 10 | +| 4 | 82.0 | 42.6 | 46.3 | 5 | +| 5 | 336.4 | 493.4 | 327.2 | 5 | +| 6 | 28.9 | 14.6 | 27.2 | 1 | +| 7 | 180.6 | 149.5 | 364.9 | 4 | +| 8 | 391.8 | 642.1 | 481.1 | 2 | +| 9 | 509.8 | 843.5 | 606.4 | 175 | +| 10 | 151.1 | 104.8 | 148.4 | 20 | +| 11 | 44.7 | 44.1 | 62.0 | 0 [1] | +| 12 | 74.1 | 49.9 | 58.3 | 2 | +| 13 | 98.7 | 58.3 | 100.4 | 30 | +| 14 | 43.0 | 27.4 | 71.1 | 1 | +| 15 | 121.5 | 65.6 | 53.8 | 1 [2] | +| 16 | 24.5 | 17.5 | 34.3 | 27840 | +| 17 | 406.7 | 285.8 | 417.8 | 1 | +| 18 | 428.8 | 370.5 | TBD | 100 | +| 19 | 58.9 | 35.5 | 68.5 | 1 | +| 20 | 105.9 | 67.6 | 93.7 | 110759 | +| 21 | 562.2 | 460.1 | 698.0 | 100 | +| 22 | 36.8 | 21.5 | 22.9 | 7 | +| **Total (excl. Q18)** | **3914.7** | **3705.0** | **3936.8** | | + +The **Total** row sums the 21 queries **excluding Q18**, because Q18 has no Ballista +figure at this commit (below), so a 22-query total would not be comparable across +engines. Spark's and Comet's Q18 times are in its row. + +Row counts agree across all three engines on every query for which Ballista produced +a result. The **AQE-off** Ballista column has been removed pending a re-run at a matched core count; the numbers above for AQE off were taken at a different executor size and are @@ -331,12 +347,20 @@ constant is tuned for SF1. engines report 1 row here. (A previous result set saw Spark report 0 for Q15 depending on which statement the harness took as the result; it does not recur in this run.) -**Q18 OOMs on Ballista at this sizing.** Q18's hash-join build side is `Partitioned` +**Q18 is not measured at this commit.** Q18's hash-join build side is `Partitioned` and does not spill; with 16 task slots sharing the 48 GB pool (3 GB per slot), the per-task build side exceeds the container limit and the executor is OOM-killed ([#2025](https://github.com/apache/datafusion-ballista/issues/2025)). At the previous 8-slot sizing (6 GB per slot) Q18 completed, so this is a direct consequence of the -denser packing. The OOM is recorded as `OOM` rather than a time. +denser packing. + +It was excluded from the suite run above with `--skip 18` so that one query could not +wedge the other 21, and it has not since been run to completion at this commit — hence +`TBD` rather than `OOM`. An earlier attempt on this branch did OOM-kill an executor, +but under a different configuration (`prefer_hash_join=false`, which sends the plan +down a conversion path that discards the join's projection and so materialises more +columns than the configuration documented above). That result does not transfer, so it +is not recorded here. The `Rows` column is the row count the query returned, recorded so a time is never read without the answer it produced. @@ -350,63 +374,6 @@ regression. did not produce an answer is recorded as `FAIL`, or `OOM` where the failure is a known memory exhaustion. -### Hash join with a per-partition build-size fallback (AQE on, 2×16 cores, 64 partitions) - -A second Ballista configuration exercises `prefer_hash_join=true` together with the -AQE hash-join safety fallback -([`ballista.optimizer.hash_join_max_build_partition_bytes`](https://github.com/apache/datafusion-ballista/pull/2084)). -The fallback lowers a Partitioned hash join to `SortMergeJoin` per join when the -build side's largest partition exceeds the configured budget, so a hash-join run does -not abort on the queries whose non-spillable hash-join build exceeds one task slot's -pool — notably Q18 -([#2025](https://github.com/apache/datafusion-ballista/issues/2025)). That fallback -is what makes `prefer_hash_join=true` usable across the whole suite instead of -failing partway. - -This run uses a **larger cluster than the reference above** and a **different -`target_partitions`**, so its numbers are not comparable to the table above: TPC-H -**SF1000**, **2 executors × 16 cores** (one per node, ~2.8 GB per task slot), -**`target_partitions=64`**, **AQE on**, `prefer_hash_join=true`, -`hash_join_max_build_partition_bytes=67108864` (64 MiB), -`enable_dynamic_filter_pushdown=false`, sort-shuffle spill uncapped. Each query ran -on a **freshly restarted cluster**. Ballista @ `afef9afc`. Times in seconds. - -| Query | Ballista (AQE on, hash join + 64 MiB fallback) | -| --------: | ---------------------------------------------: | -| 1 | 72.1 | -| 2 | 52.9 | -| 3 | 154.2 | -| 4 | 72.9 | -| 5 | 355.4 | -| 6 | 65.9 | -| 7 | 371.7 | -| 8 | 452.4 | -| 9 | 559.8 | -| 10 | 152.3 | -| 11 | 74.6 | -| 12 | 88.7 | -| 13 | 99.2 | -| 14 | 38.6 | -| 15 | 62.5 | -| 16 | 29.0 | -| 17 | 446.3 | -| 18 | 744.6 | -| 19 | 80.8 | -| 20 | 117.0 | -| 21 | 605.6 | -| 22 | 20.5 | -| **Total** | **4716.9** | - -**All 22 queries completed — no OOMs, no failures.** At the 64 MiB threshold the -large-build joins (Q5, Q7–Q9, Q17, Q18, Q21) fall back to `SortMergeJoin` while -smaller-build joins stay hash, so the heavy queries run at roughly sort-merge speed -but the suite runs end to end. Q18 completes at **744.6 s** where the same -configuration with the fallback disabled (`hash_join_max_build_partition_bytes=0`) -exhausts the per-slot pool and fails the job -([#2025](https://github.com/apache/datafusion-ballista/issues/2025)). The fallback is -opt-in and off by default (a `0` budget); the 64 MiB value here is tuned to this -cluster's ~2.8 GB task-slot pool. - ### Recording a result - Pin the **exact commit** the numbers came from, not a branch name. @@ -414,9 +381,10 @@ cluster's ~2.8 GB task-slot pool. configuration this page measures — from the same commit and cluster. - Prefer the figure from a **full suite run**, not a standalone single-query run: a long-lived executor deep into a suite is not in the same state as a freshly started - one. When a query cannot complete in-suite (e.g. Q18's OOM currently wedges the - run), note that the remaining queries were run individually, as done here for - Q19–Q22. + one, and a single-query job pays a cold-cache penalty that is large enough to swamp + the effect being measured. When a query cannot complete in-suite (e.g. Q18's OOM + currently wedges the run), exclude it with `--skip ` so the rest of the column + still comes from one continuous run, and measure that query separately. - Note the **row count** each query returned. A fast wrong answer is not a result, and distributed execution has produced silently wrong row counts before. - Flag any stage whose runtime is dominated by a few partitions — that is the