From 7484f9cc642884f6d7eef60271abbf654c16da1b Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 6 Aug 2026 19:48:20 +0000 Subject: [PATCH 1/3] feat: add query-scoped dynamic filter store Add a registry in the query coordinator for collecting dynamic filter expressions by expression ID. These can be collected during planning once `apply_expressions` becomes available. This store counts received expressions (so we know when we've collected all the filters from workers in a task) and supports the merge() operation which ORs them, following the [RFC](https://github.com/datafusion-contrib/datafusion-distributed/pull/553). Remaining work - implmement dynamic filter update collection and propagation from worker <-> coordinator --- src/coordinator/distributed.rs | 2 + src/coordinator/dynamic_filters.rs | 197 +++++++++++++++++++++++++++ src/coordinator/mod.rs | 2 + src/coordinator/query_coordinator.rs | 5 + 4 files changed, 206 insertions(+) create mode 100644 src/coordinator/dynamic_filters.rs diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index cb126a884..3abb096e5 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -191,6 +191,8 @@ impl ExecutionPlan for DistributedExec { Arc::clone(&context), &self.metrics, self.metrics_store.clone(), + // Dynamic-filter discovery will supply the query's expression IDs here. + std::iter::empty(), ); let mut builder = RecordBatchReceiverStreamBuilder::new(self.schema(), 1); diff --git a/src/coordinator/dynamic_filters.rs b/src/coordinator/dynamic_filters.rs new file mode 100644 index 000000000..40bcfb615 --- /dev/null +++ b/src/coordinator/dynamic_filters.rs @@ -0,0 +1,197 @@ +use datafusion::common::{HashMap, Result, internal_err}; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::BinaryExpr; +use std::sync::{Arc, Mutex}; + +/// Identifies all instances of the same dynamic filter within a query. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) struct DynamicFilterId(pub(crate) u64); + +/// Query-scoped collection of completed dynamic filter expressions received from producers. +/// +/// The number of expected producers is deliberately not tracked here yet: under dynamic +/// planning it is not known until the producer stage has been finalized. Until then, [`Self::len`] +/// only reports how many expressions have arrived; it does not indicate completion. +pub(super) struct DynamicFilterStore { + expressions: Mutex>>>, +} + +impl DynamicFilterStore { + /// Creates an empty entry for every dynamic filter known to the query. + pub(super) fn new(ids: impl IntoIterator) -> Self { + Self { + expressions: Mutex::new(ids.into_iter().map(|id| (id, vec![])).collect()), + } + } + + /// Adds a completed producer expression and returns the number received for its filter. + pub(super) fn add( + &self, + id: DynamicFilterId, + expression: Arc, + ) -> Result { + let mut expressions = self.expressions.lock().expect("poisoned lock"); + let Some(expressions) = expressions.get_mut(&id) else { + return internal_err!("Unknown dynamic filter id: {}", id.0); + }; + expressions.push(expression); + Ok(expressions.len()) + } + + /// Returns the number of producer expressions received for a filter. + pub(super) fn len(&self, id: DynamicFilterId) -> Result { + let expressions = self.expressions.lock().expect("poisoned lock"); + let Some(expressions) = expressions.get(&id) else { + return internal_err!("Unknown dynamic filter id: {}", id.0); + }; + Ok(expressions.len()) + } + + /// Returns a snapshot of all expressions for `id`, combined with boolean OR. + /// + /// The merge is non-destructive: subsequent calls see the same expressions plus any that have + /// arrived since the previous call. The expression tree is balanced to avoid creating a deep + /// left- or right-associated tree when a filter has many producers. + pub(super) fn merge(&self, id: DynamicFilterId) -> Result>> { + let expressions = { + let all_expressions = self.expressions.lock().expect("poisoned lock"); + let Some(expressions) = all_expressions.get(&id) else { + return internal_err!("Unknown dynamic filter id: {}", id.0); + }; + expressions.clone() + }; + + Ok(merge_with_or(&expressions)) + } +} + +fn merge_with_or(expressions: &[Arc]) -> Option> { + match expressions { + [] => None, + [expression] => Some(Arc::clone(expression)), + expressions => { + let middle = expressions.len() / 2; + let left = merge_with_or(&expressions[..middle])?; + let right = merge_with_or(&expressions[middle..])?; + Some(Arc::new(BinaryExpr::new(left, Operator::Or, right))) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::physical_expr::expressions::lit; + + const FILTER_ID: DynamicFilterId = DynamicFilterId(1); + const OTHER_FILTER_ID: DynamicFilterId = DynamicFilterId(2); + + #[test] + fn new_deduplicates_filter_ids() -> Result<()> { + let store = DynamicFilterStore::new([FILTER_ID, FILTER_ID]); + + assert_eq!(store.len(FILTER_ID)?, 0); + assert!(store.add(FILTER_ID, lit(true)).is_ok()); + assert_eq!(store.len(FILTER_ID)?, 1); + Ok(()) + } + + #[test] + fn unknown_filter_ids_are_rejected() { + let store = DynamicFilterStore::new([FILTER_ID]); + + let add_error = store.add(OTHER_FILTER_ID, lit(true)).unwrap_err(); + let len_error = store.len(OTHER_FILTER_ID).unwrap_err(); + let merge_error = store.merge(OTHER_FILTER_ID).unwrap_err(); + + assert!( + add_error + .to_string() + .contains("Unknown dynamic filter id: 2") + ); + assert!( + len_error + .to_string() + .contains("Unknown dynamic filter id: 2") + ); + assert!( + merge_error + .to_string() + .contains("Unknown dynamic filter id: 2") + ); + } + + #[test] + fn merge_empty_filter_returns_none() -> Result<()> { + let store = DynamicFilterStore::new([FILTER_ID]); + + assert!(store.merge(FILTER_ID)?.is_none()); + Ok(()) + } + + #[test] + fn merge_single_filter_preserves_expression() -> Result<()> { + let store = DynamicFilterStore::new([FILTER_ID]); + let expression = lit(true); + store.add(FILTER_ID, Arc::clone(&expression))?; + + let merged = store.merge(FILTER_ID)?.unwrap(); + + assert!(Arc::ptr_eq(&expression, &merged)); + Ok(()) + } + + #[test] + fn merge_multiple_filters_is_balanced_and_deterministic() -> Result<()> { + let store = DynamicFilterStore::new([FILTER_ID]); + for value in [true, false, true, false] { + store.add(FILTER_ID, lit(value))?; + } + + let merged = store.merge(FILTER_ID)?.unwrap(); + + assert_eq!(merged.to_string(), "true OR false OR true OR false"); + let root = merged.downcast_ref::().unwrap(); + assert!(root.left().is::()); + assert!(root.right().is::()); + Ok(()) + } + + #[test] + fn merge_is_non_destructive_and_includes_later_additions() -> Result<()> { + let store = DynamicFilterStore::new([FILTER_ID]); + store.add(FILTER_ID, lit(true))?; + + assert_eq!(store.merge(FILTER_ID)?.unwrap().to_string(), "true"); + assert_eq!(store.len(FILTER_ID)?, 1); + + store.add(FILTER_ID, lit(false))?; + + assert_eq!( + store.merge(FILTER_ID)?.unwrap().to_string(), + "true OR false" + ); + assert_eq!(store.len(FILTER_ID)?, 2); + Ok(()) + } + + #[test] + fn concurrent_adds_are_counted() -> Result<()> { + let store = Arc::new(DynamicFilterStore::new([FILTER_ID])); + let threads = (0..8) + .map(|_| { + let store = Arc::clone(&store); + std::thread::spawn(move || store.add(FILTER_ID, lit(true))) + }) + .collect::>(); + + for thread in threads { + thread.join().unwrap()?; + } + + assert_eq!(store.len(FILTER_ID)?, 8); + assert!(store.merge(FILTER_ID)?.is_some()); + Ok(()) + } +} diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index c1a8a8dd2..6fa3665e2 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,4 +1,6 @@ mod distributed; +#[allow(dead_code)] // Scaffolding for distributed dynamic filtering. +mod dynamic_filters; mod latency_metric; mod metrics_store; mod prepare_dynamic_plan; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 3acced181..a2ff7efe4 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,6 +1,7 @@ use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; use crate::coordinator::MetricsStore; +use crate::coordinator::dynamic_filters::{DynamicFilterId, DynamicFilterStore}; use crate::coordinator::latency_metric::LatencyMetric; use crate::events::{RouteTasksEvent, RouteTasksHandlers}; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; @@ -46,6 +47,8 @@ const WORK_UNIT_FEED_CHUNK_SIZE: usize = 256; /// [StageCoordinator] scoped to each individual stage. pub(super) struct QueryCoordinator { task_ctx: Arc, + #[allow(dead_code)] // Populated once dynamic-filter discovery is wired in. + dynamic_filters: Arc, coordinator_to_worker_metrics: CoordinatorToWorkerMetrics, metrics_store: Option>, end_stream_notifier: Arc, @@ -58,9 +61,11 @@ impl QueryCoordinator { task_ctx: Arc, metrics_set: &ExecutionPlanMetricsSet, metrics_store: Option>, + dynamic_filter_ids: impl IntoIterator, ) -> Self { Self { task_ctx, + dynamic_filters: Arc::new(DynamicFilterStore::new(dynamic_filter_ids)), metrics_store, coordinator_to_worker_metrics: CoordinatorToWorkerMetrics::new(metrics_set), end_stream_notifier: Arc::new(Notify::new()), From 226fe9fa8339d84d928d6300f04774c0027b7ae9 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 6 Aug 2026 19:59:02 +0000 Subject: [PATCH 2/3] tests --- src/coordinator/dynamic_filters.rs | 33 +++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/src/coordinator/dynamic_filters.rs b/src/coordinator/dynamic_filters.rs index 40bcfb615..9ab49c0b8 100644 --- a/src/coordinator/dynamic_filters.rs +++ b/src/coordinator/dynamic_filters.rs @@ -53,6 +53,9 @@ impl DynamicFilterStore { /// The merge is non-destructive: subsequent calls see the same expressions plus any that have /// arrived since the previous call. The expression tree is balanced to avoid creating a deep /// left- or right-associated tree when a filter has many producers. + /// + /// DataFusion may eventually merge dynamic filter expressions natively; see + /// [apache/datafusion#23817](https://github.com/apache/datafusion/issues/23817). pub(super) fn merge(&self, id: DynamicFilterId) -> Result>> { let expressions = { let all_expressions = self.expressions.lock().expect("poisoned lock"); @@ -82,7 +85,7 @@ fn merge_with_or(expressions: &[Arc]) -> Option Result<()> { + let store = DynamicFilterStore::new([FILTER_ID]); + store.add(FILTER_ID, task_case(0)?)?; + store.add(FILTER_ID, task_case(1)?)?; + + let merged = store.merge(FILTER_ID)?.unwrap(); + + assert_eq!( + merged.to_string(), + "CASE WHEN task@0 = 0 THEN true ELSE false END OR CASE WHEN task@0 = 1 THEN true ELSE false END" + ); + Ok(()) + } + #[test] fn merge_is_non_destructive_and_includes_later_additions() -> Result<()> { let store = DynamicFilterStore::new([FILTER_ID]); @@ -194,4 +212,17 @@ mod tests { assert!(store.merge(FILTER_ID)?.is_some()); Ok(()) } + + fn task_case(task: i32) -> Result> { + let matches_task = Arc::new(BinaryExpr::new( + Arc::new(Column::new("task", 0)), + Operator::Eq, + lit(task), + )); + Ok(Arc::new(CaseExpr::try_new( + None, + vec![(matches_task, lit(true))], + Some(lit(false)), + )?)) + } } From 67f26ee1b65ffae4359eb3ac2a7e1d0998ff67c3 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 6 Aug 2026 21:30:50 +0000 Subject: [PATCH 3/3] cleanup --- src/coordinator/dynamic_filters.rs | 87 +++++++--------------------- src/coordinator/query_coordinator.rs | 4 +- 2 files changed, 23 insertions(+), 68 deletions(-) diff --git a/src/coordinator/dynamic_filters.rs b/src/coordinator/dynamic_filters.rs index 9ab49c0b8..b012a3f49 100644 --- a/src/coordinator/dynamic_filters.rs +++ b/src/coordinator/dynamic_filters.rs @@ -5,32 +5,26 @@ use datafusion::physical_expr::expressions::BinaryExpr; use std::sync::{Arc, Mutex}; /// Identifies all instances of the same dynamic filter within a query. +/// +/// Equivalent to [`PhysicalExpr::expression_id`]. #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] -pub(crate) struct DynamicFilterId(pub(crate) u64); +pub(crate) struct ExpressionId(pub(crate) u64); -/// Query-scoped collection of completed dynamic filter expressions received from producers. -/// -/// The number of expected producers is deliberately not tracked here yet: under dynamic -/// planning it is not known until the producer stage has been finalized. Until then, [`Self::len`] -/// only reports how many expressions have arrived; it does not indicate completion. +/// Collection of completed dynamic filter expressions received from producers. pub(super) struct DynamicFilterStore { - expressions: Mutex>>>, + expressions: Mutex>>>, } impl DynamicFilterStore { /// Creates an empty entry for every dynamic filter known to the query. - pub(super) fn new(ids: impl IntoIterator) -> Self { + pub(super) fn new(ids: impl IntoIterator) -> Self { Self { expressions: Mutex::new(ids.into_iter().map(|id| (id, vec![])).collect()), } } - /// Adds a completed producer expression and returns the number received for its filter. - pub(super) fn add( - &self, - id: DynamicFilterId, - expression: Arc, - ) -> Result { + /// Adds an expression and returns the number received so far for the given id. + pub(super) fn add(&self, id: ExpressionId, expression: Arc) -> Result { let mut expressions = self.expressions.lock().expect("poisoned lock"); let Some(expressions) = expressions.get_mut(&id) else { return internal_err!("Unknown dynamic filter id: {}", id.0); @@ -39,8 +33,8 @@ impl DynamicFilterStore { Ok(expressions.len()) } - /// Returns the number of producer expressions received for a filter. - pub(super) fn len(&self, id: DynamicFilterId) -> Result { + /// Returns the number of producer expressions received for a given id. + pub(super) fn count(&self, id: ExpressionId) -> Result { let expressions = self.expressions.lock().expect("poisoned lock"); let Some(expressions) = expressions.get(&id) else { return internal_err!("Unknown dynamic filter id: {}", id.0); @@ -48,15 +42,11 @@ impl DynamicFilterStore { Ok(expressions.len()) } - /// Returns a snapshot of all expressions for `id`, combined with boolean OR. - /// - /// The merge is non-destructive: subsequent calls see the same expressions plus any that have - /// arrived since the previous call. The expression tree is balanced to avoid creating a deep - /// left- or right-associated tree when a filter has many producers. + /// Returns a snapshot of all expressions for `id`, combined with an OR binary expression. /// - /// DataFusion may eventually merge dynamic filter expressions natively; see + /// DataFusion may eventually merge dynamic filter expressions natively. See /// [apache/datafusion#23817](https://github.com/apache/datafusion/issues/23817). - pub(super) fn merge(&self, id: DynamicFilterId) -> Result>> { + pub(super) fn merge(&self, id: ExpressionId) -> Result>> { let expressions = { let all_expressions = self.expressions.lock().expect("poisoned lock"); let Some(expressions) = all_expressions.get(&id) else { @@ -69,6 +59,8 @@ impl DynamicFilterStore { } } +/// Merges the provided expressions by ORing them together. The expression tree is balanced +/// to avoid creating a deep left- or right-associated tree. fn merge_with_or(expressions: &[Arc]) -> Option> { match expressions { [] => None, @@ -87,16 +79,16 @@ mod tests { use super::*; use datafusion::physical_expr::expressions::{CaseExpr, Column, lit}; - const FILTER_ID: DynamicFilterId = DynamicFilterId(1); - const OTHER_FILTER_ID: DynamicFilterId = DynamicFilterId(2); + const FILTER_ID: ExpressionId = ExpressionId(1); + const OTHER_FILTER_ID: ExpressionId = ExpressionId(2); #[test] fn new_deduplicates_filter_ids() -> Result<()> { let store = DynamicFilterStore::new([FILTER_ID, FILTER_ID]); - assert_eq!(store.len(FILTER_ID)?, 0); + assert_eq!(store.count(FILTER_ID)?, 0); assert!(store.add(FILTER_ID, lit(true)).is_ok()); - assert_eq!(store.len(FILTER_ID)?, 1); + assert_eq!(store.count(FILTER_ID)?, 1); Ok(()) } @@ -105,7 +97,7 @@ mod tests { let store = DynamicFilterStore::new([FILTER_ID]); let add_error = store.add(OTHER_FILTER_ID, lit(true)).unwrap_err(); - let len_error = store.len(OTHER_FILTER_ID).unwrap_err(); + let count_error = store.count(OTHER_FILTER_ID).unwrap_err(); let merge_error = store.merge(OTHER_FILTER_ID).unwrap_err(); assert!( @@ -114,7 +106,7 @@ mod tests { .contains("Unknown dynamic filter id: 2") ); assert!( - len_error + count_error .to_string() .contains("Unknown dynamic filter id: 2") ); @@ -176,43 +168,6 @@ mod tests { Ok(()) } - #[test] - fn merge_is_non_destructive_and_includes_later_additions() -> Result<()> { - let store = DynamicFilterStore::new([FILTER_ID]); - store.add(FILTER_ID, lit(true))?; - - assert_eq!(store.merge(FILTER_ID)?.unwrap().to_string(), "true"); - assert_eq!(store.len(FILTER_ID)?, 1); - - store.add(FILTER_ID, lit(false))?; - - assert_eq!( - store.merge(FILTER_ID)?.unwrap().to_string(), - "true OR false" - ); - assert_eq!(store.len(FILTER_ID)?, 2); - Ok(()) - } - - #[test] - fn concurrent_adds_are_counted() -> Result<()> { - let store = Arc::new(DynamicFilterStore::new([FILTER_ID])); - let threads = (0..8) - .map(|_| { - let store = Arc::clone(&store); - std::thread::spawn(move || store.add(FILTER_ID, lit(true))) - }) - .collect::>(); - - for thread in threads { - thread.join().unwrap()?; - } - - assert_eq!(store.len(FILTER_ID)?, 8); - assert!(store.merge(FILTER_ID)?.is_some()); - Ok(()) - } - fn task_case(task: i32) -> Result> { let matches_task = Arc::new(BinaryExpr::new( Arc::new(Column::new("task", 0)), diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index a2ff7efe4..cb6ef7e98 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,7 +1,7 @@ use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; use crate::coordinator::MetricsStore; -use crate::coordinator::dynamic_filters::{DynamicFilterId, DynamicFilterStore}; +use crate::coordinator::dynamic_filters::{DynamicFilterStore, ExpressionId}; use crate::coordinator::latency_metric::LatencyMetric; use crate::events::{RouteTasksEvent, RouteTasksHandlers}; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; @@ -61,7 +61,7 @@ impl QueryCoordinator { task_ctx: Arc, metrics_set: &ExecutionPlanMetricsSet, metrics_store: Option>, - dynamic_filter_ids: impl IntoIterator, + dynamic_filter_ids: impl IntoIterator, ) -> Self { Self { task_ctx,