diff --git a/docs/upgrade/3.0.0.md b/docs/upgrade/3.0.0.md index 38f21d90a..9dbd42dee 100644 --- a/docs/upgrade/3.0.0.md +++ b/docs/upgrade/3.0.0.md @@ -55,6 +55,14 @@ use datafusion_distributed::{WorkerChannel, grpc}; let client: Box = grpc::create_worker_client(channel); ``` +Custom `WorkerChannel` transports must also carry the dynamic-filter display +fields added to the coordinator protocol. Include `SetPlanRequest::dynamic_filter_ids` +when constructing or encoding a plan request, and encode/decode the new +`WorkerToCoordinatorMsg::TaskCompletedDynamicFilters` variant. The built-in gRPC +transport handles both automatically. `TaskDynamicFilter::expression` contains a typed +`datafusion_proto::protobuf::PhysicalExprNode`; custom transports should encode it only at their +wire boundary. + ## 2. Move `TaskEstimator` methods to event handlers `TaskEstimator`, `TaskEstimation`, `TaskRoutingContext`, and diff --git a/src/common/dynamic_filtering.rs b/src/common/dynamic_filtering.rs new file mode 100644 index 000000000..96c381abf --- /dev/null +++ b/src/common/dynamic_filtering.rs @@ -0,0 +1,210 @@ +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion::common::{HashMap, HashSet, Result}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; +use datafusion::physical_plan::ExecutionPlan; +use std::sync::Arc; + +/// A dynamic-filter consumer discovered in an execution plan. +#[derive(Clone)] +pub(crate) struct DiscoveredDynamicFilter { + pub(crate) id: u64, + pub(crate) expression: Arc, + pub(crate) input_schema: SchemaRef, +} + +/// Finds dynamic-filter consumers in `plan`, optionally restricting the result to `allowed_ids`. +/// +/// Producer and consumer occurrences intentionally share expression IDs. Producer occurrences are +/// therefore removed only from the node that reports them through +/// [`ExecutionPlan::dynamic_expressions_produced`], rather than subtracting producer IDs from the +/// whole plan. +pub(crate) fn discover_dynamic_filter_consumers( + plan: &Arc, + allowed_ids: Option<&HashSet>, +) -> Result> { + let mut consumers = HashMap::new(); + + plan.apply(|node| { + let produced = node.dynamic_expressions_produced(); + let input_schema = node + .children() + .first() + .map(|child| child.schema()) + .unwrap_or_else(|| node.schema()); + + node.apply_expressions(&mut |root| { + root.apply(|expression| { + let Some(_) = expression.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + + let id = expression + .expression_id() + .expect("DynamicFilterPhysicalExpr always has an expression ID"); + let is_producer_occurrence = produced + .iter() + .any(|produced| Arc::ptr_eq(produced, expression)); + let is_allowed = allowed_ids.is_none_or(|ids| ids.contains(&id)); + if !is_producer_occurrence && is_allowed { + consumers + .entry(id) + .or_insert_with(|| DiscoveredDynamicFilter { + id, + expression: Arc::clone(expression), + input_schema: Arc::clone(&input_schema), + }); + } + + Ok(TreeNodeRecursion::Continue) + }) + })?; + Ok(TreeNodeRecursion::Continue) + })?; + + let mut consumers: Vec<_> = consumers.into_values().collect(); + consumers.sort_unstable_by_key(|consumer| consumer.id); + Ok(consumers) +} + +pub(crate) fn dynamic_filter_consumer_ids(plan: &Arc) -> Result> { + Ok(discover_dynamic_filter_consumers(plan, None)? + .into_iter() + .map(|consumer| consumer.id) + .collect()) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::Result; + use datafusion::execution::{SendableRecordBatchStream, TaskContext}; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Column, lit}; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, PlanProperties, apply_expression_roots, + }; + use std::fmt::Formatter; + + #[tokio::test] + async fn discovers_nested_consumer_but_not_its_producer_occurrence() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc; + let column = Arc::new(Column::new("a", 0)) as Arc; + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&column)], + lit(true), + )) as Arc; + let nested = Arc::new(BinaryExpr::new( + Arc::clone(&dynamic_filter), + Operator::And, + lit(true), + )) as Arc; + + let consumer = + Arc::new(ExpressionExec::new(input, nested, false)) as Arc; + let plan = Arc::new(ExpressionExec::new( + consumer, + Arc::clone(&dynamic_filter), + true, + )) as Arc; + + let discovered = discover_dynamic_filter_consumers(&plan, None)?; + assert_eq!(discovered.len(), 1); + assert_eq!(discovered[0].id, dynamic_filter.expression_id().unwrap()); + + dynamic_filter + .downcast_ref::() + .unwrap() + .update(Arc::new(BinaryExpr::new(column, Operator::Gt, lit(10_i32))))?; + dynamic_filter + .downcast_ref::() + .unwrap() + .mark_complete(); + + let current = discovered[0] + .expression + .downcast_ref::() + .unwrap() + .current()?; + assert_eq!(current.to_string(), "a@0 > 10"); + Ok(()) + } + + #[derive(Debug)] + struct ExpressionExec { + input: Arc, + expression: Arc, + produces_expression: bool, + } + + impl ExpressionExec { + fn new( + input: Arc, + expression: Arc, + produces_expression: bool, + ) -> Self { + Self { + input, + expression, + produces_expression, + } + } + } + + impl DisplayAs for ExpressionExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ExpressionExec") + } + } + + impl ExecutionPlan for ExpressionExec { + fn name(&self) -> &str { + "ExpressionExec" + } + + fn properties(&self) -> &Arc { + self.input.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn dynamic_expressions_produced(&self) -> Vec> { + self.produces_expression + .then(|| Arc::clone(&self.expression)) + .into_iter() + .collect() + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + apply_expression_roots([&self.expression], f) + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self::new( + children.remove(0), + Arc::clone(&self.expression), + self.produces_expression, + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.input.execute(partition, context) + } + } +} diff --git a/src/common/mod.rs b/src/common/mod.rs index 07463ca0b..8e8b34a34 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,4 +1,5 @@ mod children_helpers; +mod dynamic_filtering; mod once_lock; mod recursion; mod task_context_helpers; @@ -7,6 +8,9 @@ mod uuid; mod vec; pub(crate) use children_helpers::require_one_child; +pub(crate) use dynamic_filtering::{ + discover_dynamic_filter_consumers, dynamic_filter_consumer_ids, +}; pub(crate) use once_lock::OnceLockResult; pub(crate) use recursion::TreeNodeExt; pub(crate) use task_context_helpers::task_ctx_with_extension; diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 1e5fc796f..2b031cad3 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,12 +1,12 @@ +use crate::DistributedConfig; use crate::common::require_one_child; -use crate::coordinator::metrics_store::MetricsStore; +use crate::coordinator::dynamic_filters::isolate_distributed_leaf_variants_for_display; use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan; use crate::coordinator::prepare_static_plan::prepare_static_plan; use crate::coordinator::query_coordinator::QueryCoordinator; -use crate::distributed_planner::NetworkBoundaryExt; -use crate::{DistributedConfig, TaskKey}; +use crate::coordinator::store::{MetricsStore, task_keys_for_plan}; use datafusion::common::internal_datafusion_err; -use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{Result, exec_err}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::PhysicalExpr; @@ -82,33 +82,13 @@ impl DistributedExec { /// /// [`rewrite_distributed_plan_with_metrics`]: crate::rewrite_distributed_plan_with_metrics pub async fn wait_for_metrics(&self) { - let mut expected_keys: Vec = Vec::new(); let Some(task_metrics) = &self.metrics_store else { return; }; let Some(plan) = self.plan_for_viz.lock().unwrap().as_ref().cloned() else { return; }; - let _ = plan.apply(|plan| { - if let Some(boundary) = plan.as_network_boundary() { - let stage = boundary.input_stage(); - for i in 0..stage.task_count() { - expected_keys.push(TaskKey { - query_id: stage.query_id(), - stage_id: stage.num(), - task_number: i, - }); - } - } - Ok(TreeNodeRecursion::Continue) - }); - if expected_keys.is_empty() { - return; - } - let mut rx = task_metrics.rx.clone(); - let _ = rx - .wait_for(|map| expected_keys.iter().all(|key| map.contains_key(key))) - .await; + task_metrics.wait_for(&task_keys_for_plan(&plan)).await; } /// Returns the plan which is lazily prepared on `execute()` and actually gets executed. @@ -124,6 +104,14 @@ impl DistributedExec { }) } + pub(crate) fn plan_for_display(&self) -> Arc { + self.plan_for_viz + .lock() + .ok() + .and_then(|plan| plan.clone()) + .unwrap_or_else(|| Arc::clone(&self.base_plan)) + } + /// Returns the head stage that was actually executed. Unlike [`Self::plan_for_viz`] (which is /// reconstructed for visualization, with `Stage::Local` boundaries and rebuilt ancestor /// `Arc`s), this returns the original `Arc` instances whose metrics were populated during @@ -223,10 +211,12 @@ impl ExecutionPlan for DistributedExec { false => prepare_static_plan(&query_coordinator, &base_plan)?, }; + let display_plan = + isolate_distributed_leaf_variants_for_display(result.plan_for_viz, &context)?; plan_for_viz .lock() .expect("poisoned lock") - .replace(result.plan_for_viz); + .replace(Arc::clone(&display_plan)); head_stage .lock() .expect("poisoned lock") @@ -237,8 +227,11 @@ impl ExecutionPlan for DistributedExec { break; // channel closed } } - drop(tx); drop(guard); + query_coordinator + .finish_dynamic_filter_display(&display_plan) + .await; + drop(tx); query_coordinator.drain_pending_tasks().await?; Ok(()) }); diff --git a/src/coordinator/dynamic_filters.rs b/src/coordinator/dynamic_filters.rs new file mode 100644 index 000000000..57bf509eb --- /dev/null +++ b/src/coordinator/dynamic_filters.rs @@ -0,0 +1,178 @@ +use crate::common::discover_dynamic_filter_consumers; +use crate::execution_plans::DistributedLeafExec; +use crate::{DistributedCodec, TaskCompletedDynamicFilters, TaskKey}; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion::common::{HashMap, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::physical_plan::from_proto::parse_physical_expr; +use datafusion_proto::physical_plan::{DeduplicatingProtoConverter, PhysicalPlanNodeExt}; +use datafusion_proto::protobuf::PhysicalPlanNode; +use std::sync::Arc; + +/// Replaces the variants in the visualization plan with independent per-task copies. +pub(super) fn isolate_distributed_leaf_variants_for_display( + plan: Arc, + task_ctx: &Arc, +) -> Result> { + let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config()); + let converter = DeduplicatingProtoConverter::default(); + plan.transform_up(|node| { + let Some(leaf) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + + let variants = leaf + .variants() + .iter() + .map(|variant| { + let proto = PhysicalPlanNode::try_from_physical_plan_with_converter( + Arc::clone(variant), + &codec, + &converter, + )?; + proto.try_into_physical_plan_with_converter(task_ctx, &codec, &converter) + }) + .collect::>>()?; + + Ok(Transformed::yes(Arc::new(DistributedLeafExec::try_new( + Arc::clone(leaf.original()), + variants, + )?) as Arc)) + }) + .map(|transformed| transformed.data) +} + +/// Applies successful worker reports only to the matching task-local visualization variants. +pub(super) fn apply_reports_to_distributed_leaves( + plan: &Arc, + reports: &HashMap, + task_ctx: &Arc, +) { + let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config()); + let _ = plan.apply(|node| { + let Some(leaf) = node.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + + for (task_key, report) in reports { + let Some(variant) = leaf.variants().get(task_key.task_number) else { + continue; + }; + let updates: HashMap<_, _> = report + .filters + .iter() + .map(|filter| (filter.expression_id, &filter.expression)) + .collect(); + let Ok(consumers) = discover_dynamic_filter_consumers(variant, None) else { + continue; + }; + for consumer in consumers { + let Some(proto) = updates.get(&consumer.id).copied() else { + continue; + }; + let Ok(reported_expression) = + parse_physical_expr(proto, task_ctx, consumer.input_schema.as_ref(), &codec) + else { + continue; + }; + let Some(reported_dynamic_filter) = + reported_expression.downcast_ref::() + else { + continue; + }; + let Ok(expression) = reported_dynamic_filter.current() else { + continue; + }; + let Some(dynamic_filter) = consumer + .expression + .downcast_ref::() + else { + continue; + }; + if dynamic_filter.update(expression).is_ok() { + dynamic_filter.mark_complete(); + } + } + } + + Ok(TreeNodeRecursion::Continue) + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_expr::expressions::{ + BinaryExpr, Column, DynamicFilterPhysicalExpr, lit, + }; + use datafusion::physical_plan::displayable; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::filter::FilterExec; + use datafusion::prelude::SessionContext; + use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; + use uuid::Uuid; + + #[test] + fn visualization_variants_do_not_share_dynamic_filter_state() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let column = Arc::new(Column::new("a", 0)) as Arc; + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&column)], + lit(true), + )) as Arc; + let input = Arc::new(EmptyExec::new(schema)) as Arc; + let variant = Arc::new(FilterExec::try_new( + Arc::clone(&dynamic_filter), + Arc::clone(&input), + )?) as Arc; + let leaf = Arc::new(DistributedLeafExec::try_new( + Arc::clone(&variant), + [Arc::clone(&variant), variant], + )?) as Arc; + + let task_ctx = SessionContext::new().task_ctx(); + let isolated = isolate_distributed_leaf_variants_for_display(leaf, &task_ctx)?; + let expression = + Arc::new(BinaryExpr::new(column, Operator::Gt, lit(10_i32))) as Arc; + dynamic_filter + .downcast_ref::() + .unwrap() + .update(expression)?; + dynamic_filter + .downcast_ref::() + .unwrap() + .mark_complete(); + let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config()); + let report = TaskCompletedDynamicFilters { + filters: vec![crate::TaskDynamicFilter { + expression_id: dynamic_filter.expression_id().unwrap(), + expression: serialize_physical_expr(&dynamic_filter, &codec)?, + }], + }; + let reports = HashMap::from_iter([( + TaskKey { + query_id: Uuid::nil(), + stage_id: 1, + task_number: 0, + }, + report, + )]); + + apply_reports_to_distributed_leaves(&isolated, &reports, &task_ctx); + let leaf = isolated.downcast_ref::().unwrap(); + let task_0 = displayable(leaf.variants()[0].as_ref()) + .one_line() + .to_string(); + let task_1 = displayable(leaf.variants()[1].as_ref()) + .one_line() + .to_string(); + assert!(task_0.contains("DynamicFilter [ a@0 > 10 ]")); + assert!(task_1.contains("DynamicFilter [ empty ]")); + Ok(()) + } +} diff --git a/src/coordinator/metrics_store.rs b/src/coordinator/metrics_store.rs deleted file mode 100644 index 8ccc1f7c0..000000000 --- a/src/coordinator/metrics_store.rs +++ /dev/null @@ -1,36 +0,0 @@ -use crate::{TaskKey, TaskMetrics}; -use datafusion::common::HashMap; -use tokio::sync::watch; - -type MetricsMap = HashMap; - -/// Stores the metrics collected from all worker tasks, and notifies waiters when new entries arrive. -#[derive(Debug, Clone)] -pub struct MetricsStore { - tx: watch::Sender, - pub(crate) rx: watch::Receiver, -} - -impl MetricsStore { - pub(crate) fn new() -> Self { - let (tx, rx) = watch::channel(HashMap::new()); - Self { tx, rx } - } - - pub(crate) fn insert(&self, key: TaskKey, metrics: TaskMetrics) { - self.tx.send_modify(|map| { - map.insert(key, metrics); - }); - } - - pub(crate) fn get(&self, key: &TaskKey) -> Option { - self.rx.borrow().get(key).cloned() - } - - #[cfg(test)] - pub(crate) fn from_entries(entries: impl IntoIterator) -> Self { - let map: HashMap<_, _> = entries.into_iter().collect(); - let (tx, rx) = watch::channel(map); - Self { tx, rx } - } -} diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index c1a8a8dd2..c08fbbd46 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,9 +1,10 @@ mod distributed; +mod dynamic_filters; mod latency_metric; -mod metrics_store; mod prepare_dynamic_plan; mod prepare_static_plan; mod query_coordinator; +mod store; pub use distributed::DistributedExec; -pub(crate) use metrics_store::MetricsStore; +pub(crate) use store::{CompletedDynamicFilterStore, MetricsStore}; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 3acced181..ee8d919c7 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,7 +1,9 @@ -use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; +use crate::common::{TreeNodeExt, dynamic_filter_consumer_ids, 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::apply_reports_to_distributed_leaves; use crate::coordinator::latency_metric::LatencyMetric; +use crate::coordinator::store::task_keys_for_plan; +use crate::coordinator::{CompletedDynamicFilterStore, MetricsStore}; use crate::events::{RouteTasksEvent, RouteTasksHandlers}; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; use crate::passthrough_headers::get_passthrough_headers; @@ -11,8 +13,8 @@ use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time}; use crate::{ BytesCounterMetric, BytesMetricExt, CoordinatorToWorkerMsg, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedTaskContext, - DistributedWorkUnitFeedContext, LoadInfo, SetPlanRequest, TaskKey, WorkUnitFeedDeclaration, - WorkerToCoordinatorMsg, get_distributed_channel_resolver, + DistributedWorkUnitFeedContext, LoadInfo, SetPlanRequest, TaskCompletedDynamicFilters, TaskKey, + WorkUnitFeedDeclaration, WorkerToCoordinatorMsg, get_distributed_channel_resolver, }; use datafusion::common::DataFusionError; use datafusion::common::instant::Instant; @@ -23,7 +25,7 @@ use datafusion::execution::TaskContext; use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; -use datafusion_proto::physical_plan::AsExecutionPlan; +use datafusion_proto::physical_plan::{DeduplicatingProtoConverter, PhysicalPlanNodeExt}; use datafusion_proto::protobuf::PhysicalPlanNode; use futures::{Stream, StreamExt, TryStreamExt}; use prost::Message; @@ -48,6 +50,7 @@ pub(super) struct QueryCoordinator { task_ctx: Arc, coordinator_to_worker_metrics: CoordinatorToWorkerMetrics, metrics_store: Option>, + completed_dynamic_filter_store: Arc, end_stream_notifier: Arc, join_set: Mutex>>, } @@ -62,6 +65,7 @@ impl QueryCoordinator { Self { task_ctx, metrics_store, + completed_dynamic_filter_store: Arc::new(CompletedDynamicFilterStore::new()), coordinator_to_worker_metrics: CoordinatorToWorkerMetrics::new(metrics_set), end_stream_notifier: Arc::new(Notify::new()), join_set: Mutex::new(JoinSet::new()), @@ -79,6 +83,7 @@ impl QueryCoordinator { task_ctx: &self.task_ctx, metrics: &self.coordinator_to_worker_metrics, metrics_store: &self.metrics_store, + completed_dynamic_filter_store: &self.completed_dynamic_filter_store, end_stream_notifier: &self.end_stream_notifier, join_set: &self.join_set, } @@ -96,6 +101,17 @@ impl QueryCoordinator { NotifyGuard(Arc::clone(&self.end_stream_notifier)) } + pub(super) async fn finish_dynamic_filter_display( + &self, + plan_for_viz: &Arc, + ) { + let reports = self + .completed_dynamic_filter_store + .wait_for(&task_keys_for_plan(plan_for_viz)) + .await; + apply_reports_to_distributed_leaves(plan_for_viz, &reports, &self.task_ctx); + } + /// Blocks until all background tasks have finished (e.g., sending WorkUnit feeds, or collecting /// metrics) pub(super) async fn drain_pending_tasks(self) -> Result<()> { @@ -123,6 +139,7 @@ pub(super) struct StageCoordinator<'a> { task_ctx: &'a Arc, metrics: &'a CoordinatorToWorkerMetrics, metrics_store: &'a Option>, + completed_dynamic_filter_store: &'a Arc, end_stream_notifier: &'a Arc, join_set: &'a Mutex>>, } @@ -142,10 +159,15 @@ impl<'a> StageCoordinator<'a> { let session_config = self.task_ctx.session_config(); let codec = DistributedCodec::new_combined_with_user(session_config); - let (specialized, work_unit_feed_declarations) = self.task_specialized_plan(task_i)?; + let (specialized, work_unit_feed_declarations, dynamic_filter_ids) = + self.task_specialized_plan(task_i)?; - let plan_proto = - PhysicalPlanNode::try_from_physical_plan(specialized, &codec)?.encode_to_vec(); + let plan_proto = PhysicalPlanNode::try_from_physical_plan_with_converter( + specialized, + &codec, + &DeduplicatingProtoConverter::default(), + )? + .encode_to_vec(); let plan_size = plan_proto.len(); let task_key = TaskKey { @@ -153,11 +175,11 @@ impl<'a> StageCoordinator<'a> { stage_id: self.stage_id, task_number: task_i, }; - let msg = CoordinatorToWorkerMsg::SetPlanRequest(SetPlanRequest { task_key, task_count: self.task_count, plan_proto, + dynamic_filter_ids, work_unit_feed_declarations, target_worker_url: url.clone(), query_start_time_ns: self.metrics.instantiation_time, @@ -184,8 +206,8 @@ impl<'a> StageCoordinator<'a> { // 3. Here, `end_stream_notifier` fires and the coordinator->worker channel is // gracefully ended. // 4. The coordinator->worker channel EOS is received in `impl_coordinator_channel.rs`. - // 5. The metrics are send back in the worker->coordinator channel, and then that - // channel is closed. + // 5. The metrics and final dynamic filters are sent back in the + // worker->coordinator channel, and then that channel is closed. .chain(keep_stream_alive(Arc::clone(self.end_stream_notifier))) .boxed(); @@ -224,6 +246,7 @@ impl<'a> StageCoordinator<'a> { task_number: task_i, }; let task_metrics = self.metrics_store.clone(); + let completed_dynamic_filter_store = Arc::clone(self.completed_dynamic_filter_store); let (load_info_tx, load_info_rx) = tokio::sync::mpsc::unbounded_channel(); let mut load_info_tx_opt = Some(load_info_tx); @@ -246,8 +269,15 @@ impl<'a> StageCoordinator<'a> { WorkerToCoordinatorMsg::LoadInfoEos => { let _ = load_info_tx_opt.take(); } + WorkerToCoordinatorMsg::TaskCompletedDynamicFilters(filters) => { + completed_dynamic_filter_store.insert(task_key, filters); + } } } + if completed_dynamic_filter_store.get(&task_key).is_none() { + completed_dynamic_filter_store + .insert(task_key, TaskCompletedDynamicFilters::default()); + } }); load_info_rx } @@ -332,13 +362,18 @@ impl<'a> StageCoordinator<'a> { fn task_specialized_plan( &self, task_i: usize, - ) -> Result<(Arc, Vec)> { + ) -> Result<( + Arc, + Vec, + Vec, + )> { let session_config = self.task_ctx.session_config(); let wuf_registry = session_config .get_extension::() .unwrap_or_default(); let mut work_unit_feed_declarations = vec![]; + let mut dynamic_filter_ids = datafusion::common::HashSet::new(); let d_ctx = DistributedTaskContext { task_index: task_i, task_count: self.task_count, @@ -360,12 +395,19 @@ impl<'a> StageCoordinator<'a> { if let Some(dle) = plan.downcast_ref::() { let specialized = dle.to_task_specialized(d_ctx.task_index); + dynamic_filter_ids.extend(dynamic_filter_consumer_ids(&specialized)?); return Ok(Transformed::yes(specialized)); } Ok(Transformed::no(plan)) })?; - Ok((transformed.data, work_unit_feed_declarations)) + let mut dynamic_filter_ids: Vec<_> = dynamic_filter_ids.into_iter().collect(); + dynamic_filter_ids.sort_unstable(); + Ok(( + transformed.data, + work_unit_feed_declarations, + dynamic_filter_ids, + )) } /// Returns as many URLs as the task count for the stage this [StageCoordinator] diff --git a/src/coordinator/store.rs b/src/coordinator/store.rs new file mode 100644 index 000000000..b903a17c9 --- /dev/null +++ b/src/coordinator/store.rs @@ -0,0 +1,78 @@ +use crate::distributed_planner::NetworkBoundaryExt; +use crate::{TaskCompletedDynamicFilters, TaskKey, TaskMetrics}; +use datafusion::common::HashMap; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion::physical_plan::ExecutionPlan; +use std::sync::Arc; +use tokio::sync::watch; + +type StoreMap = HashMap; + +/// Stores task-scoped values and notifies waiters when entries change. +#[derive(Debug, Clone)] +pub(crate) struct Store { + tx: watch::Sender>, + rx: watch::Receiver>, +} + +impl Store { + pub(crate) fn new() -> Self { + let (tx, rx) = watch::channel(HashMap::new()); + Self { tx, rx } + } + + pub(crate) fn insert(&self, key: TaskKey, value: T) { + self.tx.send_modify(|map| { + map.insert(key, value); + }); + } + + pub(crate) fn get(&self, key: &TaskKey) -> Option + where + T: Clone, + { + self.rx.borrow().get(key).cloned() + } + + pub(crate) async fn wait_for(&self, expected_keys: &[TaskKey]) -> StoreMap + where + T: Clone, + { + let mut rx = self.rx.clone(); + if !expected_keys.is_empty() { + let _ = rx + .wait_for(|map| expected_keys.iter().all(|key| map.contains_key(key))) + .await; + } + let values = rx.borrow().clone(); + values + } + + #[cfg(test)] + pub(crate) fn from_entries(entries: impl IntoIterator) -> Self { + let map: HashMap<_, _> = entries.into_iter().collect(); + let (tx, rx) = watch::channel(map); + Self { tx, rx } + } +} + +pub(crate) type MetricsStore = Store; +pub(crate) type CompletedDynamicFilterStore = Store; + +pub(crate) fn task_keys_for_plan(plan: &Arc) -> Vec { + let mut task_keys = Vec::new(); + let _ = plan.apply(|plan| { + if let Some(boundary) = plan.as_network_boundary() { + let stage = boundary.input_stage(); + for task_number in 0..stage.task_count() { + task_keys.push(TaskKey { + query_id: stage.query_id(), + stage_id: stage.num(), + task_number, + }); + } + } + Ok(TreeNodeRecursion::Continue) + }); + task_keys +} diff --git a/src/lib.rs b/src/lib.rs index 561799e46..b66451e5b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -53,9 +53,9 @@ pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; pub use protocol::{ ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, - GetWorkerInfoResponse, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, - WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, - get_distributed_channel_resolver, + GetWorkerInfoResponse, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskCompletedDynamicFilters, + TaskDynamicFilter, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, + WorkerChannel, WorkerToCoordinatorMsg, get_distributed_channel_resolver, }; pub use stage::{ DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze, diff --git a/src/protocol/grpc/generated/worker.rs b/src/protocol/grpc/generated/worker.rs index d1f0fb4a9..e13e52880 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/grpc/generated/worker.rs @@ -24,7 +24,7 @@ pub mod coordinator_to_worker_msg { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct WorkerToCoordinatorMsg { - #[prost(oneof = "worker_to_coordinator_msg::Inner", tags = "1, 2, 3")] + #[prost(oneof = "worker_to_coordinator_msg::Inner", tags = "1, 2, 3, 4")] pub inner: ::core::option::Option, } /// Nested message and enum types in `WorkerToCoordinatorMsg`. @@ -43,6 +43,25 @@ pub mod worker_to_coordinator_msg { LoadInfo(super::LoadInfo), #[prost(bool, tag = "3")] LoadInfoEos(bool), + /// Final task-local dynamic filters used by distributed leaf variants. + #[prost(message, tag = "4")] + TaskCompletedDynamicFilters(super::TaskCompletedDynamicFilters), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct TaskCompletedDynamicFilters { + #[prost(message, repeated, tag = "1")] + pub filters: ::prost::alloc::vec::Vec, +} +/// Nested message and enum types in `TaskCompletedDynamicFilters`. +pub mod task_completed_dynamic_filters { + #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] + pub struct DynamicFilter { + #[prost(uint64, tag = "1")] + pub expression_id: u64, + /// Serialized datafusion.proto.PhysicalExprNode. + #[prost(bytes = "vec", tag = "2")] + pub expression_proto: ::prost::alloc::vec::Vec, } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -80,8 +99,8 @@ pub struct LoadInfo { /// The amount of rows that were pulled from leaf nodes while this partition was sampling data. #[prost(uint64, tag = "8")] pub rows_pulled_from_leaf: u64, - /// Whether the sampled partition stream reached end-of-stream by the time this LoadInfo was - /// captured. + /// Whether the sampled partition stream reached end-of-stream (i.e. the partition finished + /// producing all of its output) by the time this LoadInfo was captured. #[prost(bool, tag = "9")] pub reached_eos: bool, } @@ -119,6 +138,9 @@ pub struct SetPlanRequest { /// relative to when the query was fired in the coordinator. #[prost(uint64, tag = "6")] pub query_start_time_ns: u64, + /// Dynamic-filter consumer IDs belonging to the selected DistributedLeafExec variants. + #[prost(uint64, repeated, tag = "7")] + pub dynamic_filter_ids: ::prost::alloc::vec::Vec, } /// Nested message and enum types in `SetPlanRequest`. pub mod set_plan_request { diff --git a/src/protocol/grpc/worker.proto b/src/protocol/grpc/worker.proto index 24f5c9871..638c086e0 100644 --- a/src/protocol/grpc/worker.proto +++ b/src/protocol/grpc/worker.proto @@ -39,9 +39,22 @@ message WorkerToCoordinatorMsg { LoadInfo load_info = 2; bool load_info_eos = 3; + + // Final task-local dynamic filters used by distributed leaf variants. + TaskCompletedDynamicFilters task_completed_dynamic_filters = 4; } } +message TaskCompletedDynamicFilters { + message DynamicFilter { + uint64 expression_id = 1; + // Serialized datafusion.proto.PhysicalExprNode. + bytes expression_proto = 2; + } + + repeated DynamicFilter filters = 1; +} + message TaskMetrics { // Metrics for a single task's plan nodes in pre-order traversal order. // The TaskKey is implicit — it is determined by the SetPlanRequest that @@ -107,6 +120,8 @@ message SetPlanRequest { // Unix nanos when the query started as reported by the coordinator. Used for collecting temporal metrics // relative to when the query was fired in the coordinator. uint64 query_start_time_ns = 6; + // Dynamic-filter consumer IDs belonging to the selected DistributedLeafExec variants. + repeated uint64 dynamic_filter_ids = 7; } message WorkUnitBatch { diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index 566126715..19650d58e 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -9,8 +9,8 @@ use crate::{ BytesMetricExt, CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, FirstLatencyMetric, GetWorkerInfoRequest, GetWorkerInfoResponse, LatencyMetricExt, LoadInfo, MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P95LatencyMetric, ProducerHeadSpec, - SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, - WorkerChannel, WorkerToCoordinatorMsg, + SetPlanRequest, TaskCompletedDynamicFilters, TaskDynamicFilter, TaskKey, TaskMetrics, + WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; use arrow_flight::decode::FlightRecordBatchStream; @@ -24,6 +24,7 @@ use datafusion::execution::TaskContext; use datafusion::execution::memory_pool::MemoryConsumer; use datafusion::physical_expr_common::metrics::{Count, MetricBuilder, MetricValue, Time}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_proto::protobuf::PhysicalExprNode; use futures::stream::BoxStream; use futures::{FutureExt, Stream, StreamExt, TryStreamExt}; use http::{Extensions, HeaderMap}; @@ -432,6 +433,7 @@ fn encode_set_plan_request(request: SetPlanRequest) -> pb::SetPlanRequest { .collect(), target_worker_url: request.target_worker_url.to_string(), query_start_time_ns: request.query_start_time_ns as u64, + dynamic_filter_ids: request.dynamic_filter_ids, } } @@ -487,10 +489,33 @@ fn decode_worker_to_coordinator_msg( pb::worker_to_coordinator_msg::Inner::LoadInfoEos(_) => { WorkerToCoordinatorMsg::LoadInfoEos } + pb::worker_to_coordinator_msg::Inner::TaskCompletedDynamicFilters(filters) => { + WorkerToCoordinatorMsg::TaskCompletedDynamicFilters( + decode_task_completed_dynamic_filters(filters)?, + ) + } }, ) } +fn decode_task_completed_dynamic_filters( + filters: pb::TaskCompletedDynamicFilters, +) -> Result { + Ok(TaskCompletedDynamicFilters { + filters: filters + .filters + .into_iter() + .map(|filter| { + Ok(TaskDynamicFilter { + expression_id: filter.expression_id, + expression: PhysicalExprNode::decode(filter.expression_proto.as_slice()) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + }) + }) + .collect::>()?, + }) +} + fn decode_task_metrics(task_metrics: pb::TaskMetrics) -> Result { Ok(TaskMetrics { pre_order_plan_metrics: task_metrics diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index 6fba14ee5..cb4ee675b 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -8,8 +8,8 @@ use crate::protocol::ProducerHeadSpec; use crate::protocol::grpc::{ObservabilityServiceImpl, ObservabilityServiceServer}; use crate::{ CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, LoadInfo, SetPlanRequest, - TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, Worker, - WorkerResolver, WorkerToCoordinatorMsg, + TaskCompletedDynamicFilters, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, + WorkUnitMsg, Worker, WorkerResolver, WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; @@ -213,6 +213,7 @@ fn decode_set_plan_request(request: pb::SetPlanRequest) -> Result>()?, target_worker_url: parse_url(&request.target_worker_url, "target_worker_url")?, query_start_time_ns: request.query_start_time_ns as usize, + dynamic_filter_ids: request.dynamic_filter_ids, }) } @@ -261,10 +262,30 @@ fn encode_worker_to_coordinator_msg( WorkerToCoordinatorMsg::LoadInfoEos => { pb::worker_to_coordinator_msg::Inner::LoadInfoEos(true) } + WorkerToCoordinatorMsg::TaskCompletedDynamicFilters(filters) => { + pb::worker_to_coordinator_msg::Inner::TaskCompletedDynamicFilters( + encode_task_completed_dynamic_filters(filters), + ) + } }), }) } +fn encode_task_completed_dynamic_filters( + filters: TaskCompletedDynamicFilters, +) -> pb::TaskCompletedDynamicFilters { + pb::TaskCompletedDynamicFilters { + filters: filters + .filters + .into_iter() + .map(|filter| pb::task_completed_dynamic_filters::DynamicFilter { + expression_id: filter.expression_id, + expression_proto: filter.expression.encode_to_vec(), + }) + .collect(), + } +} + fn encode_task_metrics(task_metrics: TaskMetrics) -> Result { Ok(pb::TaskMetrics { pre_order_plan_metrics: task_metrics diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 6df0e11b1..9e241b701 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -9,6 +9,7 @@ pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, - LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, - WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, + LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskCompletedDynamicFilters, TaskDynamicFilter, + TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, + WorkerToCoordinatorMsg, }; diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index 3ad3816cd..9694b6a72 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -3,6 +3,7 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::Result; use datafusion::execution::TaskContext; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion_proto::protobuf::PhysicalExprNode; use futures::stream::BoxStream; use http::HeaderMap; use std::sync::Arc; @@ -79,6 +80,9 @@ pub struct SetPlanRequest { // might want to omit the serde step, so there should be a way to pass here a normal plan, and // pass the serializer/deserialized separately instead of being coupled to protobuf serialization pub plan_proto: Vec, + /// Expression IDs for dynamic-filter consumers that came from the selected + /// `DistributedLeafExec` variants. Workers report only these filters for visualization. + pub dynamic_filter_ids: Vec, /// Information about all the work unit feeds that will be streamed from coordinator to worker. /// This information is needed here because at the moment of setting the plan, all the appropriate /// channels for the incoming work unit feeds need to be constructed. @@ -121,12 +125,29 @@ pub enum WorkerToCoordinatorMsg { /// ensuring metrics are never lost due to early stream termination. /// metrics[i] is the set of metrics for plan node i in pre-order traversal order. TaskMetrics(TaskMetrics), + /// Sends the final dynamic filters used by dynamic filter consumers back to the coorindator + /// for displaying. + TaskCompletedDynamicFilters(TaskCompletedDynamicFilters), /// Load information reported by a task. This information is used for dynamically /// sizing the number of workers involved in a query. LoadInfo(LoadInfo), LoadInfoEos, } +#[derive(Clone, Debug, Default)] +pub struct TaskCompletedDynamicFilters { + /// Final expressions keyed by their DataFusion physical-expression ID. The TaskKey is + /// implicit from the coordinator channel that carried this message. + pub filters: Vec, +} + +#[derive(Clone, Debug)] +pub struct TaskDynamicFilter { + pub expression_id: u64, + /// A `DynamicFilterPhysicalExpr` proto containing its final predicate and completion state. + pub expression: PhysicalExprNode, +} + #[derive(Clone, Debug)] pub struct TaskMetrics { /// Metrics for a single task's plan nodes in pre-order traversal order. diff --git a/src/stage.rs b/src/stage.rs index cd054c575..967135309 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -288,8 +288,12 @@ fn display_ascii( show_metrics: bool, f: &mut String, ) -> std::fmt::Result { + let prepared_plan = match stage { + Either::Left(distributed_exec) => Some(distributed_exec.plan_for_display()), + Either::Right(_) => None, + }; let plan = match stage { - Either::Left(distributed_exec) => distributed_exec.children().first().unwrap(), + Either::Left(_) => prepared_plan.as_ref().unwrap(), Either::Right(stage) => { let Some(plan) = stage.local_plan() else { return write!(f, "StageExec: encoded input plan"); diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index 020bce948..e9641fdfa 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,4 +1,4 @@ -use crate::common::TreeNodeExt; +use crate::common::{TreeNodeExt, discover_dynamic_filter_consumers}; use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers}; use crate::execution_plans::SamplerExec; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; @@ -6,16 +6,20 @@ use crate::worker::LocalWorkerContext; use crate::worker::task_data::TaskDataMetrics; use crate::{ CoordinatorToWorkerMsg, DistributedCodec, DistributedConfig, DistributedExt, - DistributedTaskContext, TaskData, TaskMetrics, Worker, WorkerQueryContext, - WorkerToCoordinatorMsg, + DistributedTaskContext, TaskCompletedDynamicFilters, TaskData, TaskDynamicFilter, TaskMetrics, + Worker, WorkerQueryContext, WorkerToCoordinatorMsg, }; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{DataFusionError, Result, exec_datafusion_err, internal_err}; +use datafusion::common::{DataFusionError, HashSet, Result, exec_datafusion_err, internal_err}; use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; -use datafusion_proto::physical_plan::AsExecutionPlan; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; +use datafusion_proto::physical_plan::{ + AsExecutionPlan, DeduplicatingProtoConverter, PhysicalPlanNodeExt, +}; use datafusion_proto::protobuf::PhysicalPlanNode; +use datafusion_proto::protobuf::physical_expr_node::ExprType; use futures::stream::{BoxStream, FuturesUnordered}; use futures::{FutureExt, StreamExt, TryStreamExt}; use http::HeaderMap; @@ -51,6 +55,7 @@ impl Worker { } let (metrics_tx, metrics_rx) = oneshot::channel(); + let (dynamic_filters_tx, dynamic_filters_rx) = oneshot::channel(); let mut load_info_rxs = vec![]; let task_data = || async { @@ -88,7 +93,11 @@ impl Worker { let task_ctx = session_state.task_ctx(); let proto_node = PhysicalPlanNode::try_decode(request.plan_proto.as_ref())?; let ev = WorkerPlanRewriteEvent { - plan: proto_node.try_into_physical_plan(&task_ctx, &codec)?, + plan: proto_node.try_into_physical_plan_with_converter( + &task_ctx, + &codec, + &DeduplicatingProtoConverter::default(), + )?, session_config: session_state.config(), }; let plan = WorkerPlanRewriteHandlers::handle(ev)?.plan; @@ -119,18 +128,20 @@ impl Worker { // Continue reading remaining messages (work unit feed data) in the background. let mut work_unit_senders = Some(remote_work_unit_feed_registry.senders); let task_data_entries = Arc::clone(&self.task_data_entries); + let dynamic_filter_ids: HashSet<_> = request.dynamic_filter_ids.iter().copied().collect(); - // This tokio task takes ownership of the `oneshot::Sender` that keeps - // alive the worker->coordinator stream. as soon as this task ends, the runtime metrics - // are send back and the worker->coordinator stream ends. The flow is the following: + // This tokio task takes ownership of the final-report senders that keep the + // worker->coordinator stream alive. As soon as this task ends, the runtime metrics and + // final dynamic filters are sent back and the worker->coordinator stream ends. The flow + // is the following: // 1. The query ends normally, as all Arrow RecordBatches are already streamed. // 2. In DistributedExec::execute(), the end query guard is dropped. // 3. In StageCoordinator::send_plan_task(), `end_stream_notifier` fires and the // coordinator->worker channel is gracefully ended. // 4. The coordinator->worker channel EOS is received by this same function, ending the // while loop inside this `tokio::spawn` below. - // 5. The metrics are send back in the worker->coordinator channel, and then that channel - // is closed. + // 5. The metrics and final dynamic filters are sent back in the worker->coordinator + // channel, and then that channel is closed. #[allow(clippy::disallowed_methods)] tokio::spawn(async move { let mut stream = stream.map_ok(set_work_unit_received_time); @@ -173,6 +184,7 @@ impl Worker { } let metrics_tx = task_data.metrics_tx.lock().unwrap().take(); + let mut dynamic_filters = TaskCompletedDynamicFilters::default(); if let Some(Ok(plan)) = task_data.final_plan.get() { let d_ctx = DistributedTaskContext { task_index: key.task_number, @@ -183,7 +195,14 @@ impl Worker { if let Some(metrics_tx) = metrics_tx { send_metrics_via_channel(metrics_tx, plan, d_ctx, task_data_metrics); } + dynamic_filters = build_task_completed_dynamic_filters( + plan, + &dynamic_filter_ids, + &task_data.task_ctx, + ) + .unwrap_or_default(); } + let _ = dynamic_filters_tx.send(dynamic_filters); task_data_entries.invalidate(&key).await }); @@ -206,10 +225,48 @@ impl Worker { Some(WorkerToCoordinatorMsg::TaskMetrics(task_metrics)) }); - Ok(futures::stream::select(load_info_stream, metrics_stream) - .map(Ok) - .boxed()) + let dynamic_filters_stream = dynamic_filters_rx.into_stream().filter_map( + async |dynamic_filters_or_channel_dropped| { + let dynamic_filters = dynamic_filters_or_channel_dropped.ok()?; + Some(WorkerToCoordinatorMsg::TaskCompletedDynamicFilters( + dynamic_filters, + )) + }, + ); + + Ok(futures::stream::select( + load_info_stream, + futures::stream::select(metrics_stream, dynamic_filters_stream), + ) + .map(Ok) + .boxed()) + } +} + +fn build_task_completed_dynamic_filters( + plan: &Arc, + allowed_ids: &HashSet, + task_ctx: &Arc, +) -> Result { + let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config()); + let mut filters = vec![]; + for consumer in discover_dynamic_filter_consumers(plan, Some(allowed_ids))? { + // Serializing the complete DynamicFilterPhysicalExpr preserves both its current + // predicate and its completion state through DataFusion's native proto hook. + let expression = serialize_physical_expr(&consumer.expression, &codec)?; + let Some(ExprType::DynamicFilter(dynamic_filter)) = expression.expr_type.as_ref() else { + return internal_err!("discovered dynamic filter did not serialize as one"); + }; + // A cancelled or short-circuited task can leave filters incomplete. Do not report those + // as final values for display. + if dynamic_filter.is_complete { + filters.push(TaskDynamicFilter { + expression_id: consumer.id, + expression, + }); + } } + Ok(TaskCompletedDynamicFilters { filters }) } /// Collects metrics from the plan in pre-order traversal order and sends them via the diff --git a/tests/dynamic_filtering.rs b/tests/dynamic_filtering.rs new file mode 100644 index 000000000..baba4fe89 --- /dev/null +++ b/tests/dynamic_filtering.rs @@ -0,0 +1,116 @@ +#[cfg(all(feature = "integration", test))] +mod tests { + use datafusion::common::Result; + use datafusion::physical_plan::collect; + use datafusion_distributed::test_utils::localhost::start_localhost_context; + use datafusion_distributed::test_utils::parquet::register_parquet_tables; + use datafusion_distributed::{ + DefaultSessionBuilder, DistributedExt, assert_snapshot, display_plan_ascii, + }; + use std::sync::Arc; + + #[tokio::test] + async fn collect_left_local_dynamic_filters() -> Result<()> { + let display = execute_local_hash_join(true).await?; + assert_snapshot!(display, @r" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 3] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 3 ── tasks=2, partitions=6 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(key@0, RainToday@0)], projection=[] + │ CoalescePartitionsExec + │ [Stage 2] => NetworkBroadcastExec: partitions_per_consumer=3, stage_partitions=6, input_tasks=2 + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ RainToday@19 >= No AND RainToday@19 <= Yes AND RainToday@19 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=RainToday_null_count@1 != row_count@2 AND RainToday_max@0 >= No AND RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND (RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND Yes <= RainToday_max@0 OR RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= No AND No <= RainToday_max@0), required_guarantees=[RainToday in (No, Yes)] + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ RainToday@19 >= No AND RainToday@19 <= Yes AND RainToday@19 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=RainToday_null_count@1 != row_count@2 AND RainToday_max@0 >= No AND RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND (RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND Yes <= RainToday_max@0 OR RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= No AND No <= RainToday_max@0), required_guarantees=[RainToday in (No, Yes)] + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=12 + │ BroadcastExec: input_partitions=3, consumer_tasks=2, output_partitions=6 + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + └────────────────────────────────────────────────── + "); + Ok(()) + } + + #[tokio::test] + async fn partitioned_local_dynamic_filters() -> Result<()> { + let display = execute_local_hash_join(false).await?; + assert_snapshot!(display, @r" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 3] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 3 ── tasks=2, partitions=3 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(key@0, RainToday@0)], projection=[] + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + │ [Stage 2] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([RainToday@0], 6), input_partitions=3 + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + └────────────────────────────────────────────────── + "); + Ok(()) + } + + async fn execute_local_hash_join(broadcast_joins: bool) -> Result { + let (ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; + let ctx = ctx.with_distributed_broadcast_joins(broadcast_joins)?; + if !broadcast_joins { + let state = ctx.state_ref(); + let mut state = state.write(); + let optimizer = &mut state.config_mut().options_mut().optimizer; + optimizer.hash_join_single_partition_threshold = 0; + optimizer.hash_join_single_partition_threshold_rows = 0; + } + register_parquet_tables(&ctx).await?; + + let plan = ctx + .sql( + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + JOIN weather probe ON build.key = probe."RainToday" + "#, + ) + .await? + .create_physical_plan() + .await?; + + let results = collect(Arc::clone(&plan), ctx.task_ctx()).await?; + assert_eq!( + results.iter().map(|batch| batch.num_rows()).sum::(), + 1 + ); + + Ok(display_plan_ascii(plan.as_ref(), false)) + } +}