From d417dc6174bea307b7bddaa5703d75cc76e0d846 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Wed, 12 Aug 2026 17:58:24 +0000 Subject: [PATCH 01/17] ensure apply_expressions is implemented for DistributedLeafExec --- src/execution_plans/distributed_leaf.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/execution_plans/distributed_leaf.rs b/src/execution_plans/distributed_leaf.rs index 9f3d9884f..b7aa2645c 100644 --- a/src/execution_plans/distributed_leaf.rs +++ b/src/execution_plans/distributed_leaf.rs @@ -164,6 +164,9 @@ impl ExecutionPlan for DistributedLeafExec { &self, f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { + // `original` is deliberately hidden from `children()` so this remains a distributed leaf. + // It is the canonical representation of the task-specific variants, so expose its + // expression roots exactly as a transparent wrapper would. self.original.apply_expressions(f) } From 99a9ab636ff285675b101c4d489e7e7f9186b2ea Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Wed, 12 Aug 2026 18:16:33 +0000 Subject: [PATCH 02/17] remove extra comment --- src/execution_plans/distributed_leaf.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/execution_plans/distributed_leaf.rs b/src/execution_plans/distributed_leaf.rs index b7aa2645c..9f3d9884f 100644 --- a/src/execution_plans/distributed_leaf.rs +++ b/src/execution_plans/distributed_leaf.rs @@ -164,9 +164,6 @@ impl ExecutionPlan for DistributedLeafExec { &self, f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // `original` is deliberately hidden from `children()` so this remains a distributed leaf. - // It is the canonical representation of the task-specific variants, so expose its - // expression roots exactly as a transparent wrapper would. self.original.apply_expressions(f) } From 429fdeccfcb449920cbab4e9557f4c27520a645b Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 11 Aug 2026 15:03:39 +0000 Subject: [PATCH 03/17] display dynamic filters during execution --- docs/upgrade/3.0.0.md | 6 + src/common/dynamic_filtering.rs | 210 ++++++++++++++++++++++ src/common/mod.rs | 4 + src/coordinator/distributed.rs | 18 +- src/coordinator/dynamic_filters.rs | 225 ++++++++++++++++++++++++ src/coordinator/mod.rs | 1 + src/coordinator/query_coordinator.rs | 53 +++++- src/execution_plans/distributed_leaf.rs | 34 ++++ src/lib.rs | 6 +- src/protocol/grpc/generated/worker.rs | 28 ++- src/protocol/grpc/worker.proto | 15 ++ src/protocol/grpc/worker_client.rs | 23 ++- src/protocol/grpc/worker_service.rs | 23 ++- src/protocol/mod.rs | 4 +- src/protocol/worker_channel.rs | 20 +++ src/stage.rs | 10 +- src/worker/impl_coordinator_channel.rs | 72 ++++++-- tests/dynamic_filtering.rs | 62 +++++++ tests/tpcds_correctness_test.rs | 41 ++++- tests/tpch_plans_test.rs | 48 ++++- 20 files changed, 862 insertions(+), 41 deletions(-) create mode 100644 src/common/dynamic_filtering.rs create mode 100644 src/coordinator/dynamic_filters.rs create mode 100644 tests/dynamic_filtering.rs diff --git a/docs/upgrade/3.0.0.md b/docs/upgrade/3.0.0.md index b690e1448..119913f98 100644 --- a/docs/upgrade/3.0.0.md +++ b/docs/upgrade/3.0.0.md @@ -55,6 +55,12 @@ 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::TaskDynamicFilters` variant. The built-in gRPC +transport handles both automatically. + ## 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 7dfd3a96d..589c9b37a 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,4 +1,5 @@ mod children_helpers; +mod dynamic_filtering; mod maybe_encoded; mod once_lock; mod recursion; @@ -8,6 +9,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 use maybe_encoded::MaybeEncoded; pub(crate) use once_lock::OnceLockResult; pub(crate) use recursion::TreeNodeExt; diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 1e5fc796f..81d794413 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,4 +1,5 @@ use crate::common::require_one_child; +use crate::coordinator::dynamic_filters::isolate_distributed_leaf_display_variants; use crate::coordinator::metrics_store::MetricsStore; use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan; use crate::coordinator::prepare_static_plan::prepare_static_plan; @@ -124,6 +125,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 +232,12 @@ impl ExecutionPlan for DistributedExec { false => prepare_static_plan(&query_coordinator, &base_plan)?, }; + let display_plan = + isolate_distributed_leaf_display_variants(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 +248,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..0aadea77f --- /dev/null +++ b/src/coordinator/dynamic_filters.rs @@ -0,0 +1,225 @@ +use crate::common::discover_dynamic_filter_consumers; +use crate::execution_plans::DistributedLeafExec; +use crate::{DistributedCodec, TaskDynamicFilters, TaskKey}; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion::common::{HashMap, HashSet, 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::{PhysicalExprNode, PhysicalPlanNode}; +use prost::Message; +use std::sync::{Arc, Mutex}; +use tokio::sync::Notify; + +#[derive(Default)] +struct ReportState { + expected: HashSet, + finished: HashSet, + reports: HashMap, +} + +/// Query-scoped storage for final dynamic-filter reports. +#[derive(Default)] +pub(super) struct DynamicFilterReports { + state: Mutex, + changed: Notify, +} + +impl DynamicFilterReports { + pub(super) fn expect(&self, key: TaskKey) { + self.state.lock().unwrap().expected.insert(key); + } + + pub(super) fn insert(&self, key: TaskKey, report: TaskDynamicFilters) { + self.state.lock().unwrap().reports.insert(key, report); + self.changed.notify_one(); + } + + pub(super) fn finish(&self, key: TaskKey) { + self.state.lock().unwrap().finished.insert(key); + self.changed.notify_one(); + } + + pub(super) async fn wait_for_finished(&self) -> HashMap { + loop { + let notified = self.changed.notified(); + { + let state = self.state.lock().unwrap(); + if state.expected.is_subset(&state.finished) { + return state.reports.clone(); + } + } + notified.await; + } + } +} + +/// Gives every distributed-leaf task an independent plan used only for rendering. +pub(super) fn isolate_distributed_leaf_display_variants( + 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 display_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(leaf.with_display_variants(display_variants)?) as Arc, + )) + }) + .map(|transformed| transformed.data) +} + +/// Applies successful worker reports only to the matching task-local display 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.display_variants.get(task_key.task_number) else { + continue; + }; + let updates: HashMap<_, _> = report + .filters + .iter() + .map(|filter| (filter.expression_id, filter.expression.as_slice())) + .collect(); + let Ok(consumers) = discover_dynamic_filter_consumers(variant, None) else { + continue; + }; + for consumer in consumers { + let Some(encoded) = updates.get(&consumer.id) else { + continue; + }; + let Ok(proto) = PhysicalExprNode::decode(*encoded) 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 task_display_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_display_variants(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 = TaskDynamicFilters { + filters: vec![crate::TaskDynamicFilter { + expression_id: dynamic_filter.expression_id().unwrap(), + expression: serialize_physical_expr(&dynamic_filter, &codec)?.encode_to_vec(), + }], + }; + 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.display_variants[0].as_ref()) + .one_line() + .to_string(); + let task_1 = displayable(leaf.display_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/mod.rs b/src/coordinator/mod.rs index c1a8a8dd2..6d0da6552 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,4 +1,5 @@ mod distributed; +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 2f9c99fb0..11519c411 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,7 +1,10 @@ use crate::codec::{decode_execution_plan, encode_execution_plan}; -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::{ + DynamicFilterReports, apply_reports_to_distributed_leaves, +}; use crate::coordinator::latency_metric::LatencyMetric; use crate::events::{RouteTasksEvent, RouteTasksHandlers}; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; @@ -47,6 +50,7 @@ pub(super) struct QueryCoordinator { metrics: ExecutionPlanMetricsSet, coordinator_to_worker_metrics: CoordinatorToWorkerMetrics, metrics_store: Option>, + dynamic_filter_reports: Arc, end_stream_notifier: Arc, join_set: Mutex>>, } @@ -62,6 +66,7 @@ impl QueryCoordinator { task_ctx, metrics: metrics_set.clone(), metrics_store, + dynamic_filter_reports: Arc::new(DynamicFilterReports::default()), coordinator_to_worker_metrics: CoordinatorToWorkerMetrics::new(metrics_set), end_stream_notifier: Arc::new(Notify::new()), join_set: Mutex::new(JoinSet::new()), @@ -80,6 +85,7 @@ impl QueryCoordinator { metrics_set: &self.metrics, metrics: &self.coordinator_to_worker_metrics, metrics_store: &self.metrics_store, + dynamic_filter_reports: &self.dynamic_filter_reports, end_stream_notifier: &self.end_stream_notifier, join_set: &self.join_set, } @@ -97,6 +103,14 @@ 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.dynamic_filter_reports.wait_for_finished().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<()> { @@ -116,6 +130,12 @@ impl QueryCoordinator { /// - Building tasks that communicate a serialized plan to multiple workers for further execution. /// - Building tasks that stream partition feeds from local [WorkUnitFeedExec] nodes to their /// remote counterparts. +type SpecializedTaskPlan = ( + Arc, + Vec, + Vec, +); + pub(super) struct StageCoordinator<'a> { plan: &'a Arc, query_id: Uuid, @@ -125,6 +145,7 @@ pub(super) struct StageCoordinator<'a> { metrics_set: &'a ExecutionPlanMetricsSet, metrics: &'a CoordinatorToWorkerMetrics, metrics_store: &'a Option>, + dynamic_filter_reports: &'a Arc, end_stream_notifier: &'a Arc, join_set: &'a Mutex>>, } @@ -143,18 +164,21 @@ impl<'a> StageCoordinator<'a> { )> { let session_config = self.task_ctx.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 task_key = TaskKey { query_id: self.query_id, stage_id: self.stage_id, task_number: task_i, }; + self.dynamic_filter_reports.expect(task_key); let set_plan_request = SetPlanRequest { task_key, task_count: self.task_count, plan: MaybeEncoded::Decoded(specialized), + dynamic_filter_ids, work_unit_feed_declarations, target_worker_url: url.clone(), query_start_time_ns: self.metrics.instantiation_time, @@ -178,8 +202,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(); @@ -236,6 +260,7 @@ impl<'a> StageCoordinator<'a> { task_number: task_i, }; let task_metrics = self.metrics_store.clone(); + let dynamic_filter_reports = Arc::clone(self.dynamic_filter_reports); let (load_info_tx, load_info_rx) = tokio::sync::mpsc::unbounded_channel(); let mut load_info_tx_opt = Some(load_info_tx); @@ -258,8 +283,13 @@ impl<'a> StageCoordinator<'a> { WorkerToCoordinatorMsg::LoadInfoEos => { let _ = load_info_tx_opt.take(); } + WorkerToCoordinatorMsg::TaskDynamicFilters(filters) => { + dynamic_filter_reports.insert(task_key, filters); + dynamic_filter_reports.finish(task_key); + } } } + dynamic_filter_reports.finish(task_key); }); load_info_rx } @@ -341,16 +371,14 @@ impl<'a> StageCoordinator<'a> { /// trimming down any unnecessary information that the specific `task_i` task is not going to /// need, like unexecuted branches in [ChildrenIsolatorUnionExec], or unexecuted variants of /// [DistributedLeafExec]. - fn task_specialized_plan( - &self, - task_i: usize, - ) -> Result<(Arc, Vec)> { + fn task_specialized_plan(&self, task_i: usize) -> Result { 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, @@ -382,12 +410,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/execution_plans/distributed_leaf.rs b/src/execution_plans/distributed_leaf.rs index 9f3d9884f..b861d43bf 100644 --- a/src/execution_plans/distributed_leaf.rs +++ b/src/execution_plans/distributed_leaf.rs @@ -74,6 +74,9 @@ pub struct DistributedLeafExec { pub(crate) original: Arc, pub(crate) properties: Arc, pub(crate) variants: Vec>, + /// Per-task variants used only for rendering. These are isolated from one another before a + /// query runs so that applying task-local runtime state cannot leak into another task's line. + pub(crate) display_variants: Vec>, } impl DistributedLeafExec { @@ -115,6 +118,7 @@ impl DistributedLeafExec { Ok(Self { original, properties, + display_variants: variants.clone(), variants, }) } @@ -138,6 +142,36 @@ impl DistributedLeafExec { pub(crate) fn to_task_specialized(&self, task_i: usize) -> Arc { Arc::clone(&self.variants[task_i]) } + + pub(crate) fn with_display_variants( + &self, + display_variants: Vec>, + ) -> Result { + if display_variants.len() != self.variants.len() { + return plan_err!( + "DistributedLeafExec received {} display variants for {} execution variants", + display_variants.len(), + self.variants.len() + ); + } + for (display, execution) in display_variants.iter().zip(&self.variants) { + if display.schema() != execution.schema() + || display.properties().partitioning.partition_count() + != execution.properties().partitioning.partition_count() + { + return plan_err!( + "DistributedLeafExec display variant properties differ from its execution variant" + ); + } + } + + Ok(Self { + original: Arc::clone(&self.original), + properties: Arc::clone(&self.properties), + variants: self.variants.clone(), + display_variants, + }) + } } impl DisplayAs for DistributedLeafExec { diff --git a/src/lib.rs b/src/lib.rs index 3ad27ddd9..fb2dc7124 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -54,9 +54,9 @@ pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; pub use protocol::{ ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, - GetWorkerInfoResponse, LoadInfo, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, - WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, - get_distributed_channel_resolver, + GetWorkerInfoResponse, LoadInfo, SetPlanRequest, TaskDynamicFilter, TaskDynamicFilters, + 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 6bf8beb89..c819e7c8b 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")] + TaskDynamicFilters(super::TaskDynamicFilters), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct TaskDynamicFilters { + #[prost(message, repeated, tag = "1")] + pub filters: ::prost::alloc::vec::Vec, +} +/// Nested message and enum types in `TaskDynamicFilters`. +pub mod task_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: ::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 2affe1187..82f842dcf 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. + TaskDynamicFilters task_dynamic_filters = 4; } } +message TaskDynamicFilters { + message DynamicFilter { + uint64 expression_id = 1; + // Serialized datafusion.proto.PhysicalExprNode. + bytes expression = 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 060d8e875..043c220ea 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -9,9 +9,9 @@ use crate::{ BytesMetricExt, CoordinatorToWorkerMsg, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedConfig, ExecuteTaskRequest, FirstLatencyMetric, GetWorkerInfoRequest, GetWorkerInfoResponse, LatencyMetricExt, LoadInfo, MaxLatencyMetric, MaybeEncoded, - MinLatencyMetric, P50LatencyMetric, P95LatencyMetric, ProducerHead, SetPlanRequest, TaskKey, - TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, - WorkerToCoordinatorMsg, + MinLatencyMetric, P50LatencyMetric, P95LatencyMetric, ProducerHead, SetPlanRequest, + TaskDynamicFilter, TaskDynamicFilters, TaskKey, TaskMetrics, WorkUnitBatch, + WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; use arrow_flight::decode::FlightRecordBatchStream; @@ -450,6 +450,7 @@ fn encode_set_plan_request( .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, }) } @@ -508,10 +509,26 @@ fn decode_worker_to_coordinator_msg( pb::worker_to_coordinator_msg::Inner::LoadInfoEos(_) => { WorkerToCoordinatorMsg::LoadInfoEos } + pb::worker_to_coordinator_msg::Inner::TaskDynamicFilters(filters) => { + WorkerToCoordinatorMsg::TaskDynamicFilters(decode_task_dynamic_filters(filters)) + } }, ) } +fn decode_task_dynamic_filters(filters: pb::TaskDynamicFilters) -> TaskDynamicFilters { + TaskDynamicFilters { + filters: filters + .filters + .into_iter() + .map(|filter| TaskDynamicFilter { + expression_id: filter.expression_id, + expression: filter.expression, + }) + .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 1f0fdf0a6..4419db210 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -7,8 +7,8 @@ use crate::common::{deserialize_uuid, now_ns}; use crate::protocol::grpc::{ObservabilityServiceImpl, ObservabilityServiceServer}; use crate::{ CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, LoadInfo, MaybeEncoded, - ProducerHead, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, - WorkUnitMsg, Worker, WorkerResolver, WorkerToCoordinatorMsg, + ProducerHead, SetPlanRequest, TaskDynamicFilters, TaskKey, TaskMetrics, WorkUnitBatch, + WorkUnitFeedDeclaration, WorkUnitMsg, Worker, WorkerResolver, WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; @@ -228,6 +228,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, }) } @@ -272,10 +273,28 @@ fn encode_worker_to_coordinator_msg( WorkerToCoordinatorMsg::LoadInfoEos => { pb::worker_to_coordinator_msg::Inner::LoadInfoEos(true) } + WorkerToCoordinatorMsg::TaskDynamicFilters(filters) => { + pb::worker_to_coordinator_msg::Inner::TaskDynamicFilters( + encode_task_dynamic_filters(filters), + ) + } }), }) } +fn encode_task_dynamic_filters(filters: TaskDynamicFilters) -> pb::TaskDynamicFilters { + pb::TaskDynamicFilters { + filters: filters + .filters + .into_iter() + .map(|filter| pb::task_dynamic_filters::DynamicFilter { + expression_id: filter.expression_id, + expression: filter.expression, + }) + .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 4bd507dfe..b6e6b3af8 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -11,6 +11,6 @@ pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; pub use in_process::LocalWorkerContext; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, - LoadInfo, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, - WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, + LoadInfo, SetPlanRequest, TaskDynamicFilter, TaskDynamicFilters, TaskKey, TaskMetrics, + WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, }; diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index 0bea0bbe2..85058de5f 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -78,6 +78,9 @@ pub struct SetPlanRequest { pub task_count: usize, /// The subplan the worker is expected to execute. pub plan: MaybeEncoded>, + /// 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. @@ -124,6 +127,23 @@ pub enum WorkerToCoordinatorMsg { /// sizing the number of workers involved in a query. LoadInfo(LoadInfo), LoadInfoEos, + /// Final task-local dynamic filters used by distributed leaf variants. + TaskDynamicFilters(TaskDynamicFilters), +} + +#[derive(Clone, Debug, Default)] +pub struct TaskDynamicFilters { + /// 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 serialized `DynamicFilterPhysicalExpr` `PhysicalExprNode`, including its final predicate + /// and completion state. + pub expression: Vec, } #[derive(Clone, Debug)] diff --git a/src/stage.rs b/src/stage.rs index cd054c575..8b20d0dab 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"); @@ -415,7 +419,7 @@ fn display_inner_distributed_leaf( && !by_task.is_empty() { writeln!(f, "{indent} DistributedLeafExec:")?; - for (task_i, variant) in leaf.variants.iter().enumerate() { + for (task_i, variant) in leaf.display_variants.iter().enumerate() { let variant = displayable(variant.as_ref()).one_line().to_string(); let metrics = match by_task.is_empty() { true => String::new(), @@ -429,7 +433,7 @@ fn display_inner_distributed_leaf( false => String::new(), }; writeln!(f, "{indent} DistributedLeafExec:{header}")?; - for (task_i, variant) in leaf.variants.iter().enumerate() { + for (task_i, variant) in leaf.display_variants.iter().enumerate() { let variant = displayable(variant.as_ref()).one_line().to_string(); writeln!(f, "{indent} t{task_i}: {}", variant.trim_end())?; } diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index efba61c28..b75d3a5d0 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,21 +1,25 @@ -use crate::common::TreeNodeExt; +use crate::common::{TreeNodeExt, discover_dynamic_filter_consumers}; use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers}; use crate::execution_plans::SamplerExec; use crate::protocol::LocalWorkerContext; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::task_data::TaskDataMetrics; use crate::{ - CoordinatorToWorkerMsg, DistributedConfig, DistributedExt, DistributedTaskContext, - SetPlanRequest, TaskData, TaskMetrics, Worker, WorkerQueryContext, WorkerToCoordinatorMsg, + CoordinatorToWorkerMsg, DistributedCodec, DistributedConfig, DistributedExt, + DistributedTaskContext, SetPlanRequest, TaskData, TaskDynamicFilter, TaskDynamicFilters, + TaskMetrics, Worker, WorkerQueryContext, WorkerToCoordinatorMsg, }; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{DataFusionError, Result, exec_datafusion_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::to_proto::serialize_physical_expr; +use datafusion_proto::protobuf::physical_expr_node::ExprType; use futures::stream::{BoxStream, FuturesUnordered}; use futures::{FutureExt, StreamExt, TryStreamExt}; use http::HeaderMap; +use prost::Message; use std::sync::{Arc, OnceLock}; use tokio::sync::oneshot; use tokio::sync::oneshot::Sender; @@ -40,6 +44,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 { @@ -108,18 +113,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); @@ -157,6 +164,7 @@ impl Worker { } let metrics_tx = task_data.metrics_tx.lock().unwrap().take(); + let mut dynamic_filters = TaskDynamicFilters::default(); if let Some(Ok(plan)) = task_data.final_plan.get() { let d_ctx = DistributedTaskContext { task_index: key.task_number, @@ -167,7 +175,11 @@ 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_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 }); @@ -190,10 +202,46 @@ 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::TaskDynamicFilters(dynamic_filters)) + }, + ); + + Ok(futures::stream::select( + load_info_stream, + futures::stream::select(metrics_stream, dynamic_filters_stream), + ) + .map(Ok) + .boxed()) + } +} + +fn build_task_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: expression.encode_to_vec(), + }); + } } + Ok(TaskDynamicFilters { 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..2302974e1 --- /dev/null +++ b/tests/dynamic_filtering.rs @@ -0,0 +1,62 @@ +#[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, display_plan_ascii}; + use std::sync::Arc; + + #[tokio::test] + async fn completed_leaf_dynamic_filters_are_displayed_per_task() -> Result<()> { + for dynamic_task_count in [false, true] { + assert_completed_leaf_filters(dynamic_task_count).await?; + } + Ok(()) + } + + async fn assert_completed_leaf_filters(dynamic_task_count: bool) -> Result<()> { + let (ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; + let ctx = ctx.with_distributed_dynamic_task_count(dynamic_task_count)?; + 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 + ); + + let display = display_plan_ascii(plan.as_ref(), false); + let task_filters: Vec<_> = display + .lines() + .filter(|line| line.contains(": DataSourceExec:")) + .filter_map(|line| line.split_once("DynamicFilter [ ")) + .map(|(_, filter)| filter.split_once(" ]").unwrap().0) + .collect(); + + assert!( + task_filters.len() >= 2, + "expected task-local leaf filters in plan:\n{display}" + ); + assert!( + task_filters.iter().all(|filter| *filter != "empty"), + "expected completed filters in plan:\n{display}" + ); + Ok(()) + } +} diff --git a/tests/tpcds_correctness_test.rs b/tests/tpcds_correctness_test.rs index e4baeba9d..fe1f4c4eb 100644 --- a/tests/tpcds_correctness_test.rs +++ b/tests/tpcds_correctness_test.rs @@ -10,7 +10,7 @@ mod tests { compare_ordering, compare_result_set, }; use datafusion_distributed::{ - DefaultSessionBuilder, DistributedExec, DistributedExt, display_plan_ascii, + DefaultSessionBuilder, DistributedExec, DistributedExt, assert_snapshot, display_plan_ascii, }; use datafusion_distributed_benchmarks::datasets::{register_tables, tpcds}; use std::fs; @@ -597,6 +597,12 @@ mod tests { } let display = display_plan_ascii(d_plan.as_ref(), false); println!("Query {query_id}:\n{display}"); + if query_id == "q1" { + assert_snapshot!(displayed_dynamic_filter_predicates(&display), @r" + DynamicFilter [ c_customer_sk@0 >= 5 AND c_customer_sk@0 <= 100000 AND true ] AND DynamicFilter [ empty ] + DynamicFilter [ sr_returned_date_sk@0 >= 2451545 AND sr_returned_date_sk@0 <= 2451910 AND true ] AND DynamicFilter [ empty ] + "); + } // The comparison functions can be computationally expensive, so we spawn them in tokio // blocking tasks so that they do not block the tokio runtime. @@ -618,4 +624,37 @@ mod tests { Ok(()) } + + fn displayed_dynamic_filter_predicates(display: &str) -> String { + let mut predicates: Vec<_> = display + .lines() + .filter(|line| line.contains(": DataSourceExec:")) + .filter_map(|line| { + line.split_once("predicate=") + .map(|(_, predicate)| predicate) + }) + .map(|predicate| { + predicate + .split(", pruning_predicate=") + .next() + .unwrap_or(predicate) + }) + .map(|predicate| { + predicate + .split(", dynamic_rg_pruning=") + .next() + .unwrap_or(predicate) + }) + .map(|predicate| { + predicate + .split(", required_guarantees=") + .next() + .unwrap_or(predicate) + }) + .filter(|predicate| predicate.contains("DynamicFilter")) + .collect(); + predicates.sort_unstable(); + predicates.dedup(); + predicates.join("\n") + } } diff --git a/tests/tpch_plans_test.rs b/tests/tpch_plans_test.rs index a815c6255..fc6909a38 100644 --- a/tests/tpch_plans_test.rs +++ b/tests/tpch_plans_test.rs @@ -1,5 +1,7 @@ #[cfg(all(feature = "integration", feature = "tpch", test))] mod tests { + use datafusion::physical_plan::{ExecutionPlan, collect}; + use datafusion::prelude::SessionContext; use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context; use datafusion_distributed::{ DefaultSessionBuilder, DistributedExt, assert_snapshot, display_plan_ascii, @@ -8,6 +10,7 @@ mod tests { use std::error::Error; use std::fs; use std::path::Path; + use std::sync::Arc; use tokio::sync::OnceCell; const NUM_WORKERS: usize = 4; @@ -219,6 +222,16 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_tpch_3_dynamic_filters_after_execution() -> Result<(), Box> { + let predicates = executed_dynamic_filter_predicates("q3").await?; + assert_snapshot!(predicates, @r" + DynamicFilter [ l_orderkey@0 >= 5 AND l_orderkey@0 <= 119972 AND true ] + DynamicFilter [ o_custkey@1 >= 1 AND o_custkey@1 <= 2993 AND true ] + "); + Ok(()) + } + #[tokio::test] async fn test_tpch_4() -> Result<(), Box> { let plan = test_tpch_query("q4").await?; @@ -1434,6 +1447,39 @@ mod tests { // test_tpch_query generates and displays a distributed plan for each TPC-H query. async fn test_tpch_query(query_id: &str) -> Result> { + let (_, plan) = prepare_tpch_query(query_id).await?; + Ok(display_plan_ascii(plan.as_ref(), false)) + } + + async fn executed_dynamic_filter_predicates(query_id: &str) -> Result> { + let (ctx, plan) = prepare_tpch_query(query_id).await?; + collect(Arc::clone(&plan), ctx.task_ctx()).await?; + let display = display_plan_ascii(plan.as_ref(), false); + let mut predicates: Vec<_> = display + .lines() + .filter(|line| line.contains(": DataSourceExec:")) + .filter_map(|line| line.find("DynamicFilter [ ").map(|start| &line[start..])) + .map(|predicate| { + predicate + .split(", dynamic_rg_pruning=") + .next() + .unwrap_or(predicate) + }) + .map(|predicate| { + predicate + .split(", pruning_predicate=") + .next() + .unwrap_or(predicate) + }) + .collect(); + predicates.sort_unstable(); + predicates.dedup(); + Ok(predicates.join("\n")) + } + + async fn prepare_tpch_query( + query_id: &str, + ) -> Result<(SessionContext, Arc), Box> { let d_ctx = start_in_memory_context(NUM_WORKERS, DefaultSessionBuilder).await; let data_dir = ensure_tpch_data(TPCH_SCALE_FACTOR, TPCH_DATA_PARTS).await; let sql = tpch::get_query(query_id)?; @@ -1473,7 +1519,7 @@ mod tests { df.create_physical_plan().await? }; - Ok(display_plan_ascii(plan.as_ref(), false)) + Ok((d_ctx, plan)) } // OnceCell to ensure TPCH tables are generated only once for tests From b8c282b24ecdef3883eba670b4ad2a312e834b1a Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 11 Aug 2026 16:54:48 +0000 Subject: [PATCH 04/17] remove viz variants --- src/coordinator/distributed.rs | 4 +-- src/coordinator/dynamic_filters.rs | 25 +++++++++--------- src/execution_plans/distributed_leaf.rs | 34 ------------------------- src/stage.rs | 4 +-- 4 files changed, 17 insertions(+), 50 deletions(-) diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 81d794413..9aac19e08 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,5 +1,5 @@ use crate::common::require_one_child; -use crate::coordinator::dynamic_filters::isolate_distributed_leaf_display_variants; +use crate::coordinator::dynamic_filters::isolate_distributed_leaf_variants_for_display; use crate::coordinator::metrics_store::MetricsStore; use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan; use crate::coordinator::prepare_static_plan::prepare_static_plan; @@ -233,7 +233,7 @@ impl ExecutionPlan for DistributedExec { }; let display_plan = - isolate_distributed_leaf_display_variants(result.plan_for_viz, &context)?; + isolate_distributed_leaf_variants_for_display(result.plan_for_viz, &context)?; plan_for_viz .lock() .expect("poisoned lock") diff --git a/src/coordinator/dynamic_filters.rs b/src/coordinator/dynamic_filters.rs index 0aadea77f..4b6c381b2 100644 --- a/src/coordinator/dynamic_filters.rs +++ b/src/coordinator/dynamic_filters.rs @@ -56,8 +56,8 @@ impl DynamicFilterReports { } } -/// Gives every distributed-leaf task an independent plan used only for rendering. -pub(super) fn isolate_distributed_leaf_display_variants( +/// 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> { @@ -68,7 +68,7 @@ pub(super) fn isolate_distributed_leaf_display_variants( return Ok(Transformed::no(node)); }; - let display_variants = leaf + let variants = leaf .variants() .iter() .map(|variant| { @@ -81,14 +81,15 @@ pub(super) fn isolate_distributed_leaf_display_variants( }) .collect::>>()?; - Ok(Transformed::yes( - Arc::new(leaf.with_display_variants(display_variants)?) as Arc, - )) + 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 display variants. +/// Applies successful worker reports only to the matching task-local visualization variants. pub(super) fn apply_reports_to_distributed_leaves( plan: &Arc, reports: &HashMap, @@ -101,7 +102,7 @@ pub(super) fn apply_reports_to_distributed_leaves( }; for (task_key, report) in reports { - let Some(variant) = leaf.display_variants.get(task_key.task_number) else { + let Some(variant) = leaf.variants().get(task_key.task_number) else { continue; }; let updates: HashMap<_, _> = report @@ -165,7 +166,7 @@ mod tests { use uuid::Uuid; #[test] - fn task_display_variants_do_not_share_dynamic_filter_state() -> Result<()> { + 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( @@ -183,7 +184,7 @@ mod tests { )?) as Arc; let task_ctx = SessionContext::new().task_ctx(); - let isolated = isolate_distributed_leaf_display_variants(leaf, &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 @@ -212,10 +213,10 @@ mod tests { apply_reports_to_distributed_leaves(&isolated, &reports, &task_ctx); let leaf = isolated.downcast_ref::().unwrap(); - let task_0 = displayable(leaf.display_variants[0].as_ref()) + let task_0 = displayable(leaf.variants()[0].as_ref()) .one_line() .to_string(); - let task_1 = displayable(leaf.display_variants[1].as_ref()) + let task_1 = displayable(leaf.variants()[1].as_ref()) .one_line() .to_string(); assert!(task_0.contains("DynamicFilter [ a@0 > 10 ]")); diff --git a/src/execution_plans/distributed_leaf.rs b/src/execution_plans/distributed_leaf.rs index b861d43bf..9f3d9884f 100644 --- a/src/execution_plans/distributed_leaf.rs +++ b/src/execution_plans/distributed_leaf.rs @@ -74,9 +74,6 @@ pub struct DistributedLeafExec { pub(crate) original: Arc, pub(crate) properties: Arc, pub(crate) variants: Vec>, - /// Per-task variants used only for rendering. These are isolated from one another before a - /// query runs so that applying task-local runtime state cannot leak into another task's line. - pub(crate) display_variants: Vec>, } impl DistributedLeafExec { @@ -118,7 +115,6 @@ impl DistributedLeafExec { Ok(Self { original, properties, - display_variants: variants.clone(), variants, }) } @@ -142,36 +138,6 @@ impl DistributedLeafExec { pub(crate) fn to_task_specialized(&self, task_i: usize) -> Arc { Arc::clone(&self.variants[task_i]) } - - pub(crate) fn with_display_variants( - &self, - display_variants: Vec>, - ) -> Result { - if display_variants.len() != self.variants.len() { - return plan_err!( - "DistributedLeafExec received {} display variants for {} execution variants", - display_variants.len(), - self.variants.len() - ); - } - for (display, execution) in display_variants.iter().zip(&self.variants) { - if display.schema() != execution.schema() - || display.properties().partitioning.partition_count() - != execution.properties().partitioning.partition_count() - { - return plan_err!( - "DistributedLeafExec display variant properties differ from its execution variant" - ); - } - } - - Ok(Self { - original: Arc::clone(&self.original), - properties: Arc::clone(&self.properties), - variants: self.variants.clone(), - display_variants, - }) - } } impl DisplayAs for DistributedLeafExec { diff --git a/src/stage.rs b/src/stage.rs index 8b20d0dab..967135309 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -419,7 +419,7 @@ fn display_inner_distributed_leaf( && !by_task.is_empty() { writeln!(f, "{indent} DistributedLeafExec:")?; - for (task_i, variant) in leaf.display_variants.iter().enumerate() { + for (task_i, variant) in leaf.variants.iter().enumerate() { let variant = displayable(variant.as_ref()).one_line().to_string(); let metrics = match by_task.is_empty() { true => String::new(), @@ -433,7 +433,7 @@ fn display_inner_distributed_leaf( false => String::new(), }; writeln!(f, "{indent} DistributedLeafExec:{header}")?; - for (task_i, variant) in leaf.display_variants.iter().enumerate() { + for (task_i, variant) in leaf.variants.iter().enumerate() { let variant = displayable(variant.as_ref()).one_line().to_string(); writeln!(f, "{indent} t{task_i}: {}", variant.trim_end())?; } From 64f9f5d5ad8773e570f597ae5ee9b6065a0a9a2d Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 11 Aug 2026 17:27:22 +0000 Subject: [PATCH 05/17] store refactor --- docs/upgrade/3.0.0.md | 2 +- src/coordinator/distributed.rs | 29 ++-------- src/coordinator/dynamic_filters.rs | 54 ++---------------- src/coordinator/metrics_store.rs | 36 ------------ src/coordinator/mod.rs | 4 +- src/coordinator/query_coordinator.rs | 37 +++++++------ src/coordinator/store.rs | 77 ++++++++++++++++++++++++++ src/lib.rs | 6 +- src/protocol/grpc/generated/worker.rs | 10 ++-- src/protocol/grpc/worker.proto | 4 +- src/protocol/grpc/worker_client.rs | 14 +++-- src/protocol/grpc/worker_service.rs | 16 +++--- src/protocol/mod.rs | 2 +- src/protocol/worker_channel.rs | 4 +- src/worker/impl_coordinator_channel.rs | 25 +++++---- 15 files changed, 155 insertions(+), 165 deletions(-) delete mode 100644 src/coordinator/metrics_store.rs create mode 100644 src/coordinator/store.rs diff --git a/docs/upgrade/3.0.0.md b/docs/upgrade/3.0.0.md index 119913f98..c6a79e199 100644 --- a/docs/upgrade/3.0.0.md +++ b/docs/upgrade/3.0.0.md @@ -58,7 +58,7 @@ 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::TaskDynamicFilters` variant. The built-in gRPC +`WorkerToCoordinatorMsg::TaskCompletedDynamicFilters` variant. The built-in gRPC transport handles both automatically. ## 2. Move `TaskEstimator` methods to event handlers diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 9aac19e08..2b031cad3 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,13 +1,12 @@ +use crate::DistributedConfig; use crate::common::require_one_child; use crate::coordinator::dynamic_filters::isolate_distributed_leaf_variants_for_display; -use crate::coordinator::metrics_store::MetricsStore; 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; @@ -83,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. diff --git a/src/coordinator/dynamic_filters.rs b/src/coordinator/dynamic_filters.rs index 4b6c381b2..84c447a35 100644 --- a/src/coordinator/dynamic_filters.rs +++ b/src/coordinator/dynamic_filters.rs @@ -1,8 +1,8 @@ use crate::common::discover_dynamic_filter_consumers; use crate::execution_plans::DistributedLeafExec; -use crate::{DistributedCodec, TaskDynamicFilters, TaskKey}; +use crate::{DistributedCodec, TaskCompletedDynamicFilters, TaskKey}; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; -use datafusion::common::{HashMap, HashSet, Result}; +use datafusion::common::{HashMap, Result}; use datafusion::execution::TaskContext; use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_plan::ExecutionPlan; @@ -10,51 +10,7 @@ use datafusion_proto::physical_plan::from_proto::parse_physical_expr; use datafusion_proto::physical_plan::{DeduplicatingProtoConverter, PhysicalPlanNodeExt}; use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode}; use prost::Message; -use std::sync::{Arc, Mutex}; -use tokio::sync::Notify; - -#[derive(Default)] -struct ReportState { - expected: HashSet, - finished: HashSet, - reports: HashMap, -} - -/// Query-scoped storage for final dynamic-filter reports. -#[derive(Default)] -pub(super) struct DynamicFilterReports { - state: Mutex, - changed: Notify, -} - -impl DynamicFilterReports { - pub(super) fn expect(&self, key: TaskKey) { - self.state.lock().unwrap().expected.insert(key); - } - - pub(super) fn insert(&self, key: TaskKey, report: TaskDynamicFilters) { - self.state.lock().unwrap().reports.insert(key, report); - self.changed.notify_one(); - } - - pub(super) fn finish(&self, key: TaskKey) { - self.state.lock().unwrap().finished.insert(key); - self.changed.notify_one(); - } - - pub(super) async fn wait_for_finished(&self) -> HashMap { - loop { - let notified = self.changed.notified(); - { - let state = self.state.lock().unwrap(); - if state.expected.is_subset(&state.finished) { - return state.reports.clone(); - } - } - notified.await; - } - } -} +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( @@ -92,7 +48,7 @@ pub(super) fn isolate_distributed_leaf_variants_for_display( /// Applies successful worker reports only to the matching task-local visualization variants. pub(super) fn apply_reports_to_distributed_leaves( plan: &Arc, - reports: &HashMap, + reports: &HashMap, task_ctx: &Arc, ) { let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config()); @@ -196,7 +152,7 @@ mod tests { .unwrap() .mark_complete(); let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config()); - let report = TaskDynamicFilters { + let report = TaskCompletedDynamicFilters { filters: vec![crate::TaskDynamicFilter { expression_id: dynamic_filter.expression_id().unwrap(), expression: serialize_physical_expr(&dynamic_filter, &codec)?.encode_to_vec(), 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 6d0da6552..c08fbbd46 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,10 +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 11519c411..c91266ae7 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,11 +1,10 @@ use crate::codec::{decode_execution_plan, encode_execution_plan}; 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::{ - DynamicFilterReports, apply_reports_to_distributed_leaves, -}; +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; @@ -15,7 +14,8 @@ use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time}; use crate::{ CoordinatorToWorkerMsg, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedTaskContext, DistributedWorkUnitFeedContext, LoadInfo, LocalWorkerContext, MaybeEncoded, SetPlanRequest, - TaskKey, WorkUnitFeedDeclaration, WorkerToCoordinatorMsg, get_distributed_channel_resolver, + TaskCompletedDynamicFilters, TaskKey, WorkUnitFeedDeclaration, WorkerToCoordinatorMsg, + get_distributed_channel_resolver, }; use datafusion::common::DataFusionError; use datafusion::common::instant::Instant; @@ -50,7 +50,7 @@ pub(super) struct QueryCoordinator { metrics: ExecutionPlanMetricsSet, coordinator_to_worker_metrics: CoordinatorToWorkerMetrics, metrics_store: Option>, - dynamic_filter_reports: Arc, + completed_dynamic_filter_store: Arc, end_stream_notifier: Arc, join_set: Mutex>>, } @@ -66,7 +66,7 @@ impl QueryCoordinator { task_ctx, metrics: metrics_set.clone(), metrics_store, - dynamic_filter_reports: Arc::new(DynamicFilterReports::default()), + 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()), @@ -85,7 +85,7 @@ impl QueryCoordinator { metrics_set: &self.metrics, metrics: &self.coordinator_to_worker_metrics, metrics_store: &self.metrics_store, - dynamic_filter_reports: &self.dynamic_filter_reports, + completed_dynamic_filter_store: &self.completed_dynamic_filter_store, end_stream_notifier: &self.end_stream_notifier, join_set: &self.join_set, } @@ -107,7 +107,10 @@ impl QueryCoordinator { &self, plan_for_viz: &Arc, ) { - let reports = self.dynamic_filter_reports.wait_for_finished().await; + 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); } @@ -145,7 +148,7 @@ pub(super) struct StageCoordinator<'a> { metrics_set: &'a ExecutionPlanMetricsSet, metrics: &'a CoordinatorToWorkerMetrics, metrics_store: &'a Option>, - dynamic_filter_reports: &'a Arc, + completed_dynamic_filter_store: &'a Arc, end_stream_notifier: &'a Arc, join_set: &'a Mutex>>, } @@ -172,8 +175,6 @@ impl<'a> StageCoordinator<'a> { stage_id: self.stage_id, task_number: task_i, }; - self.dynamic_filter_reports.expect(task_key); - let set_plan_request = SetPlanRequest { task_key, task_count: self.task_count, @@ -260,7 +261,7 @@ impl<'a> StageCoordinator<'a> { task_number: task_i, }; let task_metrics = self.metrics_store.clone(); - let dynamic_filter_reports = Arc::clone(self.dynamic_filter_reports); + 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); @@ -283,13 +284,15 @@ impl<'a> StageCoordinator<'a> { WorkerToCoordinatorMsg::LoadInfoEos => { let _ = load_info_tx_opt.take(); } - WorkerToCoordinatorMsg::TaskDynamicFilters(filters) => { - dynamic_filter_reports.insert(task_key, filters); - dynamic_filter_reports.finish(task_key); + WorkerToCoordinatorMsg::TaskCompletedDynamicFilters(filters) => { + completed_dynamic_filter_store.insert(task_key, filters); } } } - dynamic_filter_reports.finish(task_key); + if completed_dynamic_filter_store.get(&task_key).is_none() { + completed_dynamic_filter_store + .insert(task_key, TaskCompletedDynamicFilters::default()); + } }); load_info_rx } diff --git a/src/coordinator/store.rs b/src/coordinator/store.rs new file mode 100644 index 000000000..d6553e0dd --- /dev/null +++ b/src/coordinator/store.rs @@ -0,0 +1,77 @@ +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; + } + rx.borrow().clone() + } + + #[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 fb2dc7124..ae881f489 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -54,9 +54,9 @@ pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; pub use protocol::{ ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, - GetWorkerInfoResponse, LoadInfo, SetPlanRequest, TaskDynamicFilter, TaskDynamicFilters, - TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, - WorkerToCoordinatorMsg, get_distributed_channel_resolver, + GetWorkerInfoResponse, LoadInfo, 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 c819e7c8b..95806d859 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/grpc/generated/worker.rs @@ -45,16 +45,16 @@ pub mod worker_to_coordinator_msg { LoadInfoEos(bool), /// Final task-local dynamic filters used by distributed leaf variants. #[prost(message, tag = "4")] - TaskDynamicFilters(super::TaskDynamicFilters), + TaskCompletedDynamicFilters(super::TaskCompletedDynamicFilters), } } #[derive(Clone, PartialEq, ::prost::Message)] -pub struct TaskDynamicFilters { +pub struct TaskCompletedDynamicFilters { #[prost(message, repeated, tag = "1")] - pub filters: ::prost::alloc::vec::Vec, + pub filters: ::prost::alloc::vec::Vec, } -/// Nested message and enum types in `TaskDynamicFilters`. -pub mod task_dynamic_filters { +/// 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")] diff --git a/src/protocol/grpc/worker.proto b/src/protocol/grpc/worker.proto index 82f842dcf..c4da0e786 100644 --- a/src/protocol/grpc/worker.proto +++ b/src/protocol/grpc/worker.proto @@ -41,11 +41,11 @@ message WorkerToCoordinatorMsg { bool load_info_eos = 3; // Final task-local dynamic filters used by distributed leaf variants. - TaskDynamicFilters task_dynamic_filters = 4; + TaskCompletedDynamicFilters task_completed_dynamic_filters = 4; } } -message TaskDynamicFilters { +message TaskCompletedDynamicFilters { message DynamicFilter { uint64 expression_id = 1; // Serialized datafusion.proto.PhysicalExprNode. diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index 043c220ea..a2eb77aa6 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -10,7 +10,7 @@ use crate::{ DistributedConfig, ExecuteTaskRequest, FirstLatencyMetric, GetWorkerInfoRequest, GetWorkerInfoResponse, LatencyMetricExt, LoadInfo, MaxLatencyMetric, MaybeEncoded, MinLatencyMetric, P50LatencyMetric, P95LatencyMetric, ProducerHead, SetPlanRequest, - TaskDynamicFilter, TaskDynamicFilters, TaskKey, TaskMetrics, WorkUnitBatch, + TaskCompletedDynamicFilters, TaskDynamicFilter, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; @@ -509,15 +509,19 @@ fn decode_worker_to_coordinator_msg( pb::worker_to_coordinator_msg::Inner::LoadInfoEos(_) => { WorkerToCoordinatorMsg::LoadInfoEos } - pb::worker_to_coordinator_msg::Inner::TaskDynamicFilters(filters) => { - WorkerToCoordinatorMsg::TaskDynamicFilters(decode_task_dynamic_filters(filters)) + pb::worker_to_coordinator_msg::Inner::TaskCompletedDynamicFilters(filters) => { + WorkerToCoordinatorMsg::TaskCompletedDynamicFilters( + decode_task_completed_dynamic_filters(filters), + ) } }, ) } -fn decode_task_dynamic_filters(filters: pb::TaskDynamicFilters) -> TaskDynamicFilters { - TaskDynamicFilters { +fn decode_task_completed_dynamic_filters( + filters: pb::TaskCompletedDynamicFilters, +) -> TaskCompletedDynamicFilters { + TaskCompletedDynamicFilters { filters: filters .filters .into_iter() diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index 4419db210..20c097ac4 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -7,7 +7,7 @@ use crate::common::{deserialize_uuid, now_ns}; use crate::protocol::grpc::{ObservabilityServiceImpl, ObservabilityServiceServer}; use crate::{ CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, LoadInfo, MaybeEncoded, - ProducerHead, SetPlanRequest, TaskDynamicFilters, TaskKey, TaskMetrics, WorkUnitBatch, + ProducerHead, SetPlanRequest, TaskCompletedDynamicFilters, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, Worker, WorkerResolver, WorkerToCoordinatorMsg, }; @@ -273,21 +273,23 @@ fn encode_worker_to_coordinator_msg( WorkerToCoordinatorMsg::LoadInfoEos => { pb::worker_to_coordinator_msg::Inner::LoadInfoEos(true) } - WorkerToCoordinatorMsg::TaskDynamicFilters(filters) => { - pb::worker_to_coordinator_msg::Inner::TaskDynamicFilters( - encode_task_dynamic_filters(filters), + WorkerToCoordinatorMsg::TaskCompletedDynamicFilters(filters) => { + pb::worker_to_coordinator_msg::Inner::TaskCompletedDynamicFilters( + encode_task_completed_dynamic_filters(filters), ) } }), }) } -fn encode_task_dynamic_filters(filters: TaskDynamicFilters) -> pb::TaskDynamicFilters { - pb::TaskDynamicFilters { +fn encode_task_completed_dynamic_filters( + filters: TaskCompletedDynamicFilters, +) -> pb::TaskCompletedDynamicFilters { + pb::TaskCompletedDynamicFilters { filters: filters .filters .into_iter() - .map(|filter| pb::task_dynamic_filters::DynamicFilter { + .map(|filter| pb::task_completed_dynamic_filters::DynamicFilter { expression_id: filter.expression_id, expression: filter.expression, }) diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index b6e6b3af8..1ee655e3a 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -11,6 +11,6 @@ pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; pub use in_process::LocalWorkerContext; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, - LoadInfo, SetPlanRequest, TaskDynamicFilter, TaskDynamicFilters, TaskKey, TaskMetrics, + LoadInfo, 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 85058de5f..f3ee7ef58 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -128,11 +128,11 @@ pub enum WorkerToCoordinatorMsg { LoadInfo(LoadInfo), LoadInfoEos, /// Final task-local dynamic filters used by distributed leaf variants. - TaskDynamicFilters(TaskDynamicFilters), + TaskCompletedDynamicFilters(TaskCompletedDynamicFilters), } #[derive(Clone, Debug, Default)] -pub struct TaskDynamicFilters { +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, diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index b75d3a5d0..b1051d278 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -6,8 +6,8 @@ use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_t use crate::worker::task_data::TaskDataMetrics; use crate::{ CoordinatorToWorkerMsg, DistributedCodec, DistributedConfig, DistributedExt, - DistributedTaskContext, SetPlanRequest, TaskData, TaskDynamicFilter, TaskDynamicFilters, - TaskMetrics, Worker, WorkerQueryContext, WorkerToCoordinatorMsg, + DistributedTaskContext, SetPlanRequest, TaskCompletedDynamicFilters, TaskData, + TaskDynamicFilter, TaskMetrics, Worker, WorkerQueryContext, WorkerToCoordinatorMsg, }; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{DataFusionError, HashSet, Result, exec_datafusion_err, internal_err}; @@ -164,7 +164,7 @@ impl Worker { } let metrics_tx = task_data.metrics_tx.lock().unwrap().take(); - let mut dynamic_filters = TaskDynamicFilters::default(); + 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, @@ -175,9 +175,12 @@ 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_dynamic_filters(plan, &dynamic_filter_ids, &task_data.task_ctx) - .unwrap_or_default(); + 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 @@ -205,7 +208,9 @@ impl Worker { 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::TaskDynamicFilters(dynamic_filters)) + Some(WorkerToCoordinatorMsg::TaskCompletedDynamicFilters( + dynamic_filters, + )) }, ); @@ -218,11 +223,11 @@ impl Worker { } } -fn build_task_dynamic_filters( +fn build_task_completed_dynamic_filters( plan: &Arc, allowed_ids: &HashSet, task_ctx: &Arc, -) -> Result { +) -> 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))? { @@ -241,7 +246,7 @@ fn build_task_dynamic_filters( }); } } - Ok(TaskDynamicFilters { filters }) + Ok(TaskCompletedDynamicFilters { filters }) } /// Collects metrics from the plan in pre-order traversal order and sends them via the From 1115fb4482b7b08c9d433475819185e72eb1ba5a Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 11 Aug 2026 18:01:04 +0000 Subject: [PATCH 06/17] dont use vec --- docs/upgrade/3.0.0.md | 4 +++- src/coordinator/dynamic_filters.rs | 14 +++++--------- src/protocol/grpc/worker_client.rs | 20 ++++++++++++-------- src/protocol/grpc/worker_service.rs | 2 +- src/protocol/worker_channel.rs | 11 ++++++----- src/worker/impl_coordinator_channel.rs | 3 +-- 6 files changed, 28 insertions(+), 26 deletions(-) diff --git a/docs/upgrade/3.0.0.md b/docs/upgrade/3.0.0.md index c6a79e199..9748e61b4 100644 --- a/docs/upgrade/3.0.0.md +++ b/docs/upgrade/3.0.0.md @@ -59,7 +59,9 @@ 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. +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 diff --git a/src/coordinator/dynamic_filters.rs b/src/coordinator/dynamic_filters.rs index 84c447a35..57bf509eb 100644 --- a/src/coordinator/dynamic_filters.rs +++ b/src/coordinator/dynamic_filters.rs @@ -8,8 +8,7 @@ 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::{PhysicalExprNode, PhysicalPlanNode}; -use prost::Message; +use datafusion_proto::protobuf::PhysicalPlanNode; use std::sync::Arc; /// Replaces the variants in the visualization plan with independent per-task copies. @@ -64,20 +63,17 @@ pub(super) fn apply_reports_to_distributed_leaves( let updates: HashMap<_, _> = report .filters .iter() - .map(|filter| (filter.expression_id, filter.expression.as_slice())) + .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(encoded) = updates.get(&consumer.id) else { - continue; - }; - let Ok(proto) = PhysicalExprNode::decode(*encoded) else { + 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) + parse_physical_expr(proto, task_ctx, consumer.input_schema.as_ref(), &codec) else { continue; }; @@ -155,7 +151,7 @@ mod tests { let report = TaskCompletedDynamicFilters { filters: vec![crate::TaskDynamicFilter { expression_id: dynamic_filter.expression_id().unwrap(), - expression: serialize_physical_expr(&dynamic_filter, &codec)?.encode_to_vec(), + expression: serialize_physical_expr(&dynamic_filter, &codec)?, }], }; let reports = HashMap::from_iter([( diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index a2eb77aa6..aaec300ae 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -25,6 +25,7 @@ use datafusion::execution::TaskContext; use datafusion::execution::memory_pool::MemoryConsumer; use datafusion::physical_expr_common::metrics::{Count, Label, 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}; @@ -511,7 +512,7 @@ fn decode_worker_to_coordinator_msg( } pb::worker_to_coordinator_msg::Inner::TaskCompletedDynamicFilters(filters) => { WorkerToCoordinatorMsg::TaskCompletedDynamicFilters( - decode_task_completed_dynamic_filters(filters), + decode_task_completed_dynamic_filters(filters)?, ) } }, @@ -520,17 +521,20 @@ fn decode_worker_to_coordinator_msg( fn decode_task_completed_dynamic_filters( filters: pb::TaskCompletedDynamicFilters, -) -> TaskCompletedDynamicFilters { - TaskCompletedDynamicFilters { +) -> Result { + Ok(TaskCompletedDynamicFilters { filters: filters .filters .into_iter() - .map(|filter| TaskDynamicFilter { - expression_id: filter.expression_id, - expression: filter.expression, + .map(|filter| { + Ok(TaskDynamicFilter { + expression_id: filter.expression_id, + expression: PhysicalExprNode::decode(filter.expression.as_slice()) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + }) }) - .collect(), - } + .collect::>()?, + }) } fn decode_task_metrics(task_metrics: pb::TaskMetrics) -> Result { diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index 20c097ac4..ed415f364 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -291,7 +291,7 @@ fn encode_task_completed_dynamic_filters( .into_iter() .map(|filter| pb::task_completed_dynamic_filters::DynamicFilter { expression_id: filter.expression_id, - expression: filter.expression, + expression: filter.expression.encode_to_vec(), }) .collect(), } diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index f3ee7ef58..87952ce4e 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -5,6 +5,7 @@ use datafusion::common::Result; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion_proto::protobuf::PhysicalExprNode; use futures::stream::BoxStream; use http::HeaderMap; use std::sync::Arc; @@ -123,12 +124,13 @@ 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, - /// Final task-local dynamic filters used by distributed leaf variants. - TaskCompletedDynamicFilters(TaskCompletedDynamicFilters), } #[derive(Clone, Debug, Default)] @@ -141,9 +143,8 @@ pub struct TaskCompletedDynamicFilters { #[derive(Clone, Debug)] pub struct TaskDynamicFilter { pub expression_id: u64, - /// A serialized `DynamicFilterPhysicalExpr` `PhysicalExprNode`, including its final predicate - /// and completion state. - pub expression: Vec, + /// A `DynamicFilterPhysicalExpr` proto containing its final predicate and completion state. + pub expression: PhysicalExprNode, } #[derive(Clone, Debug)] diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index b1051d278..ea36ea99a 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -19,7 +19,6 @@ use datafusion_proto::protobuf::physical_expr_node::ExprType; use futures::stream::{BoxStream, FuturesUnordered}; use futures::{FutureExt, StreamExt, TryStreamExt}; use http::HeaderMap; -use prost::Message; use std::sync::{Arc, OnceLock}; use tokio::sync::oneshot; use tokio::sync::oneshot::Sender; @@ -242,7 +241,7 @@ fn build_task_completed_dynamic_filters( if dynamic_filter.is_complete { filters.push(TaskDynamicFilter { expression_id: consumer.id, - expression: expression.encode_to_vec(), + expression, }); } } From 00e3f5b04776cc7ac75fb97148aeb8da0095ffb4 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 11 Aug 2026 18:14:36 +0000 Subject: [PATCH 07/17] rename --- src/protocol/grpc/generated/worker.rs | 2 +- src/protocol/grpc/worker.proto | 2 +- src/protocol/grpc/worker_client.rs | 2 +- src/protocol/grpc/worker_service.rs | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/protocol/grpc/generated/worker.rs b/src/protocol/grpc/generated/worker.rs index 95806d859..0793b2f7b 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/grpc/generated/worker.rs @@ -61,7 +61,7 @@ pub mod task_completed_dynamic_filters { pub expression_id: u64, /// Serialized datafusion.proto.PhysicalExprNode. #[prost(bytes = "vec", tag = "2")] - pub expression: ::prost::alloc::vec::Vec, + pub expression_proto: ::prost::alloc::vec::Vec, } } #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/src/protocol/grpc/worker.proto b/src/protocol/grpc/worker.proto index c4da0e786..b4f4331d9 100644 --- a/src/protocol/grpc/worker.proto +++ b/src/protocol/grpc/worker.proto @@ -49,7 +49,7 @@ message TaskCompletedDynamicFilters { message DynamicFilter { uint64 expression_id = 1; // Serialized datafusion.proto.PhysicalExprNode. - bytes expression = 2; + bytes expression_proto = 2; } repeated DynamicFilter filters = 1; diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index aaec300ae..69ebbe703 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -529,7 +529,7 @@ fn decode_task_completed_dynamic_filters( .map(|filter| { Ok(TaskDynamicFilter { expression_id: filter.expression_id, - expression: PhysicalExprNode::decode(filter.expression.as_slice()) + expression: PhysicalExprNode::decode(filter.expression_proto.as_slice()) .map_err(|error| DataFusionError::External(Box::new(error)))?, }) }) diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index ed415f364..1b210288d 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -291,7 +291,7 @@ fn encode_task_completed_dynamic_filters( .into_iter() .map(|filter| pb::task_completed_dynamic_filters::DynamicFilter { expression_id: filter.expression_id, - expression: filter.expression.encode_to_vec(), + expression_proto: filter.expression.encode_to_vec(), }) .collect(), } From 72a8e6c62804a3a89d94ae5910e6417a721e3e4e Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 11 Aug 2026 19:27:14 +0000 Subject: [PATCH 08/17] refactor tests --- tests/dynamic_filtering.rs | 102 ++++++++++++++++++++++++-------- tests/tpcds_correctness_test.rs | 41 +------------ tests/tpch_plans_test.rs | 48 +-------------- 3 files changed, 80 insertions(+), 111 deletions(-) diff --git a/tests/dynamic_filtering.rs b/tests/dynamic_filtering.rs index 2302974e1..baba4fe89 100644 --- a/tests/dynamic_filtering.rs +++ b/tests/dynamic_filtering.rs @@ -4,20 +4,90 @@ mod tests { 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, display_plan_ascii}; + use datafusion_distributed::{ + DefaultSessionBuilder, DistributedExt, assert_snapshot, display_plan_ascii, + }; use std::sync::Arc; #[tokio::test] - async fn completed_leaf_dynamic_filters_are_displayed_per_task() -> Result<()> { - for dynamic_task_count in [false, true] { - assert_completed_leaf_filters(dynamic_task_count).await?; - } + 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 assert_completed_leaf_filters(dynamic_task_count: bool) -> Result<()> { + async fn execute_local_hash_join(broadcast_joins: bool) -> Result { let (ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; - let ctx = ctx.with_distributed_dynamic_task_count(dynamic_task_count)?; + 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 @@ -41,22 +111,6 @@ mod tests { 1 ); - let display = display_plan_ascii(plan.as_ref(), false); - let task_filters: Vec<_> = display - .lines() - .filter(|line| line.contains(": DataSourceExec:")) - .filter_map(|line| line.split_once("DynamicFilter [ ")) - .map(|(_, filter)| filter.split_once(" ]").unwrap().0) - .collect(); - - assert!( - task_filters.len() >= 2, - "expected task-local leaf filters in plan:\n{display}" - ); - assert!( - task_filters.iter().all(|filter| *filter != "empty"), - "expected completed filters in plan:\n{display}" - ); - Ok(()) + Ok(display_plan_ascii(plan.as_ref(), false)) } } diff --git a/tests/tpcds_correctness_test.rs b/tests/tpcds_correctness_test.rs index fe1f4c4eb..e4baeba9d 100644 --- a/tests/tpcds_correctness_test.rs +++ b/tests/tpcds_correctness_test.rs @@ -10,7 +10,7 @@ mod tests { compare_ordering, compare_result_set, }; use datafusion_distributed::{ - DefaultSessionBuilder, DistributedExec, DistributedExt, assert_snapshot, display_plan_ascii, + DefaultSessionBuilder, DistributedExec, DistributedExt, display_plan_ascii, }; use datafusion_distributed_benchmarks::datasets::{register_tables, tpcds}; use std::fs; @@ -597,12 +597,6 @@ mod tests { } let display = display_plan_ascii(d_plan.as_ref(), false); println!("Query {query_id}:\n{display}"); - if query_id == "q1" { - assert_snapshot!(displayed_dynamic_filter_predicates(&display), @r" - DynamicFilter [ c_customer_sk@0 >= 5 AND c_customer_sk@0 <= 100000 AND true ] AND DynamicFilter [ empty ] - DynamicFilter [ sr_returned_date_sk@0 >= 2451545 AND sr_returned_date_sk@0 <= 2451910 AND true ] AND DynamicFilter [ empty ] - "); - } // The comparison functions can be computationally expensive, so we spawn them in tokio // blocking tasks so that they do not block the tokio runtime. @@ -624,37 +618,4 @@ mod tests { Ok(()) } - - fn displayed_dynamic_filter_predicates(display: &str) -> String { - let mut predicates: Vec<_> = display - .lines() - .filter(|line| line.contains(": DataSourceExec:")) - .filter_map(|line| { - line.split_once("predicate=") - .map(|(_, predicate)| predicate) - }) - .map(|predicate| { - predicate - .split(", pruning_predicate=") - .next() - .unwrap_or(predicate) - }) - .map(|predicate| { - predicate - .split(", dynamic_rg_pruning=") - .next() - .unwrap_or(predicate) - }) - .map(|predicate| { - predicate - .split(", required_guarantees=") - .next() - .unwrap_or(predicate) - }) - .filter(|predicate| predicate.contains("DynamicFilter")) - .collect(); - predicates.sort_unstable(); - predicates.dedup(); - predicates.join("\n") - } } diff --git a/tests/tpch_plans_test.rs b/tests/tpch_plans_test.rs index fc6909a38..a815c6255 100644 --- a/tests/tpch_plans_test.rs +++ b/tests/tpch_plans_test.rs @@ -1,7 +1,5 @@ #[cfg(all(feature = "integration", feature = "tpch", test))] mod tests { - use datafusion::physical_plan::{ExecutionPlan, collect}; - use datafusion::prelude::SessionContext; use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context; use datafusion_distributed::{ DefaultSessionBuilder, DistributedExt, assert_snapshot, display_plan_ascii, @@ -10,7 +8,6 @@ mod tests { use std::error::Error; use std::fs; use std::path::Path; - use std::sync::Arc; use tokio::sync::OnceCell; const NUM_WORKERS: usize = 4; @@ -222,16 +219,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn test_tpch_3_dynamic_filters_after_execution() -> Result<(), Box> { - let predicates = executed_dynamic_filter_predicates("q3").await?; - assert_snapshot!(predicates, @r" - DynamicFilter [ l_orderkey@0 >= 5 AND l_orderkey@0 <= 119972 AND true ] - DynamicFilter [ o_custkey@1 >= 1 AND o_custkey@1 <= 2993 AND true ] - "); - Ok(()) - } - #[tokio::test] async fn test_tpch_4() -> Result<(), Box> { let plan = test_tpch_query("q4").await?; @@ -1447,39 +1434,6 @@ mod tests { // test_tpch_query generates and displays a distributed plan for each TPC-H query. async fn test_tpch_query(query_id: &str) -> Result> { - let (_, plan) = prepare_tpch_query(query_id).await?; - Ok(display_plan_ascii(plan.as_ref(), false)) - } - - async fn executed_dynamic_filter_predicates(query_id: &str) -> Result> { - let (ctx, plan) = prepare_tpch_query(query_id).await?; - collect(Arc::clone(&plan), ctx.task_ctx()).await?; - let display = display_plan_ascii(plan.as_ref(), false); - let mut predicates: Vec<_> = display - .lines() - .filter(|line| line.contains(": DataSourceExec:")) - .filter_map(|line| line.find("DynamicFilter [ ").map(|start| &line[start..])) - .map(|predicate| { - predicate - .split(", dynamic_rg_pruning=") - .next() - .unwrap_or(predicate) - }) - .map(|predicate| { - predicate - .split(", pruning_predicate=") - .next() - .unwrap_or(predicate) - }) - .collect(); - predicates.sort_unstable(); - predicates.dedup(); - Ok(predicates.join("\n")) - } - - async fn prepare_tpch_query( - query_id: &str, - ) -> Result<(SessionContext, Arc), Box> { let d_ctx = start_in_memory_context(NUM_WORKERS, DefaultSessionBuilder).await; let data_dir = ensure_tpch_data(TPCH_SCALE_FACTOR, TPCH_DATA_PARTS).await; let sql = tpch::get_query(query_id)?; @@ -1519,7 +1473,7 @@ mod tests { df.create_physical_plan().await? }; - Ok((d_ctx, plan)) + Ok(display_plan_ascii(plan.as_ref(), false)) } // OnceCell to ensure TPCH tables are generated only once for tests From 730266f6a6ee7790d70f01fa4fa267b957e8c580 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Fri, 21 Aug 2026 18:39:03 +0000 Subject: [PATCH 09/17] remove allowlist from set plan request --- docs/upgrade/3.0.0.md | 8 -- src/codec/mod.rs | 3 +- src/codec/physical_plan.rs | 28 ++++++- src/common/dynamic_filtering.rs | 78 ++++++++++++------ src/common/mod.rs | 4 +- src/coordinator/dynamic_filters.rs | 11 +-- src/coordinator/query_coordinator.rs | 22 +----- src/protocol/grpc/generated/worker.rs | 16 +++- src/protocol/grpc/worker.proto | 15 +++- src/protocol/grpc/worker_client.rs | 1 - src/protocol/grpc/worker_service.rs | 1 - src/protocol/worker_channel.rs | 3 - src/worker/impl_coordinator_channel.rs | 25 +++--- tests/dynamic_filtering.rs | 105 +++++++++++++++++++++---- 14 files changed, 211 insertions(+), 109 deletions(-) diff --git a/docs/upgrade/3.0.0.md b/docs/upgrade/3.0.0.md index 9748e61b4..b690e1448 100644 --- a/docs/upgrade/3.0.0.md +++ b/docs/upgrade/3.0.0.md @@ -55,14 +55,6 @@ 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/codec/mod.rs b/src/codec/mod.rs index 020937ded..eb970d35f 100644 --- a/src/codec/mod.rs +++ b/src/codec/mod.rs @@ -4,7 +4,8 @@ mod user_codec; pub use distributed_codec::DistributedCodec; pub(crate) use physical_plan::{ - decode_execution_plan, decode_partitioning, encode_execution_plan, encode_partitioning, + decode_execution_plan, decode_partitioning, decode_physical_expr, encode_execution_plan, + encode_partitioning, encode_physical_expr, }; pub(crate) use user_codec::{ get_distributed_user_codecs, set_distributed_user_codec, set_distributed_user_codec_arc, diff --git a/src/codec/physical_plan.rs b/src/codec/physical_plan.rs index 6ff9836c6..43fd60ba0 100644 --- a/src/codec/physical_plan.rs +++ b/src/codec/physical_plan.rs @@ -1,15 +1,17 @@ use super::DistributedCodec; -use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::datatypes::{Schema, SchemaRef}; use datafusion::common::Result; use datafusion::execution::TaskContext; -use datafusion::physical_expr::Partitioning; +use datafusion::physical_expr::{Partitioning, PhysicalExpr}; use datafusion::physical_plan::ExecutionPlan; use datafusion_proto::bytes::{ physical_plan_from_bytes_with_proto_converter, physical_plan_to_bytes_with_proto_converter, }; use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning; use datafusion_proto::physical_plan::to_proto::serialize_partitioning; -use datafusion_proto::physical_plan::{DeduplicatingProtoConverter, PhysicalPlanDecodeContext}; +use datafusion_proto::physical_plan::{ + DeduplicatingProtoConverter, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, +}; use datafusion_proto::protobuf; use datafusion_proto::protobuf::proto_error; use prost::Message; @@ -42,6 +44,26 @@ pub(crate) fn decode_execution_plan( physical_plan_from_bytes_with_proto_converter(encoded, task_ctx, &codec, &converter) } +pub(crate) fn encode_physical_expr( + expression: &Arc, + task_ctx: &TaskContext, +) -> Result { + let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config()); + let converter = new_proto_converter(); + converter.physical_expr_to_proto(expression, &codec) +} + +pub(crate) fn decode_physical_expr( + proto: &protobuf::PhysicalExprNode, + input_schema: &Schema, + task_ctx: &TaskContext, +) -> Result> { + let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config()); + let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx, &codec); + let converter = new_proto_converter(); + converter.proto_to_physical_expr(proto, input_schema, &decode_ctx) +} + pub(crate) fn encode_partitioning( partitioning: &Partitioning, task_ctx: &TaskContext, diff --git a/src/common/dynamic_filtering.rs b/src/common/dynamic_filtering.rs index 96c381abf..7582fa252 100644 --- a/src/common/dynamic_filtering.rs +++ b/src/common/dynamic_filtering.rs @@ -1,33 +1,41 @@ use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; -use datafusion::common::{HashMap, HashSet, Result}; +use datafusion::common::{HashMap, HashSet, Result, internal_err}; 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. +/// A dynamic-filter consumer discovered in an execution plan along with the schema its evaluated +/// against. #[derive(Clone)] pub(crate) struct DiscoveredDynamicFilter { pub(crate) id: u64, pub(crate) expression: Arc, + // The schema the filter is evaluated against. 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. +/// Finds dynamic-filter consumers in `plan`, deduplicated by expression ID. 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 produced_ids: HashSet<_> = node + .dynamic_expressions_produced() + .into_iter() + .map(|produced| { + let Some(id) = produced.expression_id() else { + return internal_err!( + "{}::dynamic_expressions_produced returned an expression without an expression ID", + node.name() + ); + }; + Ok(id) + }) + .collect::>()?; let input_schema = node .children() .first() @@ -40,14 +48,13 @@ pub(crate) fn discover_dynamic_filter_consumers( 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 { + let Some(id) = expression.expression_id() else { + return internal_err!( + "DynamicFilterPhysicalExpr did not have an expression ID" + ); + }; + let is_producer_occurrence = produced_ids.contains(&id); + if !is_producer_occurrence { consumers .entry(id) .or_insert_with(|| DiscoveredDynamicFilter { @@ -68,13 +75,6 @@ pub(crate) fn discover_dynamic_filter_consumers( 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::*; @@ -84,6 +84,7 @@ mod tests { use datafusion::logical_expr::Operator; use datafusion::physical_expr::expressions::{BinaryExpr, Column, lit}; use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, PlanProperties, apply_expression_roots, }; @@ -112,7 +113,7 @@ mod tests { true, )) as Arc; - let discovered = discover_dynamic_filter_consumers(&plan, None)?; + let discovered = discover_dynamic_filter_consumers(&plan)?; assert_eq!(discovered.len(), 1); assert_eq!(discovered[0].id, dynamic_filter.expression_id().unwrap()); @@ -134,6 +135,31 @@ mod tests { Ok(()) } + #[test] + fn deduplicates_consumers_with_the_same_expression_id() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("a", 0))], + lit(true), + )) as Arc; + let consumers = (0..2) + .map(|_| { + Arc::new(ExpressionExec::new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + Arc::clone(&dynamic_filter), + false, + )) as Arc + }) + .collect(); + let plan = UnionExec::try_new(consumers)?; + + let discovered = discover_dynamic_filter_consumers(&plan)?; + + assert_eq!(discovered.len(), 1); + assert_eq!(discovered[0].id, dynamic_filter.expression_id().unwrap()); + Ok(()) + } + #[derive(Debug)] struct ExpressionExec { input: Arc, diff --git a/src/common/mod.rs b/src/common/mod.rs index 589c9b37a..97cca7394 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -9,9 +9,7 @@ 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 dynamic_filtering::discover_dynamic_filter_consumers; pub use maybe_encoded::MaybeEncoded; pub(crate) use once_lock::OnceLockResult; pub(crate) use recursion::TreeNodeExt; diff --git a/src/coordinator/dynamic_filters.rs b/src/coordinator/dynamic_filters.rs index 57bf509eb..0985e36d1 100644 --- a/src/coordinator/dynamic_filters.rs +++ b/src/coordinator/dynamic_filters.rs @@ -1,3 +1,4 @@ +use crate::codec::decode_physical_expr; use crate::common::discover_dynamic_filter_consumers; use crate::execution_plans::DistributedLeafExec; use crate::{DistributedCodec, TaskCompletedDynamicFilters, TaskKey}; @@ -6,7 +7,6 @@ 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; @@ -50,7 +50,6 @@ pub(super) fn apply_reports_to_distributed_leaves( 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); @@ -65,7 +64,7 @@ pub(super) fn apply_reports_to_distributed_leaves( .iter() .map(|filter| (filter.expression_id, &filter.expression)) .collect(); - let Ok(consumers) = discover_dynamic_filter_consumers(variant, None) else { + let Ok(consumers) = discover_dynamic_filter_consumers(variant) else { continue; }; for consumer in consumers { @@ -73,7 +72,7 @@ pub(super) fn apply_reports_to_distributed_leaves( continue; }; let Ok(reported_expression) = - parse_physical_expr(proto, task_ctx, consumer.input_schema.as_ref(), &codec) + decode_physical_expr(proto, consumer.input_schema.as_ref(), task_ctx) else { continue; }; @@ -114,7 +113,6 @@ mod tests { 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] @@ -147,11 +145,10 @@ mod tests { .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)?, + expression: crate::codec::encode_physical_expr(&dynamic_filter, &task_ctx)?, }], }; let reports = HashMap::from_iter([( diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index c91266ae7..c635c2238 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,5 +1,5 @@ use crate::codec::{decode_execution_plan, encode_execution_plan}; -use crate::common::{TreeNodeExt, dynamic_filter_consumer_ids, now_ns, task_ctx_with_extension}; +use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; use crate::coordinator::dynamic_filters::apply_reports_to_distributed_leaves; use crate::coordinator::latency_metric::LatencyMetric; @@ -133,11 +133,7 @@ impl QueryCoordinator { /// - Building tasks that communicate a serialized plan to multiple workers for further execution. /// - Building tasks that stream partition feeds from local [WorkUnitFeedExec] nodes to their /// remote counterparts. -type SpecializedTaskPlan = ( - Arc, - Vec, - Vec, -); +type SpecializedTaskPlan = (Arc, Vec); pub(super) struct StageCoordinator<'a> { plan: &'a Arc, @@ -167,8 +163,7 @@ impl<'a> StageCoordinator<'a> { )> { let session_config = self.task_ctx.session_config(); - let (specialized, work_unit_feed_declarations, dynamic_filter_ids) = - self.task_specialized_plan(task_i)?; + let (specialized, work_unit_feed_declarations) = self.task_specialized_plan(task_i)?; let task_key = TaskKey { query_id: self.query_id, @@ -179,7 +174,6 @@ impl<'a> StageCoordinator<'a> { task_key, task_count: self.task_count, plan: MaybeEncoded::Decoded(specialized), - dynamic_filter_ids, work_unit_feed_declarations, target_worker_url: url.clone(), query_start_time_ns: self.metrics.instantiation_time, @@ -381,7 +375,6 @@ impl<'a> StageCoordinator<'a> { .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, @@ -413,19 +406,12 @@ 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)) })?; - 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, - )) + Ok((transformed.data, work_unit_feed_declarations)) } /// Returns as many URLs as the task count for the stage this [StageCoordinator] diff --git a/src/protocol/grpc/generated/worker.rs b/src/protocol/grpc/generated/worker.rs index 0793b2f7b..73368b080 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/grpc/generated/worker.rs @@ -43,7 +43,18 @@ 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. + /// Final dynamic filters used by dynamic-filter consumer execution-plan nodes. + /// + /// Filters are deduplicated by expression_id because consumers with the same ID share + /// logical filter state within a task. For example, this plan includes one entry: + /// + /// HashJoin producer: expression_id=10 + /// ├── DataSourceExec build side + /// └── UnionExec probe side + /// ├── DataSourceExec A consumer: expression_id=10 + /// └── DataSourceExec B consumer: expression_id=10 + /// + /// Another task in the same stage may report a different value for expression_id=10. #[prost(message, tag = "4")] TaskCompletedDynamicFilters(super::TaskCompletedDynamicFilters), } @@ -138,9 +149,6 @@ 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 b4f4331d9..52b668bb9 100644 --- a/src/protocol/grpc/worker.proto +++ b/src/protocol/grpc/worker.proto @@ -40,7 +40,18 @@ message WorkerToCoordinatorMsg { bool load_info_eos = 3; - // Final task-local dynamic filters used by distributed leaf variants. + // Final dynamic filters used by dynamic-filter consumer execution-plan nodes. + // + // Filters are deduplicated by expression_id because consumers with the same ID share + // logical filter state within a task. For example, this plan includes one entry: + // + // HashJoin producer: expression_id=10 + // ├── DataSourceExec build side + // └── UnionExec probe side + // ├── DataSourceExec A consumer: expression_id=10 + // └── DataSourceExec B consumer: expression_id=10 + // + // Another task in the same stage may report a different value for expression_id=10. TaskCompletedDynamicFilters task_completed_dynamic_filters = 4; } } @@ -120,8 +131,6 @@ 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 69ebbe703..7196561eb 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -451,7 +451,6 @@ fn encode_set_plan_request( .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, }) } diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index 1b210288d..4b32f91fa 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -228,7 +228,6 @@ 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, }) } diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index 87952ce4e..03f0ebb82 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -79,9 +79,6 @@ pub struct SetPlanRequest { pub task_count: usize, /// The subplan the worker is expected to execute. pub plan: MaybeEncoded>, - /// 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. diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index ea36ea99a..838cb3f4b 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,3 +1,4 @@ +use crate::codec::encode_physical_expr; use crate::common::{TreeNodeExt, discover_dynamic_filter_consumers}; use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers}; use crate::execution_plans::SamplerExec; @@ -5,16 +6,15 @@ use crate::protocol::LocalWorkerContext; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::task_data::TaskDataMetrics; use crate::{ - CoordinatorToWorkerMsg, DistributedCodec, DistributedConfig, DistributedExt, - DistributedTaskContext, SetPlanRequest, TaskCompletedDynamicFilters, TaskData, - TaskDynamicFilter, TaskMetrics, Worker, WorkerQueryContext, WorkerToCoordinatorMsg, + CoordinatorToWorkerMsg, DistributedConfig, DistributedExt, DistributedTaskContext, + SetPlanRequest, TaskCompletedDynamicFilters, TaskData, TaskDynamicFilter, TaskMetrics, Worker, + WorkerQueryContext, WorkerToCoordinatorMsg, }; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{DataFusionError, HashSet, Result, exec_datafusion_err, internal_err}; +use datafusion::common::{DataFusionError, Result, exec_datafusion_err, internal_err}; use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; use datafusion_proto::protobuf::physical_expr_node::ExprType; use futures::stream::{BoxStream, FuturesUnordered}; use futures::{FutureExt, StreamExt, TryStreamExt}; @@ -112,7 +112,6 @@ 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 final-report senders that keep the // worker->coordinator stream alive. As soon as this task ends, the runtime metrics and @@ -174,12 +173,8 @@ 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(); + dynamic_filters = build_task_completed_dynamic_filters(plan, &task_data.task_ctx) + .unwrap_or_default(); } let _ = dynamic_filters_tx.send(dynamic_filters); task_data_entries.invalidate(&key).await @@ -224,15 +219,13 @@ impl Worker { 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))? { + for consumer in discover_dynamic_filter_consumers(plan)? { // 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 expression = encode_physical_expr(&consumer.expression, task_ctx)?; let Some(ExprType::DynamicFilter(dynamic_filter)) = expression.expr_type.as_ref() else { return internal_err!("discovered dynamic filter did not serialize as one"); }; diff --git a/tests/dynamic_filtering.rs b/tests/dynamic_filtering.rs index baba4fe89..49771039c 100644 --- a/tests/dynamic_filtering.rs +++ b/tests/dynamic_filtering.rs @@ -78,9 +78,97 @@ mod tests { Ok(()) } + #[tokio::test] + async fn collect_left_union_probe_deduplicates_dynamic_filters() -> Result<()> { + let display = execute_local_union_probe_hash_join().await?; + assert_snapshot!(display, @r" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 2] => NetworkCoalesceExec: output_partitions=14, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=14 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(key@0, MinTemp@0)], projection=[] + │ CoalescePartitionsExec + │ [Stage 1] => NetworkBroadcastExec: partitions_per_consumer=3, stage_partitions=6, input_tasks=1 + │ DistributedUnionExec: t0:[c0, c2, c4] t1:[c1, c3] + │ 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:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp], file_type=parquet, predicate=DynamicFilter [ MinTemp@0 >= -5.3 AND MinTemp@0 <= 20.9 AND true ], dynamic_rg_pruning=eligible, pruning_predicate=MinTemp_null_count@1 != row_count@2 AND MinTemp_max@0 >= -5.3 AND MinTemp_null_count@1 != row_count@2 AND MinTemp_min@3 <= 20.9, required_guarantees=[] + │ ProjectionExec: expr=[CAST(-1000 AS Float64) as MinTemp] + │ PlaceholderRowExec + │ 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:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp], file_type=parquet, predicate=DynamicFilter [ MinTemp@0 >= -5.3 AND MinTemp@0 <= 20.9 AND true ], dynamic_rg_pruning=eligible, pruning_predicate=MinTemp_null_count@1 != row_count@2 AND MinTemp_max@0 >= -5.3 AND MinTemp_null_count@1 != row_count@2 AND MinTemp_min@3 <= 20.9, required_guarantees=[] + │ ProjectionExec: expr=[CAST(-1001 AS Float64) as MinTemp] + │ PlaceholderRowExec + │ ProjectionExec: expr=[CAST(-1002 AS Float64) as MinTemp] + │ PlaceholderRowExec + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=1, partitions=6 + │ BroadcastExec: input_partitions=3, consumer_tasks=2, output_partitions=6 + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ RepartitionExec: partitioning=Hash([key@0], 3), 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:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp@0 as key], file_type=parquet + └────────────────────────────────────────────────── + "); + Ok(()) + } + async fn execute_local_hash_join(broadcast_joins: bool) -> Result { + execute_local_query( + broadcast_joins, + false, + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + JOIN weather probe ON build.key = probe."RainToday" + "#, + ) + .await + } + + async fn execute_local_union_probe_hash_join() -> Result { + execute_local_query( + true, + true, + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "MinTemp" AS key + FROM weather + ) build + JOIN ( + SELECT "MinTemp" FROM weather + UNION ALL + SELECT CAST(-1000.0 AS DOUBLE) AS "MinTemp" + UNION ALL + SELECT "MinTemp" FROM weather + UNION ALL + SELECT CAST(-1001.0 AS DOUBLE) AS "MinTemp" + UNION ALL + SELECT CAST(-1002.0 AS DOUBLE) AS "MinTemp" + ) probe ON build.key = probe."MinTemp" + "#, + ) + .await + } + + async fn execute_local_query( + broadcast_joins: bool, + one_task_per_leaf: bool, + sql: &str, + ) -> Result { let (ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; - let ctx = ctx.with_distributed_broadcast_joins(broadcast_joins)?; + let mut ctx = ctx.with_distributed_broadcast_joins(broadcast_joins)?; + if one_task_per_leaf { + ctx = ctx.with_distributed_desired_task_count_handler(1usize); + } if !broadcast_joins { let state = ctx.state_ref(); let mut state = state.write(); @@ -90,20 +178,7 @@ mod tests { } 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 plan = ctx.sql(sql).await?.create_physical_plan().await?; let results = collect(Arc::clone(&plan), ctx.task_ctx()).await?; assert_eq!( From f69728c5abdea83b41842ba84146261f380215f1 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Fri, 21 Aug 2026 19:11:24 +0000 Subject: [PATCH 10/17] remove plan_for_display --- src/common/dynamic_filtering.rs | 1 - src/coordinator/distributed.rs | 27 ++++++++------------------- src/metrics/task_metrics_rewriter.rs | 2 +- src/stage.rs | 2 +- 4 files changed, 10 insertions(+), 22 deletions(-) diff --git a/src/common/dynamic_filtering.rs b/src/common/dynamic_filtering.rs index 7582fa252..0d2fba113 100644 --- a/src/common/dynamic_filtering.rs +++ b/src/common/dynamic_filtering.rs @@ -12,7 +12,6 @@ use std::sync::Arc; pub(crate) struct DiscoveredDynamicFilter { pub(crate) id: u64, pub(crate) expression: Arc, - // The schema the filter is evaluated against. pub(crate) input_schema: SchemaRef, } diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 2b031cad3..8e2c63996 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -91,20 +91,9 @@ impl DistributedExec { task_metrics.wait_for(&task_keys_for_plan(&plan)).await; } - /// Returns the plan which is lazily prepared on `execute()` and actually gets executed. - /// It is updated on every call to `execute()`. Returns an error if `.execute()` has not been - /// called. - pub(crate) fn plan_for_viz(&self) -> Result> { - self.plan_for_viz - .lock() - .map_err(|e| internal_datafusion_err!("Failed to lock prepared plan: {}", e))? - .clone() - .ok_or_else(|| { - internal_datafusion_err!("No prepared plan found. Was execute() called?") - }) - } - - pub(crate) fn plan_for_display(&self) -> Arc { + /// Returns the execution-derived plan used for visualization, falling back to the initial plan + /// until `execute()` has prepared it. It is updated on every call to `execute()`. + pub(crate) fn plan_for_viz(&self) -> Arc { self.plan_for_viz .lock() .ok() @@ -180,7 +169,7 @@ impl ExecutionPlan for DistributedExec { } let base_plan = Arc::clone(&self.base_plan); - let plan_for_viz = Arc::clone(&self.plan_for_viz); + let plan_for_viz_slot = Arc::clone(&self.plan_for_viz); let head_stage = Arc::clone(&self.head_stage); let query_coordinator = QueryCoordinator::new( @@ -211,12 +200,12 @@ impl ExecutionPlan for DistributedExec { false => prepare_static_plan(&query_coordinator, &base_plan)?, }; - let display_plan = + let plan_for_viz = isolate_distributed_leaf_variants_for_display(result.plan_for_viz, &context)?; - plan_for_viz + plan_for_viz_slot .lock() .expect("poisoned lock") - .replace(Arc::clone(&display_plan)); + .replace(Arc::clone(&plan_for_viz)); head_stage .lock() .expect("poisoned lock") @@ -229,7 +218,7 @@ impl ExecutionPlan for DistributedExec { } drop(guard); query_coordinator - .finish_dynamic_filter_display(&display_plan) + .finish_dynamic_filter_display(&plan_for_viz) .await; drop(tx); query_coordinator.drain_pending_tasks().await?; diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index a3633203f..cfb12f41d 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -61,7 +61,7 @@ pub async fn rewrite_distributed_plan_with_metrics( // Rewrite the DistributedExec's child plan with metrics. let dist_exec_plan_with_metrics = rewrite_local_plan_with_metrics( format.to_rewrite_ctx(0), // Task id is 0 for the DistributedExec plan - distributed_exec.plan_for_viz()?, + distributed_exec.plan_for_viz(), task_metrics, )?; diff --git a/src/stage.rs b/src/stage.rs index 967135309..b1625c04a 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -289,7 +289,7 @@ fn display_ascii( f: &mut String, ) -> std::fmt::Result { let prepared_plan = match stage { - Either::Left(distributed_exec) => Some(distributed_exec.plan_for_display()), + Either::Left(distributed_exec) => Some(distributed_exec.plan_for_viz()), Either::Right(_) => None, }; let plan = match stage { From 92fb135697076479bf8a19ecceb4bbd92a26bb5d Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Fri, 21 Aug 2026 20:24:20 +0000 Subject: [PATCH 11/17] make displaying align with metrics --- docs/source/user-guide/05-metrics.md | 13 ++- src/coordinator/distributed.rs | 115 ++++++++++++++++++--------- src/coordinator/dynamic_filters.rs | 21 +++++ src/coordinator/mod.rs | 1 + src/coordinator/query_coordinator.rs | 16 +--- src/lib.rs | 2 +- src/metrics/task_metrics_rewriter.rs | 9 +-- src/stage.rs | 6 +- tests/dynamic_filtering.rs | 17 +++- 9 files changed, 130 insertions(+), 70 deletions(-) diff --git a/docs/source/user-guide/05-metrics.md b/docs/source/user-guide/05-metrics.md index 98d4f1440..60a52b3d9 100644 --- a/docs/source/user-guide/05-metrics.md +++ b/docs/source/user-guide/05-metrics.md @@ -26,8 +26,11 @@ channel, so they are not lost even if the result stream is dropped early (for ex ## Rendering a plan with metrics -Two functions, both exported from the crate root, do the work: +These functions, all exported from the crate root, do the work: +- `rewrite_distributed_plan_with_dynamic_filters(plan)` — folds the completed dynamic filters + reported by each worker task into an isolated copy of the plan. When displaying both dynamic + filters and metrics, apply the dynamic-filter rewrite first. - `rewrite_distributed_plan_with_metrics(plan, format)` — folds every task's metrics back into the coordinator's copy of the plan. It waits for all worker metrics to arrive, so the result is always complete. The `format` is a `DistributedMetricsFormat`: @@ -54,11 +57,15 @@ execute_stream(plan.clone(), ctx.task_ctx())? .try_collect::>() .await?; -// 3. Fold the per-task metrics back into the plan... +// 3. Fold the completed per-task dynamic filters back into the plan... +let plan = + rewrite_distributed_plan_with_dynamic_filters(plan).await?; + +// 4. Fold the per-task metrics back into the plan... let plan = rewrite_distributed_plan_with_metrics(plan, DistributedMetricsFormat::Aggregated).await?; -// 4. ...and render it. +// 5. ...and render it. println!("{}", display_plan_ascii(plan.as_ref(), true)); ``` diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 8e2c63996..461dd07ba 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,13 +1,13 @@ -use crate::DistributedConfig; use crate::common::require_one_child; 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::coordinator::store::{MetricsStore, task_keys_for_plan}; +use crate::coordinator::store::{CompletedDynamicFilterStore, MetricsStore, task_keys_for_plan}; +use crate::{DistributedConfig, TaskCompletedDynamicFilters, TaskKey}; use datafusion::common::internal_datafusion_err; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{Result, exec_err}; +use datafusion::common::{HashMap, Result, exec_err}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr_common::metrics::MetricsSet; @@ -33,18 +33,26 @@ pub struct DistributedExec { /// - If the plan is going to be distributed dynamically during execution, this is the initial /// non-distributed plan. base_plan: Arc, - /// Resulting [ExecutionPlan] after execution ready for visualization purposes. - /// - If the plan was distributed statically, this is equal to the base plan. - /// - If the plan is going to be distributed dynamically during execution, this is the resulting - /// plan re-calculated based on runtime statistics. - plan_for_viz: Arc>>>, - /// The head stage meant to be executed locally on [DistributedExec::execute]. - head_stage: Arc>>>, + prepared_execution: Arc>>, /// DataFusion metrics. metrics: ExecutionPlanMetricsSet, /// Storage where metrics collected from workers at runtime will place their results as they /// finish their respective remote tasks. pub(crate) metrics_store: Option>, + /// Storage for the completed dynamic filters reported by each worker task. + pub(crate) completed_dynamic_filter_store: Arc, +} + +/// Execution state produced by distributed planning (static or dynamic) retained +/// for post-execution work such as plan rewrites to display metrics and dynamic filters. +#[derive(Debug, Clone)] +struct PreparedExecution { + /// Resulting plan reconstructed after static or dynamic planning. + plan_for_viz: Arc, + /// The head stage actually executed locally by the coordinator. + head_stage: Arc, + /// The task context of the [`DistributedExec`]. Useful for decoding protobufs. + task_ctx: Arc, } pub(super) struct PreparedPlan { @@ -58,10 +66,10 @@ impl DistributedExec { pub fn new(base_plan: Arc) -> Self { Self { base_plan, - plan_for_viz: Arc::new(Mutex::new(None)), - head_stage: Arc::new(Mutex::new(None)), + prepared_execution: Arc::new(Mutex::new(None)), metrics: ExecutionPlanMetricsSet::new(), metrics_store: None, + completed_dynamic_filter_store: Arc::new(CompletedDynamicFilterStore::new()), } } @@ -85,20 +93,35 @@ impl DistributedExec { let Some(task_metrics) = &self.metrics_store else { return; }; - let Some(plan) = self.plan_for_viz.lock().unwrap().as_ref().cloned() else { + let Ok(plan) = self.plan_for_viz() else { return; }; task_metrics.wait_for(&task_keys_for_plan(&plan)).await; } - /// Returns the execution-derived plan used for visualization, falling back to the initial plan - /// until `execute()` has prepared it. It is updated on every call to `execute()`. - pub(crate) fn plan_for_viz(&self) -> Arc { - self.plan_for_viz + pub(crate) async fn wait_for_dynamic_filters( + &self, + ) -> Result> { + let plan = self.plan_for_viz()?; + Ok(self + .completed_dynamic_filter_store + .wait_for(&task_keys_for_plan(&plan)) + .await) + } + + fn prepared_execution(&self) -> Result { + self.prepared_execution .lock() - .ok() - .and_then(|plan| plan.clone()) - .unwrap_or_else(|| Arc::clone(&self.base_plan)) + .map_err(|e| internal_datafusion_err!("Failed to lock prepared execution: {e}"))? + .clone() + .ok_or_else(|| { + internal_datafusion_err!("No prepared execution found. Was execute() called?") + }) + } + + /// Returns the plan reconstructed from the execution for visualization and rewriting. + pub(crate) fn plan_for_viz(&self) -> Result> { + Ok(self.prepared_execution()?.plan_for_viz) } /// Returns the head stage that was actually executed. Unlike [`Self::plan_for_viz`] (which is @@ -106,11 +129,28 @@ impl DistributedExec { /// `Arc`s), this returns the original `Arc` instances whose metrics were populated during /// execution. pub(crate) fn head_stage(&self) -> Result> { - self.head_stage - .lock() - .map_err(|e| internal_datafusion_err!("Failed to lock head stage: {}", e))? - .clone() - .ok_or_else(|| internal_datafusion_err!("No head stage found. Was execute() called?")) + Ok(self.prepared_execution()?.head_stage) + } + + pub(crate) fn task_ctx(&self) -> Result> { + Ok(self.prepared_execution()?.task_ctx) + } + + /// Builds a non-executable visualization result while preserving the state needed by a + /// subsequent rewrite. Dynamic filters must be rewritten before metrics. + pub(crate) fn with_rewritten_plan( + &self, + plan_for_viz: Arc, + ) -> Result> { + let mut prepared_execution = self.prepared_execution()?; + prepared_execution.plan_for_viz = Arc::clone(&plan_for_viz); + Ok(Arc::new(Self { + base_plan: plan_for_viz, + prepared_execution: Arc::new(Mutex::new(Some(prepared_execution))), + metrics: self.metrics.clone(), + metrics_store: self.metrics_store.clone(), + completed_dynamic_filter_store: Arc::clone(&self.completed_dynamic_filter_store), + })) } } @@ -146,10 +186,10 @@ impl ExecutionPlan for DistributedExec { ) -> Result> { Ok(Arc::new(DistributedExec { base_plan: require_one_child(&children)?, - plan_for_viz: Arc::new(Mutex::new(None)), - head_stage: Arc::new(Mutex::new(None)), + prepared_execution: Arc::new(Mutex::new(None)), metrics: self.metrics.clone(), metrics_store: self.metrics_store.clone(), + completed_dynamic_filter_store: Arc::clone(&self.completed_dynamic_filter_store), })) } @@ -169,13 +209,13 @@ impl ExecutionPlan for DistributedExec { } let base_plan = Arc::clone(&self.base_plan); - let plan_for_viz_slot = Arc::clone(&self.plan_for_viz); - let head_stage = Arc::clone(&self.head_stage); + let prepared_execution = Arc::clone(&self.prepared_execution); let query_coordinator = QueryCoordinator::new( Arc::clone(&context), &self.metrics, self.metrics_store.clone(), + Arc::clone(&self.completed_dynamic_filter_store), ); let mut builder = RecordBatchReceiverStreamBuilder::new(self.schema(), 1); @@ -202,14 +242,14 @@ impl ExecutionPlan for DistributedExec { let plan_for_viz = isolate_distributed_leaf_variants_for_display(result.plan_for_viz, &context)?; - plan_for_viz_slot - .lock() - .expect("poisoned lock") - .replace(Arc::clone(&plan_for_viz)); - head_stage + prepared_execution .lock() - .expect("poisoned lock") - .replace(Arc::clone(&result.head_stage)); + .map_err(|e| internal_datafusion_err!("Failed to lock prepared execution: {e}"))? + .replace(PreparedExecution { + plan_for_viz, + head_stage: Arc::clone(&result.head_stage), + task_ctx: Arc::clone(&context), + }); let mut stream = result.head_stage.execute(partition, context)?; while let Some(msg) = stream.next().await { if tx.send(msg).await.is_err() { @@ -217,9 +257,6 @@ impl ExecutionPlan for DistributedExec { } } drop(guard); - query_coordinator - .finish_dynamic_filter_display(&plan_for_viz) - .await; drop(tx); query_coordinator.drain_pending_tasks().await?; Ok(()) diff --git a/src/coordinator/dynamic_filters.rs b/src/coordinator/dynamic_filters.rs index 0985e36d1..bc32852af 100644 --- a/src/coordinator/dynamic_filters.rs +++ b/src/coordinator/dynamic_filters.rs @@ -1,5 +1,6 @@ use crate::codec::decode_physical_expr; use crate::common::discover_dynamic_filter_consumers; +use crate::coordinator::DistributedExec; use crate::execution_plans::DistributedLeafExec; use crate::{DistributedCodec, TaskCompletedDynamicFilters, TaskKey}; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; @@ -11,6 +12,26 @@ use datafusion_proto::physical_plan::{DeduplicatingProtoConverter, PhysicalPlanN use datafusion_proto::protobuf::PhysicalPlanNode; use std::sync::Arc; +/// Rewrites an executed distributed plan with the completed dynamic filters reported by its +/// worker tasks. +/// +/// When composing this with [`crate::rewrite_distributed_plan_with_metrics`], dynamic filters must +/// be rewritten first. +pub async fn rewrite_distributed_plan_with_dynamic_filters( + plan: Arc, +) -> Result> { + let Some(distributed_exec) = plan.downcast_ref::() else { + return Ok(plan); + }; + + let plan_for_viz = distributed_exec.plan_for_viz()?; + let task_ctx = distributed_exec.task_ctx()?; + let reports = distributed_exec.wait_for_dynamic_filters().await?; + let plan_for_viz = isolate_distributed_leaf_variants_for_display(plan_for_viz, &task_ctx)?; + apply_reports_to_distributed_leaves(&plan_for_viz, &reports, &task_ctx); + distributed_exec.with_rewritten_plan(plan_for_viz) +} + /// Replaces the variants in the visualization plan with independent per-task copies. pub(super) fn isolate_distributed_leaf_variants_for_display( plan: Arc, diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index c08fbbd46..dc1b4e520 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -7,4 +7,5 @@ mod query_coordinator; mod store; pub use distributed::DistributedExec; +pub use dynamic_filters::rewrite_distributed_plan_with_dynamic_filters; pub(crate) use store::{CompletedDynamicFilterStore, MetricsStore}; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index c635c2238..95f04d660 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,9 +1,7 @@ use crate::codec::{decode_execution_plan, encode_execution_plan}; use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; -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}; @@ -61,12 +59,13 @@ impl QueryCoordinator { task_ctx: Arc, metrics_set: &ExecutionPlanMetricsSet, metrics_store: Option>, + completed_dynamic_filter_store: Arc, ) -> Self { Self { task_ctx, metrics: metrics_set.clone(), metrics_store, - completed_dynamic_filter_store: Arc::new(CompletedDynamicFilterStore::new()), + completed_dynamic_filter_store, coordinator_to_worker_metrics: CoordinatorToWorkerMetrics::new(metrics_set), end_stream_notifier: Arc::new(Notify::new()), join_set: Mutex::new(JoinSet::new()), @@ -103,17 +102,6 @@ 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<()> { diff --git a/src/lib.rs b/src/lib.rs index ae881f489..bf4bf4693 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,7 +18,7 @@ mod worker_resolver; #[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; -pub use coordinator::DistributedExec; +pub use coordinator::{DistributedExec, rewrite_distributed_plan_with_dynamic_filters}; pub use distributed_ext::{DistributedExt, DistributedGetterExt}; pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, ProducerHead, SessionStateBuilderExt, diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index cfb12f41d..99ca5c024 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -12,9 +12,9 @@ use datafusion::common::tree_node::Transformed; use datafusion::common::tree_node::TreeNode; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::Result; +use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::internal_err; use datafusion::physical_plan::metrics::{Label, Metric, MetricsSet}; -use datafusion::physical_plan::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions}; use std::sync::Arc; /// Format to use when displaying metrics for a distributed plan. @@ -61,7 +61,7 @@ pub async fn rewrite_distributed_plan_with_metrics( // Rewrite the DistributedExec's child plan with metrics. let dist_exec_plan_with_metrics = rewrite_local_plan_with_metrics( format.to_rewrite_ctx(0), // Task id is 0 for the DistributedExec plan - distributed_exec.plan_for_viz(), + distributed_exec.plan_for_viz()?, task_metrics, )?; @@ -89,10 +89,7 @@ pub async fn rewrite_distributed_plan_with_metrics( Ok(Transformed::no(plan)) })?; - plan.replace_children( - vec![transformed.data], - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) + distributed_exec.with_rewritten_plan(transformed.data) } /// Extra information for rewriting local plans. diff --git a/src/stage.rs b/src/stage.rs index b1625c04a..cd054c575 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -288,12 +288,8 @@ 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_viz()), - Either::Right(_) => None, - }; let plan = match stage { - Either::Left(_) => prepared_plan.as_ref().unwrap(), + Either::Left(distributed_exec) => distributed_exec.children().first().unwrap(), Either::Right(stage) => { let Some(plan) = stage.local_plan() else { return write!(f, "StageExec: encoded input plan"); diff --git a/tests/dynamic_filtering.rs b/tests/dynamic_filtering.rs index 49771039c..b090fb790 100644 --- a/tests/dynamic_filtering.rs +++ b/tests/dynamic_filtering.rs @@ -5,7 +5,9 @@ mod tests { 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, + DefaultSessionBuilder, DistributedExt, DistributedMetricsFormat, assert_snapshot, + display_plan_ascii, rewrite_distributed_plan_with_dynamic_filters, + rewrite_distributed_plan_with_metrics, }; use std::sync::Arc; @@ -186,6 +188,17 @@ mod tests { 1 ); - Ok(display_plan_ascii(plan.as_ref(), false)) + let original_display = display_plan_ascii(plan.as_ref(), false); + + let plan_with_dynamic_filters = + rewrite_distributed_plan_with_dynamic_filters(Arc::clone(&plan)).await?; + assert_eq!(display_plan_ascii(plan.as_ref(), false), original_display); + + let plan_with_metrics = rewrite_distributed_plan_with_metrics( + plan_with_dynamic_filters, + DistributedMetricsFormat::Aggregated, + ) + .await?; + Ok(display_plan_ascii(plan_with_metrics.as_ref(), false)) } } From cdd153ea4823aad91a9bba22ade35bb763e3b39b Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 25 Aug 2026 22:01:37 +0000 Subject: [PATCH 12/17] remove type aliases and use select_all --- src/coordinator/distributed.rs | 12 ++++++------ src/coordinator/mod.rs | 2 +- src/coordinator/query_coordinator.rs | 25 +++++++++++++------------ src/coordinator/store.rs | 5 +---- src/metrics/task_metrics_rewriter.rs | 10 +++++----- src/stage.rs | 7 ++++--- src/worker/impl_coordinator_channel.rs | 11 ++++++----- 7 files changed, 36 insertions(+), 36 deletions(-) diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 461dd07ba..86a6f5922 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -3,8 +3,8 @@ use crate::coordinator::dynamic_filters::isolate_distributed_leaf_variants_for_d 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::coordinator::store::{CompletedDynamicFilterStore, MetricsStore, task_keys_for_plan}; -use crate::{DistributedConfig, TaskCompletedDynamicFilters, TaskKey}; +use crate::coordinator::store::{Store, task_keys_for_plan}; +use crate::{DistributedConfig, TaskCompletedDynamicFilters, TaskKey, TaskMetrics}; use datafusion::common::internal_datafusion_err; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{HashMap, Result, exec_err}; @@ -38,9 +38,9 @@ pub struct DistributedExec { metrics: ExecutionPlanMetricsSet, /// Storage where metrics collected from workers at runtime will place their results as they /// finish their respective remote tasks. - pub(crate) metrics_store: Option>, + pub(crate) metrics_store: Option>>, /// Storage for the completed dynamic filters reported by each worker task. - pub(crate) completed_dynamic_filter_store: Arc, + pub(crate) completed_dynamic_filter_store: Arc>, } /// Execution state produced by distributed planning (static or dynamic) retained @@ -69,14 +69,14 @@ impl DistributedExec { prepared_execution: Arc::new(Mutex::new(None)), metrics: ExecutionPlanMetricsSet::new(), metrics_store: None, - completed_dynamic_filter_store: Arc::new(CompletedDynamicFilterStore::new()), + completed_dynamic_filter_store: Arc::new(Store::new()), } } /// Enables task metrics collection from remote workers. pub fn with_metrics_collection(mut self, enabled: bool) -> Self { self.metrics_store = match enabled { - true => Some(Arc::new(MetricsStore::new())), + true => Some(Arc::new(Store::new())), false => None, }; self diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index dc1b4e520..7eb6dbef7 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -8,4 +8,4 @@ mod store; pub use distributed::DistributedExec; pub use dynamic_filters::rewrite_distributed_plan_with_dynamic_filters; -pub(crate) use store::{CompletedDynamicFilterStore, MetricsStore}; +pub(crate) use store::Store; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 95f04d660..ff818ad0b 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,8 +1,8 @@ use crate::codec::{decode_execution_plan, encode_execution_plan}; use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; +use crate::coordinator::Store; use crate::coordinator::latency_metric::LatencyMetric; -use crate::coordinator::{CompletedDynamicFilterStore, MetricsStore}; use crate::events::{RouteTasksEvent, RouteTasksHandlers}; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; use crate::passthrough_headers::get_passthrough_headers; @@ -12,8 +12,8 @@ use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time}; use crate::{ CoordinatorToWorkerMsg, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedTaskContext, DistributedWorkUnitFeedContext, LoadInfo, LocalWorkerContext, MaybeEncoded, SetPlanRequest, - TaskCompletedDynamicFilters, TaskKey, WorkUnitFeedDeclaration, WorkerToCoordinatorMsg, - get_distributed_channel_resolver, + TaskCompletedDynamicFilters, TaskKey, TaskMetrics, WorkUnitFeedDeclaration, + WorkerToCoordinatorMsg, get_distributed_channel_resolver, }; use datafusion::common::DataFusionError; use datafusion::common::instant::Instant; @@ -47,8 +47,8 @@ pub(super) struct QueryCoordinator { task_ctx: Arc, metrics: ExecutionPlanMetricsSet, coordinator_to_worker_metrics: CoordinatorToWorkerMetrics, - metrics_store: Option>, - completed_dynamic_filter_store: Arc, + metrics_store: Option>>, + completed_dynamic_filter_store: Arc>, end_stream_notifier: Arc, join_set: Mutex>>, } @@ -58,8 +58,8 @@ impl QueryCoordinator { pub(super) fn new( task_ctx: Arc, metrics_set: &ExecutionPlanMetricsSet, - metrics_store: Option>, - completed_dynamic_filter_store: Arc, + metrics_store: Option>>, + completed_dynamic_filter_store: Arc>, ) -> Self { Self { task_ctx, @@ -121,8 +121,6 @@ impl QueryCoordinator { /// - Building tasks that communicate a serialized plan to multiple workers for further execution. /// - Building tasks that stream partition feeds from local [WorkUnitFeedExec] nodes to their /// remote counterparts. -type SpecializedTaskPlan = (Arc, Vec); - pub(super) struct StageCoordinator<'a> { plan: &'a Arc, query_id: Uuid, @@ -131,8 +129,8 @@ pub(super) struct StageCoordinator<'a> { task_ctx: &'a Arc, metrics_set: &'a ExecutionPlanMetricsSet, metrics: &'a CoordinatorToWorkerMetrics, - metrics_store: &'a Option>, - completed_dynamic_filter_store: &'a Arc, + metrics_store: &'a Option>>, + completed_dynamic_filter_store: &'a Arc>, end_stream_notifier: &'a Arc, join_set: &'a Mutex>>, } @@ -356,7 +354,10 @@ impl<'a> StageCoordinator<'a> { /// trimming down any unnecessary information that the specific `task_i` task is not going to /// need, like unexecuted branches in [ChildrenIsolatorUnionExec], or unexecuted variants of /// [DistributedLeafExec]. - fn task_specialized_plan(&self, task_i: usize) -> Result { + fn task_specialized_plan( + &self, + task_i: usize, + ) -> Result<(Arc, Vec)> { let session_config = self.task_ctx.session_config(); let wuf_registry = session_config .get_extension::() diff --git a/src/coordinator/store.rs b/src/coordinator/store.rs index d6553e0dd..677bcfc4c 100644 --- a/src/coordinator/store.rs +++ b/src/coordinator/store.rs @@ -1,5 +1,5 @@ +use crate::TaskKey; 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; @@ -55,9 +55,6 @@ impl Store { } } -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| { diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index 99ca5c024..f072f4a76 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -1,11 +1,11 @@ use crate::common::TreeNodeExt; -use crate::coordinator::{DistributedExec, MetricsStore}; +use crate::coordinator::{DistributedExec, Store}; use crate::distributed_planner::NetworkBoundaryExt; use crate::execution_plans::MetricsWrapperExec; use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL; use crate::metrics::collect_plan_metrics; use crate::stage::{LocalStage, Stage}; -use crate::{DistributedTaskContext, TaskKey}; +use crate::{DistributedTaskContext, TaskKey, TaskMetrics}; use datafusion::common::HashMap; use datafusion::common::plan_err; use datafusion::common::tree_node::Transformed; @@ -207,7 +207,7 @@ pub fn rewrite_local_plan_with_metrics( /// Note: Metrics may be aggregated by name (ex. output_rows) automatically by various datafusion utils. pub fn stage_metrics_rewriter( stage: &LocalStage, - metrics_collection: Arc, + metrics_collection: Arc>, format: DistributedMetricsFormat, ) -> Result> { // Phase 1 — accumulate per-task metrics into a map keyed by node identity. @@ -283,7 +283,7 @@ pub fn stage_metrics_rewriter( #[cfg(test)] mod tests { use crate::DistributedExt; - use crate::coordinator::MetricsStore; + use crate::coordinator::Store; use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL; use crate::metrics::task_metrics_rewriter::MetricsWrapperExec; use crate::metrics::task_metrics_rewriter::{ @@ -449,7 +449,7 @@ mod tests { let num_metrics_per_task_per_node = 4; // Generate metrics for each task and store them in the map. - let metrics_collection = MetricsStore::from_entries((0..stage.tasks).map(|task_id| { + let metrics_collection = Store::from_entries((0..stage.tasks).map(|task_id| { let task_key = TaskKey { query_id: stage.query_id, stage_id: stage.num, diff --git a/src/stage.rs b/src/stage.rs index cd054c575..37220f044 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -1,4 +1,4 @@ -use crate::coordinator::{DistributedExec, MetricsStore}; +use crate::coordinator::{DistributedExec, Store}; use crate::execution_plans::{DistributedLeafExec, NetworkCoalesceExec}; use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL; use datafusion::common::{HashMap, Statistics, config_err}; @@ -224,7 +224,8 @@ impl DistributedTaskContext { } use crate::{ - DistributedMetricsFormat, NetworkShuffleExec, TaskKey, rewrite_distributed_plan_with_metrics, + DistributedMetricsFormat, NetworkShuffleExec, TaskKey, TaskMetrics, + rewrite_distributed_plan_with_metrics, }; use crate::{NetworkBoundary, NetworkBoundaryExt}; use datafusion::arrow::datatypes::SchemaRef; @@ -439,7 +440,7 @@ fn display_inner_distributed_leaf( /// Gathers the metrics global to a stage. These metrics are not specific to any plan node, and /// are instead global to a whole stage. -fn gather_stage_header_metrics(stage: &Stage, metrics_store: &MetricsStore) -> MetricsSet { +fn gather_stage_header_metrics(stage: &Stage, metrics_store: &Store) -> MetricsSet { let mut task_key = TaskKey { query_id: stage.query_id(), stage_id: stage.num(), diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index 838cb3f4b..10ce4da32 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -16,7 +16,7 @@ use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use datafusion_proto::protobuf::physical_expr_node::ExprType; -use futures::stream::{BoxStream, FuturesUnordered}; +use futures::stream::{BoxStream, FuturesUnordered, select_all}; use futures::{FutureExt, StreamExt, TryStreamExt}; use http::HeaderMap; use std::sync::{Arc, OnceLock}; @@ -208,10 +208,11 @@ impl Worker { }, ); - Ok(futures::stream::select( - load_info_stream, - futures::stream::select(metrics_stream, dynamic_filters_stream), - ) + Ok(select_all([ + load_info_stream.boxed(), + metrics_stream.boxed(), + dynamic_filters_stream.boxed(), + ]) .map(Ok) .boxed()) } From 886cf131eaed032104dd3eb4061498c47f2189fc Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 25 Aug 2026 22:07:21 +0000 Subject: [PATCH 13/17] qualified import --- src/worker/impl_coordinator_channel.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index 10ce4da32..ac7412ce2 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -12,7 +12,7 @@ use crate::{ }; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{DataFusionError, Result, exec_datafusion_err, internal_err}; -use datafusion::execution::SessionStateBuilder; +use datafusion::execution::{SessionStateBuilder, TaskContext}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use datafusion_proto::protobuf::physical_expr_node::ExprType; @@ -220,7 +220,7 @@ impl Worker { fn build_task_completed_dynamic_filters( plan: &Arc, - task_ctx: &Arc, + task_ctx: &Arc, ) -> Result { let mut filters = vec![]; for consumer in discover_dynamic_filter_consumers(plan)? { From 7a14f8d0f22c270ee7a655889130e3accc0f8717 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Mon, 31 Aug 2026 18:59:56 +0000 Subject: [PATCH 14/17] refactors --- Cargo.toml | 5 + src/common/mod.rs | 2 - src/coordinator/distributed.rs | 52 ++- src/coordinator/dynamic_filters.rs | 196 ---------- src/coordinator/mod.rs | 2 - src/coordinator/query_coordinator.rs | 31 +- src/distributed_ext.rs | 42 +++ src/distributed_planner/distributed_config.rs | 4 + .../distributed_query_planner.rs | 20 +- .../discovery.rs} | 48 ++- src/dynamic_filtering/display.rs | 356 ++++++++++++++++++ src/dynamic_filtering/mod.rs | 50 +++ src/lib.rs | 4 +- src/metrics/task_metrics_rewriter.rs | 2 +- src/stage.rs | 8 +- src/worker/impl_coordinator_channel.rs | 48 ++- src/worker/task_data.rs | 6 +- src/worker/test_utils/worker_handles.rs | 2 + .../.partitioned_join.rs.pending-snap | 13 + tests/dynamic_filtering/aggregates.rs | 71 ++++ .../collect_left_join.rs} | 195 ++++------ tests/dynamic_filtering/common.rs | 161 ++++++++ tests/dynamic_filtering/config.rs | 23 ++ tests/dynamic_filtering/main.rs | 6 + tests/dynamic_filtering/partitioned_join.rs | 92 +++++ tests/dynamic_filtering/sorts.rs | 79 ++++ 26 files changed, 1148 insertions(+), 370 deletions(-) delete mode 100644 src/coordinator/dynamic_filters.rs rename src/{common/dynamic_filtering.rs => dynamic_filtering/discovery.rs} (82%) create mode 100644 src/dynamic_filtering/display.rs create mode 100644 src/dynamic_filtering/mod.rs create mode 100644 tests/dynamic_filtering/.partitioned_join.rs.pending-snap create mode 100644 tests/dynamic_filtering/aggregates.rs rename tests/{dynamic_filtering.rs => dynamic_filtering/collect_left_join.rs} (73%) create mode 100644 tests/dynamic_filtering/common.rs create mode 100644 tests/dynamic_filtering/config.rs create mode 100644 tests/dynamic_filtering/main.rs create mode 100644 tests/dynamic_filtering/partitioned_join.rs create mode 100644 tests/dynamic_filtering/sorts.rs diff --git a/Cargo.toml b/Cargo.toml index 5b8b06972..8a6f9a648 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,11 @@ license = "Apache-2.0" documentation = "https://datafusion-contrib.github.io/datafusion-distributed/" repository = "https://github.com/datafusion-contrib/datafusion-distributed" +[[test]] +name = "dynamic_filtering" +path = "tests/dynamic_filtering/main.rs" +required-features = ["integration"] + [dependencies] chrono = { version = "0.4.44" } datafusion = { workspace = true, features = [ diff --git a/src/common/mod.rs b/src/common/mod.rs index 97cca7394..7dfd3a96d 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,5 +1,4 @@ mod children_helpers; -mod dynamic_filtering; mod maybe_encoded; mod once_lock; mod recursion; @@ -9,7 +8,6 @@ mod uuid; mod vec; pub(crate) use children_helpers::require_one_child; -pub(crate) use dynamic_filtering::discover_dynamic_filter_consumers; pub use maybe_encoded::MaybeEncoded; pub(crate) use once_lock::OnceLockResult; pub(crate) use recursion::TreeNodeExt; diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 86a6f5922..533f7eea1 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,9 +1,9 @@ use crate::common::require_one_child; -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::coordinator::store::{Store, task_keys_for_plan}; +use crate::dynamic_filtering::sever_dynamic_filter_relationships_in_plan_for_display; use crate::{DistributedConfig, TaskCompletedDynamicFilters, TaskKey, TaskMetrics}; use datafusion::common::internal_datafusion_err; use datafusion::common::tree_node::TreeNodeRecursion; @@ -40,7 +40,7 @@ pub struct DistributedExec { /// finish their respective remote tasks. pub(crate) metrics_store: Option>>, /// Storage for the completed dynamic filters reported by each worker task. - pub(crate) completed_dynamic_filter_store: Arc>, + pub(crate) completed_dynamic_filter_store: Option>>, } /// Execution state produced by distributed planning (static or dynamic) retained @@ -69,7 +69,7 @@ impl DistributedExec { prepared_execution: Arc::new(Mutex::new(None)), metrics: ExecutionPlanMetricsSet::new(), metrics_store: None, - completed_dynamic_filter_store: Arc::new(Store::new()), + completed_dynamic_filter_store: None, } } @@ -82,6 +82,15 @@ impl DistributedExec { self } + /// Enables collection of completed dynamic filters from remote workers for display. + pub fn with_dynamic_filter_collection(mut self, enabled: bool) -> Self { + self.completed_dynamic_filter_store = match enabled { + true => Some(Arc::new(Store::new())), + false => None, + }; + self + } + /// Waits until all worker tasks have reported their metrics back via the coordinator channel. /// /// Metrics are delivered asynchronously after query execution completes, so callers that need @@ -101,12 +110,12 @@ impl DistributedExec { pub(crate) async fn wait_for_dynamic_filters( &self, - ) -> Result> { + ) -> Result>> { + let Some(store) = &self.completed_dynamic_filter_store else { + return Ok(None); + }; let plan = self.plan_for_viz()?; - Ok(self - .completed_dynamic_filter_store - .wait_for(&task_keys_for_plan(&plan)) - .await) + Ok(Some(store.wait_for(&task_keys_for_plan(&plan)).await)) } fn prepared_execution(&self) -> Result { @@ -124,6 +133,13 @@ impl DistributedExec { Ok(self.prepared_execution()?.plan_for_viz) } + /// Returns the prepared visualization plan when available, or the original optimized plan + /// before execution has prepared one. + pub(crate) fn plan_for_viz_or_base_plan(&self) -> Arc { + self.plan_for_viz() + .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 @@ -138,18 +154,18 @@ impl DistributedExec { /// Builds a non-executable visualization result while preserving the state needed by a /// subsequent rewrite. Dynamic filters must be rewritten before metrics. - pub(crate) fn with_rewritten_plan( + pub(crate) fn with_plan_for_viz( &self, plan_for_viz: Arc, ) -> Result> { let mut prepared_execution = self.prepared_execution()?; prepared_execution.plan_for_viz = Arc::clone(&plan_for_viz); Ok(Arc::new(Self { - base_plan: plan_for_viz, + base_plan: Arc::clone(&self.base_plan), prepared_execution: Arc::new(Mutex::new(Some(prepared_execution))), metrics: self.metrics.clone(), metrics_store: self.metrics_store.clone(), - completed_dynamic_filter_store: Arc::clone(&self.completed_dynamic_filter_store), + completed_dynamic_filter_store: self.completed_dynamic_filter_store.clone(), })) } } @@ -189,7 +205,7 @@ impl ExecutionPlan for DistributedExec { prepared_execution: Arc::new(Mutex::new(None)), metrics: self.metrics.clone(), metrics_store: self.metrics_store.clone(), - completed_dynamic_filter_store: Arc::clone(&self.completed_dynamic_filter_store), + completed_dynamic_filter_store: self.completed_dynamic_filter_store.clone(), })) } @@ -210,12 +226,13 @@ impl ExecutionPlan for DistributedExec { let base_plan = Arc::clone(&self.base_plan); let prepared_execution = Arc::clone(&self.prepared_execution); + let collect_dynamic_filters = self.completed_dynamic_filter_store.is_some(); let query_coordinator = QueryCoordinator::new( Arc::clone(&context), &self.metrics, self.metrics_store.clone(), - Arc::clone(&self.completed_dynamic_filter_store), + self.completed_dynamic_filter_store.clone(), ); let mut builder = RecordBatchReceiverStreamBuilder::new(self.schema(), 1); @@ -240,8 +257,13 @@ impl ExecutionPlan for DistributedExec { false => prepare_static_plan(&query_coordinator, &base_plan)?, }; - let plan_for_viz = - isolate_distributed_leaf_variants_for_display(result.plan_for_viz, &context)?; + let plan_for_viz = match collect_dynamic_filters { + true => sever_dynamic_filter_relationships_in_plan_for_display( + result.plan_for_viz, + &context, + )?, + false => result.plan_for_viz, + }; prepared_execution .lock() .map_err(|e| internal_datafusion_err!("Failed to lock prepared execution: {e}"))? diff --git a/src/coordinator/dynamic_filters.rs b/src/coordinator/dynamic_filters.rs deleted file mode 100644 index bc32852af..000000000 --- a/src/coordinator/dynamic_filters.rs +++ /dev/null @@ -1,196 +0,0 @@ -use crate::codec::decode_physical_expr; -use crate::common::discover_dynamic_filter_consumers; -use crate::coordinator::DistributedExec; -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::{DeduplicatingProtoConverter, PhysicalPlanNodeExt}; -use datafusion_proto::protobuf::PhysicalPlanNode; -use std::sync::Arc; - -/// Rewrites an executed distributed plan with the completed dynamic filters reported by its -/// worker tasks. -/// -/// When composing this with [`crate::rewrite_distributed_plan_with_metrics`], dynamic filters must -/// be rewritten first. -pub async fn rewrite_distributed_plan_with_dynamic_filters( - plan: Arc, -) -> Result> { - let Some(distributed_exec) = plan.downcast_ref::() else { - return Ok(plan); - }; - - let plan_for_viz = distributed_exec.plan_for_viz()?; - let task_ctx = distributed_exec.task_ctx()?; - let reports = distributed_exec.wait_for_dynamic_filters().await?; - let plan_for_viz = isolate_distributed_leaf_variants_for_display(plan_for_viz, &task_ctx)?; - apply_reports_to_distributed_leaves(&plan_for_viz, &reports, &task_ctx); - distributed_exec.with_rewritten_plan(plan_for_viz) -} - -/// 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 _ = 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) else { - continue; - }; - for consumer in consumers { - let Some(proto) = updates.get(&consumer.id).copied() else { - continue; - }; - let Ok(reported_expression) = - decode_physical_expr(proto, consumer.input_schema.as_ref(), task_ctx) - 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 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 report = TaskCompletedDynamicFilters { - filters: vec![crate::TaskDynamicFilter { - expression_id: dynamic_filter.expression_id().unwrap(), - expression: crate::codec::encode_physical_expr(&dynamic_filter, &task_ctx)?, - }], - }; - 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/mod.rs b/src/coordinator/mod.rs index 7eb6dbef7..fd7706fda 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,5 +1,4 @@ mod distributed; -mod dynamic_filters; mod latency_metric; mod prepare_dynamic_plan; mod prepare_static_plan; @@ -7,5 +6,4 @@ mod query_coordinator; mod store; pub use distributed::DistributedExec; -pub use dynamic_filters::rewrite_distributed_plan_with_dynamic_filters; pub(crate) use store::Store; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index ff818ad0b..a53809b6a 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -3,6 +3,7 @@ use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; use crate::coordinator::Store; use crate::coordinator::latency_metric::LatencyMetric; +use crate::dynamic_filtering::maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships; use crate::events::{RouteTasksEvent, RouteTasksHandlers}; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; use crate::passthrough_headers::get_passthrough_headers; @@ -48,7 +49,7 @@ pub(super) struct QueryCoordinator { metrics: ExecutionPlanMetricsSet, coordinator_to_worker_metrics: CoordinatorToWorkerMetrics, metrics_store: Option>>, - completed_dynamic_filter_store: Arc>, + completed_dynamic_filter_store: Option>>, end_stream_notifier: Arc, join_set: Mutex>>, } @@ -59,7 +60,7 @@ impl QueryCoordinator { task_ctx: Arc, metrics_set: &ExecutionPlanMetricsSet, metrics_store: Option>>, - completed_dynamic_filter_store: Arc>, + completed_dynamic_filter_store: Option>>, ) -> Self { Self { task_ctx, @@ -130,7 +131,7 @@ pub(super) struct StageCoordinator<'a> { metrics_set: &'a ExecutionPlanMetricsSet, metrics: &'a CoordinatorToWorkerMetrics, metrics_store: &'a Option>>, - completed_dynamic_filter_store: &'a Arc>, + completed_dynamic_filter_store: &'a Option>>, end_stream_notifier: &'a Arc, join_set: &'a Mutex>>, } @@ -241,7 +242,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 completed_dynamic_filter_store = self.completed_dynamic_filter_store.clone(); let (load_info_tx, load_info_rx) = tokio::sync::mpsc::unbounded_channel(); let mut load_info_tx_opt = Some(load_info_tx); @@ -265,13 +266,21 @@ impl<'a> StageCoordinator<'a> { let _ = load_info_tx_opt.take(); } WorkerToCoordinatorMsg::TaskCompletedDynamicFilters(filters) => { - completed_dynamic_filter_store.insert(task_key, filters); + if let Some(store) = &completed_dynamic_filter_store { + store.insert(task_key, filters); + } } } } - if completed_dynamic_filter_store.get(&task_key).is_none() { - completed_dynamic_filter_store - .insert(task_key, TaskCompletedDynamicFilters::default()); + // The store abstraction relies on one entry being present for each task to mark + // completition. Since not all tasks report dynamic filters, add placeholders here. + // + // Also, note that completed_dynamic_filter_store will be None if collection is + // disabled. + if let Some(store) = &completed_dynamic_filter_store + && store.get(&task_key).is_none() + { + store.insert(task_key, TaskCompletedDynamicFilters::default()); } }); load_info_rx @@ -400,7 +409,11 @@ impl<'a> StageCoordinator<'a> { Ok(Transformed::no(plan)) })?; - Ok((transformed.data, work_unit_feed_declarations)) + let plan = maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships( + Arc::clone(&transformed.data), + self.task_ctx, + )?; + Ok((plan, work_unit_feed_declarations)) } /// Returns as many URLs as the task count for the stage this [StageCoordinator] diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index d89f4b4e7..04a600ac4 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -375,6 +375,20 @@ pub trait DistributedExt: Sized { /// Same as [DistributedExt::with_distributed_metrics_collection] but with an in-place mutation. fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>; + /// Collects completed dynamic filters from worker tasks so they can be displayed in the + /// distributed plan. This does not enable or disable dynamic filtering during execution. + fn with_distributed_dynamic_filter_collection( + self, + enabled: bool, + ) -> Result; + + /// Same as [DistributedExt::with_distributed_dynamic_filter_collection] but with an in-place + /// mutation. + fn set_distributed_dynamic_filter_collection( + &mut self, + enabled: bool, + ) -> Result<(), DataFusionError>; + /// Enables children isolator unions for distributing UNION operations across as many tasks as /// the sum of all the tasks required for each child. /// @@ -803,6 +817,15 @@ impl DistributedExt for SessionConfig { Ok(()) } + fn set_distributed_dynamic_filter_collection( + &mut self, + enabled: bool, + ) -> Result<(), DataFusionError> { + let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?; + d_cfg.collect_dynamic_filters = enabled; + Ok(()) + } + fn set_distributed_children_isolator_unions( &mut self, enabled: bool, @@ -957,6 +980,10 @@ impl DistributedExt for SessionConfig { #[expr($?;Ok(self))] fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result; + #[call(set_distributed_dynamic_filter_collection)] + #[expr($?;Ok(self))] + fn with_distributed_dynamic_filter_collection(mut self, enabled: bool) -> Result; + #[call(set_distributed_children_isolator_unions)] #[expr($?;Ok(self))] fn with_distributed_children_isolator_unions(mut self, enabled: bool) -> Result; @@ -1083,6 +1110,11 @@ impl DistributedExt for SessionStateBuilder { #[expr($?;Ok(self))] fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result; + fn set_distributed_dynamic_filter_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[call(set_distributed_dynamic_filter_collection)] + #[expr($?;Ok(self))] + fn with_distributed_dynamic_filter_collection(mut self, enabled: bool) -> Result; + fn set_distributed_children_isolator_unions(&mut self, enabled: bool) -> Result<(), DataFusionError>; #[call(set_distributed_children_isolator_unions)] #[expr($?;Ok(self))] @@ -1233,6 +1265,11 @@ impl DistributedExt for SessionState { #[expr($?;Ok(self))] fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result; + fn set_distributed_dynamic_filter_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[call(set_distributed_dynamic_filter_collection)] + #[expr($?;Ok(self))] + fn with_distributed_dynamic_filter_collection(mut self, enabled: bool) -> Result; + fn set_distributed_children_isolator_unions(&mut self, enabled: bool) -> Result<(), DataFusionError>; #[call(set_distributed_children_isolator_unions)] #[expr($?;Ok(self))] @@ -1376,6 +1413,11 @@ impl DistributedExt for SessionContext { #[expr($?;Ok(self))] fn with_distributed_metrics_collection(self, enabled: bool) -> Result; + fn set_distributed_dynamic_filter_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[call(set_distributed_dynamic_filter_collection)] + #[expr($?;Ok(self))] + fn with_distributed_dynamic_filter_collection(self, enabled: bool) -> Result; + fn set_distributed_children_isolator_unions(&mut self, enabled: bool) -> Result<(), DataFusionError>; #[call(set_distributed_children_isolator_unions)] #[expr($?;Ok(self))] diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index fa330456b..e6f3583b5 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -30,6 +30,10 @@ extensions_options! { /// Propagate collected metrics from all nodes in the plan across network boundaries /// so that they can be reconstructed on the head node of the plan. pub collect_metrics: bool, default = true + /// Collect completed dynamic filters from worker tasks so that they can be displayed in + /// the distributed plan. This does not control whether dynamic filtering is used during + /// query execution. + pub collect_dynamic_filters: bool, default = true /// Enable broadcast joins for CollectLeft hash joins. When enabled, the build side of /// a CollectLeft join is broadcast to all consumer tasks. pub broadcast_joins: bool, default = true diff --git a/src/distributed_planner/distributed_query_planner.rs b/src/distributed_planner/distributed_query_planner.rs index d312b3293..c6c63472f 100644 --- a/src/distributed_planner/distributed_query_planner.rs +++ b/src/distributed_planner/distributed_query_planner.rs @@ -131,9 +131,7 @@ fn create_distributed_plan( return Ok(plan); } let plan = push_fetch_into_network_coalesce(plan)?; - return Ok(Arc::new( - DistributedExec::new(plan).with_metrics_collection(d_cfg.collect_metrics), - )); + return Ok(create_distributed_exec(Arc::clone(&plan), d_cfg)); } let mut plan = Arc::clone(&original_plan); @@ -150,9 +148,7 @@ fn create_distributed_plan( if d_cfg.dynamic_task_count { // The task count will be decided dynamically at execution time. - return Ok(Arc::new( - DistributedExec::new(plan).with_metrics_collection(d_cfg.collect_metrics), - )); + return Ok(create_distributed_exec(Arc::clone(&plan), d_cfg)); } // Compute per-node task counts and inject `Network*Exec` nodes at the stage boundaries. @@ -167,12 +163,18 @@ fn create_distributed_plan( let plan = partial_reduce_below_network_shuffles(plan, cfg)?; let plan = push_fetch_into_network_coalesce(plan)?; - Ok(Arc::new( - DistributedExec::new(plan).with_metrics_collection(d_cfg.collect_metrics), - )) + Ok(create_distributed_exec(Arc::clone(&plan), d_cfg)) }) } +fn create_distributed_exec(plan: Arc, d_cfg: &DistributedConfig) -> Arc { + Arc::new( + DistributedExec::new(plan) + .with_metrics_collection(d_cfg.collect_metrics) + .with_dynamic_filter_collection(d_cfg.collect_dynamic_filters), + ) +} + #[cfg(test)] mod tests { use crate::assert_snapshot; diff --git a/src/common/dynamic_filtering.rs b/src/dynamic_filtering/discovery.rs similarity index 82% rename from src/common/dynamic_filtering.rs rename to src/dynamic_filtering/discovery.rs index 0d2fba113..2c14c1d2f 100644 --- a/src/common/dynamic_filtering.rs +++ b/src/dynamic_filtering/discovery.rs @@ -6,7 +6,7 @@ use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use std::sync::Arc; -/// A dynamic-filter consumer discovered in an execution plan along with the schema its evaluated +/// A dynamic-filter consumer discovered in an execution plan along with the schema it is evaluated /// against. #[derive(Clone)] pub(crate) struct DiscoveredDynamicFilter { @@ -74,6 +74,33 @@ pub(crate) fn discover_dynamic_filter_consumers( Ok(consumers) } +/// Returns whether `plan` contains only the consumer side of a dynamic filter +/// relationship. +pub(crate) fn has_nonlocal_dynamic_filter_relationships( + plan: &Arc, +) -> Result { + let consumer_ids: HashSet<_> = discover_dynamic_filter_consumers(plan)? + .into_iter() + .map(|consumer| consumer.id) + .collect(); + + let mut producer_ids = HashSet::new(); + plan.apply(|node| { + for produced in node.dynamic_expressions_produced() { + let Some(id) = produced.expression_id() else { + return internal_err!( + "{}::dynamic_expressions_produced returned an expression without an expression ID", + node.name() + ); + }; + producer_ids.insert(id); + } + Ok(TreeNodeRecursion::Continue) + })?; + + Ok(consumer_ids != producer_ids) +} + #[cfg(test)] mod tests { use super::*; @@ -115,6 +142,7 @@ mod tests { let discovered = discover_dynamic_filter_consumers(&plan)?; assert_eq!(discovered.len(), 1); assert_eq!(discovered[0].id, dynamic_filter.expression_id().unwrap()); + assert!(!has_nonlocal_dynamic_filter_relationships(&plan)?); dynamic_filter .downcast_ref::() @@ -156,6 +184,24 @@ mod tests { assert_eq!(discovered.len(), 1); assert_eq!(discovered[0].id, dynamic_filter.expression_id().unwrap()); + assert!(has_nonlocal_dynamic_filter_relationships(&plan)?); + Ok(()) + } + + #[test] + fn identifies_a_producer_without_a_local_consumer() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("a", 0))], + lit(true), + )) as Arc; + let plan = Arc::new(ExpressionExec::new( + Arc::new(EmptyExec::new(schema)), + dynamic_filter, + true, + )) as Arc; + + assert!(has_nonlocal_dynamic_filter_relationships(&plan)?); Ok(()) } diff --git a/src/dynamic_filtering/display.rs b/src/dynamic_filtering/display.rs new file mode 100644 index 000000000..522313e83 --- /dev/null +++ b/src/dynamic_filtering/display.rs @@ -0,0 +1,356 @@ +use crate::codec::decode_physical_expr; +use crate::coordinator::DistributedExec; +use crate::dynamic_filtering::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::empty::EmptyExec; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::{ + ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, ReplaceChildrenOptions, +}; +use datafusion_proto::physical_plan::{DeduplicatingProtoConverter, PhysicalPlanNodeExt}; +use datafusion_proto::protobuf::PhysicalPlanNode; +use datafusion_proto::protobuf::physical_expr_node::ExprType; +use std::sync::Arc; + +/// Rewrites an executed distributed plan with the dynamic filters reported by its completed +/// worker tasks. +/// +/// When composing this with [`crate::rewrite_distributed_plan_with_metrics`], dynamic filters must +/// be rewritten first. +pub async fn rewrite_distributed_plan_with_dynamic_filters( + plan: Arc, +) -> Result> { + let Some(distributed_exec) = plan.downcast_ref::() else { + return Ok(plan); + }; + + let Some(reports) = distributed_exec.wait_for_dynamic_filters().await? else { + return Ok(plan); + }; + let plan_for_viz = distributed_exec.plan_for_viz()?; + let task_ctx = distributed_exec.task_ctx()?; + // Avoids mutating the `plan_for_viz` of the incoming DistributedExec. + let plan_for_viz = + sever_dynamic_filter_relationships_in_plan_for_display(plan_for_viz, &task_ctx)?; + apply_reports_to_distributed_leaves(&plan_for_viz, &reports, &task_ctx); + distributed_exec.with_plan_for_viz(plan_for_viz) +} + +/// The purpose of this function is to sever dynamic filter connections so we can update +/// filter values for display purposes without having an update in one node propagate to another. +/// +/// For example, in this plan, we would like to be able to [`update()`] every variant independently +/// without mutating the producer or other variants (some producers like [`SortExec`] display their +/// dynamic filters). +/// +/// ```text +/// RepartitionExec: +/// AggregateExec: mode=Partial +/// HashJoinExec: mode=Partitioned +/// DistributedLeafExec: +/// t0: DataSourceExec: ... +/// t1: DataSourceExec: ... +/// DistributedLeafExec: +/// t0: DataSourceExec: predicate=DynamicFilter [ f_dkey@2 >= A AND f_dkey@2 <= A AND f_dkey@2 IN (SET) ([]) ] <- unique filter +/// t1: DataSourceExec: predicate=DynamicFilter [ f_dkey@2 >= B AND f_dkey@2 <= B AND f_dkey@2 IN (SET) ([]) ] <- unique filter +/// ``` +/// +/// This is done by deep-copying every leaf variant and every [`SortExec`] so we don't have to +/// worry about any shared state. +/// +/// [`update()`]: DynamicFilterPhysicalExpr::update() +pub(crate) fn sever_dynamic_filter_relationships_in_plan_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, + )?; + let variant = + proto.try_into_physical_plan_with_converter(task_ctx, &codec, &converter)?; + isolate_sort_dynamic_filters_for_display(variant, task_ctx) + }) + .collect::>>()?; + + Ok(Transformed::yes(Arc::new(DistributedLeafExec::try_new( + Arc::clone(leaf.original()), + variants, + )?) as Arc)) + }) + .and_then(|transformed| isolate_sort_dynamic_filters_for_display(transformed.data, task_ctx)) +} + +/// Deep-copies dynamic-filter-producing sorts without serializing their input subtrees. +/// +/// A `SortExec` requires one child, so each sort is temporarily given an `EmptyExec` with the +/// input's schema and partition count. Only that two-node plan is round-tripped through protobuf; +/// the real child is restored afterwards: +/// +/// ```text +/// SortExec SortExec isolated SortExec +/// real child -> EmptyExec -> real child +/// protobuf roundtrip restore + recompute properties +/// ``` +/// +/// The roundtrip preserves the filter value and expression ID while severing its runtime state +/// from the executable plan. Recomputing properties after restoring the real child is required +/// because the placeholder does not preserve its ordering, equivalences, boundedness, or +/// partitioning kind. +/// +/// See . +fn isolate_sort_dynamic_filters_for_display( + plan: Arc, + task_ctx: &Arc, +) -> Result> { + let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config()); + plan.transform_up(|node| { + let Some(sort) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + if node.dynamic_expressions_produced().is_empty() { + return Ok(Transformed::no(node)); + } + + let input = Arc::clone(sort.input()); + let placeholder = Arc::new( + EmptyExec::new(input.schema()) + .with_partitions(input.output_partitioning().partition_count()), + ) as Arc; + let recompute = ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute); + let sort_with_placeholder = node.replace_children(vec![placeholder], recompute)?; + + let converter = DeduplicatingProtoConverter::default(); + let proto = PhysicalPlanNode::try_from_physical_plan_with_converter( + sort_with_placeholder, + &codec, + &converter, + )?; + let isolated = proto.try_into_physical_plan_with_converter(task_ctx, &codec, &converter)?; + let isolated = isolated.replace_children(vec![input], recompute)?; + + Ok(Transformed::yes(isolated)) + }) + .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 _ = 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) else { + continue; + }; + for consumer in consumers { + let Some(proto) = updates.get(&consumer.id).copied() else { + continue; + }; + let Some(ExprType::DynamicFilter(dynamic_filter_proto)) = proto.expr_type.as_ref() + else { + continue; + }; + let Ok(reported_expression) = + decode_physical_expr(proto, consumer.input_schema.as_ref(), task_ctx) + 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_proto.generation > 1 { + let _ = dynamic_filter.update(expression); + } + } + } + + Ok(TreeNodeRecursion::Continue) + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::mock_exec::MockExec; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{ + BinaryExpr, Column, DynamicFilterPhysicalExpr, lit, + }; + use datafusion::physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr}; + use datafusion::physical_plan::displayable; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::filter::FilterExec; + use datafusion::physical_plan::sorts::sort::SortExec; + use datafusion::prelude::SessionContext; + 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 = sever_dynamic_filter_relationships_in_plan_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 report = TaskCompletedDynamicFilters { + filters: vec![crate::TaskDynamicFilter { + expression_id: dynamic_filter.expression_id().unwrap(), + expression: crate::codec::encode_physical_expr(&dynamic_filter, &task_ctx)?, + }], + }; + 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(()) + } + + #[test] + fn visualization_isolates_sort_producer_filter_without_mutating_original() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + // MockExec has no protobuf representation. A successful isolation therefore proves the + // real child was replaced before serializing the SortExec. + let input = Arc::new(MockExec::new_partitioned(vec![vec![], vec![]], schema)) + as Arc; + let ordering = + LexOrdering::new([PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))]) + .unwrap(); + let sort = Arc::new( + SortExec::new(ordering, Arc::clone(&input)) + .with_fetch(Some(10)) + .with_preserve_partitioning(true), + ) as Arc; + let produced = sort.dynamic_expressions_produced(); + let dynamic_filter = produced[0] + .downcast_ref::() + .unwrap(); + let expression_id = dynamic_filter.expression_id(); + dynamic_filter.update(Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + lit(10_i32), + )))?; + + let task_ctx = SessionContext::new().task_ctx(); + let isolated = isolate_sort_dynamic_filters_for_display(Arc::clone(&sort), &task_ctx)?; + let isolated_sort = isolated.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(isolated_sort.input(), &input)); + assert_eq!(isolated_sort.fetch(), Some(10)); + assert!(isolated_sort.preserve_partitioning()); + assert_eq!(isolated.schema(), sort.schema()); + assert_eq!( + isolated.output_partitioning().partition_count(), + sort.output_partitioning().partition_count() + ); + assert_eq!( + isolated.dynamic_expressions_produced()[0].expression_id(), + expression_id + ); + assert!( + displayable(isolated.as_ref()) + .one_line() + .to_string() + .contains("filter=[a@0 > 10]") + ); + + dynamic_filter.update(Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + lit(20_i32), + )))?; + assert!( + displayable(sort.as_ref()) + .one_line() + .to_string() + .contains("filter=[a@0 > 20]") + ); + assert!( + displayable(isolated.as_ref()) + .one_line() + .to_string() + .contains("filter=[a@0 > 10]") + ); + Ok(()) + } +} diff --git a/src/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs new file mode 100644 index 000000000..8e5d07996 --- /dev/null +++ b/src/dynamic_filtering/mod.rs @@ -0,0 +1,50 @@ +mod discovery; +mod display; + +use crate::codec::{decode_execution_plan, encode_execution_plan}; +use datafusion::common::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use std::sync::Arc; + +pub(crate) use discovery::*; +pub use display::rewrite_distributed_plan_with_dynamic_filters; +pub(crate) use display::sever_dynamic_filter_relationships_in_plan_for_display; + +// We must take care to avoid partial dynamic filter updates when sending an +// in-memory plan. +// +// Consider this partitioned hash join topology where the consumer task is +// collocated with one producer on worker A: +// ```text +// Worker A +// +// Stage 2 Task 0 +// HashJoinExec <- Dynamic Filter Produced: (foo > 100) +// +// Stage 1 Task 0 +// DataSourceExec <- consumer +// +// Worker B +// Stage 2 Task 1 +// HashJoinExec <- Dynamic Filter Produced: (foo != 150) +// ``` +// +// The in-process transport allows the Worker A join to propagate its filter to +// the consumer and mark it as completed, so the consumer incorrectly applies +// (foo > 100) instead of (foo > 100 OR foo != 150). +// +// In this situation, we roundtrip Stage 1 Task 0 to sever the in-memory +// relationship. The dynamic filter update from the producer must reach the +// coordinator for merging prior to being forwarded to the consumer. +pub(crate) fn maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships( + plan: Arc, + task_ctx: &Arc, +) -> Result> { + if has_nonlocal_dynamic_filter_relationships(&plan)? { + let encoded = encode_execution_plan(plan, task_ctx)?; + decode_execution_plan(&encoded, task_ctx) + } else { + Ok(plan) + } +} diff --git a/src/lib.rs b/src/lib.rs index bf4bf4693..87fd144a4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,7 @@ mod config_extension_ext; mod coordinator; mod distributed_ext; mod distributed_planner; +mod dynamic_filtering; mod execution_plans; mod explain_analyze; mod metrics; @@ -18,11 +19,12 @@ mod worker_resolver; #[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; -pub use coordinator::{DistributedExec, rewrite_distributed_plan_with_dynamic_filters}; +pub use coordinator::DistributedExec; pub use distributed_ext::{DistributedExt, DistributedGetterExt}; pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, ProducerHead, SessionStateBuilderExt, }; +pub use dynamic_filtering::rewrite_distributed_plan_with_dynamic_filters; pub use events::{ DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, RouteTasksEvent, RouteTasksEventResponse, RouteTasksHandler, ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index f072f4a76..987f5ec65 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -89,7 +89,7 @@ pub async fn rewrite_distributed_plan_with_metrics( Ok(Transformed::no(plan)) })?; - distributed_exec.with_rewritten_plan(transformed.data) + distributed_exec.with_plan_for_viz(transformed.data) } /// Extra information for rewriting local plans. diff --git a/src/stage.rs b/src/stage.rs index 37220f044..0e2d4bf1d 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -290,12 +290,12 @@ fn display_ascii( f: &mut String, ) -> std::fmt::Result { let plan = match stage { - Either::Left(distributed_exec) => distributed_exec.children().first().unwrap(), + Either::Left(distributed_exec) => distributed_exec.plan_for_viz_or_base_plan(), Either::Right(stage) => { let Some(plan) = stage.local_plan() else { return write!(f, "StageExec: encoded input plan"); }; - plan + Arc::clone(plan) } }; match stage { @@ -329,7 +329,7 @@ fn display_ascii( HORIZONTAL.repeat(5), stage.num(), HORIZONTAL.repeat(2), - format_tasks_for_stage(stage.task_count(), plan) + format_tasks_for_stage(stage.task_count(), &plan) )?; if show_metrics && let Some(metrics_store) = &root.metrics_store { let metrics = gather_stage_header_metrics(stage, metrics_store); @@ -342,7 +342,7 @@ fn display_ascii( } let mut plan_str = String::new(); - display_inner_ascii(plan, 0, show_metrics, &mut plan_str)?; + display_inner_ascii(&plan, 0, show_metrics, &mut plan_str)?; let plan_str = plan_str .split('\n') .filter(|v| !v.is_empty()) diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index ac7412ce2..a26a5049e 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,5 +1,6 @@ use crate::codec::encode_physical_expr; -use crate::common::{TreeNodeExt, discover_dynamic_filter_consumers}; +use crate::common::TreeNodeExt; +use crate::dynamic_filtering::discover_dynamic_filter_consumers; use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers}; use crate::execution_plans::SamplerExec; use crate::protocol::LocalWorkerContext; @@ -62,6 +63,7 @@ impl Worker { let d_cfg = DistributedConfig::from_config_options(cfg.options())?; let shuffle_batch_size = d_cfg.shuffle_batch_size; let collect_metrics = d_cfg.collect_metrics; + let collect_dynamic_filters = d_cfg.collect_dynamic_filters; if shuffle_batch_size != 0 { cfg = cfg.with_batch_size(shuffle_batch_size); } @@ -97,6 +99,10 @@ impl Worker { true => Arc::new(std::sync::Mutex::new(Some(metrics_tx))), false => Arc::new(std::sync::Mutex::new(None)), }, + completed_dynamic_filters_tx: match collect_dynamic_filters { + true => Arc::new(std::sync::Mutex::new(Some(dynamic_filters_tx))), + false => Arc::new(std::sync::Mutex::new(None)), + }, task_data_metrics: Arc::new(TaskDataMetrics::new(request.query_start_time_ns)), }) }; @@ -161,8 +167,10 @@ impl Worker { } } + // Send metrics and completed dynamic filters if enabled. let metrics_tx = task_data.metrics_tx.lock().unwrap().take(); - let mut dynamic_filters = TaskCompletedDynamicFilters::default(); + let dynamic_filters_tx = task_data.completed_dynamic_filters_tx.lock().unwrap().take(); + let mut dynamic_filters = None; if let Some(Ok(plan)) = task_data.final_plan.get() { let d_ctx = DistributedTaskContext { task_index: key.task_number, @@ -173,10 +181,16 @@ 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, &task_data.task_ctx) - .unwrap_or_default(); + if dynamic_filters_tx.is_some() { + dynamic_filters = Some( + build_task_completed_dynamic_filters(plan, &task_data.task_ctx) + .unwrap_or_default(), + ); + } + } + if let Some(dynamic_filters_tx) = dynamic_filters_tx { + let _ = dynamic_filters_tx.send(dynamic_filters.unwrap_or_default()); } - let _ = dynamic_filters_tx.send(dynamic_filters); task_data_entries.invalidate(&key).await }); @@ -218,26 +232,28 @@ impl Worker { } } +/// Finds all consumed dynamic filters and serializes them. +/// +/// Note that it's possible that a dynamic filter is consumed by the leaf, updated by +/// the producer, then read here, meaning the observed dynamic filter was not +/// necessarily the one applied. This may happen in upstream datafusion as well. +/// Generally this happens because dynamic filter updates happen asynchronously to execution, +/// meaning consumers do not necessarily have to wait for dynamic filters to update / complete +/// before executing. fn build_task_completed_dynamic_filters( plan: &Arc, task_ctx: &Arc, ) -> Result { let mut filters = vec![]; for consumer in discover_dynamic_filter_consumers(plan)? { - // Serializing the complete DynamicFilterPhysicalExpr preserves both its current - // predicate and its completion state through DataFusion's native proto hook. let expression = encode_physical_expr(&consumer.expression, task_ctx)?; - let Some(ExprType::DynamicFilter(dynamic_filter)) = expression.expr_type.as_ref() else { + let Some(ExprType::DynamicFilter(_)) = 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, - }); - } + filters.push(TaskDynamicFilter { + expression_id: consumer.id, + expression, + }); } Ok(TaskCompletedDynamicFilters { filters }) } diff --git a/src/worker/task_data.rs b/src/worker/task_data.rs index 2ff377173..81ad44617 100644 --- a/src/worker/task_data.rs +++ b/src/worker/task_data.rs @@ -1,6 +1,6 @@ use crate::common::OnceLockResult; use crate::common::now_ns; -use crate::{MaxLatencyMetric, ProducerHead, TaskMetrics}; +use crate::{MaxLatencyMetric, ProducerHead, TaskCompletedDynamicFilters, TaskMetrics}; use datafusion::common::{DataFusionError, Result}; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; @@ -22,6 +22,10 @@ pub struct TaskData { /// `Option::take`) when the coordinator channel reaches EOS, sending the collected metrics /// back to the coordinator through the `CoordinatorChannel` side channel. pub(super) metrics_tx: Arc>>>, + /// Sender half of the completed dynamic-filter channel. It is absent when the user does not + /// want to display dynamic filters. + pub(super) completed_dynamic_filters_tx: + Arc>>>, /// Metrics related to the execution of a task within a stage. This metrics, instead of being /// associated to a specific node, they are global to the task, like the time at which the plan /// was fed by the coordinator to the worker. diff --git a/src/worker/test_utils/worker_handles.rs b/src/worker/test_utils/worker_handles.rs index 4fe19c5c8..13e8b7eb7 100644 --- a/src/worker/test_utils/worker_handles.rs +++ b/src/worker/test_utils/worker_handles.rs @@ -207,12 +207,14 @@ pub async fn register_plan_on_worker( .get_with(task_key, async { Default::default() }) .await; let (metrics_tx, _metrics_rx) = tokio::sync::oneshot::channel(); + let (dynamic_filters_tx, _dynamic_filters_rx) = tokio::sync::oneshot::channel(); swmr_task_data .write(Ok(TaskData { task_ctx, base_plan: plan, final_plan: Default::default(), metrics_tx: Arc::new(std::sync::Mutex::new(Some(metrics_tx))), + completed_dynamic_filters_tx: Arc::new(std::sync::Mutex::new(Some(dynamic_filters_tx))), task_data_metrics: Arc::new(TaskDataMetrics::new(0)), })) .expect("failed to write to task data"); diff --git a/tests/dynamic_filtering/.partitioned_join.rs.pending-snap b/tests/dynamic_filtering/.partitioned_join.rs.pending-snap new file mode 100644 index 000000000..27a9158bc --- /dev/null +++ b/tests/dynamic_filtering/.partitioned_join.rs.pending-snap @@ -0,0 +1,13 @@ +{"run_id":"1787768243-897618719","line":21,"new":{"module_name":"dynamic_filtering__partitioned_join__tests","snapshot_name":"local_dynamic_filters","metadata":{"source":"tests/dynamic_filtering/partitioned_join.rs","assertion_line":21,"expression":"display"},"snapshot":"┌───── DistributedExec\n│ CoalescePartitionsExec\n│ [Stage 2] => NetworkCoalesceExec: output_partitions=8, input_tasks=2\n└──────────────────────────────────────────────────\n ┌───── Stage 2 ── tasks=2, partitions=4\n │ ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as n]\n │ AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=2\n └──────────────────────────────────────────────────\n ┌───── Stage 1 ── tasks=2, partitions=8\n │ RepartitionExec: partitioning=Hash([env@0], 8), input_partitions=4\n │ AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@0)], projection=[env@0]\n │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], [], []]}, projection=[env, service, d_dkey], output_partitioning=Hash([d_dkey@2], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ t1: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet], [], []]}, projection=[env, service, d_dkey], output_partitioning=Hash([d_dkey@2], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], [], []]}, projection=[f_dkey], output_ordering=[f_dkey@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@0], 4), file_type=parquet\n │ t1: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet], [], []]}, projection=[f_dkey], output_ordering=[f_dkey@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@0], 4), file_type=parquet\n └──────────────────────────────────────────────────"},"old":{"module_name":"dynamic_filtering__partitioned_join__tests","metadata":{},"snapshot":"┌───── DistributedExec\n│ CoalescePartitionsExec\n│ [Stage 2] => NetworkCoalesceExec: output_partitions=8, input_tasks=2\n└──────────────────────────────────────────────────\n ┌───── Stage 2 ── tasks=2, partitions=4\n │ ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as n]\n │ AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=1\n └──────────────────────────────────────────────────\n ┌───── Stage 1 ── tasks=1, partitions=8\n │ RepartitionExec: partitioning=Hash([env@0], 8), input_partitions=4\n │ AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@0)], projection=[env@0]\n │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(B), (C), (D)], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(B), (C), (D)], 4), file_type=parquet, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN f_dkey@2 >= A AND f_dkey@2 <= A AND f_dkey@2 IN (SET) ([]) WHEN 1 THEN f_dkey@2 >= B AND f_dkey@2 <= B AND f_dkey@2 IN (SET) ([B]) WHEN 2 THEN false ELSE false END ], dynamic_rg_pruning=eligible\n └──────────────────────────────────────────────────"}} +{"run_id":"1787769218-789252811","line":21,"new":{"module_name":"dynamic_filtering__partitioned_join__tests","snapshot_name":"local_dynamic_filters","metadata":{"source":"tests/dynamic_filtering/partitioned_join.rs","assertion_line":21,"expression":"display"},"snapshot":"┌───── DistributedExec\n│ CoalescePartitionsExec\n│ [Stage 2] => NetworkCoalesceExec: output_partitions=4, input_tasks=2\n└──────────────────────────────────────────────────\n ┌───── Stage 2 ── tasks=2, partitions=2\n │ ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as n]\n │ AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ [Stage 1] => NetworkShuffleExec: output_partitions=2, input_tasks=2\n └──────────────────────────────────────────────────\n ┌───── Stage 1 ── tasks=2, partitions=4\n │ RepartitionExec: partitioning=Hash([env@0], 4), input_partitions=2\n │ AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@0)], projection=[env@0]\n │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ t1: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=DynamicFilter [ f_dkey@2 >= A AND f_dkey@2 <= A AND f_dkey@2 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=f_dkey_null_count@1 != row_count@2 AND f_dkey_max@0 >= A AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= A AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= A AND A <= f_dkey_max@0, required_guarantees=[f_dkey in (A)]\n │ t1: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=DynamicFilter [ f_dkey@2 >= B AND f_dkey@2 <= B AND f_dkey@2 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=f_dkey_null_count@1 != row_count@2 AND f_dkey_max@0 >= B AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= B AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= B AND B <= f_dkey_max@0, required_guarantees=[f_dkey in (B)]\n └──────────────────────────────────────────────────"},"old":{"module_name":"dynamic_filtering__partitioned_join__tests","metadata":{},"snapshot":"┌───── DistributedExec\n│ CoalescePartitionsExec\n│ [Stage 2] => NetworkCoalesceExec: output_partitions=8, input_tasks=2\n└──────────────────────────────────────────────────\n ┌───── Stage 2 ── tasks=2, partitions=4\n │ ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as n]\n │ AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=1\n └──────────────────────────────────────────────────\n ┌───── Stage 1 ── tasks=1, partitions=8\n │ RepartitionExec: partitioning=Hash([env@0], 8), input_partitions=4\n │ AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@0)], projection=[env@0]\n │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(B), (C), (D)], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(B), (C), (D)], 4), file_type=parquet, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN f_dkey@2 >= A AND f_dkey@2 <= A AND f_dkey@2 IN (SET) ([]) WHEN 1 THEN f_dkey@2 >= B AND f_dkey@2 <= B AND f_dkey@2 IN (SET) ([B]) WHEN 2 THEN false ELSE false END ], dynamic_rg_pruning=eligible\n └──────────────────────────────────────────────────"}} +{"run_id":"1787769306-743977842","line":21,"new":null,"old":null} +{"run_id":"1787769316-595683950","line":21,"new":null,"old":null} +{"run_id":"1787769316-595683950","line":62,"new":null,"old":null} +{"run_id":"1787771166-861678916","line":21,"new":null,"old":null} +{"run_id":"1787771166-861678916","line":62,"new":null,"old":null} +{"run_id":"1787777444-237809113","line":21,"new":null,"old":null} +{"run_id":"1787777444-237809113","line":62,"new":null,"old":null} +{"run_id":"1788198409-484379595","line":21,"new":null,"old":null} +{"run_id":"1788198409-484379595","line":62,"new":null,"old":null} +{"run_id":"1788202690-432432755","line":21,"new":null,"old":null} +{"run_id":"1788202690-432432755","line":62,"new":null,"old":null} diff --git a/tests/dynamic_filtering/aggregates.rs b/tests/dynamic_filtering/aggregates.rs new file mode 100644 index 000000000..41c3cc4f6 --- /dev/null +++ b/tests/dynamic_filtering/aggregates.rs @@ -0,0 +1,71 @@ +#[cfg(test)] +mod tests { + use crate::common::TestQuery; + use datafusion::common::Result; + use datafusion_distributed::assert_snapshot; + + /// A partial aggregate updates a data source in the same stage. + #[tokio::test] + async fn local_dynamic_filters() -> Result<()> { + let display = TestQuery::new( + r#" + SELECT MIN("MinTemp") + FROM weather + WHERE "RainToday" = 'Yes' + "#, + ) + .execute() + .await?; + assert_snapshot!(display, @" + ┌───── DistributedExec + │ AggregateExec: mode=Final, gby=[], aggr=[min(weather.MinTemp)] + │ CoalescePartitionsExec + │ [Stage 1] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=6 + │ AggregateExec: mode=Partial, gby=[], aggr=[min(weather.MinTemp)] + │ FilterExec: RainToday@1 = Yes, projection=[MinTemp@0] + │ 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=[MinTemp, RainToday], file_type=parquet, predicate=RainToday@19 = Yes AND DynamicFilter [ MinTemp@0 < 4.3 ], dynamic_rg_pruning=eligible, pruning_predicate=RainToday_null_count@2 != row_count@3 AND RainToday_min@0 <= Yes AND Yes <= RainToday_max@1 AND MinTemp_null_count@5 != row_count@3 AND MinTemp_min@4 < 4.3, required_guarantees=[RainToday in (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=[MinTemp, RainToday], file_type=parquet, predicate=RainToday@19 = Yes AND DynamicFilter [ MinTemp@0 < -1.6 ], dynamic_rg_pruning=eligible, pruning_predicate=RainToday_null_count@2 != row_count@3 AND RainToday_min@0 <= Yes AND Yes <= RainToday_max@1 AND MinTemp_null_count@5 != row_count@3 AND MinTemp_min@4 < -1.6, required_guarantees=[RainToday in (Yes)] + └────────────────────────────────────────────────── + "); + Ok(()) + } + + /// A partial aggregate does not yet update a data source across a shuffle. + #[tokio::test] + async fn remote_dynamic_filters() -> Result<()> { + let display = TestQuery::new( + r#" + SELECT MIN(key) + FROM ( + SELECT DISTINCT "MinTemp" AS key + FROM weather + ) + "#, + ) + .execute() + .await?; + assert_snapshot!(display, @" + ┌───── DistributedExec + │ AggregateExec: mode=Final, gby=[], aggr=[min(key)] + │ CoalescePartitionsExec + │ [Stage 2] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=3 + │ AggregateExec: mode=Partial, gby=[], aggr=[min(key)] + │ 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=[MinTemp@0 as key], 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=[MinTemp@0 as key], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + └────────────────────────────────────────────────── + "); + Ok(()) + } +} diff --git a/tests/dynamic_filtering.rs b/tests/dynamic_filtering/collect_left_join.rs similarity index 73% rename from tests/dynamic_filtering.rs rename to tests/dynamic_filtering/collect_left_join.rs index b090fb790..15a1c55eb 100644 --- a/tests/dynamic_filtering.rs +++ b/tests/dynamic_filtering/collect_left_join.rs @@ -1,19 +1,25 @@ -#[cfg(all(feature = "integration", test))] +#[cfg(test)] mod tests { + use crate::common::TestQuery; 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, DistributedMetricsFormat, assert_snapshot, - display_plan_ascii, rewrite_distributed_plan_with_dynamic_filters, - rewrite_distributed_plan_with_metrics, - }; - use std::sync::Arc; + use datafusion_distributed::assert_snapshot; + /// A CollectLeft HashJoinExec producer propagates identical updates to local consumers. #[tokio::test] - async fn collect_left_local_dynamic_filters() -> Result<()> { - let display = execute_local_hash_join(true).await?; + async fn local_dynamic_filters() -> Result<()> { + let display = TestQuery::new( + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + JOIN weather probe ON build.key = probe."RainToday" + "#, + ) + .with_broadcast_joins() + .execute() + .await?; assert_snapshot!(display, @r" ┌───── DistributedExec │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] @@ -46,43 +52,89 @@ mod tests { Ok(()) } + /// A CollectLeft HashJoinExec does not propagate dynamic filters to a remote consumer. #[tokio::test] - async fn partitioned_local_dynamic_filters() -> Result<()> { - let display = execute_local_hash_join(false).await?; + async fn remote_dynamic_filters() -> Result<()> { + let display = TestQuery::new( + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + RIGHT SEMI JOIN ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) probe ON build.key = probe.key + "#, + ) + .with_broadcast_joins() + .execute() + .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 4] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 └────────────────────────────────────────────────── - ┌───── Stage 3 ── tasks=2, partitions=3 + ┌───── Stage 4 ── tasks=2, partitions=3 │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] - │ HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(key@0, RainToday@0)], projection=[] + │ HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(key@0, key@0)], projection=[] + │ CoalescePartitionsExec + │ [Stage 2] => NetworkBroadcastExec: partitions_per_consumer=3, stage_partitions=6, input_tasks=2 │ 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 3] => NetworkShuffleExec: output_partitions=3, input_tasks=2 └────────────────────────────────────────────────── - ┌───── Stage 1 ── tasks=2, partitions=6 + ┌───── 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 + └────────────────────────────────────────────────── + ┌───── Stage 3 ── 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 + │ 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, 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@19 as key], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── "); Ok(()) } #[tokio::test] - async fn collect_left_union_probe_deduplicates_dynamic_filters() -> Result<()> { - let display = execute_local_union_probe_hash_join().await?; + async fn union_probe_deduplicates_dynamic_filters() -> Result<()> { + let display = TestQuery::new( + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "MinTemp" AS key + FROM weather + ) build + JOIN ( + SELECT "MinTemp" FROM weather + UNION ALL + SELECT CAST(-1000.0 AS DOUBLE) AS "MinTemp" + UNION ALL + SELECT "MinTemp" FROM weather + UNION ALL + SELECT CAST(-1001.0 AS DOUBLE) AS "MinTemp" + UNION ALL + SELECT CAST(-1002.0 AS DOUBLE) AS "MinTemp" + ) probe ON build.key = probe."MinTemp" + "#, + ) + .with_broadcast_joins() + .with_one_task_per_leaf() + .execute() + .await?; assert_snapshot!(display, @r" ┌───── DistributedExec │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] @@ -118,87 +170,4 @@ mod tests { "); Ok(()) } - - async fn execute_local_hash_join(broadcast_joins: bool) -> Result { - execute_local_query( - broadcast_joins, - false, - r#" - SELECT COUNT(*) - FROM ( - SELECT DISTINCT "RainToday" AS key - FROM weather - ) build - JOIN weather probe ON build.key = probe."RainToday" - "#, - ) - .await - } - - async fn execute_local_union_probe_hash_join() -> Result { - execute_local_query( - true, - true, - r#" - SELECT COUNT(*) - FROM ( - SELECT DISTINCT "MinTemp" AS key - FROM weather - ) build - JOIN ( - SELECT "MinTemp" FROM weather - UNION ALL - SELECT CAST(-1000.0 AS DOUBLE) AS "MinTemp" - UNION ALL - SELECT "MinTemp" FROM weather - UNION ALL - SELECT CAST(-1001.0 AS DOUBLE) AS "MinTemp" - UNION ALL - SELECT CAST(-1002.0 AS DOUBLE) AS "MinTemp" - ) probe ON build.key = probe."MinTemp" - "#, - ) - .await - } - - async fn execute_local_query( - broadcast_joins: bool, - one_task_per_leaf: bool, - sql: &str, - ) -> Result { - let (ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; - let mut ctx = ctx.with_distributed_broadcast_joins(broadcast_joins)?; - if one_task_per_leaf { - ctx = ctx.with_distributed_desired_task_count_handler(1usize); - } - 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(sql).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 - ); - - let original_display = display_plan_ascii(plan.as_ref(), false); - - let plan_with_dynamic_filters = - rewrite_distributed_plan_with_dynamic_filters(Arc::clone(&plan)).await?; - assert_eq!(display_plan_ascii(plan.as_ref(), false), original_display); - - let plan_with_metrics = rewrite_distributed_plan_with_metrics( - plan_with_dynamic_filters, - DistributedMetricsFormat::Aggregated, - ) - .await?; - Ok(display_plan_ascii(plan_with_metrics.as_ref(), false)) - } } diff --git a/tests/dynamic_filtering/common.rs b/tests/dynamic_filtering/common.rs new file mode 100644 index 000000000..cfbc32985 --- /dev/null +++ b/tests/dynamic_filtering/common.rs @@ -0,0 +1,161 @@ +use datafusion::arrow::datatypes::DataType; +use datafusion::common::{Result, ScalarValue, SplitPoint}; +use datafusion::datasource::file_format::parquet::ParquetFormat; +use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, +}; +use datafusion::logical_expr::{Partitioning, RangePartitioning}; +use datafusion::physical_plan::collect; +use datafusion::prelude::{SessionContext, col}; +use datafusion_distributed::test_utils::localhost::start_localhost_context; +use datafusion_distributed::test_utils::parquet::register_parquet_tables; +use datafusion_distributed::test_utils::routing::url_emitter_route_tasks; +use datafusion_distributed::{ + DefaultSessionBuilder, DistributedExt, DistributedMetricsFormat, display_plan_ascii, + rewrite_distributed_plan_with_dynamic_filters, rewrite_distributed_plan_with_metrics, +}; +use std::sync::Arc; + +pub(crate) struct TestQuery<'a> { + sql: &'a str, + expected_rows: usize, + broadcast_joins: bool, + one_task_per_leaf: bool, + collect_dynamic_filters: bool, +} + +impl<'a> TestQuery<'a> { + pub(crate) fn new(sql: &'a str) -> Self { + Self { + sql, + expected_rows: 1, + broadcast_joins: false, + one_task_per_leaf: false, + collect_dynamic_filters: true, + } + } + + pub(crate) fn with_expected_rows(mut self, expected_rows: usize) -> Self { + self.expected_rows = expected_rows; + self + } + + pub(crate) fn with_broadcast_joins(mut self) -> Self { + self.broadcast_joins = true; + self + } + + pub(crate) fn with_one_task_per_leaf(mut self) -> Self { + self.one_task_per_leaf = true; + self + } + + pub(crate) fn without_dynamic_filter_collection(mut self) -> Self { + self.collect_dynamic_filters = false; + self + } + + pub(crate) async fn execute(self) -> Result { + let (ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; + let mut ctx = ctx + .with_distributed_broadcast_joins(self.broadcast_joins)? + .with_distributed_dynamic_filter_collection(self.collect_dynamic_filters)?; + if self.one_task_per_leaf { + ctx = ctx.with_distributed_desired_task_count_handler(1usize); + } + if !self.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?; + execute_query_and_display( + &ctx, + self.sql, + self.expected_rows, + self.collect_dynamic_filters, + ) + .await + } +} + +pub(crate) async fn execute_range_partitioned_query( + sql: &str, + expected_rows: usize, +) -> Result { + let (ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await; + let ctx = ctx + .with_distributed_broadcast_joins(false)? + .with_distributed_desired_task_count_handler(2usize) + .with_distributed_route_tasks_handler(url_emitter_route_tasks); + { + let state = ctx.state_ref(); + let mut state = state.write(); + let options = state.config_mut().options_mut(); + options.execution.target_partitions = 2; + options.optimizer.hash_join_single_partition_threshold = 0; + options.optimizer.hash_join_single_partition_threshold_rows = 0; + } + + register_range_partitioned_table(&ctx, "dim", "testdata/join/parquet/dim", "d_dkey").await?; + register_range_partitioned_table(&ctx, "fact", "testdata/join/parquet/fact", "f_dkey").await?; + + execute_query_and_display(&ctx, sql, expected_rows, true).await +} + +async fn register_range_partitioned_table( + ctx: &SessionContext, + name: &str, + path: &str, + partition_column: &str, +) -> Result<()> { + let table_url = ListingTableUrl::parse(path)?; + let output_partitioning = Partitioning::Range(RangePartitioning::try_new( + vec![col(partition_column).sort(true, false)], + vec![SplitPoint::new(vec![ScalarValue::Utf8(Some( + "C".to_string(), + ))])], + )?); + let options = ListingOptions::new(Arc::new(ParquetFormat::default())) + .with_table_partition_cols(vec![(partition_column.to_string(), DataType::Utf8)]) + .with_output_partitioning(Some(output_partitioning)); + let config = ListingTableConfig::new(table_url) + .with_listing_options(options) + .infer_schema(&ctx.state()) + .await?; + ctx.register_table(name, Arc::new(ListingTable::try_new(config)?))?; + Ok(()) +} + +async fn execute_query_and_display( + ctx: &SessionContext, + sql: &str, + expected_rows: usize, + collect_dynamic_filters: bool, +) -> Result { + let plan = ctx.sql(sql).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::(), + expected_rows + ); + + let original_display = display_plan_ascii(plan.as_ref(), false); + let plan_with_dynamic_filters = + rewrite_distributed_plan_with_dynamic_filters(Arc::clone(&plan)).await?; + assert_eq!( + Arc::ptr_eq(&plan, &plan_with_dynamic_filters), + !collect_dynamic_filters + ); + assert_eq!(display_plan_ascii(plan.as_ref(), false), original_display); + + let plan_with_metrics = rewrite_distributed_plan_with_metrics( + plan_with_dynamic_filters, + DistributedMetricsFormat::Aggregated, + ) + .await?; + Ok(display_plan_ascii(plan_with_metrics.as_ref(), false)) +} diff --git a/tests/dynamic_filtering/config.rs b/tests/dynamic_filtering/config.rs new file mode 100644 index 000000000..cbe55030f --- /dev/null +++ b/tests/dynamic_filtering/config.rs @@ -0,0 +1,23 @@ +#[cfg(test)] +mod tests { + use crate::common::TestQuery; + use datafusion::common::Result; + + #[tokio::test] + async fn completed_filter_collection_can_be_disabled() -> Result<()> { + TestQuery::new( + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + JOIN weather probe ON build.key = probe."RainToday" + "#, + ) + .without_dynamic_filter_collection() + .execute() + .await?; + Ok(()) + } +} diff --git a/tests/dynamic_filtering/main.rs b/tests/dynamic_filtering/main.rs new file mode 100644 index 000000000..f7bdf8839 --- /dev/null +++ b/tests/dynamic_filtering/main.rs @@ -0,0 +1,6 @@ +mod aggregates; +mod collect_left_join; +mod common; +mod config; +mod partitioned_join; +mod sorts; diff --git a/tests/dynamic_filtering/partitioned_join.rs b/tests/dynamic_filtering/partitioned_join.rs new file mode 100644 index 000000000..b0c530c4c --- /dev/null +++ b/tests/dynamic_filtering/partitioned_join.rs @@ -0,0 +1,92 @@ +#[cfg(test)] +mod tests { + use crate::common::{TestQuery, execute_range_partitioned_query}; + use datafusion::common::Result; + use datafusion_distributed::assert_snapshot; + + /// A Partitioned HashJoinExec propagates dynamic filters to local consumers. + #[tokio::test] + async fn local_dynamic_filters() -> Result<()> { + let display = execute_range_partitioned_query( + r#" + SELECT d.env, COUNT(*) AS n + FROM dim d + JOIN fact f ON d.d_dkey = f.f_dkey + WHERE d.service = 'log' + GROUP BY d.env + "#, + 2, + ) + .await?; + assert_snapshot!(display, @r" + ┌───── DistributedExec + │ CoalescePartitionsExec + │ [Stage 2] => NetworkCoalesceExec: output_partitions=4, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=2 + │ ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as n] + │ AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))] + │ [Stage 1] => NetworkShuffleExec: output_partitions=2, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=4 + │ RepartitionExec: partitioning=Hash([env@0], 4), input_partitions=2 + │ AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))] + │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@0)], projection=[env@0] + │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] + │ t1: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=DynamicFilter [ f_dkey@2 >= A AND f_dkey@2 <= A AND f_dkey@2 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=f_dkey_null_count@1 != row_count@2 AND f_dkey_max@0 >= A AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= A AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= A AND A <= f_dkey_max@0, required_guarantees=[f_dkey in (A)] + │ t1: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=DynamicFilter [ f_dkey@2 >= B AND f_dkey@2 <= B AND f_dkey@2 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=f_dkey_null_count@1 != row_count@2 AND f_dkey_max@0 >= B AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= B AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= B AND B <= f_dkey_max@0, required_guarantees=[f_dkey in (B)] + └────────────────────────────────────────────────── + "); + Ok(()) + } + + /// A Partitioned HashJoinExec does not propagate dynamic filters to a remote consumer. + #[tokio::test] + async fn remote_dynamic_filters() -> Result<()> { + let display = TestQuery::new( + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + JOIN weather probe ON build.key = probe."RainToday" + "#, + ) + .execute() + .await?; + assert_snapshot!(display, @" + ┌───── 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(()) + } +} diff --git a/tests/dynamic_filtering/sorts.rs b/tests/dynamic_filtering/sorts.rs new file mode 100644 index 000000000..48ad97772 --- /dev/null +++ b/tests/dynamic_filtering/sorts.rs @@ -0,0 +1,79 @@ +#[cfg(test)] +mod tests { + use crate::common::TestQuery; + use datafusion::common::Result; + use datafusion_distributed::{assert_snapshot, test_utils::insta::settings}; + + /// A TopK SortExec applies dynamic filters to local data sources. + #[tokio::test] + async fn local_dynamic_filters() -> Result<()> { + let display = TestQuery::new( + r#" + SELECT "MinTemp" + FROM weather + ORDER BY "MinTemp" DESC + LIMIT 10 + "#, + ) + .with_expected_rows(10) + .execute() + .await?; + let mut settings = settings(); + settings.add_filter( + r"(DynamicFilter \[[^\]\n]*? > )-?\d+(?:\.\d+)?( \])", + "${1}${2}", + ); + settings.add_filter(r"(_max@\d+ > )-?\d+(?:\.\d+)?", "${1}"); + settings.bind(|| assert_snapshot!(display, @" + ┌───── DistributedExec + │ SortPreservingMergeExec: [MinTemp@0 DESC], fetch=10 + │ [Stage 1] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=6 + │ SortExec: TopK(fetch=10), expr=[MinTemp@0 DESC], preserve_partitioning=[true] + │ 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=[MinTemp], file_type=parquet, predicate=DynamicFilter [ MinTemp@0 IS NULL OR MinTemp@0 > ], dynamic_rg_pruning=eligible, pruning_predicate=MinTemp_null_count@0 > 0 OR MinTemp_null_count@0 != row_count@2 AND MinTemp_max@1 > , required_guarantees=[] + │ 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=[MinTemp], file_type=parquet, predicate=DynamicFilter [ MinTemp@0 IS NULL OR MinTemp@0 > ], dynamic_rg_pruning=eligible, pruning_predicate=MinTemp_null_count@0 > 0 OR MinTemp_null_count@0 != row_count@2 AND MinTemp_max@1 > , required_guarantees=[] + └────────────────────────────────────────────────── + ")); + Ok(()) + } + + /// A TopK sort does not yet update dynamic filters to remote consumers. + #[tokio::test] + async fn remote_dynamic_filters() -> Result<()> { + let display = TestQuery::new( + r#" + SELECT key + FROM ( + SELECT DISTINCT "MinTemp" AS key + FROM weather + ) + ORDER BY key DESC + LIMIT 10 + "#, + ) + .with_expected_rows(10) + .execute() + .await?; + assert_snapshot!(display, @" + ┌───── DistributedExec + │ SortPreservingMergeExec: [key@0 DESC], fetch=10 + │ [Stage 2] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=3 + │ SortExec: TopK(fetch=10), expr=[key@0 DESC], preserve_partitioning=[true] + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[], lim=[10] + │ [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=[], lim=[10] + │ 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=[MinTemp@0 as key], 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=[MinTemp@0 as key], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + └────────────────────────────────────────────────── + "); + Ok(()) + } +} From 96ec2c2410538ec43496b474eff40e9a930463c4 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Mon, 31 Aug 2026 20:14:44 +0000 Subject: [PATCH 15/17] handle error --- src/worker/impl_coordinator_channel.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index a26a5049e..f6cabb13b 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -168,6 +168,7 @@ impl Worker { } // Send metrics and completed dynamic filters if enabled. + // TODO(#686): handle errors let metrics_tx = task_data.metrics_tx.lock().unwrap().take(); let dynamic_filters_tx = task_data.completed_dynamic_filters_tx.lock().unwrap().take(); let mut dynamic_filters = None; @@ -182,6 +183,7 @@ impl Worker { send_metrics_via_channel(metrics_tx, plan, d_ctx, task_data_metrics); } if dynamic_filters_tx.is_some() { + // TODO(#686): handle error dynamic_filters = Some( build_task_completed_dynamic_filters(plan, &task_data.task_ctx) .unwrap_or_default(), @@ -189,6 +191,7 @@ impl Worker { } } if let Some(dynamic_filters_tx) = dynamic_filters_tx { + // TODO(#686): handle error let _ = dynamic_filters_tx.send(dynamic_filters.unwrap_or_default()); } task_data_entries.invalidate(&key).await From 0f3fe492cd1e0e64dde2c937ac34f46e2c9b4373 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Mon, 31 Aug 2026 20:57:44 +0000 Subject: [PATCH 16/17] unnest proto --- src/common/maybe_encoded.rs | 50 +++++++++++++++++++++++--- src/dynamic_filtering/display.rs | 45 +++++++++++------------ src/protocol/grpc/generated/worker.rs | 21 +++++------ src/protocol/grpc/worker.proto | 12 +++---- src/protocol/grpc/worker_client.rs | 12 +++---- src/protocol/grpc/worker_service.rs | 45 ++++++++++++++++------- src/protocol/worker_channel.rs | 6 ++-- src/worker/impl_coordinator_channel.rs | 31 +++++++--------- 8 files changed, 137 insertions(+), 85 deletions(-) diff --git a/src/common/maybe_encoded.rs b/src/common/maybe_encoded.rs index faef314c1..7e01ad2b3 100644 --- a/src/common/maybe_encoded.rs +++ b/src/common/maybe_encoded.rs @@ -1,17 +1,21 @@ use crate::codec::{ - decode_execution_plan, decode_partitioning, encode_execution_plan, encode_partitioning, + decode_execution_plan, decode_partitioning, decode_physical_expr, encode_execution_plan, + encode_partitioning, encode_physical_expr, }; -use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::datatypes::{Schema, SchemaRef}; use datafusion::common::{Result, internal_err}; use datafusion::execution::TaskContext; -use datafusion::physical_expr::Partitioning; +use datafusion::physical_expr::{Partitioning, PhysicalExpr}; use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::protobuf::PhysicalExprNode; +use datafusion_proto::protobuf::proto_error; +use prost::Message; use std::sync::Arc; /// A value that a transport may either leave encoded or materialize in memory. /// Users are free to pass [MaybeEncoded::Encoded] or [MaybeEncoded::Decoded] at any /// moment and Distributed DataFusion's code will internally know how to handle it. -#[derive(Clone)] +#[derive(Clone, Debug)] pub enum MaybeEncoded { Encoded(Vec), Decoded(T), @@ -82,6 +86,44 @@ impl MaybeEncoded { } } +impl MaybeEncoded> { + /// Returns the encoded [`PhysicalExpr`] as protobuf bytes: + /// - If in `Decoded` state, it encodes it using the codecs registered in the [`TaskContext`]. + /// - If in `Encoded` state, it passes through the existing bytes. + pub fn encode(self, ctx: &Arc) -> Result> { + match self { + Self::Encoded(encoded) => Ok(encoded), + Self::Decoded(expression) => { + Ok(encode_physical_expr(&expression, ctx)?.encode_to_vec()) + } + } + } + + /// Returns the decoded [`PhysicalExpr`]. + /// - If in `Decoded` state, it passes through the expression. + /// - If in `Encoded` state, it decodes it using the provided schema and task context. + pub fn decode( + self, + input_schema: &Schema, + task_ctx: &TaskContext, + ) -> Result> { + self.decode_with(|encoded| { + let proto = PhysicalExprNode::decode(encoded.as_slice()) + .map_err(|error| proto_error(error.to_string()))?; + decode_physical_expr(&proto, input_schema, task_ctx) + }) + } + + /// Materializes the expression's protobuf representation without changing the stored form. + pub(crate) fn to_proto(&self, task_ctx: &TaskContext) -> Result { + match self { + Self::Encoded(encoded) => PhysicalExprNode::decode(encoded.as_slice()) + .map_err(|error| proto_error(error.to_string())), + Self::Decoded(expression) => encode_physical_expr(expression, task_ctx), + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/dynamic_filtering/display.rs b/src/dynamic_filtering/display.rs index 522313e83..2826b3919 100644 --- a/src/dynamic_filtering/display.rs +++ b/src/dynamic_filtering/display.rs @@ -41,12 +41,11 @@ pub async fn rewrite_distributed_plan_with_dynamic_filters( distributed_exec.with_plan_for_viz(plan_for_viz) } -/// The purpose of this function is to sever dynamic filter connections so we can update -/// filter values for display purposes without having an update in one node propagate to another. +/// Severs dynamic filter connections so we can update filter values for +/// display purposes without having an update in one node propagate to another. /// /// For example, in this plan, we would like to be able to [`update()`] every variant independently -/// without mutating the producer or other variants (some producers like [`SortExec`] display their -/// dynamic filters). +/// without mutating the producer or other variants /// /// ```text /// RepartitionExec: @@ -60,7 +59,7 @@ pub async fn rewrite_distributed_plan_with_dynamic_filters( /// t1: DataSourceExec: predicate=DynamicFilter [ f_dkey@2 >= B AND f_dkey@2 <= B AND f_dkey@2 IN (SET) ([]) ] <- unique filter /// ``` /// -/// This is done by deep-copying every leaf variant and every [`SortExec`] so we don't have to +/// This is done by deep-copying every leaf variant so we don't have to /// worry about any shared state. /// /// [`update()`]: DynamicFilterPhysicalExpr::update() @@ -79,6 +78,7 @@ pub(crate) fn sever_dynamic_filter_relationships_in_plan_for_display( .variants() .iter() .map(|variant| { + // Proto roundtrip is used to deep copy the variant. let proto = PhysicalPlanNode::try_from_physical_plan_with_converter( Arc::clone(variant), &codec, @@ -86,6 +86,7 @@ pub(crate) fn sever_dynamic_filter_relationships_in_plan_for_display( )?; let variant = proto.try_into_physical_plan_with_converter(task_ctx, &codec, &converter)?; + // The variant can have a SortExec isolate_sort_dynamic_filters_for_display(variant, task_ctx) }) .collect::>>()?; @@ -95,27 +96,23 @@ pub(crate) fn sever_dynamic_filter_relationships_in_plan_for_display( variants, )?) as Arc)) }) + // Handle SortExec nodes not inside variants. .and_then(|transformed| isolate_sort_dynamic_filters_for_display(transformed.data, task_ctx)) } -/// Deep-copies dynamic-filter-producing sorts without serializing their input subtrees. +/// Deep-copies dynamic-filter-producing [`SortExec`]s by doing a proto roundtrip. Some producers +/// like [`SortExec`] display their dynamic filters, so we need to explicitly handle +/// displaying different dynamic filters for each task containing a [`SortExec`]. For now, we +/// just clear the dynamic filter and don't display it. /// -/// A `SortExec` requires one child, so each sort is temporarily given an `EmptyExec` with the -/// input's schema and partition count. Only that two-node plan is round-tripped through protobuf; -/// the real child is restored afterwards: +/// To avoid serializing an entire subtree, we swap in an [`EmptyExec`]: /// /// ```text -/// SortExec SortExec isolated SortExec -/// real child -> EmptyExec -> real child -/// protobuf roundtrip restore + recompute properties +/// SortExec SortExec SortExec +/// ...children -> EmptyExec -> ...children /// ``` /// -/// The roundtrip preserves the filter value and expression ID while severing its runtime state -/// from the executable plan. Recomputing properties after restoring the real child is required -/// because the placeholder does not preserve its ordering, equivalences, boundedness, or -/// partitioning kind. -/// -/// See . +/// TODO(#677): display producer dynamic filters fn isolate_sort_dynamic_filters_for_display( plan: Arc, task_ctx: &Arc, @@ -175,7 +172,10 @@ pub(super) fn apply_reports_to_distributed_leaves( continue; }; for consumer in consumers { - let Some(proto) = updates.get(&consumer.id).copied() else { + let Some(expression) = updates.get(&consumer.id).copied() else { + continue; + }; + let Ok(proto) = expression.to_proto(task_ctx) else { continue; }; let Some(ExprType::DynamicFilter(dynamic_filter_proto)) = proto.expr_type.as_ref() @@ -183,7 +183,7 @@ pub(super) fn apply_reports_to_distributed_leaves( continue; }; let Ok(reported_expression) = - decode_physical_expr(proto, consumer.input_schema.as_ref(), task_ctx) + decode_physical_expr(&proto, consumer.input_schema.as_ref(), task_ctx) else { continue; }; @@ -215,6 +215,7 @@ pub(super) fn apply_reports_to_distributed_leaves( mod tests { use super::*; use crate::test_utils::mock_exec::MockExec; + use crate::{MaybeEncoded, TaskDynamicFilter}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::logical_expr::Operator; use datafusion::physical_expr::expressions::{ @@ -259,9 +260,9 @@ mod tests { .unwrap() .mark_complete(); let report = TaskCompletedDynamicFilters { - filters: vec![crate::TaskDynamicFilter { + filters: vec![TaskDynamicFilter { expression_id: dynamic_filter.expression_id().unwrap(), - expression: crate::codec::encode_physical_expr(&dynamic_filter, &task_ctx)?, + expression: MaybeEncoded::Decoded(Arc::clone(&dynamic_filter)), }], }; let reports = HashMap::from_iter([( diff --git a/src/protocol/grpc/generated/worker.rs b/src/protocol/grpc/generated/worker.rs index 73368b080..78bbd1058 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/grpc/generated/worker.rs @@ -59,21 +59,18 @@ pub mod worker_to_coordinator_msg { TaskCompletedDynamicFilters(super::TaskCompletedDynamicFilters), } } +#[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)] 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, - } + pub filters: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct TaskMetrics { diff --git a/src/protocol/grpc/worker.proto b/src/protocol/grpc/worker.proto index 52b668bb9..61319f5f8 100644 --- a/src/protocol/grpc/worker.proto +++ b/src/protocol/grpc/worker.proto @@ -56,13 +56,13 @@ message WorkerToCoordinatorMsg { } } -message TaskCompletedDynamicFilters { - message DynamicFilter { - uint64 expression_id = 1; - // Serialized datafusion.proto.PhysicalExprNode. - bytes expression_proto = 2; - } +message DynamicFilter { + uint64 expression_id = 1; + // Serialized datafusion.proto.PhysicalExprNode. + bytes expression_proto = 2; +} +message TaskCompletedDynamicFilters { repeated DynamicFilter filters = 1; } diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index 7196561eb..81af88ab7 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -25,7 +25,6 @@ use datafusion::execution::TaskContext; use datafusion::execution::memory_pool::MemoryConsumer; use datafusion::physical_expr_common::metrics::{Count, Label, 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}; @@ -525,14 +524,11 @@ fn decode_task_completed_dynamic_filters( 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)))?, - }) + .map(|filter| TaskDynamicFilter { + expression_id: filter.expression_id, + expression: MaybeEncoded::Encoded(filter.expression_proto), }) - .collect::>()?, + .collect(), }) } diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index 4b32f91fa..f0231a635 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -20,7 +20,7 @@ use datafusion::arrow::array::{Array, AsArray, RecordBatch, RecordBatchOptions}; use datafusion::arrow::ipc::CompressionType; use datafusion::arrow::ipc::writer::IpcWriteOptions; use datafusion::common::DataFusionError; -use datafusion::execution::SendableRecordBatchStream; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; use prost::Message; @@ -107,6 +107,14 @@ impl pb::worker_service_server::WorkerService for Worker { }; let set_plan_request = decode_set_plan_request(set_plan_request)?; + let task_key = set_plan_request.task_key; + // Dynamic-filter reports may carry decoded physical expressions. Retain the worker's + // task data so the gRPC boundary can encode them with its configured codecs, even after + // the completed task has been removed from the worker cache. + let task_data_entry = self + .task_data_entries + .get_with(task_key, async { Default::default() }) + .await; let input_stream = body .map_err(map_status_to_datafusion_error) @@ -118,9 +126,15 @@ impl pb::worker_service_server::WorkerService for Worker { let output_stream = self .coordinator_channel(metadata.into_headers(), set_plan_request, input_stream) .await - .map_err(datafusion_error_to_tonic_status)? - .map(|msg| match msg { - Ok(msg) => encode_worker_to_coordinator_msg(msg), + .map_err(datafusion_error_to_tonic_status)?; + let task_data = task_data_entry + .read_now() + .ok_or_else(|| Status::internal("worker task data was not initialized"))? + .map_err(|error| datafusion_error_to_tonic_status(DataFusionError::Shared(error)))?; + let task_ctx = Arc::clone(&task_data.task_ctx); + let output_stream = output_stream + .map(move |msg| match msg { + Ok(msg) => encode_worker_to_coordinator_msg(msg, &task_ctx), Err(err) => Err(datafusion_error_to_tonic_status(err)), }) .boxed(); @@ -258,6 +272,7 @@ pub(super) fn decode_producer_head(proto: pb::execute_task_request::ProducerHead fn encode_worker_to_coordinator_msg( msg: WorkerToCoordinatorMsg, + task_ctx: &Arc, ) -> Result { Ok(pb::WorkerToCoordinatorMsg { inner: Some(match msg { @@ -274,7 +289,7 @@ fn encode_worker_to_coordinator_msg( } WorkerToCoordinatorMsg::TaskCompletedDynamicFilters(filters) => { pb::worker_to_coordinator_msg::Inner::TaskCompletedDynamicFilters( - encode_task_completed_dynamic_filters(filters), + encode_task_completed_dynamic_filters(filters, task_ctx)?, ) } }), @@ -283,17 +298,23 @@ fn encode_worker_to_coordinator_msg( fn encode_task_completed_dynamic_filters( filters: TaskCompletedDynamicFilters, -) -> pb::TaskCompletedDynamicFilters { - pb::TaskCompletedDynamicFilters { + task_ctx: &Arc, +) -> Result { + Ok(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(), + .map(|filter| { + Ok(pb::DynamicFilter { + expression_id: filter.expression_id, + expression_proto: filter + .expression + .encode(task_ctx) + .map_err(datafusion_error_to_tonic_status)?, + }) }) - .collect(), - } + .collect::>()?, + }) } fn encode_task_metrics(task_metrics: TaskMetrics) -> Result { diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index 03f0ebb82..d702818f2 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -3,9 +3,9 @@ use async_trait::async_trait; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::Result; use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use datafusion_proto::protobuf::PhysicalExprNode; use futures::stream::BoxStream; use http::HeaderMap; use std::sync::Arc; @@ -140,8 +140,8 @@ pub struct TaskCompletedDynamicFilters { #[derive(Clone, Debug)] pub struct TaskDynamicFilter { pub expression_id: u64, - /// A `DynamicFilterPhysicalExpr` proto containing its final predicate and completion state. - pub expression: PhysicalExprNode, + /// A `DynamicFilterPhysicalExpr` containing its final predicate and completion state. + pub expression: MaybeEncoded>, } #[derive(Clone, Debug)] diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index f6cabb13b..b795a1adf 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,4 +1,3 @@ -use crate::codec::encode_physical_expr; use crate::common::TreeNodeExt; use crate::dynamic_filtering::discover_dynamic_filter_consumers; use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers}; @@ -8,15 +7,14 @@ use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_t use crate::worker::task_data::TaskDataMetrics; use crate::{ CoordinatorToWorkerMsg, DistributedConfig, DistributedExt, DistributedTaskContext, - SetPlanRequest, TaskCompletedDynamicFilters, TaskData, TaskDynamicFilter, TaskMetrics, Worker, - WorkerQueryContext, WorkerToCoordinatorMsg, + MaybeEncoded, SetPlanRequest, 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::execution::{SessionStateBuilder, TaskContext}; +use datafusion::common::{DataFusionError, Result, exec_datafusion_err}; +use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; -use datafusion_proto::protobuf::physical_expr_node::ExprType; use futures::stream::{BoxStream, FuturesUnordered, select_all}; use futures::{FutureExt, StreamExt, TryStreamExt}; use http::HeaderMap; @@ -170,7 +168,11 @@ impl Worker { // Send metrics and completed dynamic filters if enabled. // TODO(#686): handle errors let metrics_tx = task_data.metrics_tx.lock().unwrap().take(); - let dynamic_filters_tx = task_data.completed_dynamic_filters_tx.lock().unwrap().take(); + let dynamic_filters_tx = task_data + .completed_dynamic_filters_tx + .lock() + .unwrap() + .take(); let mut dynamic_filters = None; if let Some(Ok(plan)) = task_data.final_plan.get() { let d_ctx = DistributedTaskContext { @@ -184,10 +186,8 @@ impl Worker { } if dynamic_filters_tx.is_some() { // TODO(#686): handle error - dynamic_filters = Some( - build_task_completed_dynamic_filters(plan, &task_data.task_ctx) - .unwrap_or_default(), - ); + dynamic_filters = + Some(build_task_completed_dynamic_filters(plan).unwrap_or_default()); } } if let Some(dynamic_filters_tx) = dynamic_filters_tx { @@ -235,7 +235,7 @@ impl Worker { } } -/// Finds all consumed dynamic filters and serializes them. +/// Finds all consumed dynamic filters for the completed task report. /// /// Note that it's possible that a dynamic filter is consumed by the leaf, updated by /// the producer, then read here, meaning the observed dynamic filter was not @@ -245,17 +245,12 @@ impl Worker { /// before executing. fn build_task_completed_dynamic_filters( plan: &Arc, - task_ctx: &Arc, ) -> Result { let mut filters = vec![]; for consumer in discover_dynamic_filter_consumers(plan)? { - let expression = encode_physical_expr(&consumer.expression, task_ctx)?; - let Some(ExprType::DynamicFilter(_)) = expression.expr_type.as_ref() else { - return internal_err!("discovered dynamic filter did not serialize as one"); - }; filters.push(TaskDynamicFilter { expression_id: consumer.id, - expression, + expression: MaybeEncoded::Decoded(consumer.expression), }); } Ok(TaskCompletedDynamicFilters { filters }) From 96d9835a1810b9b5c9cc65317b41597bda153ff4 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Mon, 31 Aug 2026 21:48:01 +0000 Subject: [PATCH 17/17] lints and comments --- src/distributed_ext.rs | 2 +- .../distributed_query_planner.rs | 5 ++++- src/metrics/task_metrics_rewriter.rs | 4 +++- .../.partitioned_join.rs.pending-snap | 13 ------------- tests/dynamic_filtering/collect_left_join.rs | 4 ++++ tests/dynamic_filtering/common.rs | 19 +++++++++++-------- 6 files changed, 23 insertions(+), 24 deletions(-) delete mode 100644 tests/dynamic_filtering/.partitioned_join.rs.pending-snap diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 04a600ac4..8c081e101 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -382,7 +382,7 @@ pub trait DistributedExt: Sized { enabled: bool, ) -> Result; - /// Same as [DistributedExt::with_distributed_dynamic_filter_collection] but with an in-place + /// Same as [`DistributedExt::with_distributed_dynamic_filter_collection`] but with an in-place /// mutation. fn set_distributed_dynamic_filter_collection( &mut self, diff --git a/src/distributed_planner/distributed_query_planner.rs b/src/distributed_planner/distributed_query_planner.rs index c6c63472f..c26f8854d 100644 --- a/src/distributed_planner/distributed_query_planner.rs +++ b/src/distributed_planner/distributed_query_planner.rs @@ -167,7 +167,10 @@ fn create_distributed_plan( }) } -fn create_distributed_exec(plan: Arc, d_cfg: &DistributedConfig) -> Arc { +fn create_distributed_exec( + plan: Arc, + d_cfg: &DistributedConfig, +) -> Arc { Arc::new( DistributedExec::new(plan) .with_metrics_collection(d_cfg.collect_metrics) diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index 987f5ec65..e87c4e674 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -602,7 +602,9 @@ mod tests { rewrite_distributed_plan_with_metrics(plan, DistributedMetricsFormat::Aggregated) .await .unwrap(); - assert_metrics_present_in_plan(&rewritten_plan); + assert!(rewritten_plan.metrics().unwrap().iter().count() > 0); + let distributed_exec = rewritten_plan.downcast_ref::().unwrap(); + assert_metrics_present_in_plan(&distributed_exec.plan_for_viz().unwrap()); } #[test] diff --git a/tests/dynamic_filtering/.partitioned_join.rs.pending-snap b/tests/dynamic_filtering/.partitioned_join.rs.pending-snap deleted file mode 100644 index 27a9158bc..000000000 --- a/tests/dynamic_filtering/.partitioned_join.rs.pending-snap +++ /dev/null @@ -1,13 +0,0 @@ -{"run_id":"1787768243-897618719","line":21,"new":{"module_name":"dynamic_filtering__partitioned_join__tests","snapshot_name":"local_dynamic_filters","metadata":{"source":"tests/dynamic_filtering/partitioned_join.rs","assertion_line":21,"expression":"display"},"snapshot":"┌───── DistributedExec\n│ CoalescePartitionsExec\n│ [Stage 2] => NetworkCoalesceExec: output_partitions=8, input_tasks=2\n└──────────────────────────────────────────────────\n ┌───── Stage 2 ── tasks=2, partitions=4\n │ ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as n]\n │ AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=2\n └──────────────────────────────────────────────────\n ┌───── Stage 1 ── tasks=2, partitions=8\n │ RepartitionExec: partitioning=Hash([env@0], 8), input_partitions=4\n │ AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@0)], projection=[env@0]\n │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], [], []]}, projection=[env, service, d_dkey], output_partitioning=Hash([d_dkey@2], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ t1: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet], [], []]}, projection=[env, service, d_dkey], output_partitioning=Hash([d_dkey@2], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], [], []]}, projection=[f_dkey], output_ordering=[f_dkey@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@0], 4), file_type=parquet\n │ t1: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet], [], []]}, projection=[f_dkey], output_ordering=[f_dkey@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@0], 4), file_type=parquet\n └──────────────────────────────────────────────────"},"old":{"module_name":"dynamic_filtering__partitioned_join__tests","metadata":{},"snapshot":"┌───── DistributedExec\n│ CoalescePartitionsExec\n│ [Stage 2] => NetworkCoalesceExec: output_partitions=8, input_tasks=2\n└──────────────────────────────────────────────────\n ┌───── Stage 2 ── tasks=2, partitions=4\n │ ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as n]\n │ AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=1\n └──────────────────────────────────────────────────\n ┌───── Stage 1 ── tasks=1, partitions=8\n │ RepartitionExec: partitioning=Hash([env@0], 8), input_partitions=4\n │ AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@0)], projection=[env@0]\n │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(B), (C), (D)], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(B), (C), (D)], 4), file_type=parquet, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN f_dkey@2 >= A AND f_dkey@2 <= A AND f_dkey@2 IN (SET) ([]) WHEN 1 THEN f_dkey@2 >= B AND f_dkey@2 <= B AND f_dkey@2 IN (SET) ([B]) WHEN 2 THEN false ELSE false END ], dynamic_rg_pruning=eligible\n └──────────────────────────────────────────────────"}} -{"run_id":"1787769218-789252811","line":21,"new":{"module_name":"dynamic_filtering__partitioned_join__tests","snapshot_name":"local_dynamic_filters","metadata":{"source":"tests/dynamic_filtering/partitioned_join.rs","assertion_line":21,"expression":"display"},"snapshot":"┌───── DistributedExec\n│ CoalescePartitionsExec\n│ [Stage 2] => NetworkCoalesceExec: output_partitions=4, input_tasks=2\n└──────────────────────────────────────────────────\n ┌───── Stage 2 ── tasks=2, partitions=2\n │ ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as n]\n │ AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ [Stage 1] => NetworkShuffleExec: output_partitions=2, input_tasks=2\n └──────────────────────────────────────────────────\n ┌───── Stage 1 ── tasks=2, partitions=4\n │ RepartitionExec: partitioning=Hash([env@0], 4), input_partitions=2\n │ AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@0)], projection=[env@0]\n │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ t1: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=DynamicFilter [ f_dkey@2 >= A AND f_dkey@2 <= A AND f_dkey@2 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=f_dkey_null_count@1 != row_count@2 AND f_dkey_max@0 >= A AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= A AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= A AND A <= f_dkey_max@0, required_guarantees=[f_dkey in (A)]\n │ t1: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=DynamicFilter [ f_dkey@2 >= B AND f_dkey@2 <= B AND f_dkey@2 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=f_dkey_null_count@1 != row_count@2 AND f_dkey_max@0 >= B AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= B AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= B AND B <= f_dkey_max@0, required_guarantees=[f_dkey in (B)]\n └──────────────────────────────────────────────────"},"old":{"module_name":"dynamic_filtering__partitioned_join__tests","metadata":{},"snapshot":"┌───── DistributedExec\n│ CoalescePartitionsExec\n│ [Stage 2] => NetworkCoalesceExec: output_partitions=8, input_tasks=2\n└──────────────────────────────────────────────────\n ┌───── Stage 2 ── tasks=2, partitions=4\n │ ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as n]\n │ AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=1\n └──────────────────────────────────────────────────\n ┌───── Stage 1 ── tasks=1, partitions=8\n │ RepartitionExec: partitioning=Hash([env@0], 8), input_partitions=4\n │ AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))]\n │ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@0)], projection=[env@0]\n │ FilterExec: service@1 = log, projection=[env@0, d_dkey@2]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(B), (C), (D)], 4), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)]\n │ DistributedLeafExec:\n │ t0: DataSourceExec: file_groups={4 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(B), (C), (D)], 4), file_type=parquet, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN f_dkey@2 >= A AND f_dkey@2 <= A AND f_dkey@2 IN (SET) ([]) WHEN 1 THEN f_dkey@2 >= B AND f_dkey@2 <= B AND f_dkey@2 IN (SET) ([B]) WHEN 2 THEN false ELSE false END ], dynamic_rg_pruning=eligible\n └──────────────────────────────────────────────────"}} -{"run_id":"1787769306-743977842","line":21,"new":null,"old":null} -{"run_id":"1787769316-595683950","line":21,"new":null,"old":null} -{"run_id":"1787769316-595683950","line":62,"new":null,"old":null} -{"run_id":"1787771166-861678916","line":21,"new":null,"old":null} -{"run_id":"1787771166-861678916","line":62,"new":null,"old":null} -{"run_id":"1787777444-237809113","line":21,"new":null,"old":null} -{"run_id":"1787777444-237809113","line":62,"new":null,"old":null} -{"run_id":"1788198409-484379595","line":21,"new":null,"old":null} -{"run_id":"1788198409-484379595","line":62,"new":null,"old":null} -{"run_id":"1788202690-432432755","line":21,"new":null,"old":null} -{"run_id":"1788202690-432432755","line":62,"new":null,"old":null} diff --git a/tests/dynamic_filtering/collect_left_join.rs b/tests/dynamic_filtering/collect_left_join.rs index 15a1c55eb..00105d363 100644 --- a/tests/dynamic_filtering/collect_left_join.rs +++ b/tests/dynamic_filtering/collect_left_join.rs @@ -132,6 +132,10 @@ mod tests { "#, ) .with_broadcast_joins() + // Indirectly forces the union to put c0 and c2 on the same task. We would like to test + // that consumers of a dynamic filter on the same node get the same filter. + // + // If we don't do this, inject_network_boundaries splits them up to spread out the load. .with_one_task_per_leaf() .execute() .await?; diff --git a/tests/dynamic_filtering/common.rs b/tests/dynamic_filtering/common.rs index cfbc32985..251056611 100644 --- a/tests/dynamic_filtering/common.rs +++ b/tests/dynamic_filtering/common.rs @@ -11,8 +11,8 @@ use datafusion_distributed::test_utils::localhost::start_localhost_context; use datafusion_distributed::test_utils::parquet::register_parquet_tables; use datafusion_distributed::test_utils::routing::url_emitter_route_tasks; use datafusion_distributed::{ - DefaultSessionBuilder, DistributedExt, DistributedMetricsFormat, display_plan_ascii, - rewrite_distributed_plan_with_dynamic_filters, rewrite_distributed_plan_with_metrics, + DefaultSessionBuilder, DistributedExt, display_plan_ascii, + rewrite_distributed_plan_with_dynamic_filters, }; use std::sync::Arc; @@ -35,21 +35,25 @@ impl<'a> TestQuery<'a> { } } + /// Assert the number of rows after the query runs. pub(crate) fn with_expected_rows(mut self, expected_rows: usize) -> Self { self.expected_rows = expected_rows; self } + /// Forces collect left joins and enables distributed broadcast joins. pub(crate) fn with_broadcast_joins(mut self) -> Self { self.broadcast_joins = true; self } + /// Sets the desired task count to 1. pub(crate) fn with_one_task_per_leaf(mut self) -> Self { self.one_task_per_leaf = true; self } + /// Disables dynamic filter collection. pub(crate) fn without_dynamic_filter_collection(mut self) -> Self { self.collect_dynamic_filters = false; self @@ -64,6 +68,7 @@ impl<'a> TestQuery<'a> { ctx = ctx.with_distributed_desired_task_count_handler(1usize); } if !self.broadcast_joins { + // Force partitioned hash joins. let state = ctx.state_ref(); let mut state = state.write(); let optimizer = &mut state.config_mut().options_mut().optimizer; @@ -152,10 +157,8 @@ async fn execute_query_and_display( ); assert_eq!(display_plan_ascii(plan.as_ref(), false), original_display); - let plan_with_metrics = rewrite_distributed_plan_with_metrics( - plan_with_dynamic_filters, - DistributedMetricsFormat::Aggregated, - ) - .await?; - Ok(display_plan_ascii(plan_with_metrics.as_ref(), false)) + Ok(display_plan_ascii( + plan_with_dynamic_filters.as_ref(), + false, + )) }