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..b012a3f49 --- /dev/null +++ b/src/coordinator/dynamic_filters.rs @@ -0,0 +1,183 @@ +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. +/// +/// Equivalent to [`PhysicalExpr::expression_id`]. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) struct ExpressionId(pub(crate) u64); + +/// Collection of completed dynamic filter expressions received from producers. +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 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); + }; + expressions.push(expression); + Ok(expressions.len()) + } + + /// 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); + }; + Ok(expressions.len()) + } + + /// Returns a snapshot of all expressions for `id`, combined with an OR binary expression. + /// + /// 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: ExpressionId) -> 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)) + } +} + +/// 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, + [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::{CaseExpr, Column, lit}; + + 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.count(FILTER_ID)?, 0); + assert!(store.add(FILTER_ID, lit(true)).is_ok()); + assert_eq!(store.count(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 count_error = store.count(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!( + count_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_ors_case_expressions() -> 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(()) + } + + 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)), + )?)) + } +} 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..cb6ef7e98 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::{DynamicFilterStore, ExpressionId}; 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()),