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/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 new file mode 100644 index 000000000..0d2fba113 --- /dev/null +++ b/src/common/dynamic_filtering.rs @@ -0,0 +1,235 @@ +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; +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 along with the schema its evaluated +/// against. +#[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`, deduplicated by expression ID. +pub(crate) fn discover_dynamic_filter_consumers( + plan: &Arc, +) -> Result> { + let mut consumers = HashMap::new(); + + plan.apply(|node| { + 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() + .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 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 { + 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) +} + +#[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::union::UnionExec; + 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)?; + 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(()) + } + + #[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, + 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..97cca7394 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,7 @@ 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 1e5fc796f..86a6f5922 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,13 +1,13 @@ use crate::common::require_one_child; -use crate::coordinator::metrics_store::MetricsStore; +use crate::coordinator::dynamic_filters::isolate_distributed_leaf_variants_for_display; use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan; use crate::coordinator::prepare_static_plan::prepare_static_plan; use crate::coordinator::query_coordinator::QueryCoordinator; -use crate::distributed_planner::NetworkBoundaryExt; -use crate::{DistributedConfig, TaskKey}; +use crate::coordinator::store::{Store, task_keys_for_plan}; +use crate::{DistributedConfig, TaskCompletedDynamicFilters, TaskKey, TaskMetrics}; use datafusion::common::internal_datafusion_err; -use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; -use datafusion::common::{Result, exec_err}; +use datafusion::common::tree_node::TreeNodeRecursion; +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>, + 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,17 +66,17 @@ 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(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 @@ -82,58 +90,67 @@ 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 { + let Ok(plan) = self.plan_for_viz() 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. - /// 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 + 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() - .map_err(|e| internal_datafusion_err!("Failed to lock prepared plan: {}", e))? + .map_err(|e| internal_datafusion_err!("Failed to lock prepared execution: {e}"))? .clone() .ok_or_else(|| { - internal_datafusion_err!("No prepared plan found. Was execute() called?") + 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 /// reconstructed for visualization, with `Stage::Local` boundaries and rebuilt ancestor /// `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), + })) } } @@ -169,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), })) } @@ -192,13 +209,13 @@ impl ExecutionPlan for DistributedExec { } let base_plan = Arc::clone(&self.base_plan); - let plan_for_viz = 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); @@ -223,22 +240,24 @@ impl ExecutionPlan for DistributedExec { false => prepare_static_plan(&query_coordinator, &base_plan)?, }; - plan_for_viz + let plan_for_viz = + isolate_distributed_leaf_variants_for_display(result.plan_for_viz, &context)?; + prepared_execution .lock() - .expect("poisoned lock") - .replace(result.plan_for_viz); - head_stage - .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() { break; // channel closed } } - drop(tx); drop(guard); + 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..bc32852af --- /dev/null +++ b/src/coordinator/dynamic_filters.rs @@ -0,0 +1,196 @@ +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/metrics_store.rs b/src/coordinator/metrics_store.rs deleted file mode 100644 index 8ccc1f7c0..000000000 --- a/src/coordinator/metrics_store.rs +++ /dev/null @@ -1,36 +0,0 @@ -use crate::{TaskKey, TaskMetrics}; -use datafusion::common::HashMap; -use tokio::sync::watch; - -type MetricsMap = HashMap; - -/// Stores the metrics collected from all worker tasks, and notifies waiters when new entries arrive. -#[derive(Debug, Clone)] -pub struct MetricsStore { - tx: watch::Sender, - pub(crate) rx: watch::Receiver, -} - -impl MetricsStore { - pub(crate) fn new() -> Self { - let (tx, rx) = watch::channel(HashMap::new()); - Self { tx, rx } - } - - pub(crate) fn insert(&self, key: TaskKey, metrics: TaskMetrics) { - self.tx.send_modify(|map| { - map.insert(key, metrics); - }); - } - - pub(crate) fn get(&self, key: &TaskKey) -> Option { - self.rx.borrow().get(key).cloned() - } - - #[cfg(test)] - pub(crate) fn from_entries(entries: impl IntoIterator) -> Self { - let map: HashMap<_, _> = entries.into_iter().collect(); - let (tx, rx) = watch::channel(map); - Self { tx, rx } - } -} diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index c1a8a8dd2..7eb6dbef7 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,9 +1,11 @@ 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 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 2f9c99fb0..ff818ad0b 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,7 +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::MetricsStore; +use crate::coordinator::Store; use crate::coordinator::latency_metric::LatencyMetric; use crate::events::{RouteTasksEvent, RouteTasksHandlers}; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; @@ -12,7 +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, - 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; @@ -46,7 +47,8 @@ pub(super) struct QueryCoordinator { task_ctx: Arc, metrics: ExecutionPlanMetricsSet, coordinator_to_worker_metrics: CoordinatorToWorkerMetrics, - metrics_store: Option>, + metrics_store: Option>>, + completed_dynamic_filter_store: Arc>, end_stream_notifier: Arc, join_set: Mutex>>, } @@ -56,12 +58,14 @@ impl QueryCoordinator { pub(super) fn new( task_ctx: Arc, metrics_set: &ExecutionPlanMetricsSet, - metrics_store: Option>, + metrics_store: Option>>, + completed_dynamic_filter_store: Arc>, ) -> Self { Self { task_ctx, metrics: metrics_set.clone(), metrics_store, + 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()), @@ -80,6 +84,7 @@ impl QueryCoordinator { metrics_set: &self.metrics, metrics: &self.coordinator_to_worker_metrics, metrics_store: &self.metrics_store, + completed_dynamic_filter_store: &self.completed_dynamic_filter_store, end_stream_notifier: &self.end_stream_notifier, join_set: &self.join_set, } @@ -124,7 +129,8 @@ pub(super) struct StageCoordinator<'a> { task_ctx: &'a Arc, metrics_set: &'a ExecutionPlanMetricsSet, metrics: &'a CoordinatorToWorkerMetrics, - metrics_store: &'a Option>, + metrics_store: &'a Option>>, + completed_dynamic_filter_store: &'a Arc>, end_stream_notifier: &'a Arc, join_set: &'a Mutex>>, } @@ -150,7 +156,6 @@ impl<'a> StageCoordinator<'a> { stage_id: self.stage_id, task_number: task_i, }; - let set_plan_request = SetPlanRequest { task_key, task_count: self.task_count, @@ -178,8 +183,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 +241,7 @@ impl<'a> StageCoordinator<'a> { task_number: task_i, }; let task_metrics = self.metrics_store.clone(); + let completed_dynamic_filter_store = Arc::clone(self.completed_dynamic_filter_store); let (load_info_tx, load_info_rx) = tokio::sync::mpsc::unbounded_channel(); let mut load_info_tx_opt = Some(load_info_tx); @@ -258,8 +264,15 @@ impl<'a> StageCoordinator<'a> { WorkerToCoordinatorMsg::LoadInfoEos => { let _ = load_info_tx_opt.take(); } + WorkerToCoordinatorMsg::TaskCompletedDynamicFilters(filters) => { + completed_dynamic_filter_store.insert(task_key, filters); + } } } + if completed_dynamic_filter_store.get(&task_key).is_none() { + completed_dynamic_filter_store + .insert(task_key, TaskCompletedDynamicFilters::default()); + } }); load_info_rx } diff --git a/src/coordinator/store.rs b/src/coordinator/store.rs new file mode 100644 index 000000000..677bcfc4c --- /dev/null +++ b/src/coordinator/store.rs @@ -0,0 +1,74 @@ +use crate::TaskKey; +use crate::distributed_planner::NetworkBoundaryExt; +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) 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 3ad27ddd9..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, @@ -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, 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/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index a3633203f..f072f4a76 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -1,20 +1,20 @@ 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; 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. @@ -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. @@ -210,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. @@ -286,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::{ @@ -452,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/protocol/grpc/generated/worker.rs b/src/protocol/grpc/generated/worker.rs index 6bf8beb89..73368b080 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,36 @@ pub mod worker_to_coordinator_msg { LoadInfo(super::LoadInfo), #[prost(bool, tag = "3")] LoadInfoEos(bool), + /// 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), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct TaskCompletedDynamicFilters { + #[prost(message, repeated, tag = "1")] + pub filters: ::prost::alloc::vec::Vec, +} +/// Nested message and enum types in `TaskCompletedDynamicFilters`. +pub mod task_completed_dynamic_filters { + #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] + pub struct DynamicFilter { + #[prost(uint64, tag = "1")] + pub expression_id: u64, + /// Serialized datafusion.proto.PhysicalExprNode. + #[prost(bytes = "vec", tag = "2")] + pub expression_proto: ::prost::alloc::vec::Vec, } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -80,8 +110,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, } diff --git a/src/protocol/grpc/worker.proto b/src/protocol/grpc/worker.proto index 2affe1187..52b668bb9 100644 --- a/src/protocol/grpc/worker.proto +++ b/src/protocol/grpc/worker.proto @@ -39,9 +39,33 @@ message WorkerToCoordinatorMsg { LoadInfo load_info = 2; bool load_info_eos = 3; + + // 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; } } +message TaskCompletedDynamicFilters { + message DynamicFilter { + uint64 expression_id = 1; + // Serialized datafusion.proto.PhysicalExprNode. + bytes expression_proto = 2; + } + + repeated DynamicFilter filters = 1; +} + message TaskMetrics { // Metrics for a single task's plan nodes in pre-order traversal order. // The TaskKey is implicit — it is determined by the SetPlanRequest that diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index 060d8e875..7196561eb 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, + TaskCompletedDynamicFilters, TaskDynamicFilter, TaskKey, TaskMetrics, WorkUnitBatch, + WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; use arrow_flight::decode::FlightRecordBatchStream; @@ -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}; @@ -508,10 +509,33 @@ fn decode_worker_to_coordinator_msg( pb::worker_to_coordinator_msg::Inner::LoadInfoEos(_) => { WorkerToCoordinatorMsg::LoadInfoEos } + pb::worker_to_coordinator_msg::Inner::TaskCompletedDynamicFilters(filters) => { + WorkerToCoordinatorMsg::TaskCompletedDynamicFilters( + decode_task_completed_dynamic_filters(filters)?, + ) + } }, ) } +fn decode_task_completed_dynamic_filters( + filters: pb::TaskCompletedDynamicFilters, +) -> Result { + Ok(TaskCompletedDynamicFilters { + filters: filters + .filters + .into_iter() + .map(|filter| { + Ok(TaskDynamicFilter { + expression_id: filter.expression_id, + expression: PhysicalExprNode::decode(filter.expression_proto.as_slice()) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + }) + }) + .collect::>()?, + }) +} + fn decode_task_metrics(task_metrics: pb::TaskMetrics) -> Result { Ok(TaskMetrics { pre_order_plan_metrics: task_metrics diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index 1f0fdf0a6..4b32f91fa 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, TaskCompletedDynamicFilters, TaskKey, TaskMetrics, WorkUnitBatch, + WorkUnitFeedDeclaration, WorkUnitMsg, Worker, WorkerResolver, WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; @@ -272,10 +272,30 @@ fn encode_worker_to_coordinator_msg( WorkerToCoordinatorMsg::LoadInfoEos => { pb::worker_to_coordinator_msg::Inner::LoadInfoEos(true) } + WorkerToCoordinatorMsg::TaskCompletedDynamicFilters(filters) => { + pb::worker_to_coordinator_msg::Inner::TaskCompletedDynamicFilters( + encode_task_completed_dynamic_filters(filters), + ) + } }), }) } +fn encode_task_completed_dynamic_filters( + filters: TaskCompletedDynamicFilters, +) -> pb::TaskCompletedDynamicFilters { + pb::TaskCompletedDynamicFilters { + filters: filters + .filters + .into_iter() + .map(|filter| pb::task_completed_dynamic_filters::DynamicFilter { + expression_id: filter.expression_id, + expression_proto: filter.expression.encode_to_vec(), + }) + .collect(), + } +} + fn encode_task_metrics(task_metrics: TaskMetrics) -> Result { Ok(pb::TaskMetrics { pre_order_plan_metrics: task_metrics diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 4bd507dfe..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, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, - WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, + 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 0bea0bbe2..03f0ebb82 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; @@ -120,12 +121,29 @@ pub enum WorkerToCoordinatorMsg { /// ensuring metrics are never lost due to early stream termination. /// metrics[i] is the set of metrics for plan node i in pre-order traversal order. TaskMetrics(TaskMetrics), + /// Sends the final dynamic filters used by dynamic filter consumers back to the coorindator + /// for displaying. + TaskCompletedDynamicFilters(TaskCompletedDynamicFilters), /// Load information reported by a task. This information is used for dynamically /// sizing the number of workers involved in a query. LoadInfo(LoadInfo), LoadInfoEos, } +#[derive(Clone, Debug, Default)] +pub struct TaskCompletedDynamicFilters { + /// Final expressions keyed by their DataFusion physical-expression ID. The TaskKey is + /// implicit from the coordinator channel that carried this message. + pub filters: Vec, +} + +#[derive(Clone, Debug)] +pub struct TaskDynamicFilter { + pub expression_id: u64, + /// A `DynamicFilterPhysicalExpr` proto containing its final predicate and completion state. + pub expression: PhysicalExprNode, +} + #[derive(Clone, Debug)] pub struct TaskMetrics { /// Metrics for a single task's plan nodes in pre-order traversal order. diff --git a/src/stage.rs b/src/stage.rs index cd054c575..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 efba61c28..ac7412ce2 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,4 +1,5 @@ -use crate::common::TreeNodeExt; +use crate::codec::encode_physical_expr; +use crate::common::{TreeNodeExt, discover_dynamic_filter_consumers}; use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers}; use crate::execution_plans::SamplerExec; use crate::protocol::LocalWorkerContext; @@ -6,14 +7,16 @@ use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_t use crate::worker::task_data::TaskDataMetrics; use crate::{ CoordinatorToWorkerMsg, DistributedConfig, DistributedExt, DistributedTaskContext, - SetPlanRequest, TaskData, TaskMetrics, Worker, WorkerQueryContext, WorkerToCoordinatorMsg, + SetPlanRequest, TaskCompletedDynamicFilters, TaskData, TaskDynamicFilter, TaskMetrics, Worker, + WorkerQueryContext, WorkerToCoordinatorMsg, }; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{DataFusionError, Result, exec_datafusion_err}; -use datafusion::execution::SessionStateBuilder; +use datafusion::common::{DataFusionError, Result, exec_datafusion_err, internal_err}; +use datafusion::execution::{SessionStateBuilder, TaskContext}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; -use futures::stream::{BoxStream, FuturesUnordered}; +use datafusion_proto::protobuf::physical_expr_node::ExprType; +use futures::stream::{BoxStream, FuturesUnordered, select_all}; use futures::{FutureExt, StreamExt, TryStreamExt}; use http::HeaderMap; use std::sync::{Arc, OnceLock}; @@ -40,6 +43,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 { @@ -109,17 +113,18 @@ impl Worker { let mut work_unit_senders = Some(remote_work_unit_feed_registry.senders); let task_data_entries = Arc::clone(&self.task_data_entries); - // 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 +162,7 @@ impl Worker { } let metrics_tx = task_data.metrics_tx.lock().unwrap().take(); + let mut dynamic_filters = TaskCompletedDynamicFilters::default(); if let Some(Ok(plan)) = task_data.final_plan.get() { let d_ctx = DistributedTaskContext { task_index: key.task_number, @@ -167,7 +173,10 @@ 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(); } + let _ = dynamic_filters_tx.send(dynamic_filters); task_data_entries.invalidate(&key).await }); @@ -190,10 +199,47 @@ impl Worker { Some(WorkerToCoordinatorMsg::TaskMetrics(task_metrics)) }); - Ok(futures::stream::select(load_info_stream, metrics_stream) - .map(Ok) - .boxed()) + let dynamic_filters_stream = dynamic_filters_rx.into_stream().filter_map( + async |dynamic_filters_or_channel_dropped| { + let dynamic_filters = dynamic_filters_or_channel_dropped.ok()?; + Some(WorkerToCoordinatorMsg::TaskCompletedDynamicFilters( + dynamic_filters, + )) + }, + ); + + Ok(select_all([ + load_info_stream.boxed(), + metrics_stream.boxed(), + dynamic_filters_stream.boxed(), + ]) + .map(Ok) + .boxed()) + } +} + +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 { + return internal_err!("discovered dynamic filter did not serialize as one"); + }; + // A cancelled or short-circuited task can leave filters incomplete. Do not report those + // as final values for display. + if dynamic_filter.is_complete { + filters.push(TaskDynamicFilter { + expression_id: consumer.id, + expression, + }); + } } + Ok(TaskCompletedDynamicFilters { filters }) } /// Collects metrics from the plan in pre-order traversal order and sends them via the diff --git a/tests/dynamic_filtering.rs b/tests/dynamic_filtering.rs new file mode 100644 index 000000000..b090fb790 --- /dev/null +++ b/tests/dynamic_filtering.rs @@ -0,0 +1,204 @@ +#[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, DistributedMetricsFormat, assert_snapshot, + display_plan_ascii, rewrite_distributed_plan_with_dynamic_filters, + rewrite_distributed_plan_with_metrics, + }; + use std::sync::Arc; + + #[tokio::test] + async fn collect_left_local_dynamic_filters() -> Result<()> { + let display = execute_local_hash_join(true).await?; + assert_snapshot!(display, @r" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 3] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 3 ── tasks=2, partitions=6 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(key@0, RainToday@0)], projection=[] + │ CoalescePartitionsExec + │ [Stage 2] => NetworkBroadcastExec: partitions_per_consumer=3, stage_partitions=6, input_tasks=2 + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ RainToday@19 >= No AND RainToday@19 <= Yes AND RainToday@19 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=RainToday_null_count@1 != row_count@2 AND RainToday_max@0 >= No AND RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND (RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND Yes <= RainToday_max@0 OR RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= No AND No <= RainToday_max@0), required_guarantees=[RainToday in (No, Yes)] + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ RainToday@19 >= No AND RainToday@19 <= Yes AND RainToday@19 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=RainToday_null_count@1 != row_count@2 AND RainToday_max@0 >= No AND RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND (RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND Yes <= RainToday_max@0 OR RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= No AND No <= RainToday_max@0), required_guarantees=[RainToday in (No, Yes)] + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=12 + │ BroadcastExec: input_partitions=3, consumer_tasks=2, output_partitions=6 + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + └────────────────────────────────────────────────── + "); + Ok(()) + } + + #[tokio::test] + async fn partitioned_local_dynamic_filters() -> Result<()> { + let display = execute_local_hash_join(false).await?; + assert_snapshot!(display, @r" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 3] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 3 ── tasks=2, partitions=3 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(key@0, RainToday@0)], projection=[] + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + │ [Stage 2] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([RainToday@0], 6), input_partitions=3 + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + └────────────────────────────────────────────────── + "); + Ok(()) + } + + #[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 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)) + } +}