diff --git a/ballista/client/tests/runtime_filter.rs b/ballista/client/tests/runtime_filter.rs new file mode 100644 index 0000000000..51cf8e571f --- /dev/null +++ b/ballista/client/tests/runtime_filter.rs @@ -0,0 +1,132 @@ +// 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(feature = "standalone")] + +use ballista::prelude::SessionContextExt; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::arrow::util::pretty::pretty_format_batches; +use datafusion::prelude::*; +use std::sync::Arc; + +async fn run(ctx: &SessionContext, sql: &str) -> String { + let batches = tokio::time::timeout(std::time::Duration::from_secs(30), async { + ctx.sql(sql).await.unwrap().collect().await.unwrap() + }) + .await + .expect("runtime filter must not deadlock the query"); + pretty_format_batches(&batches).unwrap().to_string() +} + +#[tokio::test] +async fn distributed_runtime_filter_matches_unfiltered_joins() { + let dir = tempfile::tempdir().unwrap(); + for name in ["build", "probe"] { + std::fs::create_dir(dir.path().join(name)).unwrap(); + } + std::fs::write(dir.path().join("build/a.csv"), "k,v\n2,1\n,2\n").unwrap(); + std::fs::write(dir.path().join("build/b.csv"), "k,v\n8,3\n2,4\n").unwrap(); + for part in 0..4 { + let mut csv = "k,v\n,0\n".to_string(); + for k in 0..25 { + csv.push_str(&format!("{},{}\n", part * 25 + k, k)); + } + std::fs::write(dir.path().join(format!("probe/{part}.csv")), csv).unwrap(); + } + let ctx = SessionContext::standalone().await.unwrap(); + for setting in [ + "ballista.planner.adaptive.enabled = true", + "datafusion.optimizer.prefer_hash_join = true", + "datafusion.execution.target_partitions = 4", + ] { + run(&ctx, &format!("SET {setting}")).await; + } + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, true), + Field::new("v", DataType::Int64, true), + ])); + // Parquet provides build cardinality, allowing AQE to choose a broadcast + // stage whose completion naturally precedes the probe scan. + let build_parquet = dir.path().join("build_parquet"); + let local = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(2)); + local + .read_csv( + dir.path().join("build").to_str().unwrap(), + CsvReadOptions::new().schema(&schema), + ) + .await + .unwrap() + .write_parquet(build_parquet.to_str().unwrap(), Default::default(), None) + .await + .unwrap(); + ctx.register_parquet( + "build", + build_parquet.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + ctx.register_csv( + "probe", + dir.path().join("probe").to_str().unwrap(), + CsvReadOptions::new().schema(&schema), + ) + .await + .unwrap(); + let queries = [ + "SELECT b.k, b.v, p.v FROM build b JOIN probe p ON b.k = p.k WHERE b.v > 0 AND p.v >= 0 ORDER BY b.k, b.v, p.v", + "SELECT b.k, b.v, p.v FROM build b JOIN probe p ON b.k IS NOT DISTINCT FROM p.k WHERE b.v > 0 AND p.v >= 0 ORDER BY b.k, b.v, p.v", + "SELECT p.k FROM build b RIGHT JOIN probe p ON b.k = p.k ORDER BY p.k", + "SELECT b.k, p.k FROM build b JOIN (SELECT k FROM probe ORDER BY k LIMIT 3) p ON b.k = p.k ORDER BY b.k, p.k", + ]; + for query in queries { + run(&ctx, "SET ballista.planner.runtime_filter.enabled = false").await; + let expected = run(&ctx, query).await; + run(&ctx, "SET ballista.planner.runtime_filter.enabled = true").await; + assert_eq!(run(&ctx, query).await, expected, "{query}"); + } + let plan = run(&ctx, &format!("EXPLAIN ANALYZE {}", queries[0])).await; + assert!(plan.contains("RuntimeFilterBuildExec"), "{plan}"); + assert!(plan.contains("RuntimeFilterApplyExec"), "{plan}"); + assert!(plan.contains("ready=true"), "{plan}"); + assert!( + plan.split("rows_pruned=").skip(1).any(|suffix| { + suffix + .chars() + .take_while(|c| c.is_ascii_digit()) + .collect::() + .parse::() + .is_ok_and(|rows| rows > 0) + }), + "runtime filter must actually prune rows: {plan}" + ); + // Independently runnable partitioned stages must preserve results even + // when the producer has not reported before the probe is dispatched. + for setting in [ + "ballista.optimizer.broadcast_join_threshold_bytes = 0", + "ballista.optimizer.broadcast_join_threshold_rows = 0", + "ballista.optimizer.stage_build_side = false", + ] { + run(&ctx, &format!("SET {setting}")).await; + } + let partitioned = run(&ctx, queries[0]).await; + let plan = run(&ctx, &format!("EXPLAIN ANALYZE {}", queries[0])).await; + assert!(plan.contains("mode=Partitioned"), "{plan}"); + run(&ctx, "SET ballista.planner.runtime_filter.enabled = false").await; + assert_eq!(partitioned, run(&ctx, queries[0]).await); +} diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index f6db57ddfa..c232d229b0 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -182,6 +182,10 @@ pub const BALLISTA_SEMI_JOIN_REDUCTION: &str = "ballista.optimizer.semi_join_red pub const BALLISTA_BALANCE_SCAN_FILE_GROUPS: &str = "ballista.optimizer.balance_scan_file_groups"; +/// Enable cross-stage hash-join runtime filters (AQE only). +pub const BALLISTA_RUNTIME_FILTER_ENABLED: &str = + "ballista.planner.runtime_filter.enabled"; + /// Configuration key to enable AQE coalesce-shuffle-partitions rule. /// Disabled by default — opt in when the workload benefits from larger /// downstream tasks more than from preserved parallelism. @@ -509,6 +513,9 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| "Wait time in milliseconds between IO retries in the Ballista client.".to_string(), DataType::UInt64, Some(3000.to_string())), + ConfigEntry::new(BALLISTA_RUNTIME_FILTER_ENABLED.to_string(), + "Enables bounded distributed runtime filters for AQE inner and semi hash joins.".to_string(), + DataType::Boolean, Some(false.to_string())), ConfigEntry::new(BALLISTA_COALESCE_ENABLED.to_string(), "Enables the AQE coalesce-shuffle-partitions rule. \ Disabled by default — opt in when fewer/larger \ @@ -1125,6 +1132,11 @@ impl BallistaConfig { self.get_bool_setting(BALLISTA_BALANCE_SCAN_FILE_GROUPS) } + /// Whether distributed runtime filtering is enabled. + pub fn runtime_filter_enabled(&self) -> bool { + self.get_bool_setting(BALLISTA_RUNTIME_FILTER_ENABLED) + } + /// Returns whether the AQE coalesce-shuffle-partitions rule is enabled. pub fn coalesce_enabled(&self) -> bool { self.get_bool_setting(BALLISTA_COALESCE_ENABLED) diff --git a/ballista/core/src/execution_plans/runtime_filter.rs b/ballista/core/src/execution_plans/runtime_filter.rs index b4e079d2bd..ecafb917e9 100644 --- a/ballista/core/src/execution_plans/runtime_filter.rs +++ b/ballista/core/src/execution_plans/runtime_filter.rs @@ -114,25 +114,46 @@ fn bloom_bits(hash: u64) -> impl Iterator { } impl RuntimeFilter { + /// Whether the summary accepts every non-NULL key. Union cannot make a + /// saturated Bloom selective again; its range bounds may still be useful. + pub fn is_saturated(&self) -> bool { + self.in_hashes.is_none() && self.bloom.iter().all(|&byte| byte == u8::MAX) + } + + fn is_nearly_saturated(&self) -> bool { + // With four probes, 99% occupancy implies about 96% false positives. + // Keep the complete payload and its bounds, but avoid probing this Bloom. + self.in_hashes.is_none() + && self + .bloom + .iter() + .map(|byte| byte.count_ones() as usize) + .sum::() + * 100 + >= BLOOM_BYTES * 8 * 99 + } + fn update(&mut self, keys: &ArrayRef) -> Result<()> { self.has_null |= keys.null_count() > 0; - with_hashes([keys], &HASH_STATE, |hashes| { - for (row, &hash) in hashes.iter().enumerate() { - if keys.is_null(row) { - continue; - } - for bit in bloom_bits(hash) { - self.bloom[bit / 8] |= 1 << (bit % 8); - } - if let Some(values) = &mut self.in_hashes { - values.insert(hash); - if values.len() > MAX_VALUES { - self.in_hashes = None; + if !self.is_saturated() { + with_hashes([keys], &HASH_STATE, |hashes| { + for (row, &hash) in hashes.iter().enumerate() { + if keys.is_null(row) { + continue; + } + for bit in bloom_bits(hash) { + self.bloom[bit / 8] |= 1 << (bit % 8); + } + if let Some(values) = &mut self.in_hashes { + values.insert(hash); + if values.len() > MAX_VALUES { + self.in_hashes = None; + } } } - } - Ok(()) - })?; + Ok(()) + })?; + } if !self.bounds_disabled && keys.null_count() != keys.len() { let mut min = MinAccumulator::try_new(keys.data_type())?; let mut max = MaxAccumulator::try_new(keys.data_type())?; @@ -406,6 +427,16 @@ impl RuntimeFilterApplyExec { metrics: ExecutionPlanMetricsSet::new(), }) } + + /// Bind a complete report while preserving the input and join key. + pub fn with_payload(&self, payload: RuntimeFilterPayload) -> Result { + Self::try_new( + self.input.clone(), + self.filter_id, + self.key.clone(), + Some(payload), + ) + } } impl DisplayAs for RuntimeFilterApplyExec { @@ -481,6 +512,12 @@ impl ExecutionPlan for RuntimeFilterApplyExec { let Some(filter) = self.filter.clone() else { return Ok(input); }; + // Preserve the payload's bounds for scan pruning, but avoid hashing + // batches against a nearly full Bloom with little filtering capacity. + if filter.is_nearly_saturated() { + return Ok(input); + } + let key = self.key.clone(); let pruned = MetricBuilder::new(&self.metrics).counter("rows_pruned", partition); let stream = input.map(move |batch| { @@ -674,4 +711,63 @@ mod tests { }) .unwrap(); } + + #[test] + fn runtime_filter_saturation_preserves_later_bounds_and_nulls() { + let mut filter = RuntimeFilter::default(); + filter + .update(&(Arc::new(Int64Array::from_iter_values(0..1_000_000)) as ArrayRef)) + .unwrap(); + assert!(filter.is_saturated()); + filter + .update( + &(Arc::new(Int64Array::from(vec![Some(-1), Some(2_000_000), None])) + as ArrayRef), + ) + .unwrap(); + let payload = filter.to_proto(1).unwrap(); + assert!(payload.has_null); + assert_eq!( + ScalarValue::try_from(payload.min.as_ref().unwrap()).unwrap(), + ScalarValue::Int64(Some(-1)) + ); + assert_eq!( + ScalarValue::try_from(payload.max.as_ref().unwrap()).unwrap(), + ScalarValue::Int64(Some(2_000_000)) + ); + } + + #[tokio::test] + async fn runtime_filter_nearly_full_bloom_bypasses_row_filtering() { + let mut filter = RuntimeFilter::default(); + filter + .update(&(Arc::new(Int64Array::from_iter_values(0..120_000)) as ArrayRef)) + .unwrap(); + assert!(!filter.is_saturated()); + let batch = RecordBatch::try_from_iter([( + "key", + Arc::new(Int64Array::from(vec![Some(2), None])) as ArrayRef, + )]) + .unwrap(); + let input = MemorySourceConfig::try_new_exec( + &[vec![batch.clone()]], + batch.schema(), + None, + ) + .unwrap(); + let apply: Arc = Arc::new( + RuntimeFilterApplyExec::try_new( + input, + 7, + Arc::new(Column::new("key", 0)), + Some(filter.to_proto(7).unwrap()), + ) + .unwrap(), + ); + let result = collect(apply, SessionContext::new().task_ctx()) + .await + .unwrap(); + // Bypass passes the batch through, including NULLs the join will check. + assert_eq!(result, vec![batch]); + } } diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index 926c7d21b8..1759316fd3 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -821,8 +821,9 @@ impl SessionConfigHelperExt for SessionConfig { // boundaries that run as independent tasks, so when the producing // operator and the consuming scan land in different stages the // filter is never populated across the boundary and the scan blocks - // forever. Disable dynamic filter pushdown until Ballista can carry - // dynamic filters across stage boundaries. + // forever. Keep these local dynamic expressions disabled. Ballista's + // opt-in runtime-filter hub transports independent completed join + // summaries; it does not populate DataFusion's dynamic expressions. // // See https://github.com/apache/datafusion-ballista/issues/1375 .set_bool("datafusion.optimizer.enable_dynamic_filter_pushdown", false) diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index f8b7123f7f..9af263b7a5 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -73,6 +73,8 @@ pub(crate) mod execution_plan; pub mod optimizer_rule; pub mod planner; pub(crate) mod reuse_exchange; +mod runtime_filter; +mod runtime_filter_planner; pub(crate) mod skew_join; #[cfg(test)] mod test; @@ -110,6 +112,7 @@ pub(crate) struct AdaptiveExecutionGraph { scheduler_id: Option, /// Adaptive Planner to be used with this execution graph planner: AdaptivePlanner, + runtime_filters: Option, /// ID for this job job_id: JobId, /// Job name, can be empty string @@ -243,8 +246,23 @@ impl AdaptiveExecutionGraph { .collect(); let stages = stages?; + let filter_config = session_config + .options() + .extensions + .get::(); + let runtime_filters = filter_config + .filter(|config| config.runtime_filter_enabled()) + .map(|_| { + let mut registry = runtime_filter::RuntimeFilterRegistry::default(); + for (&id, stage) in &stages { + registry.register(id, stage.plan()); + } + registry + }); + Ok(Self { planner, + runtime_filters, scheduler_id: Some(scheduler_id.to_string()), job_id: job_id.to_owned(), job_name: job_name.to_owned(), @@ -577,6 +595,9 @@ impl AdaptiveExecutionGraph { self.session_config.clone(), stage, )?; + if let Some(registry) = &mut self.runtime_filters { + registry.register(stage_id, stage.plan()); + } self.stages.insert(stage_id, stage); } } @@ -657,6 +678,24 @@ impl AdaptiveExecutionGraph { } }); + if let (Some(id), Some(registry)) = (running_stage_id, &mut self.runtime_filters) + { + let plan = self.stages[&id].plan_arc(); + let prepared = match registry.prepare(id, plan.clone(), &self.stages) { + Ok(plan) => plan, + Err(error) => { + log::warn!( + "Runtime filter unavailable for stage {id}; bypassing: {error}" + ); + plan + } + }; + let ExecutionStage::Running(stage) = self.stages.get_mut(&id).unwrap() else { + unreachable!() + }; + stage.plan = prepared; + } + // If no available tasks found in the running stage, // try to find a resolved stage and convert it to the running stage if running_stage_id.is_none() { @@ -1635,13 +1674,8 @@ impl ExecutionGraph for AdaptiveExecutionGraph { let job_id = self.job_id.clone(); let session_id = self.session_id.clone(); - let mut next_task = self.stages.iter_mut().find(|(_stage_id, stage)| { - if let ExecutionStage::Running(stage) = stage { - stage.available_tasks() > 0 - } else { - false - } - }).map(|(stage_id, stage)| { + let ready_id = self.get_running_stage_id(&[]); + let mut next_task = ready_id.and_then(|id| self.stages.get_mut(&id).map(|stage| (id, stage))).map(|(stage_id, stage)| { if let ExecutionStage::Running(stage) = stage { use ballista_core::serde::scheduler::TaskKey; @@ -1680,7 +1714,7 @@ impl ExecutionGraph for AdaptiveExecutionGraph { let key = TaskKey { job_id, - stage_id: *stage_id, + stage_id, task_id, }; diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index dac30cf578..a3fb723344 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -104,6 +104,7 @@ pub struct AdaptivePlanner { /// instead of running a stage of its own, by its key (see /// [`crate::state::aqe::reuse_exchange`]). Spark's AQE `stageCache`. reusable_stages: HashMap>, + runtime_filters: Option, } impl Debug for AdaptivePlanner { @@ -136,7 +137,15 @@ impl AdaptivePlanner { let planner = DefaultPhysicalPlanner::default(); let plan = planner.optimize_physical_plan(plan, &session_state, |_, _| {})?; + let runtime_filters = session_state + .config() + .options() + .extensions + .get::() + .is_some_and(|c| c.runtime_filter_enabled()) + .then(super::runtime_filter_planner::RuntimeFilterPlanner::default); Ok(Self { + runtime_filters, stage_id_generator: 0, // FIXME: compatibility issue with static where stages start from 1 session_state, physical_planner: planner.into(), @@ -486,6 +495,9 @@ impl AdaptivePlanner { Ok((None, stages_to_cancel)) } Some(stages) => { + if let Some(filters) = &mut self.runtime_filters { + filters.discover(&self.plan)?; + } let config = self.session_state.config().options(); let (stage_ids, shuffle_writers) = stages .into_iter() @@ -498,6 +510,10 @@ impl AdaptivePlanner { // plan in `default_optimizers()`. let plan = CoalescePartitionsRule.optimize(plan, config)?; let plan = OptimizeSkewedJoinRule.optimize(plan, config)?; + let plan = match &mut self.runtime_filters { + Some(filters) => filters.instrument(plan)?, + None => plan, + }; BallistaAdapter::adapt_to_ballista(plan, &self.job_id, config) .map(|w| (w.plan.stage_id(), w)) }) diff --git a/ballista/scheduler/src/state/aqe/runtime_filter.rs b/ballista/scheduler/src/state/aqe/runtime_filter.rs new file mode 100644 index 0000000000..062f8f7e53 --- /dev/null +++ b/ballista/scheduler/src/state/aqe/runtime_filter.rs @@ -0,0 +1,225 @@ +// 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. + +//! Job-scoped runtime filter hub. Keeping it in the execution graph gives +//! `(job_id, filter_id)` isolation and ties cleanup to the job's existing lifetime. +//! Reports stay in TaskInfo: task/stage retry and executor-loss handling already +//! invalidate those statuses, so a separate partial-report cache cannot go stale. + +use crate::state::execution_stage::ExecutionStage; +use ballista_core::execution_plans::runtime_filter::RuntimeFilter; +use ballista_core::execution_plans::{RuntimeFilterApplyExec, RuntimeFilterBuildExec}; +use ballista_core::serde::protobuf::{RuntimeFilterPayload, task_status}; +use datafusion::common::Result; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::physical_plan::ExecutionPlan; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +#[derive(Debug, Clone, Default)] +pub(super) struct RuntimeFilterRegistry { + producers: HashMap, + prepared: HashSet, +} + +impl RuntimeFilterRegistry { + pub fn register(&mut self, stage_id: usize, plan: &dyn ExecutionPlan) { + if let Some(build) = plan.downcast_ref::() { + self.producers.insert(build.filter_id, stage_id); + } + for child in plan.children() { + self.register(stage_id, child.as_ref()); + } + } + + /// Freeze available complete filters before this stage's first dispatch. + /// Missing or unfinished producers leave their consumers as pass-throughs. + pub fn prepare( + &mut self, + stage_id: usize, + plan: Arc, + stages: &HashMap, + ) -> Result> { + // Binding is attempted once, including when the caller falls back + // after an invalid report. Later tasks must keep the same stage plan. + if !self.prepared.insert(stage_id) { + return Ok(plan); + } + let plan = plan + .transform_up(|node| { + let Some(apply) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(payload) = completed_filter( + apply.filter_id, + self.producers + .get(&apply.filter_id) + .and_then(|id| stages.get(id)), + )? + else { + return Ok(Transformed::yes(apply.children()[0].clone())); + }; + Ok(Transformed::yes(Arc::new(apply.with_payload(payload)?))) + })? + .data; + Ok(plan) + } +} + +fn completed_filter( + id: u64, + stage: Option<&ExecutionStage>, +) -> Result> { + let tasks = match stage { + Some(ExecutionStage::Successful(stage)) => &stage.task_infos, + Some(ExecutionStage::Running(stage)) if stage.is_successful() => { + &stage.task_infos + } + _ => return Ok(None), + }; + let mut merged = RuntimeFilter::default(); + for task in tasks { + // Existing retry/loss handling invalidates old successful slices. + if let task_status::Status::Successful(success) = &task.task_status { + let Some(report) = success.runtime_filters.iter().find(|p| p.filter_id == id) + else { + return Ok(None); + }; + merged.merge(&RuntimeFilter::from_proto(report)?); + } + } + merged.to_proto(id).map(Some) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::state::execution_graph::create_task_info; + use crate::state::execution_stage::{RunningStage, TaskInfo}; + use ballista_core::execution_plans::runtime_filter::collect_reports; + use ballista_core::serde::protobuf::SuccessfulTask; + use datafusion::arrow::array::{Int64Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_plan::collect; + use datafusion::prelude::{SessionConfig, SessionContext}; + + fn source(values: Vec) -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from(values))], + ) + .unwrap(); + MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + async fn report(value: i64) -> RuntimeFilterPayload { + let build: Arc = Arc::new(RuntimeFilterBuildExec::new( + source(vec![value]), + 3, + Arc::new(Column::new("k", 0)), + )); + collect(build.clone(), SessionContext::new().task_ctx()) + .await + .unwrap(); + collect_reports(&build).unwrap().pop().unwrap() + } + + fn task( + id: usize, + partition: usize, + payload: Option, + ) -> TaskInfo { + let mut task = create_task_info("e".into(), id); + task.global_input_partition_ids = vec![partition]; + if let Some(payload) = payload { + task.task_status = task_status::Status::Successful(SuccessfulTask { + runtime_filters: vec![payload], + ..Default::default() + }); + } + task + } + + async fn keys(plan: Arc) -> Vec { + let batches = collect(plan, SessionContext::new().task_ctx()) + .await + .unwrap(); + batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect() + } + + #[tokio::test] + async fn runtime_filter_hub_uses_only_complete_current_attempts() { + let build: Arc = Arc::new(RuntimeFilterBuildExec::new( + source(vec![]), + 3, + Arc::new(Column::new("k", 0)), + )); + let mut stage = RunningStage::new( + 1, + 0, + build, + 2, + vec![], + HashMap::new(), + Arc::new(SessionConfig::new()), + ); + stage.pending.next_slice(2); + stage.task_infos = vec![task(0, 0, Some(report(1).await)), task(1, 1, None)]; + let mut stages = HashMap::from([(1, ExecutionStage::Running(stage))]); + let probe: Arc = Arc::new( + RuntimeFilterApplyExec::try_new( + source(vec![1, 2, 3, 4]), + 3, + Arc::new(Column::new("k", 0)), + None, + ) + .unwrap(), + ); + let mut hub = RuntimeFilterRegistry::default(); + hub.register(1, stages[&1].plan()); + let bypass = hub.prepare(2, probe.clone(), &stages).unwrap(); + assert_eq!(keys(bypass.clone()).await, [1, 2, 3, 4]); + let ExecutionStage::Running(stage) = stages.get_mut(&1).unwrap() else { + unreachable!() + }; + // Drop the lost attempt's key 1; only replacement slices may contribute. + stage.reset_task_info(0); + stage.task_infos[1] = task(1, 1, Some(report(2).await)); + stage.pending.next_slice(1); + stage.task_infos.push(task(2, 0, Some(report(3).await))); + // A stage already dispatched without a filter stays unchanged. + let frozen = hub.prepare(2, bypass, &stages).unwrap(); + assert_eq!(keys(frozen).await, [1, 2, 3, 4]); + let plan = hub.prepare(3, probe, &stages).unwrap(); + assert_eq!(keys(plan).await, [2, 3]); + } +} diff --git a/ballista/scheduler/src/state/aqe/runtime_filter_planner.rs b/ballista/scheduler/src/state/aqe/runtime_filter_planner.rs new file mode 100644 index 0000000000..2779e2a77d --- /dev/null +++ b/ballista/scheduler/src/state/aqe/runtime_filter_planner.rs @@ -0,0 +1,562 @@ +// 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. + +//! Add runtime filters only to executable stage copies. The AQE tree stays +//! untouched so scan recognition, join selection and stage boundaries do not +//! depend on whether runtime filtering is enabled. + +use super::execution_plan::{DynamicJoinSelectionExec, ExchangeExec}; +use super::optimizer_rule::SelectJoinRule; +use ballista_core::execution_plans::runtime_filter::supports_key; +use ballista_core::execution_plans::{RuntimeFilterApplyExec, RuntimeFilterBuildExec}; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion::common::{JoinType, Result}; +use datafusion::datasource::source::DataSourceExec; +use datafusion::physical_expr::{PhysicalExpr, expressions::Column}; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::joins::HashJoinExec; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::{ExecutionPlan, replace_children_if_necessary}; +use std::collections::{BTreeMap, HashSet}; +use std::sync::Arc; + +#[derive(Clone)] +struct Filter { + id: u64, + key: Arc, +} + +#[derive(Clone, Default)] +pub(super) struct RuntimeFilterPlanner { + builds: BTreeMap<(usize, usize), Filter>, + probes: BTreeMap>, + planned: HashSet, +} + +impl RuntimeFilterPlanner { + /// Discover summaries before their producer stages are submitted. A pending + /// join may still swap sides, so only the final join binds a scan consumer. + /// Independent shuffle consumers can be bound now: filtering either input + /// is safe for the supported join types, even if the join later swaps. + pub fn discover(&mut self, plan: &Arc) -> Result<()> { + plan.apply(|node| { + if node + .downcast_ref::() + .is_some_and(|e| e.shuffle_created()) + { + return Ok(TreeNodeRecursion::Jump); + } + let (left, right, on, kind, null_aware) = if let Some(join) = + node.downcast_ref::() + { + ( + join.left(), + join.right(), + join.on(), + *join.join_type(), + join.null_aware, + ) + } else if let Some(join) = node.downcast_ref::() { + ( + &join.left, + &join.right, + join.on.as_slice(), + join.join_type, + join.null_aware, + ) + } else { + return Ok(TreeNodeRecursion::Continue); + }; + if supported_join(kind, null_aware) + && let Some((left_key, right_key)) = on.first() + { + if left.is::() { + self.discover_pair(left, right, left_key, right_key)?; + } else if node.is::() { + self.discover_pair(right, left, right_key, left_key)?; + } + } + Ok(TreeNodeRecursion::Continue) + })?; + Ok(()) + } + + fn discover_pair( + &mut self, + build: &Arc, + probe: &Arc, + build_key: &Arc, + probe_key: &Arc, + ) -> Result<()> { + let Some(exchange) = build.downcast_ref::() else { + return Ok(()); + }; + let Some(stage) = exchange.stage_id() else { + return Ok(()); + }; + let Some(column) = build_key.downcast_ref::() else { + return Ok(()); + }; + let data_type = build_key.data_type(build.schema().as_ref())?; + if !supports_key(&data_type) + || data_type != probe_key.data_type(probe.schema().as_ref())? + { + return Ok(()); + } + let probe_exchange = probe.downcast_ref::(); + let (probe_input, probe_stage) = if let Some(probe) = probe_exchange { + let Some(id) = probe.stage_id() else { + return Ok(()); + }; + // Both boundaries have been made runnable in this pass. Never + // wait for a future stage, or create a boundary to obtain a filter. + if probe.shuffle_created() || self.planned.contains(&id) { + return Ok(()); + } + (probe.input().clone(), Some(id)) + } else { + (probe.clone(), None) + }; + if probe_column( + probe_input.as_ref(), + probe_key.as_ref(), + probe_stage.is_some(), + ) + .is_none() + { + return Ok(()); + } + let producer_key = (stage, column.index()); + // A fallback probe can be created after its producer completed. Reuse + // that producer's existing summary; only new taps need an unplanned stage. + if !self.builds.contains_key(&producer_key) { + if exchange.shuffle_created() || self.planned.contains(&stage) { + return Ok(()); + } + // Unfinished exchanges have unknown statistics. Compare their inputs + // to avoid collecting a known larger prospective build side. + if SelectJoinRule::supports_swap_join_order( + exchange.input().as_ref(), + probe_input.as_ref(), + )? { + return Ok(()); + } + self.builds.insert( + producer_key, + Filter { + id: self.builds.len() as u64, + key: build_key.clone(), + }, + ); + } + let filter = &self.builds[&producer_key]; + if let Some(probe_stage) = probe_stage { + let probes = self.probes.entry(probe_stage).or_default(); + if !probes.iter().any(|p| p.id == filter.id) { + probes.push(Filter { + id: filter.id, + key: probe_key.clone(), + }); + } + } + Ok(()) + } + + pub fn instrument( + &mut self, + boundary: Arc, + ) -> Result> { + let stage = if let Some(exchange) = boundary.downcast_ref::() { + exchange.stage_id().unwrap() + } else { + boundary + .downcast_ref::() + .unwrap() + .stage_id() + .unwrap() + }; + self.planned.insert(stage); + let mut input = boundary.children()[0] + .clone() + .transform_down(|node| { + // A resolved exchange is a reader, not part of this executable stage. + if node.is::() { + return Ok(Transformed::new(node, false, TreeNodeRecursion::Jump)); + } + let Some(join) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + if !supported_join(*join.join_type(), join.null_aware) { + return Ok(Transformed::no(node)); + } + let Some((left_key, right_key)) = join.on().first() else { + return Ok(Transformed::no(node)); + }; + let Some(exchange) = join.left().downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(column) = left_key.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(producer) = exchange.stage_id() else { + return Ok(Transformed::no(node)); + }; + let Some(filter) = self.builds.get(&(producer, column.index())) else { + return Ok(Transformed::no(node)); + }; + let Some(probe) = push_apply( + join.right().clone(), + filter.id, + right_key.clone(), + false, + )? + else { + return Ok(Transformed::no(node)); + }; + let left = join.left().clone(); + Ok(Transformed::yes(replace_children_if_necessary( + node, + vec![left, probe], + )?)) + })? + .data; + if let Some(probes) = self.probes.get(&stage) { + for filter in probes { + if let Some(probe) = + push_apply(input.clone(), filter.id, filter.key.clone(), true)? + { + input = probe; + } + } + } + for (_, filter) in self.builds.range((stage, 0)..=(stage, usize::MAX)) { + input = Arc::new(RuntimeFilterBuildExec::new( + input, + filter.id, + filter.key.clone(), + )); + } + replace_children_if_necessary(boundary, vec![input]) + } +} + +fn supported_join(kind: JoinType, null_aware: bool) -> bool { + !null_aware + && matches!( + kind, + JoinType::Inner | JoinType::LeftSemi | JoinType::RightSemi + ) +} + +/// Find a scan key without constructing a speculative plan. Only cross +/// value-preserving unary operators; a predicate or shuffle must benefit. +fn probe_column( + mut plan: &dyn ExecutionPlan, + key: &dyn PhysicalExpr, + mut useful: bool, +) -> Option { + let mut column = key.downcast_ref::()?.index(); + loop { + if plan.is::() { + return useful.then_some(column); + } + // LIMIT may be fused into a filter or coalesce. Never change which + // rows it selects by inserting a runtime filter below it. + if plan.fetch().is_some() { + return None; + } + if let Some(projection) = plan.downcast_ref::() { + column = projection.expr()[column] + .expr + .downcast_ref::()? + .index(); + plan = projection.input().as_ref(); + } else if let Some(filter) = plan.downcast_ref::() { + column = filter.projection().as_ref().map_or(column, |p| p[column]); + useful = true; + plan = filter.input().as_ref(); + } else if let Some(repartition) = plan.downcast_ref::() { + plan = repartition.input().as_ref(); + } else { + let coalesce = plan.downcast_ref::()?; + plan = coalesce.input().as_ref(); + } + } +} + +fn push_apply( + plan: Arc, + id: u64, + key: Arc, + useful: bool, +) -> Result>> { + let Some(column) = probe_column(plan.as_ref(), key.as_ref(), useful) else { + return Ok(None); + }; + // probe_column established a single scan below a unary pipeline. + plan.transform_up(|node| { + if !node.is::() { + return Ok(Transformed::no(node)); + } + let key = Arc::new(Column::new(node.schema().field(column).name(), column)); + Ok(Transformed::yes(Arc::new(RuntimeFilterApplyExec::try_new( + node, id, key, None, + )?))) + }) + .map(|plan| Some(plan.data)) +} + +#[cfg(test)] +mod tests { + use super::*; + use ballista_core::execution_plans::runtime_filter::collect_reports; + use ballista_core::serde::protobuf::RuntimeFilterPayload; + use datafusion::arrow::array::{ArrayRef, Int64Array, RecordBatch}; + use datafusion::arrow::compute::concat_batches; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::NullEquality; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::physical_expr::expressions::{col, lit}; + use datafusion::physical_plan::filter::FilterExecBuilder; + use datafusion::physical_plan::{collect, displayable, joins::PartitionMode}; + use datafusion::prelude::SessionContext; + + fn source(values: impl IntoIterator) -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from_iter_values(values))], + ) + .unwrap(); + MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + async fn report(values: impl IntoIterator) -> RuntimeFilterPayload { + let build: Arc = Arc::new(RuntimeFilterBuildExec::new( + source(values), + 0, + Arc::new(Column::new("k", 0)), + )); + collect(build.clone(), SessionContext::new().task_ctx()) + .await + .unwrap(); + collect_reports(&build).unwrap().pop().unwrap() + } + + fn bind( + plan: Arc, + payload: &RuntimeFilterPayload, + ) -> Arc { + plan.transform_up(|node| { + let Some(apply) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + Ok(Transformed::yes(Arc::new( + apply.with_payload(payload.clone())?, + ))) + }) + .unwrap() + .data + } + + #[test] + fn runtime_filter_skips_larger_producer_behind_unfinished_exchange() { + let exchange = Arc::new(ExchangeExec::new(source(0..1000), None, 0)); + exchange.set_stage_id(0); + let probe = Arc::new(FilterExec::try_new(lit(true), source(0..2)).unwrap()); + let key: Arc = Arc::new(Column::new("k", 0)); + let join: Arc = Arc::new( + HashJoinExec::try_new( + exchange.clone(), + probe, + vec![(key.clone(), key)], + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ); + let mut filters = RuntimeFilterPlanner::default(); + filters.discover(&join).unwrap(); + let stage = filters.instrument(exchange).unwrap(); + assert!( + !displayable(stage.as_ref()) + .indent(false) + .to_string() + .contains("RuntimeFilterBuildExec"), + "the exchange has unknown output statistics, but its input is already known to exceed the probe" + ); + } + + #[tokio::test] + async fn runtime_filter_reuses_completed_build_for_partitioned_fallback() { + use crate::state::aqe::planner::AdaptivePlanner; + use crate::state::aqe::test::{ballista_ctx, mock_partitions_with_size}; + use datafusion::datasource::MemTable; + + let ctx = ballista_ctx(); + for setting in [ + "ballista.planner.runtime_filter.enabled = true", + "ballista.optimizer.broadcast_join_threshold_bytes = 500", + ] { + ctx.sql(&format!("SET {setting}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + } + for (name, rows) in [("build", 64), ("probe", 4096)] { + let batch = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int64Array::from_iter_values(0..rows)) as ArrayRef, + ), + ( + "val", + Arc::new(Int64Array::from_iter_values((0..rows).map(|n| n % 2))) + as ArrayRef, + ), + ]) + .unwrap(); + let table = MemTable::try_new(batch.schema(), vec![vec![batch]; 4]).unwrap(); + ctx.register_table(name, Arc::new(table)).unwrap(); + } + let mut planner = AdaptivePlanner::from_logical_sql( + &ctx, + "SELECT p.id FROM build b JOIN probe p ON b.id = p.id WHERE b.val > 0 AND p.val > 0", + None, + ) + .await + .unwrap(); + let first = planner.actionable_stages().unwrap().0.unwrap(); + assert_eq!(first.len(), 1); + let build = displayable(first[0].plan.as_ref()) + .indent(false) + .to_string(); + assert!( + build.contains("RuntimeFilterBuildExec: id=0, key=id@0"), + "{build}" + ); + planner + .finalise_stage_internal(0, mock_partitions_with_size(128, 1024)) + .unwrap(); + let next = planner.actionable_stages().unwrap().0.unwrap(); + let probe = displayable(next[0].plan.as_ref()).indent(false).to_string(); + assert!( + probe.contains("RuntimeFilterApplyExec: id=0, key=id@0"), + "{probe}" + ); + } + + #[tokio::test] + async fn runtime_filter_maps_keys_through_scan_projections() { + let batch = RecordBatch::try_from_iter([ + ( + "value", + Arc::new(Int64Array::from(vec![10, 20, 30])) as ArrayRef, + ), + ("k", Arc::new(Int64Array::from(vec![0, 1, 2])) as ArrayRef), + ]) + .unwrap(); + let scan = MemorySourceConfig::try_new_exec( + &[vec![batch.clone()]], + batch.schema(), + None, + ) + .unwrap(); + let filter = FilterExecBuilder::new(lit(true), scan) + .apply_projection(Some(vec![1, 0])) + .unwrap() + .build() + .unwrap(); + let schema = filter.schema(); + let probe = Arc::new( + ProjectionExec::try_new( + vec![ + (col("value", &schema).unwrap(), "value".into()), + (col("k", &schema).unwrap(), "join_key".into()), + ], + Arc::new(filter), + ) + .unwrap(), + ); + let plan = push_apply(probe, 0, Arc::new(Column::new("join_key", 1)), false) + .unwrap() + .unwrap(); + let plan = bind(plan, &report(0..2).await); + let schema = plan.schema(); + let batches = collect(plan, SessionContext::new().task_ctx()) + .await + .unwrap(); + let batch = concat_batches(&schema, &batches).unwrap(); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[10, 20] + ); + } + + #[tokio::test] + async fn runtime_filter_preserves_limits_fused_into_unary_operators() { + let payload = report([2]).await; + let key: Arc = Arc::new(Column::new("k", 0)); + let probes: [Arc; 2] = [ + Arc::new( + FilterExecBuilder::new(lit(true), source(0..4)) + .with_fetch(Some(2)) + .build() + .unwrap(), + ), + Arc::new(CoalescePartitionsExec::new(source(0..4)).with_fetch(Some(2))), + ]; + for probe in probes { + // The limited input is [0, 1]. Filtering before its limit would + // instead retain key 2 and incorrectly add a row to the join. + let probe = push_apply(probe.clone(), 0, key.clone(), true) + .unwrap() + .unwrap_or(probe); + let join = Arc::new( + HashJoinExec::try_new( + source([2]), + bind(probe, &payload), + vec![(key.clone(), key.clone())], + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ); + let rows = collect(join, SessionContext::new().task_ctx()) + .await + .unwrap(); + assert_eq!(rows.iter().map(RecordBatch::num_rows).sum::(), 0); + } + } +} diff --git a/ballista/scheduler/src/state/aqe/test/stage_build_side.rs b/ballista/scheduler/src/state/aqe/test/stage_build_side.rs index c0b3e3eebf..357f2e2980 100644 --- a/ballista/scheduler/src/state/aqe/test/stage_build_side.rs +++ b/ballista/scheduler/src/state/aqe/test/stage_build_side.rs @@ -164,3 +164,44 @@ async fn reuses_the_staged_exchange_when_the_build_side_measures_too_large() { StatsExec: partitions=4, rows=Exact(1000000000), bytes=Exact(42949672960) "); } + +/// Enabling filters must not interfere with AQE's measured join decisions. +#[tokio::test] +async fn runtime_filters_preserve_adaptive_planning() { + use datafusion::physical_plan::displayable; + + let mut plans = Vec::new(); + for enabled in [false, true] { + let ctx = q8_shaped_ctx(); + ctx.sql("SET datafusion.optimizer.prefer_hash_join = true") + .await + .unwrap() + .collect() + .await + .unwrap(); + ctx.sql(&format!( + "SET ballista.planner.runtime_filter.enabled = {enabled}" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + let mut planner = planner_for(&ctx, Q8_SHAPED_JOIN).await; + let initial = displayable(planner.current_plan()) + .indent(false) + .to_string(); + let stages = planner.actionable_stages().unwrap().0.unwrap(); + assert_eq!(stages.len(), 1); + planner + .finalise_stage_internal(0, mock_partitions_with_statistics()) + .unwrap(); + plans.push(( + initial, + displayable(planner.current_plan()) + .indent(false) + .to_string(), + )); + } + assert_eq!(plans[0], plans[1]); +} diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index e349fce918..b4f833a96d 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -138,6 +138,7 @@ standard DataFusion settings. | ballista.planner.parallel_window.enabled | Boolean | false | Enables the AQE parallel-window rule (ParallelWindowRule), which rewrites bounded-RANGE-frame windows into a distributed range-shuffle so BoundedWindowAggExec's single-partition constraint is not a serial bottleneck. Disabled by default — opt in when the workload contains matching window shapes. | | ballista.planner.propagate_empty.enabled | Boolean | true | Enables the AQE propagate-empty-relation rule. Injects EmptyExec into the plan where an input is known to be empty, such as one side of a join, allowing downstream work to be skipped. | | ballista.planner.reuse_exchange.enabled | Boolean | true | With adaptive query planning, runs the stage of a subplan that appears more than once in a query, such as a common table expression read twice, only once. Every reader of the subplan reads the same shuffle output. | +| ballista.planner.runtime_filter.enabled | Boolean | false | Enables bounded distributed runtime filters for AQE inner and semi hash joins. | | ballista.planner.skew_join.advisory_partition_bytes | UInt64 | 67108864 | Advisory target sub-shard byte size used by the bin-packer when splitting a skewed partition. | | ballista.planner.skew_join.enabled | Boolean | false | Enables the AQE optimize-skewed-join rule (Spark's OptimizeSkewedJoin). Disabled by default — opt in for workloads where per-key skew creates straggler tasks on join inputs. Requires datafusion.optimizer.enable_dynamic_filter_pushdown=false. | | ballista.planner.skew_join.partition_threshold_bytes | UInt64 | 268435456 | Absolute byte threshold below which a partition is never split (Spark's skewedPartitionThresholdInBytes). | @@ -163,6 +164,21 @@ that do not repartition their output. Partial aggregation above the join remains eligible because the final aggregate merges its groups after the shuffle. Coalescing takes priority when it has already grouped the join inputs. +## Distributed Runtime Filters + +Enable runtime filters for AQE inner/semi hash joins: + +```sql +SET ballista.planner.adaptive.enabled = true; +SET datafusion.optimizer.prefer_hash_join = true; +SET datafusion.optimizer.enable_dynamic_filter_pushdown = false; +SET ballista.planner.runtime_filter.enabled = true; +``` + +Filters use the first join key when it is a supported column, including NULL-safe +equality. Probe stages never wait: complete reports are used, otherwise bypassed. +Filters do not cross LIMIT or aggregates. `EXPLAIN ANALYZE` shows `rows_pruned`. + ## Shuffle Settings The following session-level keys control Ballista's shuffle behavior. See