From fa8db4faa56773eeaa61d3f596d06f7d9bd5293b Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Mon, 17 Aug 2026 19:17:21 +0000 Subject: [PATCH 1/6] feat: plan distributed dynamic filters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Before stage split Producer worker plan HashJoin producer F HashJoin producer F └── NetworkShuffle └── NetworkShuffle └── Local stage ├── Remote stage └── consumer F └── anchor F | apply_expressions() finds F Register task-specialized producer and consumer topology independently from completed-consumer display state. Preserve consumers crossing a stage boundary as non-evaluating network-boundary anchors so both static and dynamic planning retain the information needed to identify remote filters. --- src/codec/distributed_codec.rs | 135 +++++++++-- src/coordinator/dynamic_filter_registry.rs | 213 ++++++++++++++++++ src/coordinator/mod.rs | 2 + src/coordinator/query_coordinator.rs | 8 + .../inject_network_boundaries.rs | 29 +-- src/dynamic_filtering/discovery.rs | 186 ++++++++++++--- src/dynamic_filtering/display.rs | 4 +- src/dynamic_filtering/mod.rs | 55 ++++- .../benchmarks/shuffle_bench.rs | 1 + .../benchmarks/transport_bench.rs | 1 + src/execution_plans/network_broadcast.rs | 20 +- src/execution_plans/network_coalesce.rs | 21 +- src/execution_plans/network_shuffle.rs | 19 +- src/test_utils/routing.rs | 35 ++- src/worker/impl_coordinator_channel.rs | 2 +- tests/dynamic_filtering/common.rs | 15 +- tests/dynamic_filtering/partitioned_join.rs | 39 ++-- 17 files changed, 685 insertions(+), 100 deletions(-) create mode 100644 src/coordinator/dynamic_filter_registry.rs diff --git a/src/codec/distributed_codec.rs b/src/codec/distributed_codec.rs index 2c3d7358d..09ad78abe 100644 --- a/src/codec/distributed_codec.rs +++ b/src/codec/distributed_codec.rs @@ -104,6 +104,7 @@ impl PhysicalExtensionCodec for DistributedCodec { partitioning, input_stage, equivalence_classes, + dynamic_filter_anchors, }) => { let schema: Schema = schema .as_ref() @@ -118,6 +119,12 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )? .ok_or(proto_error("NetworkShuffleExec is missing partitioning"))?; + let dynamic_filter_anchors = dynamic_filter_anchors + .iter() + .map(|expression| { + proto_converter.proto_to_physical_expr(expression, &schema, &decode_ctx) + }) + .collect::>>()?; let schema = Arc::new(schema); let equivalence_properties = parse_equivalence_properties( equivalence_classes, @@ -126,17 +133,21 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )?; - Ok(Arc::new(new_network_hash_shuffle_exec( - partitioning, - equivalence_properties, - parse_stage_proto(input_stage, inputs)?, - ))) + Ok(Arc::new( + new_network_hash_shuffle_exec( + partitioning, + equivalence_properties, + parse_stage_proto(input_stage, inputs)?, + ) + .with_dynamic_filter_anchors(dynamic_filter_anchors), + )) } DistributedExecNode::NetworkCoalesceTasks(NetworkCoalesceExecProto { schema, partitioning, input_stage, equivalence_classes, + dynamic_filter_anchors, }) => { let schema: Schema = schema .as_ref() @@ -151,6 +162,12 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )? .ok_or(proto_error("NetworkCoalesceExec is missing partitioning"))?; + let dynamic_filter_anchors = dynamic_filter_anchors + .iter() + .map(|expression| { + proto_converter.proto_to_physical_expr(expression, &schema, &decode_ctx) + }) + .collect::>>()?; let schema = Arc::new(schema); let equivalence_properties = parse_equivalence_properties( equivalence_classes, @@ -159,17 +176,21 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )?; - Ok(Arc::new(new_network_coalesce_tasks_exec( - partitioning, - equivalence_properties, - parse_stage_proto(input_stage, inputs)?, - ))) + Ok(Arc::new( + new_network_coalesce_tasks_exec( + partitioning, + equivalence_properties, + parse_stage_proto(input_stage, inputs)?, + ) + .with_dynamic_filter_anchors(dynamic_filter_anchors), + )) } DistributedExecNode::NetworkBroadcast(NetworkBroadcastExecProto { schema, partitioning, input_stage, equivalence_classes, + dynamic_filter_anchors, }) => { let schema: Schema = schema .as_ref() @@ -184,6 +205,12 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )? .ok_or(proto_error("NetworkBroadcastExec is missing partitioning"))?; + let dynamic_filter_anchors = dynamic_filter_anchors + .iter() + .map(|expression| { + proto_converter.proto_to_physical_expr(expression, &schema, &decode_ctx) + }) + .collect::>>()?; let schema = Arc::new(schema); let equivalence_properties = parse_equivalence_properties( equivalence_classes, @@ -192,11 +219,14 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )?; - Ok(Arc::new(new_network_broadcast_exec( - partitioning, - equivalence_properties, - parse_stage_proto(input_stage, inputs)?, - ))) + Ok(Arc::new( + new_network_broadcast_exec( + partitioning, + equivalence_properties, + parse_stage_proto(input_stage, inputs)?, + ) + .with_dynamic_filter_anchors(dynamic_filter_anchors), + )) } DistributedExecNode::Broadcast(BroadcastExecProto { consumer_task_count, @@ -310,6 +340,11 @@ impl PhysicalExtensionCodec for DistributedCodec { self, proto_converter, )?, + dynamic_filter_anchors: node + .dynamic_filter_anchors() + .iter() + .map(|expression| proto_converter.physical_expr_to_proto(expression, self)) + .collect::>>()?, }; let wrapper = DistributedExecProto { @@ -331,6 +366,11 @@ impl PhysicalExtensionCodec for DistributedCodec { self, proto_converter, )?, + dynamic_filter_anchors: node + .dynamic_filter_anchors() + .iter() + .map(|expression| proto_converter.physical_expr_to_proto(expression, self)) + .collect::>>()?, }; let wrapper = DistributedExecProto { @@ -352,6 +392,11 @@ impl PhysicalExtensionCodec for DistributedCodec { self, proto_converter, )?, + dynamic_filter_anchors: node + .dynamic_filter_anchors() + .iter() + .map(|expression| proto_converter.physical_expr_to_proto(expression, self)) + .collect::>>()?, }; let wrapper = DistributedExecProto { @@ -516,6 +561,8 @@ pub struct NetworkShuffleExecProto { input_stage: Option, #[prost(message, repeated, tag = "4")] equivalence_classes: Vec, + #[prost(message, repeated, tag = "5")] + dynamic_filter_anchors: Vec, } #[derive(Clone, PartialEq, ::prost::Message)] @@ -574,6 +621,7 @@ fn new_network_hash_shuffle_exec( )), worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, + dynamic_filter_anchors: vec![], } } @@ -590,6 +638,8 @@ pub struct NetworkCoalesceExecProto { input_stage: Option, #[prost(message, repeated, tag = "4")] equivalence_classes: Vec, + #[prost(message, repeated, tag = "5")] + dynamic_filter_anchors: Vec, } fn new_network_coalesce_tasks_exec( @@ -606,6 +656,7 @@ fn new_network_coalesce_tasks_exec( )), worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, + dynamic_filter_anchors: vec![], } } @@ -619,6 +670,8 @@ pub struct NetworkBroadcastExecProto { input_stage: Option, #[prost(message, repeated, tag = "4")] equivalence_classes: Vec, + #[prost(message, repeated, tag = "5")] + dynamic_filter_anchors: Vec, } #[derive(Clone, PartialEq, ::prost::Message)] @@ -644,19 +697,26 @@ fn new_network_broadcast_exec( )), worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, + dynamic_filter_anchors: vec![], } } #[cfg(test)] mod tests { - use super::super::physical_plan::new_proto_converter as default_proto_converter; + use super::super::physical_plan::{ + new_proto_converter as default_proto_converter, roundtrip_pb, + }; use super::*; use datafusion::arrow::datatypes::{DataType, Field}; use datafusion::physical_expr::{LexOrdering, PhysicalExpr}; use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::filter::FilterExec; use datafusion::prelude::SessionContext; use datafusion::{ - physical_expr::{Partitioning, PhysicalSortExpr, expressions::Column, expressions::col}, + physical_expr::{ + Partitioning, PhysicalSortExpr, + expressions::{Column, DynamicFilterPhysicalExpr, col, lit}, + }, physical_plan::{ExecutionPlan, displayable, sorts::sort::SortExec, union::UnionExec}, }; @@ -717,6 +777,47 @@ mod tests { Ok(()) } + #[test] + fn test_roundtrip_network_dynamic_filter_anchor() -> datafusion::common::Result<()> { + let ctx = create_context(); + let schema = schema_i32("a"); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("a", 0))], + lit(true), + )) as Arc; + let expected_id = dynamic_filter.expression_id(); + let network: Arc = Arc::new( + new_network_hash_shuffle_exec( + Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 4), + EquivalenceProperties::new(schema), + dummy_stage(), + ) + .with_dynamic_filter_anchors(vec![Arc::clone(&dynamic_filter)]), + ); + let plan: Arc = Arc::new(FilterExec::try_new(dynamic_filter, network)?); + + let decoded = roundtrip_pb(plan, &ctx)?; + let filter = decoded.downcast_ref::().unwrap(); + let predicate = filter + .predicate() + .downcast_ref::() + .unwrap(); + let network = filter.input().downcast_ref::().unwrap(); + let anchor = network.dynamic_filter_anchors()[0] + .downcast_ref::() + .unwrap(); + + assert_eq!(predicate.expression_id(), expected_id); + assert_eq!(anchor.expression_id(), expected_id); + predicate.update(lit(false))?; + assert_eq!( + anchor.current()?.to_string(), + "false", + "the filter predicate and network anchor should share state", + ); + Ok(()) + } + #[test] fn test_roundtrip_union() -> datafusion::common::Result<()> { let codec = DistributedCodec; diff --git a/src/coordinator/dynamic_filter_registry.rs b/src/coordinator/dynamic_filter_registry.rs new file mode 100644 index 000000000..d58c4421f --- /dev/null +++ b/src/coordinator/dynamic_filter_registry.rs @@ -0,0 +1,213 @@ +use crate::TaskKey; +use crate::dynamic_filtering::{ + discover_dynamic_filter_consumers, discover_dynamic_filter_producers, +}; +use datafusion::common::{HashMap, HashSet, Result}; +use datafusion::physical_plan::ExecutionPlan; +use std::sync::{Arc, Mutex}; + +#[derive(Default)] +pub(super) struct PlannedDynamicFilter { + // Producer and consumer tasks for a dynamic filter. + // + // Note that it is not guaranteed that every task within a stage produces / consumes dynamic filters. For + // example, a distributed union may prevent a dynamic filter from appearing in all tasks. So, we + // store task keys rather than stage ids. + pub(super) producer_tasks: HashSet, + pub(super) consumer_tasks: HashSet, +} + +#[derive(Default)] +pub(super) struct DynamicFilterRegistryState { + pub(super) filters: HashMap, +} + +/// Query-scoped hub for distributed dynamic filtering. It stores the locations +/// of dynamic filters and runtime state, informing the coordinator where dynamic filter +/// updates are coming from, how/if they should be merged, where they need to be forwarded. +#[derive(Default)] +pub(crate) struct DynamicFilterRegistry { + pub(super) state: Mutex, +} + +impl DynamicFilterRegistry { + pub(crate) fn new() -> Self { + Self::default() + } + + /// Adds any dynamic filter producers and consumers found in `plan` to the registry. + pub(crate) fn register_task( + &self, + plan: &Arc, + task_key: TaskKey, + ) -> Result<()> { + let producers = discover_dynamic_filter_producers(plan)?; + // We can safely ignore anchors because they are not evaluated by network boundaries. This + // means they do not need updates forwarded to them. + let consumers = discover_dynamic_filter_consumers(plan)?.consumers; + + let mut state = self.state.lock().expect("dynamic filter registry poisoned"); + for producer in producers { + state + .filters + .entry(producer.id) + .or_default() + .producer_tasks + .insert(task_key); + } + for consumer in consumers { + state + .filters + .entry(consumer.id) + .or_default() + .consumer_tasks + .insert(task_key); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::execution_plans::NetworkShuffleExec; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::tree_node::TreeNodeRecursion; + use datafusion::execution::{SendableRecordBatchStream, TaskContext}; + use datafusion::physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, lit}; + use datafusion::physical_expr::{Partitioning, PhysicalExpr}; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::repartition::RepartitionExec; + use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, PlanProperties, apply_expression_roots, + }; + use std::fmt::Formatter; + use uuid::Uuid; + + // Test that we correctly register producers and consumers while ignoring anchors. + #[test] + fn registers_dynamic_filters_by_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 id = dynamic_filter.expression_id().unwrap(); + + let producer_occurrence = + Arc::clone(&dynamic_filter).with_new_children(vec![Arc::new(Column::new("a", 0))])?; + assert!(!Arc::ptr_eq(&dynamic_filter, &producer_occurrence)); + + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc; + let repartition = Arc::new(RepartitionExec::try_new( + input, + Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 1), + )?) as Arc; + let boundary = Arc::new( + NetworkShuffleExec::try_new(repartition, 1)? + .with_dynamic_filter_anchors(vec![Arc::clone(&dynamic_filter)]), + ) as Arc; + let producer = Arc::new(ExpressionExec::new( + boundary, + producer_occurrence, + Some(Arc::clone(&dynamic_filter)), + )) as Arc; + let producer_task = task_key(0); + + let consumer = Arc::new(ExpressionExec::new( + Arc::new(EmptyExec::new(schema)), + dynamic_filter, + None, + )) as Arc; + let consumer_task = task_key(1); + + let registry = DynamicFilterRegistry::new(); + registry.register_task(&producer, producer_task)?; + registry.register_task(&consumer, consumer_task)?; + + let state = registry.state.lock().unwrap(); + let filter = state.filters.get(&id).unwrap(); + assert_eq!(filter.producer_tasks, HashSet::from([producer_task])); + assert_eq!(filter.consumer_tasks, HashSet::from([consumer_task])); + Ok(()) + } + + fn task_key(task_number: usize) -> TaskKey { + TaskKey { + query_id: Uuid::nil(), + stage_id: 3, + task_number, + } + } + + #[derive(Debug)] + struct ExpressionExec { + input: Arc, + expression: Arc, + produced_expression: Option>, + } + + impl ExpressionExec { + fn new( + input: Arc, + expression: Arc, + produced_expression: Option>, + ) -> Self { + Self { + input, + expression, + produced_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.produced_expression.iter().cloned().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.produced_expression.clone(), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.input.execute(partition, context) + } + } +} diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index fd7706fda..746052a6f 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,4 +1,5 @@ mod distributed; +mod dynamic_filter_registry; mod latency_metric; mod prepare_dynamic_plan; mod prepare_static_plan; @@ -6,4 +7,5 @@ mod query_coordinator; mod store; pub use distributed::DistributedExec; +pub(crate) use dynamic_filter_registry::DynamicFilterRegistry; pub(crate) use store::Store; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 6bb6faddb..16f9976a5 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,6 +1,7 @@ use crate::codec::roundtrip_pb; use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; +use crate::coordinator::DynamicFilterRegistry; use crate::coordinator::Store; use crate::coordinator::latency_metric::LatencyMetric; use crate::dynamic_filtering::maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships; @@ -52,6 +53,7 @@ pub(super) struct QueryCoordinator { coordinator_to_worker_metrics: CoordinatorToWorkerMetrics, metrics_store: Option>>, completed_dynamic_filter_store: Option>>, + dynamic_filter_registry: Arc, end_stream_notifier: Arc, join_set: Mutex>>, } @@ -69,6 +71,7 @@ impl QueryCoordinator { metrics: metrics_set.clone(), metrics_store, completed_dynamic_filter_store, + dynamic_filter_registry: Arc::new(DynamicFilterRegistry::new()), coordinator_to_worker_metrics: CoordinatorToWorkerMetrics::new(metrics_set), end_stream_notifier: Arc::new(Notify::new()), join_set: Mutex::new(JoinSet::new()), @@ -88,6 +91,7 @@ impl QueryCoordinator { metrics: &self.coordinator_to_worker_metrics, metrics_store: &self.metrics_store, completed_dynamic_filter_store: &self.completed_dynamic_filter_store, + dynamic_filter_registry: &self.dynamic_filter_registry, end_stream_notifier: &self.end_stream_notifier, join_set: &self.join_set, } @@ -134,6 +138,7 @@ pub(super) struct StageCoordinator<'a> { metrics: &'a CoordinatorToWorkerMetrics, metrics_store: &'a Option>>, completed_dynamic_filter_store: &'a Option>>, + dynamic_filter_registry: &'a Arc, end_stream_notifier: &'a Arc, join_set: &'a Mutex>>, } @@ -163,6 +168,9 @@ impl<'a> StageCoordinator<'a> { task_number: task_i, }; + self.dynamic_filter_registry + .register_task(&specialized, task_key)?; + let mut headers = get_config_extension_propagation_headers(session_config)?; headers.extend(get_passthrough_headers(session_config)); diff --git a/src/distributed_planner/inject_network_boundaries.rs b/src/distributed_planner/inject_network_boundaries.rs index 1452397bd..9f1d61dd5 100644 --- a/src/distributed_planner/inject_network_boundaries.rs +++ b/src/distributed_planner/inject_network_boundaries.rs @@ -1,4 +1,5 @@ use crate::distributed_planner::insert_broadcast::is_left_broadcast_safe; +use crate::dynamic_filtering::orphan_dynamic_filter_consumers; use crate::events::TaskCountAnnotation::{Desired, Maximum}; use crate::events::{ DesiredTaskCountEvent, DesiredTaskCountHandlers, ScaleUpLeafNodeEvent, ScaleUpLeafNodeHandlers, @@ -331,14 +332,15 @@ async fn _inject_network_boundaries( tasks: task_count.as_usize(), metrics_set: Default::default(), }; + let dynamic_filter_anchors = orphan_dynamic_filter_consumers(&input_stage.plan)?; let result = nb_ctx .nb_builder .build(input_stage, TypeId::of::(), nb_ctx) .await?; - let nb = Arc::new(NetworkShuffleExec::from_stage( - result.input_stage, - result.input_properties, - )); + let nb = Arc::new( + NetworkShuffleExec::from_stage(result.input_stage, result.input_properties) + .with_dynamic_filter_anchors(dynamic_filter_anchors), + ); Ok(nb_ctx.plan_with_task_count(nb, result.consumer_task_count)) } // Upon reaching a broadcast, we need to introduce a network broadcast right above it. @@ -350,14 +352,15 @@ async fn _inject_network_boundaries( tasks: task_count.as_usize(), metrics_set: Default::default(), }; + let dynamic_filter_anchors = orphan_dynamic_filter_consumers(&input_stage.plan)?; let result = nb_ctx .nb_builder .build(input_stage, TypeId::of::(), nb_ctx) .await?; - let nb = Arc::new(NetworkBroadcastExec::from_stage( - result.input_stage, - result.input_properties, - )); + let nb = Arc::new( + NetworkBroadcastExec::from_stage(result.input_stage, result.input_properties) + .with_dynamic_filter_anchors(dynamic_filter_anchors), + ); Ok(nb_ctx.plan_with_task_count(nb, result.consumer_task_count)) } // If the parent of the current node is either a `CoalescePartitionsExec` or a @@ -372,6 +375,7 @@ async fn _inject_network_boundaries( tasks: task_count.as_usize(), metrics_set: Default::default(), }; + let dynamic_filter_anchors = orphan_dynamic_filter_consumers(&input_stage.plan)?; let result = nb_ctx .nb_builder .build(input_stage, TypeId::of::(), nb_ctx) @@ -384,11 +388,10 @@ async fn _inject_network_boundaries( // The parent that triggered this branch is a `CoalescePartitionsExec` or // `SortPreservingMergeExec`, both of which fold all partitions into one — so the // stage above this boundary must run in exactly one task. - let nb = Arc::new(NetworkCoalesceExec::try_from_stage( - result.input_stage, - result.input_properties, - 1, - )?); + let nb = Arc::new( + NetworkCoalesceExec::try_from_stage(result.input_stage, result.input_properties, 1)? + .with_dynamic_filter_anchors(dynamic_filter_anchors), + ); Ok(nb_ctx.plan_with_task_count(nb, result.consumer_task_count)) } else if parent.is_none() { // We've just finished walking the head stage's subplan. Run a final propagation so diff --git a/src/dynamic_filtering/discovery.rs b/src/dynamic_filtering/discovery.rs index e07b052ad..79af4dd07 100644 --- a/src/dynamic_filtering/discovery.rs +++ b/src/dynamic_filtering/discovery.rs @@ -1,3 +1,4 @@ +use crate::NetworkBoundaryExt; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::common::{HashMap, HashSet, Result, internal_err}; @@ -6,6 +7,12 @@ use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use std::sync::Arc; +/// A dynamic filter produced by an [`ExecutionPlan`]. +#[derive(Clone)] +pub(crate) struct DiscoveredDynamicFilterProducer { + pub(crate) id: u64, +} + /// A dynamic-filter consumer discovered in an execution plan along with the schema it is evaluated /// against. #[derive(Clone)] @@ -15,11 +22,30 @@ pub(crate) struct DiscoveredDynamicFilter { pub(crate) input_schema: SchemaRef, } -/// Finds dynamic-filter consumers in `plan`, deduplicated by expression ID. +/// An anchor is an artificial dynamic filter consumer injected into network boundaries +/// to keep consumer references alive when they are moved across network boundaries. +/// +/// TODO(#697): remove anchors in df-56. +#[derive(Clone)] +pub(crate) struct DiscoveredDynamicFilterAnchor { + pub(crate) id: u64, + pub(crate) expression: Arc, +} + +pub(crate) struct DiscoveredDynamicFilterConsumers { + // Real consumers, ordered by expression id. + pub(crate) consumers: Vec, + // Artificial consumers. Dynamic filters in network boundaries. Also ordered by expression id. + pub(crate) anchors: Vec, +} + +/// Finds dynamic-filter consumers and network-boundary anchors in `plan`, deduplicated by +/// expression ID within each category. pub(crate) fn discover_dynamic_filter_consumers( plan: &Arc, -) -> Result> { +) -> Result { let mut consumers = HashMap::new(); + let mut anchors = HashMap::new(); plan.apply(|node| { let produced_ids: HashSet<_> = node @@ -40,6 +66,7 @@ pub(crate) fn discover_dynamic_filter_consumers( .first() .map(|child| child.schema()) .unwrap_or_else(|| node.schema()); + let is_network_boundary = node.is_network_boundary(); node.apply_expressions(&mut |root| { root.apply(|expression| { @@ -53,8 +80,16 @@ pub(crate) fn discover_dynamic_filter_consumers( "DynamicFilterPhysicalExpr did not have an expression ID" ); }; - let is_producer_occurrence = produced_ids.contains(&id); - if !is_producer_occurrence { + if is_network_boundary { + // Network-boundary expressions are metadata-only dependencies, not expressions + // evaluated by the node. + anchors + .entry(id) + .or_insert_with(|| DiscoveredDynamicFilterAnchor { + id, + expression: expression.clone(), + }); + } else if !produced_ids.contains(&id) { consumers .entry(id) .or_insert_with(|| DiscoveredDynamicFilter { @@ -72,45 +107,88 @@ pub(crate) fn discover_dynamic_filter_consumers( let mut consumers: Vec<_> = consumers.into_values().collect(); consumers.sort_unstable_by_key(|consumer| consumer.id); - Ok(consumers) + let mut anchors: Vec<_> = anchors.into_values().collect(); + anchors.sort_unstable_by_key(|anchor| anchor.id); + Ok(DiscoveredDynamicFilterConsumers { consumers, anchors }) } -/// Returns whether `plan` contains only the consumer side of a dynamic filter -/// relationship. -pub(crate) fn has_nonlocal_dynamic_filter_relationships( +/// Finds dynamic-filter producers in `plan`, deduplicated and ordered by expression ID. +pub(crate) fn discover_dynamic_filter_producers( 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(); +) -> Result> { + let mut producers = HashMap::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() - ); + for expression in node.dynamic_expressions_produced() { + if expression + .downcast_ref::() + .is_none() + { + continue; + } + let Some(id) = expression.expression_id() else { + return internal_err!("DynamicFilterPhysicalExpr did not have an expression ID"); }; - producer_ids.insert(id); + producers + .entry(id) + .or_insert(DiscoveredDynamicFilterProducer { id }); } Ok(TreeNodeRecursion::Continue) })?; - Ok(consumer_ids != producer_ids) + let mut producers: Vec<_> = producers.into_values().collect(); + producers.sort_unstable_by_key(|producer| producer.id); + Ok(producers) +} + +/// Finds consumers whose producer does not occur in `plan`. These consumers become orphaned +/// from their producer when the producer is moved behind a remote network boundary. These +/// orphans become network boundary anchors, artificially keeping the producers alive. +/// +/// TODO(697): remove anchors in df-56 +pub(crate) fn orphan_dynamic_filter_consumers( + plan: &Arc, +) -> Result>> { + let produced_here: HashSet<_> = discover_dynamic_filter_producers(plan)? + .into_iter() + .map(|producer| producer.id) + .collect(); + let discovered = discover_dynamic_filter_consumers(plan)?; + // Include anchors here because we want anchors to work recursively. For example, + // if a producer is in stage 4 and its consumer is in stage 1, an + // anchor should exist in stage 4. The easiest way to guarantee that is to ensure + // the anchor exists in stages 2, 3, and 4 recursively via this function. + let orphaned: HashMap<_, _> = discovered + .consumers + .into_iter() + .map(|consumer| (consumer.id, consumer.expression as Arc)) + .chain( + discovered + .anchors + .into_iter() + .map(|anchor| (anchor.id, anchor.expression)), + ) + .filter(|(id, _)| !produced_here.contains(id)) + .collect(); + let mut orphaned: Vec<_> = orphaned.into_iter().collect(); + orphaned.sort_unstable_by_key(|(id, _)| *id); + Ok(orphaned + .into_iter() + .map(|(_, expression)| expression) + .collect()) } #[cfg(test)] mod tests { use super::*; + use crate::NetworkShuffleExec; 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::Partitioning; use datafusion::physical_expr::expressions::{BinaryExpr, Column, lit}; use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, PlanProperties, apply_expression_roots, @@ -141,9 +219,13 @@ mod tests { )) as Arc; 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)?); + assert_eq!(discovered.consumers.len(), 1); + assert!(discovered.anchors.is_empty()); + assert_eq!( + discovered.consumers[0].id, + dynamic_filter.expression_id().unwrap() + ); + assert!(orphan_dynamic_filter_consumers(&plan)?.is_empty()); dynamic_filter .downcast_ref::() @@ -154,7 +236,7 @@ mod tests { .unwrap() .mark_complete(); - let current = discovered[0].expression.current()?; + let current = discovered.consumers[0].expression.current()?; assert_eq!(current.to_string(), "a@0 > 10"); Ok(()) } @@ -179,26 +261,58 @@ 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)?); + assert_eq!(discovered.consumers.len(), 1); + assert!(discovered.anchors.is_empty()); + assert_eq!( + discovered.consumers[0].id, + dynamic_filter.expression_id().unwrap() + ); Ok(()) } #[test] - fn identifies_a_producer_without_a_local_consumer() -> Result<()> { + fn distinguishes_consumers_from_network_boundary_anchors() -> 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::new(Column::new("a", 0))], + vec![Arc::clone(&column)], lit(true), )) as Arc; + let id = dynamic_filter.expression_id().unwrap(); + + let input = Arc::new(EmptyExec::new(schema)) as Arc; + let repartition = Arc::new(RepartitionExec::try_new( + input, + Partitioning::Hash(vec![column], 1), + )?) as Arc; + let boundary = Arc::new( + NetworkShuffleExec::try_new(repartition, 1)? + .with_dynamic_filter_anchors(vec![Arc::clone(&dynamic_filter)]), + ) as Arc; let plan = Arc::new(ExpressionExec::new( - Arc::new(EmptyExec::new(schema)), - dynamic_filter, - true, + boundary, + Arc::clone(&dynamic_filter), + false, )) as Arc; - assert!(has_nonlocal_dynamic_filter_relationships(&plan)?); + let discovered = discover_dynamic_filter_consumers(&plan)?; + assert_eq!( + discovered + .consumers + .iter() + .map(|consumer| consumer.id) + .collect::>(), + vec![id] + ); + assert_eq!( + discovered + .anchors + .iter() + .map(|anchor| anchor.id) + .collect::>(), + vec![id] + ); + assert_eq!(orphan_dynamic_filter_consumers(&plan)?.len(), 1); Ok(()) } diff --git a/src/dynamic_filtering/display.rs b/src/dynamic_filtering/display.rs index 8075020d6..66e5965b1 100644 --- a/src/dynamic_filtering/display.rs +++ b/src/dynamic_filtering/display.rs @@ -159,10 +159,10 @@ fn apply_reports_to_distributed_leaves( .iter() .map(|filter| (filter.expression_id, &filter.expression)) .collect(); - let Ok(consumers) = discover_dynamic_filter_consumers(variant) else { + let Ok(discovered) = discover_dynamic_filter_consumers(variant) else { continue; }; - for consumer in consumers { + for consumer in discovered.consumers { let Some(expression) = updates.get(&consumer.id).copied() else { continue; }; diff --git a/src/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs index 5fa38c2e9..d45db07bf 100644 --- a/src/dynamic_filtering/mod.rs +++ b/src/dynamic_filtering/mod.rs @@ -12,7 +12,13 @@ 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. +// in-memory plan. Because there's many failure modes, the safest option is +// to roundtrip the plan when any dynamic filter is found. +// +// The purpose of this function is to sever any possible cross-task in-memory +// dynamic filter relationships. +// +// Example 1: Producer-Consumer // // Consider this partitioned hash join topology where the consumer task is // collocated with one producer on worker A: @@ -34,14 +40,53 @@ pub(crate) use display::sever_dynamic_filter_relationships_in_plan_for_display; // 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. +// Example 2: Producer-Producer +// +// ```text +// Worker A +// +// Stage 2 Task 0 +// HashJoinExec <- Dynamic Filter Produced: (foo > 100) +// +// Stage 2 Task 1 +// HashJoinExec <- Dynamic Filter Produced: (foo != 150) +// +// Stage 1 Task 0 +// DataSourceExec <- consumer +// ``` +// +// Both producers are collocated on worker A. In this situation, they both race to +// update the dynamic filter, meaning the final expression will either be foo > 100 +// or foo != 150. The correct expression is (foo > 100 OR foo != 150). +// +// Example 3: Local Producer-Consumer +// +// ```text +// Worker A +// +// Stage 2 Task 0 +// HashJoinExec <- Dynamic Filter Produced: (foo > 100) +// DataSourceExec <- consumer +// +// Stage 2 Task 1 +// HashJoinExec <- Dynamic Filter Produced: (foo != 150) +// DataSourceExec <- consumer +// ``` +// +// Since both producers and both consumers are located on the same worker, they all share +// one in-memory dynamic filter. This ends up being a race between two writers and two readers. +// For a partitioned hash join, a producer may update its task-local consumer in memory, but +// updates must not cross task boundaries. 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 has_producers = !discover_dynamic_filter_producers(&plan)?.is_empty(); + let has_consumers = !discover_dynamic_filter_consumers(&plan)? + .consumers + .is_empty(); + + if has_producers || has_consumers { roundtrip_pb(plan, task_ctx) } else { Ok(plan) diff --git a/src/execution_plans/benchmarks/shuffle_bench.rs b/src/execution_plans/benchmarks/shuffle_bench.rs index 4f46f1be7..4a636b188 100644 --- a/src/execution_plans/benchmarks/shuffle_bench.rs +++ b/src/execution_plans/benchmarks/shuffle_bench.rs @@ -230,6 +230,7 @@ impl ShuffleFixture { )), input_stage: input_stage.clone(), worker_connections: WorkerConnectionPool::new(self.bench.producer_tasks), + dynamic_filter_anchors: vec![], }; let task_ctx = Arc::new(task_ctx_with_extension( &self.task_ctx, diff --git a/src/execution_plans/benchmarks/transport_bench.rs b/src/execution_plans/benchmarks/transport_bench.rs index ad159cd59..57c5ddae8 100644 --- a/src/execution_plans/benchmarks/transport_bench.rs +++ b/src/execution_plans/benchmarks/transport_bench.rs @@ -280,6 +280,7 @@ impl TransportFixture { worker_connections: crate::worker::WorkerConnectionPool::new( self.bench.producer_tasks, ), + dynamic_filter_anchors: vec![], }; let task_ctx = Arc::new(task_ctx_with_extension( &self.task_ctx, diff --git a/src/execution_plans/network_broadcast.rs b/src/execution_plans/network_broadcast.rs index 94bca6de3..0d3290d9e 100644 --- a/src/execution_plans/network_broadcast.rs +++ b/src/execution_plans/network_broadcast.rs @@ -12,7 +12,7 @@ use datafusion::physical_expr_common::metrics::MetricsSet; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics, - StatisticsArgs, + StatisticsArgs, apply_expression_roots, }; use std::fmt::Formatter; use std::sync::Arc; @@ -122,6 +122,7 @@ pub struct NetworkBroadcastExec { pub(crate) properties: Arc, pub(crate) input_stage: Stage, pub(crate) worker_connections: WorkerConnectionPool, + pub(crate) dynamic_filter_anchors: Vec>, } impl NetworkBroadcastExec { @@ -136,9 +137,22 @@ impl NetworkBroadcastExec { properties, worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, + dynamic_filter_anchors: vec![], } } + pub(crate) fn with_dynamic_filter_anchors( + mut self, + dynamic_filter_anchors: Vec>, + ) -> Self { + self.dynamic_filter_anchors = dynamic_filter_anchors; + self + } + + pub(crate) fn dynamic_filter_anchors(&self) -> &[Arc] { + &self.dynamic_filter_anchors + } + /// Creates a new [NetworkBroadcastExec] fed by the provided [BroadcastExec]. The input plan /// will be executed in a remote worker in `producer_tasks` number of tasks. pub fn try_new(input: Arc, producer_tasks: usize) -> Result { @@ -219,9 +233,9 @@ impl ExecutionPlan for NetworkBroadcastExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - Ok(TreeNodeRecursion::Continue) + apply_expression_roots(self.dynamic_filter_anchors.iter(), f) } fn with_new_children( diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index 4cbaec0ce..a6768f60c 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -15,7 +15,8 @@ use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan, - PlanProperties, ReplaceChildrenOptions, Statistics, StatisticsArgs, internal_err, + PlanProperties, ReplaceChildrenOptions, Statistics, StatisticsArgs, apply_expression_roots, + internal_err, }; use std::fmt::{Debug, Formatter}; use std::sync::Arc; @@ -80,6 +81,7 @@ pub struct NetworkCoalesceExec { pub(crate) properties: Arc, pub(crate) input_stage: Stage, pub(crate) worker_connections: WorkerConnectionPool, + pub(crate) dynamic_filter_anchors: Vec>, } impl NetworkCoalesceExec { @@ -98,9 +100,22 @@ impl NetworkCoalesceExec { properties: props, worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, + dynamic_filter_anchors: vec![], }) } + pub(crate) fn with_dynamic_filter_anchors( + mut self, + dynamic_filter_anchors: Vec>, + ) -> Self { + self.dynamic_filter_anchors = dynamic_filter_anchors; + self + } + + pub(crate) fn dynamic_filter_anchors(&self) -> &[Arc] { + &self.dynamic_filter_anchors + } + /// Creates a new [NetworkCoalesceExec] fed by the provided `input` plan. /// /// The `input` plan will be remotely executed in `producer_tasks` tasks, while the @@ -241,9 +256,9 @@ impl ExecutionPlan for NetworkCoalesceExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - Ok(TreeNodeRecursion::Continue) + apply_expression_roots(self.dynamic_filter_anchors.iter(), f) } fn with_new_children( diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index fbb4200a5..8172848d2 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -14,6 +14,7 @@ use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, Statistics, StatisticsArgs, + apply_expression_roots, }; use std::fmt::Formatter; use std::sync::Arc; @@ -105,6 +106,7 @@ pub struct NetworkShuffleExec { pub(crate) properties: Arc, pub(crate) input_stage: Stage, pub(crate) worker_connections: WorkerConnectionPool, + pub(crate) dynamic_filter_anchors: Vec>, } impl NetworkShuffleExec { @@ -113,9 +115,22 @@ impl NetworkShuffleExec { properties: input_properties, worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, + dynamic_filter_anchors: vec![], } } + pub(crate) fn with_dynamic_filter_anchors( + mut self, + dynamic_filter_anchors: Vec>, + ) -> Self { + self.dynamic_filter_anchors = dynamic_filter_anchors; + self + } + + pub(crate) fn dynamic_filter_anchors(&self) -> &[Arc] { + &self.dynamic_filter_anchors + } + /// Creates a new [NetworkShuffleExec] fed by the provided [RepartitionExec]. The input plan /// will be executed in a remote worker in `producer_tasks` number of tasks. pub fn try_new(input: Arc, producer_tasks: usize) -> Result { @@ -195,9 +210,9 @@ impl ExecutionPlan for NetworkShuffleExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - Ok(TreeNodeRecursion::Continue) + apply_expression_roots(self.dynamic_filter_anchors.iter(), f) } fn with_new_children( diff --git a/src/test_utils/routing.rs b/src/test_utils/routing.rs index 0c726e65c..3f65ccd98 100644 --- a/src/test_utils/routing.rs +++ b/src/test_utils/routing.rs @@ -5,7 +5,8 @@ use arrow::{ use datafusion::{ catalog::{Session, TableFunctionImpl, TableProvider}, common::{ - Result, ScalarValue, Statistics, internal_err, plan_err, tree_node::TreeNodeRecursion, + Result, ScalarValue, Statistics, exec_err, internal_err, plan_err, + tree_node::TreeNodeRecursion, }, datasource::TableType, execution::TaskContext, @@ -23,7 +24,9 @@ use datafusion_proto::{ use futures::stream; use prost::Message; use std::{fmt::Formatter, sync::Arc}; +use tokio::sync::Mutex; use tonic::async_trait; +use url::Url; use crate::{ DesiredTaskCountEvent, DesiredTaskCountEventResponse, DistributedLeafExec, @@ -386,3 +389,33 @@ impl PhysicalExtensionCodec for URLEmitterExtensionCodec { .map_err(|e| proto_error(format!("Failed to encode URLEmitterExec: {e}"))) } } + +/// Colocates all tasks on the same worker by choosing a URL once and caching it. +#[derive(Default)] +pub struct ColocateAllTasksHandler { + cached: Mutex>, +} + +#[async_trait] +impl RouteTaskHandler for ColocateAllTasksHandler { + async fn handle(&self, ev: RouteTaskEvent<'_>) -> Option> { + let url = { + let mut cached = self.cached.lock().await; + if let Some(url) = cached.as_ref() { + url.clone() + } else { + let urls = match ev.worker_resolver.get_urls() { + Ok(urls) => urls, + Err(error) => return Some(Err(error)), + }; + let Some(url) = urls.into_iter().next() else { + return Some(exec_err!("expected at least one worker URL")); + }; + *cached = Some(url.clone()); + url + } + }; + + Some(ev.dialer.dial(url).await) + } +} diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index 5072e30e2..115808d40 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -256,7 +256,7 @@ fn build_task_completed_dynamic_filters( plan: &Arc, ) -> Result { let mut filters = vec![]; - for consumer in discover_dynamic_filter_consumers(plan)? { + for consumer in discover_dynamic_filter_consumers(plan)?.consumers { filters.push(TaskDynamicFilter { expression_id: consumer.id, expression: MaybeEncoded::Decoded(consumer.expression), diff --git a/tests/dynamic_filtering/common.rs b/tests/dynamic_filtering/common.rs index 87966952f..1833da626 100644 --- a/tests/dynamic_filtering/common.rs +++ b/tests/dynamic_filtering/common.rs @@ -9,7 +9,9 @@ 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::UrlEmitterRouteTaskHandler; +use datafusion_distributed::test_utils::routing::{ + ColocateAllTasksHandler, UrlEmitterRouteTaskHandler, +}; use datafusion_distributed::{ DefaultSessionBuilder, DistributedExt, display_plan_ascii, rewrite_distributed_plan_with_dynamic_filters, @@ -89,12 +91,17 @@ impl<'a> TestQuery<'a> { pub(crate) async fn execute_range_partitioned_query( sql: &str, expected_rows: usize, + colocate_tasks: bool, ) -> Result { let (ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await; - let ctx = ctx + let mut ctx = ctx .with_distributed_broadcast_joins(false)? - .with_distributed_desired_task_count_handler(2usize) - .with_distributed_route_task_handler(UrlEmitterRouteTaskHandler); + .with_distributed_desired_task_count_handler(2usize); + ctx = if colocate_tasks { + ctx.with_distributed_route_task_handler(ColocateAllTasksHandler::default()) + } else { + ctx.with_distributed_route_task_handler(UrlEmitterRouteTaskHandler) + }; { let state = ctx.state_ref(); let mut state = state.write(); diff --git a/tests/dynamic_filtering/partitioned_join.rs b/tests/dynamic_filtering/partitioned_join.rs index b0c530c4c..a291db9f7 100644 --- a/tests/dynamic_filtering/partitioned_join.rs +++ b/tests/dynamic_filtering/partitioned_join.rs @@ -3,22 +3,35 @@ mod tests { use crate::common::{TestQuery, execute_range_partitioned_query}; use datafusion::common::Result; use datafusion_distributed::assert_snapshot; + use datafusion_distributed::test_utils::insta::insta::allow_duplicates; /// 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" + let display = execute_range_partitioned_query(LOCAL_DYNAMIC_FILTER_QUERY, 2, false).await?; + assert_local_dynamic_filters_plan(display); + Ok(()) + } + + /// Colocated tasks must not share their task-local dynamic filters. + #[tokio::test] + async fn colocated_local_dynamic_filters() -> Result<()> { + let display = execute_range_partitioned_query(LOCAL_DYNAMIC_FILTER_QUERY, 2, true).await?; + assert_local_dynamic_filters_plan(display); + Ok(()) + } + + const LOCAL_DYNAMIC_FILTER_QUERY: &str = 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 + "#; + + fn assert_local_dynamic_filters_plan(display: String) { + allow_duplicates! { + assert_snapshot!(display, @r" ┌───── DistributedExec │ CoalescePartitionsExec │ [Stage 2] => NetworkCoalesceExec: output_partitions=4, input_tasks=2 @@ -41,7 +54,7 @@ mod tests { │ 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. From 9866970af84970cbae66a24fe8dac0a7e7d3949e Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Fri, 11 Sep 2026 16:51:43 +0000 Subject: [PATCH 2/6] delete registers_dynamic_filters_by_expression_id test --- src/coordinator/dynamic_filter_registry.rs | 154 +-------------------- 1 file changed, 6 insertions(+), 148 deletions(-) diff --git a/src/coordinator/dynamic_filter_registry.rs b/src/coordinator/dynamic_filter_registry.rs index d58c4421f..6642fdf65 100644 --- a/src/coordinator/dynamic_filter_registry.rs +++ b/src/coordinator/dynamic_filter_registry.rs @@ -22,9 +22,12 @@ pub(super) struct DynamicFilterRegistryState { pub(super) filters: HashMap, } -/// Query-scoped hub for distributed dynamic filtering. It stores the locations -/// of dynamic filters and runtime state, informing the coordinator where dynamic filter -/// updates are coming from, how/if they should be merged, where they need to be forwarded. +/// Query-scoped hub for distributed dynamic filtering. +/// +/// It stores the locations of dynamic filters and their runtime state. Informs the coordinator +/// - where dynamic filter updates are coming from +/// - how/if dynamic filter updates should be merged +/// - where dynamic filter updates should be forwarded #[derive(Default)] pub(crate) struct DynamicFilterRegistry { pub(super) state: Mutex, @@ -66,148 +69,3 @@ impl DynamicFilterRegistry { Ok(()) } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::execution_plans::NetworkShuffleExec; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; - use datafusion::common::tree_node::TreeNodeRecursion; - use datafusion::execution::{SendableRecordBatchStream, TaskContext}; - use datafusion::physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, lit}; - use datafusion::physical_expr::{Partitioning, PhysicalExpr}; - use datafusion::physical_plan::empty::EmptyExec; - use datafusion::physical_plan::repartition::RepartitionExec; - use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, PlanProperties, apply_expression_roots, - }; - use std::fmt::Formatter; - use uuid::Uuid; - - // Test that we correctly register producers and consumers while ignoring anchors. - #[test] - fn registers_dynamic_filters_by_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 id = dynamic_filter.expression_id().unwrap(); - - let producer_occurrence = - Arc::clone(&dynamic_filter).with_new_children(vec![Arc::new(Column::new("a", 0))])?; - assert!(!Arc::ptr_eq(&dynamic_filter, &producer_occurrence)); - - let input = Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc; - let repartition = Arc::new(RepartitionExec::try_new( - input, - Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 1), - )?) as Arc; - let boundary = Arc::new( - NetworkShuffleExec::try_new(repartition, 1)? - .with_dynamic_filter_anchors(vec![Arc::clone(&dynamic_filter)]), - ) as Arc; - let producer = Arc::new(ExpressionExec::new( - boundary, - producer_occurrence, - Some(Arc::clone(&dynamic_filter)), - )) as Arc; - let producer_task = task_key(0); - - let consumer = Arc::new(ExpressionExec::new( - Arc::new(EmptyExec::new(schema)), - dynamic_filter, - None, - )) as Arc; - let consumer_task = task_key(1); - - let registry = DynamicFilterRegistry::new(); - registry.register_task(&producer, producer_task)?; - registry.register_task(&consumer, consumer_task)?; - - let state = registry.state.lock().unwrap(); - let filter = state.filters.get(&id).unwrap(); - assert_eq!(filter.producer_tasks, HashSet::from([producer_task])); - assert_eq!(filter.consumer_tasks, HashSet::from([consumer_task])); - Ok(()) - } - - fn task_key(task_number: usize) -> TaskKey { - TaskKey { - query_id: Uuid::nil(), - stage_id: 3, - task_number, - } - } - - #[derive(Debug)] - struct ExpressionExec { - input: Arc, - expression: Arc, - produced_expression: Option>, - } - - impl ExpressionExec { - fn new( - input: Arc, - expression: Arc, - produced_expression: Option>, - ) -> Self { - Self { - input, - expression, - produced_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.produced_expression.iter().cloned().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.produced_expression.clone(), - ))) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - self.input.execute(partition, context) - } - } -} From a1e2ed4eee03ff57b6ec10bde74f74851cca5061 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Fri, 11 Sep 2026 18:01:03 +0000 Subject: [PATCH 3/6] better comments around dynamic filter "sever" functions --- src/dynamic_filtering/display.rs | 8 +- src/dynamic_filtering/mod.rs | 173 +++++++++++++++++++------------ 2 files changed, 113 insertions(+), 68 deletions(-) diff --git a/src/dynamic_filtering/display.rs b/src/dynamic_filtering/display.rs index 66e5965b1..abcfa5537 100644 --- a/src/dynamic_filtering/display.rs +++ b/src/dynamic_filtering/display.rs @@ -51,6 +51,9 @@ pub async fn rewrite_distributed_plan_with_dynamic_filters( /// Severs dynamic filter connections so we can update filter values for /// display purposes without having an update in one node propagate to another. /// +/// Unlike [`maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships()`], this +/// isolates every dynamic filter found in the plan. +/// /// For example, in this plan, we would like to be able to [`update()`] every variant independently /// without mutating the producer or other variants /// @@ -62,14 +65,15 @@ pub async fn rewrite_distributed_plan_with_dynamic_filters( /// 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 +/// t0: DataSourceExec: predicate=DynamicFilter [ f_dkey@2 > A ] <- unique filter +/// t1: DataSourceExec: predicate=DynamicFilter [ f_dkey@2 < B ] <- unique filter /// ``` /// /// This is done by deep-copying every leaf variant so we don't have to /// worry about any shared state. /// /// [`update()`]: DynamicFilterPhysicalExpr::update() +/// [`maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships()`]: super::maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships() pub(crate) fn sever_dynamic_filter_relationships_in_plan_for_display( plan: Arc, task_ctx: &Arc, diff --git a/src/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs index d45db07bf..09b8ba5d8 100644 --- a/src/dynamic_filtering/mod.rs +++ b/src/dynamic_filtering/mod.rs @@ -11,72 +11,113 @@ 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. Because there's many failure modes, the safest option is -// to roundtrip the plan when any dynamic filter is found. -// -// The purpose of this function is to sever any possible cross-task in-memory -// dynamic filter relationships. -// -// Example 1: Producer-Consumer -// -// 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). -// -// Example 2: Producer-Producer -// -// ```text -// Worker A -// -// Stage 2 Task 0 -// HashJoinExec <- Dynamic Filter Produced: (foo > 100) -// -// Stage 2 Task 1 -// HashJoinExec <- Dynamic Filter Produced: (foo != 150) -// -// Stage 1 Task 0 -// DataSourceExec <- consumer -// ``` -// -// Both producers are collocated on worker A. In this situation, they both race to -// update the dynamic filter, meaning the final expression will either be foo > 100 -// or foo != 150. The correct expression is (foo > 100 OR foo != 150). -// -// Example 3: Local Producer-Consumer -// -// ```text -// Worker A -// -// Stage 2 Task 0 -// HashJoinExec <- Dynamic Filter Produced: (foo > 100) -// DataSourceExec <- consumer -// -// Stage 2 Task 1 -// HashJoinExec <- Dynamic Filter Produced: (foo != 150) -// DataSourceExec <- consumer -// ``` -// -// Since both producers and both consumers are located on the same worker, they all share -// one in-memory dynamic filter. This ends up being a race between two writers and two readers. -// For a partitioned hash join, a producer may update its task-local consumer in memory, but -// updates must not cross task boundaries. +/// Isolates all shared, in-memory dynamic filter state from this plan if it contains +/// any dynamic filter producers or consumers. +/// +/// Plan nodes *within* this plan will share in-memory dynamic filter state. However, they +/// will not share state with plan nodes outside of this plan, such as in parent stages. +/// +/// # Correctness: Task-Local Dynamic Filters are Safe to Apply +/// +/// Claim: It is correct for a producer to *always* update consumers within the same task. +/// +/// Proof: +/// +/// Invariant: If the task-local consumer filters out a row, removing that row must not change the output +/// of the task-local producer. +/// +/// DataFusion-Distributed dynamic-filter producers satisfy this invariant: +/// +/// 1. TopK dynamic filters are distributed into N TopK operations for N tasks followed by a global +/// TopK sort preserving TopK across those tasks. A task-local TopK dynamic filter does not +/// change any outputs in this scenario. +/// 2. A partial aggregate without grouping may push `MIN` and `MAX` bounds to consumers. In a task, a row +/// rejected by the aggregate's current bound will not change that aggregate's output. So, a local +/// min/max dynamic filter is safe to apply. +/// 3. A `CollectLeft` hash join has the complete build side in every task, so any task can apply +/// a predicate representing the build side. +/// 4. A partitioned hash join builds its predicate from the task-local hash table. There are three +/// cases for its build and probe side: +/// - Their partitioning matches. Thus, the build and probe execute the same partitions in the +/// same task, so the local dynamic filter filters rows corresponding to the local hash table. +/// - Their partitioning differs and the required repartition becomes a network shuffle. There is +/// local dynamic filter to push down anymore. +/// - Their partitioning differs but the repartition remains local. This is a single-node plan, +/// where the task executes all partitions and the producer's predicate covers the complete +/// build input. +/// +/// This reasoning applies only within one task. A consumer in another task may need the union of +/// several task-local predicates and must receive that predicate through the coordinator instead +/// of sharing a producer's in-memory state. +/// +/// # Cases this Function Avoids +/// +/// It's possible for any two tasks to be collocated and share memory because +/// - a user to implement a custom transport layer and skip all proto serialization +/// - the coordinator may send plans to it's local worker via an in-memory channel without serializing +/// +/// This causes weird scenarios that fall outside the well-defined cases above. This function is +/// responsible for isolating task plans to prevent these scenarios from occurring. +/// +/// ## Example 1: Producer-Consumer +/// +/// 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). +/// +/// ## Example 2: Producer-Producer +/// +/// ```text +/// Worker A +/// +/// Stage 2 Task 0 +/// HashJoinExec <- Dynamic Filter Produced: (foo > 100) +/// +/// Stage 2 Task 1 +/// HashJoinExec <- Dynamic Filter Produced: (foo != 150) +/// +/// Stage 1 Task 0 +/// DataSourceExec <- consumer +/// ``` +/// +/// Both producers are collocated on worker A. In this situation, they both race to +/// update the dynamic filter, meaning the final expression will either be foo > 100 +/// or foo != 150. The correct expression is (foo > 100 OR foo != 150). +/// +/// Example 3: Local Producer-Consumer +/// +/// ```text +/// Worker A +/// +/// Stage 2 Task 0 +/// HashJoinExec <- Dynamic Filter Produced: (foo > 100) +/// DataSourceExec <- consumer +/// +/// Stage 2 Task 1 +/// HashJoinExec <- Dynamic Filter Produced: (foo != 150) +/// DataSourceExec <- consumer +/// ``` +/// +/// Since both producers and both consumers are located on the same worker, they all share +/// one in-memory dynamic filter. This ends up being a race between two writers and two readers. +/// For a partitioned hash join, a producer may update its task-local consumer in memory, but +/// updates must not cross task boundaries. pub(crate) fn maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships( plan: Arc, task_ctx: &Arc, From bd4cb92b07369842d60ce89c4a7ecfeeeb029aa4 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Fri, 11 Sep 2026 18:49:45 +0000 Subject: [PATCH 4/6] migrate dynamic filter discovery tests to integration tests --- src/dynamic_filtering/discovery.rs | 242 ++--------------------- src/dynamic_filtering/mod.rs | 2 +- src/lib.rs | 6 + tests/dynamic_filtering/discovery.rs | 284 +++++++++++++++++++++++++++ tests/dynamic_filtering/main.rs | 1 + 5 files changed, 306 insertions(+), 229 deletions(-) create mode 100644 tests/dynamic_filtering/discovery.rs diff --git a/src/dynamic_filtering/discovery.rs b/src/dynamic_filtering/discovery.rs index 79af4dd07..068aa02ce 100644 --- a/src/dynamic_filtering/discovery.rs +++ b/src/dynamic_filtering/discovery.rs @@ -9,17 +9,17 @@ use std::sync::Arc; /// A dynamic filter produced by an [`ExecutionPlan`]. #[derive(Clone)] -pub(crate) struct DiscoveredDynamicFilterProducer { - pub(crate) id: u64, +pub struct DiscoveredDynamicFilterProducer { + pub id: u64, } /// A dynamic-filter consumer discovered in an execution plan along with the schema it is evaluated /// against. #[derive(Clone)] -pub(crate) struct DiscoveredDynamicFilter { - pub(crate) id: u64, - pub(crate) expression: Arc, - pub(crate) input_schema: SchemaRef, +pub struct DiscoveredDynamicFilter { + pub id: u64, + pub expression: Arc, + pub input_schema: SchemaRef, } /// An anchor is an artificial dynamic filter consumer injected into network boundaries @@ -27,21 +27,21 @@ pub(crate) struct DiscoveredDynamicFilter { /// /// TODO(#697): remove anchors in df-56. #[derive(Clone)] -pub(crate) struct DiscoveredDynamicFilterAnchor { - pub(crate) id: u64, - pub(crate) expression: Arc, +pub struct DiscoveredDynamicFilterAnchor { + pub id: u64, + pub expression: Arc, } -pub(crate) struct DiscoveredDynamicFilterConsumers { +pub struct DiscoveredDynamicFilterConsumers { // Real consumers, ordered by expression id. - pub(crate) consumers: Vec, + pub consumers: Vec, // Artificial consumers. Dynamic filters in network boundaries. Also ordered by expression id. - pub(crate) anchors: Vec, + pub anchors: Vec, } /// Finds dynamic-filter consumers and network-boundary anchors in `plan`, deduplicated by /// expression ID within each category. -pub(crate) fn discover_dynamic_filter_consumers( +pub fn discover_dynamic_filter_consumers( plan: &Arc, ) -> Result { let mut consumers = HashMap::new(); @@ -113,7 +113,7 @@ pub(crate) fn discover_dynamic_filter_consumers( } /// Finds dynamic-filter producers in `plan`, deduplicated and ordered by expression ID. -pub(crate) fn discover_dynamic_filter_producers( +pub fn discover_dynamic_filter_producers( plan: &Arc, ) -> Result> { let mut producers = HashMap::new(); @@ -176,217 +176,3 @@ pub(crate) fn orphan_dynamic_filter_consumers( .map(|(_, expression)| expression) .collect()) } - -#[cfg(test)] -mod tests { - use super::*; - use crate::NetworkShuffleExec; - 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::Partitioning; - use datafusion::physical_expr::expressions::{BinaryExpr, Column, lit}; - use datafusion::physical_plan::empty::EmptyExec; - use datafusion::physical_plan::repartition::RepartitionExec; - 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.consumers.len(), 1); - assert!(discovered.anchors.is_empty()); - assert_eq!( - discovered.consumers[0].id, - dynamic_filter.expression_id().unwrap() - ); - assert!(orphan_dynamic_filter_consumers(&plan)?.is_empty()); - - 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.consumers[0].expression.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.consumers.len(), 1); - assert!(discovered.anchors.is_empty()); - assert_eq!( - discovered.consumers[0].id, - dynamic_filter.expression_id().unwrap() - ); - Ok(()) - } - - #[test] - fn distinguishes_consumers_from_network_boundary_anchors() -> 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 id = dynamic_filter.expression_id().unwrap(); - - let input = Arc::new(EmptyExec::new(schema)) as Arc; - let repartition = Arc::new(RepartitionExec::try_new( - input, - Partitioning::Hash(vec![column], 1), - )?) as Arc; - let boundary = Arc::new( - NetworkShuffleExec::try_new(repartition, 1)? - .with_dynamic_filter_anchors(vec![Arc::clone(&dynamic_filter)]), - ) as Arc; - let plan = Arc::new(ExpressionExec::new( - boundary, - Arc::clone(&dynamic_filter), - false, - )) as Arc; - - let discovered = discover_dynamic_filter_consumers(&plan)?; - assert_eq!( - discovered - .consumers - .iter() - .map(|consumer| consumer.id) - .collect::>(), - vec![id] - ); - assert_eq!( - discovered - .anchors - .iter() - .map(|anchor| anchor.id) - .collect::>(), - vec![id] - ); - assert_eq!(orphan_dynamic_filter_consumers(&plan)?.len(), 1); - 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/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs index 09b8ba5d8..4aca4801f 100644 --- a/src/dynamic_filtering/mod.rs +++ b/src/dynamic_filtering/mod.rs @@ -7,7 +7,7 @@ use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; use std::sync::Arc; -pub(crate) use discovery::*; +pub use discovery::*; pub use display::rewrite_distributed_plan_with_dynamic_filters; pub(crate) use display::sever_dynamic_filter_relationships_in_plan_for_display; diff --git a/src/lib.rs b/src/lib.rs index d8c9a7896..17c1159e0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,6 +25,12 @@ pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, ProducerHead, SessionStateBuilderExt, }; pub use dynamic_filtering::rewrite_distributed_plan_with_dynamic_filters; +#[cfg(any(feature = "integration", test))] +pub use dynamic_filtering::{ + DiscoveredDynamicFilter, DiscoveredDynamicFilterAnchor, DiscoveredDynamicFilterConsumers, + DiscoveredDynamicFilterProducer, discover_dynamic_filter_consumers, + discover_dynamic_filter_producers, +}; pub use events::{ CoordinatorToWorkerDialer, DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, RouteTaskEvent, RouteTaskEventResponse, RouteTaskHandler, diff --git a/tests/dynamic_filtering/discovery.rs b/tests/dynamic_filtering/discovery.rs new file mode 100644 index 000000000..1a8ae7c82 --- /dev/null +++ b/tests/dynamic_filtering/discovery.rs @@ -0,0 +1,284 @@ +#[cfg(test)] +mod tests { + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion::common::{HashMap, Result}; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; + use datafusion::physical_plan::{ExecutionPlan, collect}; + use datafusion_distributed::test_utils::localhost::start_localhost_context; + use datafusion_distributed::test_utils::parquet::register_parquet_tables; + use datafusion_distributed::{ + DefaultSessionBuilder, DistributedExt, NetworkBoundaryExt, RouteTaskEvent, + RouteTaskEventResponse, RouteTaskHandler, assert_snapshot, + discover_dynamic_filter_consumers, discover_dynamic_filter_producers, + }; + use itertools::Itertools; + use std::collections::BTreeSet; + use std::fmt::Write; + use std::sync::Arc; + use tokio::sync::Mutex; + use tonic::async_trait; + + #[tokio::test] + async fn discovers_dynamic_filters_in_sql_plan() -> Result<()> { + let display = display_query( + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + JOIN weather probe ON build.key = probe."RainToday" + JOIN ( + SELECT DISTINCT "RainTomorrow" AS key + FROM weather + ) other_build ON other_build.key = probe."RainTomorrow" + WHERE probe."MinTemp" > 0 + "#, + ) + .await?; + assert_snapshot!(display, @r" + Stage 5 + AggregateExec + HashJoinExec producers=[1] + NetworkShuffleExec + AggregateExec + NetworkShuffleExec anchors=[1] + Stage 4 + RepartitionExec + AggregateExec + DataSourceExec consumers=[1] + Stage 3 + RepartitionExec + HashJoinExec producers=[2] + NetworkShuffleExec + AggregateExec + NetworkShuffleExec anchors=[2] + Stage 2 + RepartitionExec + AggregateExec + DataSourceExec consumers=[2] + Stage 1 + RepartitionExec + FilterExec + DataSourceExec + "); + Ok(()) + } + + #[tokio::test] + async fn passes_anchor_through_two_shuffles() -> Result<()> { + let display = display_query( + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + JOIN ( + SELECT "RainTomorrow" AS key, SUM(n) AS total + FROM ( + SELECT "RainTomorrow", "RainToday", COUNT(*) AS n + FROM weather + GROUP BY "RainTomorrow", "RainToday" + ) grouped + GROUP BY "RainTomorrow" + ) probe ON build.key = probe.key + "#, + ) + .await?; + assert_snapshot!(display, @r" + Stage 4 + AggregateExec + HashJoinExec producers=[1] + AggregateExec + NetworkShuffleExec + ProjectionExec + AggregateExec + NetworkShuffleExec anchors=[1] + Stage 3 + RepartitionExec + AggregateExec + ProjectionExec + AggregateExec + NetworkShuffleExec anchors=[1] + Stage 2 + RepartitionExec + AggregateExec + DataSourceExec consumers=[1] + Stage 1 + RepartitionExec + AggregateExec + DataSourceExec + "); + Ok(()) + } + + async fn display_query(sql: &str) -> Result { + let captured_plans = CapturePlans::default(); + let (ctx, _guard, _) = start_localhost_context(2, DefaultSessionBuilder).await; + let ctx = ctx + .with_distributed_broadcast_joins(false)? + .with_distributed_route_task_handler(captured_plans.clone()); + { + 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?; + collect(plan, ctx.task_ctx()).await?; + let captured_plans = captured_plans.0.lock().await; + display_dynamic_filter_discovery(&captured_plans) + } + + /// Captures the first task of each stage for displaying purposes. + #[derive(Clone, Default)] + struct CapturePlans(Arc>>>); + + #[async_trait] + impl RouteTaskHandler for CapturePlans { + async fn handle( + &self, + event: RouteTaskEvent<'_>, + ) -> Option> { + if event.task_key.task_number == 0 { + self.0.lock().await.insert( + event.task_key.stage_id, + Arc::clone(event.task_specialized_plan), + ); + } + None + } + } + + /// Map random dynamic filter expression ids to monotonic numbers 1, 2, 3... + /// for stable snapshots. + #[derive(Default)] + struct IdNormalizer(HashMap); + + impl IdNormalizer { + fn annotation(&mut self, name: &str, ids: BTreeSet) -> Option { + (!ids.is_empty()).then(|| { + let ids = ids + .into_iter() + .map(|id| { + let next = self.0.len() + 1; + self.0.entry(id).or_insert(next).to_string() + }) + .join(", "); + format!("{name}=[{ids}]") + }) + } + } + + struct DynamicFilterIds { + consumers: BTreeSet, + anchors: BTreeSet, + producers: BTreeSet, + } + + fn dynamic_filter_annotations( + node: &dyn ExecutionPlan, + discovered: &DynamicFilterIds, + normalizer: &mut IdNormalizer, + ) -> Result { + let producers = node + .dynamic_expressions_produced() + .iter() + .filter_map(dynamic_filter_id) + .filter(|id| discovered.producers.contains(id)) + .collect::>(); + let is_network_boundary = node.is_network_boundary(); + let mut anchors = BTreeSet::new(); + let mut consumers = BTreeSet::new(); + node.apply_expressions(&mut |root| { + root.apply(|expression| { + if let Some(id) = dynamic_filter_id(expression) { + if is_network_boundary && discovered.anchors.contains(&id) { + anchors.insert(id); + } else if discovered.consumers.contains(&id) && !producers.contains(&id) { + consumers.insert(id); + } + } + Ok(TreeNodeRecursion::Continue) + }) + })?; + + let annotations = [ + ("anchors", anchors), + ("consumers", consumers), + ("producers", producers), + ] + .into_iter() + .filter_map(|(name, ids)| normalizer.annotation(name, ids)) + .join(" "); + Ok(if annotations.is_empty() { + String::new() + } else { + format!(" {annotations}") + }) + } + + fn dynamic_filter_id(expression: &Arc) -> Option { + expression.downcast_ref::()?; + Some( + expression + .expression_id() + .expect("dynamic filters always have an expression ID"), + ) + } + + fn display_dynamic_filter_discovery( + plans: &HashMap>, + ) -> Result { + fn render( + node: &dyn ExecutionPlan, + depth: usize, + discovered: &DynamicFilterIds, + normalizer: &mut IdNormalizer, + output: &mut String, + ) -> Result<()> { + writeln!( + output, + "{}{}{}", + " ".repeat(depth), + node.name(), + dynamic_filter_annotations(node, discovered, normalizer)?, + ) + .expect("writing to String cannot fail"); + for child in node.children() { + render(child.as_ref(), depth + 1, discovered, normalizer, output)?; + } + Ok(()) + } + + let mut output = String::new(); + let mut normalizer = IdNormalizer::default(); + for stage_id in plans.keys().sorted().rev() { + writeln!(output, "Stage {stage_id}").expect("writing to String cannot fail"); + let plan = &plans[stage_id]; + let consumers = discover_dynamic_filter_consumers(plan)?; + let discovered = DynamicFilterIds { + consumers: consumers + .consumers + .into_iter() + .map(|consumer| consumer.id) + .collect(), + anchors: consumers + .anchors + .into_iter() + .map(|anchor| anchor.id) + .collect(), + producers: discover_dynamic_filter_producers(plan)? + .into_iter() + .map(|producer| producer.id) + .collect(), + }; + render(plan.as_ref(), 1, &discovered, &mut normalizer, &mut output)?; + } + Ok(output) + } +} diff --git a/tests/dynamic_filtering/main.rs b/tests/dynamic_filtering/main.rs index f7bdf8839..c00ae5966 100644 --- a/tests/dynamic_filtering/main.rs +++ b/tests/dynamic_filtering/main.rs @@ -2,5 +2,6 @@ mod aggregates; mod collect_left_join; mod common; mod config; +mod discovery; mod partitioned_join; mod sorts; From 41032c9b5c6ea0a1febff2de68d49444063f08f9 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Fri, 11 Sep 2026 20:39:34 +0000 Subject: [PATCH 5/6] only do dynamic filtering display and plan severing when enabled --- src/coordinator/distributed.rs | 7 +++++-- src/coordinator/query_coordinator.rs | 18 +++++++++++++----- src/dynamic_filtering/mod.rs | 8 ++++++++ 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 3d9481f90..e2b2508b7 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -3,7 +3,9 @@ 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::dynamic_filtering::{ + is_dynamic_filtering_enabled, 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; @@ -249,7 +251,8 @@ impl ExecutionPlan for DistributedExec { false => prepare_static_plan(&query_coordinator, &base_plan).await?, }; - prepared.plan_for_viz = match collect_dynamic_filters { + let dynamic_filtering_enabled = is_dynamic_filtering_enabled(context.session_config()); + prepared.plan_for_viz = match dynamic_filtering_enabled && collect_dynamic_filters { true => sever_dynamic_filter_relationships_in_plan_for_display( prepared.plan_for_viz, &context, diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 16f9976a5..499dab7ad 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -4,7 +4,10 @@ use crate::config_extension_ext::get_config_extension_propagation_headers; use crate::coordinator::DynamicFilterRegistry; 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::dynamic_filtering::{ + is_dynamic_filtering_enabled, + maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships, +}; use crate::events::{ RouteTaskEvent, RouteTaskEventResponse, RouteTaskHandlers, new_coordinator_to_worker_dialer, }; @@ -423,6 +426,7 @@ impl<'a> StageCoordinator<'a> { let wuf_registry = session_config .get_extension::() .unwrap_or_default(); + let dynamic_filtering_enabled = is_dynamic_filtering_enabled(session_config); let mut work_unit_feed_declarations = vec![]; let d_ctx = DistributedTaskContext { @@ -460,10 +464,14 @@ impl<'a> StageCoordinator<'a> { Ok(Transformed::no(plan)) })?; - let plan = maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships( - Arc::clone(&transformed.data), - self.task_ctx, - )?; + let plan = if dynamic_filtering_enabled { + maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships( + Arc::clone(&transformed.data), + self.task_ctx, + )? + } else { + transformed.data + }; Ok((plan, work_unit_feed_declarations)) } } diff --git a/src/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs index 4aca4801f..4d0f3206e 100644 --- a/src/dynamic_filtering/mod.rs +++ b/src/dynamic_filtering/mod.rs @@ -4,6 +4,7 @@ mod display; use crate::codec::roundtrip_pb; use datafusion::common::Result; use datafusion::execution::TaskContext; +use datafusion::execution::config::SessionConfig; use datafusion::physical_plan::ExecutionPlan; use std::sync::Arc; @@ -11,6 +12,13 @@ pub use discovery::*; pub use display::rewrite_distributed_plan_with_dynamic_filters; pub(crate) use display::sever_dynamic_filter_relationships_in_plan_for_display; +pub(crate) fn is_dynamic_filtering_enabled(session_config: &SessionConfig) -> bool { + session_config + .options() + .optimizer + .enable_dynamic_filter_pushdown +} + /// Isolates all shared, in-memory dynamic filter state from this plan if it contains /// any dynamic filter producers or consumers. /// From 2b1213f12c0db669a45b476010930f3f0fc236c8 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Fri, 11 Sep 2026 18:49:33 +0000 Subject: [PATCH 6/6] move dynamic filter anchors to remote stages --- src/codec/distributed_codec.rs | 165 ++++++++++-------- src/common/recursion.rs | 1 + src/coordinator/prepare_dynamic_plan.rs | 3 + src/coordinator/prepare_static_plan.rs | 3 + .../inject_network_boundaries.rs | 29 ++- .../benchmarks/shuffle_bench.rs | 2 +- .../benchmarks/transport_bench.rs | 2 +- src/execution_plans/network_broadcast.rs | 16 +- src/execution_plans/network_coalesce.rs | 16 +- src/execution_plans/network_shuffle.rs | 16 +- src/stage.rs | 10 ++ 11 files changed, 130 insertions(+), 133 deletions(-) diff --git a/src/codec/distributed_codec.rs b/src/codec/distributed_codec.rs index 09ad78abe..96277668a 100644 --- a/src/codec/distributed_codec.rs +++ b/src/codec/distributed_codec.rs @@ -14,8 +14,8 @@ use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::Result; use datafusion::error::DataFusionError; use datafusion::execution::TaskContext; -use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_expr::equivalence::{EquivalenceClass, EquivalenceGroup}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ExecutionPlan, Partitioning, PlanProperties}; @@ -66,11 +66,17 @@ impl PhysicalExtensionCodec for DistributedCodec { fn parse_stage_proto( proto: Option, inputs: &[Arc], + dynamic_filter_anchors: Vec>, ) -> Result { let Some(proto) = proto else { return Err(proto_error("Empty StageProto")); }; if let Some(input) = inputs.first().cloned() { + if !dynamic_filter_anchors.is_empty() { + return Err(proto_error( + "Dynamic filter anchors require a remote input stage", + )); + } Ok(Stage::Local(LocalStage { query_id: deserialize_uuid(proto.query_id.as_ref())?, num: proto.num as usize, @@ -94,6 +100,7 @@ impl PhysicalExtensionCodec for DistributedCodec { num: proto.num as usize, workers: worker_urls, runtime_stats: None, + dynamic_filter_anchors, })) } } @@ -104,7 +111,6 @@ impl PhysicalExtensionCodec for DistributedCodec { partitioning, input_stage, equivalence_classes, - dynamic_filter_anchors, }) => { let schema: Schema = schema .as_ref() @@ -119,8 +125,10 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )? .ok_or(proto_error("NetworkShuffleExec is missing partitioning"))?; - let dynamic_filter_anchors = dynamic_filter_anchors - .iter() + let dynamic_filter_anchors = input_stage + .as_ref() + .into_iter() + .flat_map(|stage| stage.dynamic_filter_anchors.iter()) .map(|expression| { proto_converter.proto_to_physical_expr(expression, &schema, &decode_ctx) }) @@ -133,21 +141,17 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )?; - Ok(Arc::new( - new_network_hash_shuffle_exec( - partitioning, - equivalence_properties, - parse_stage_proto(input_stage, inputs)?, - ) - .with_dynamic_filter_anchors(dynamic_filter_anchors), - )) + Ok(Arc::new(new_network_hash_shuffle_exec( + partitioning, + equivalence_properties, + parse_stage_proto(input_stage, inputs, dynamic_filter_anchors)?, + ))) } DistributedExecNode::NetworkCoalesceTasks(NetworkCoalesceExecProto { schema, partitioning, input_stage, equivalence_classes, - dynamic_filter_anchors, }) => { let schema: Schema = schema .as_ref() @@ -162,8 +166,10 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )? .ok_or(proto_error("NetworkCoalesceExec is missing partitioning"))?; - let dynamic_filter_anchors = dynamic_filter_anchors - .iter() + let dynamic_filter_anchors = input_stage + .as_ref() + .into_iter() + .flat_map(|stage| stage.dynamic_filter_anchors.iter()) .map(|expression| { proto_converter.proto_to_physical_expr(expression, &schema, &decode_ctx) }) @@ -176,21 +182,17 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )?; - Ok(Arc::new( - new_network_coalesce_tasks_exec( - partitioning, - equivalence_properties, - parse_stage_proto(input_stage, inputs)?, - ) - .with_dynamic_filter_anchors(dynamic_filter_anchors), - )) + Ok(Arc::new(new_network_coalesce_tasks_exec( + partitioning, + equivalence_properties, + parse_stage_proto(input_stage, inputs, dynamic_filter_anchors)?, + ))) } DistributedExecNode::NetworkBroadcast(NetworkBroadcastExecProto { schema, partitioning, input_stage, equivalence_classes, - dynamic_filter_anchors, }) => { let schema: Schema = schema .as_ref() @@ -205,8 +207,10 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )? .ok_or(proto_error("NetworkBroadcastExec is missing partitioning"))?; - let dynamic_filter_anchors = dynamic_filter_anchors - .iter() + let dynamic_filter_anchors = input_stage + .as_ref() + .into_iter() + .flat_map(|stage| stage.dynamic_filter_anchors.iter()) .map(|expression| { proto_converter.proto_to_physical_expr(expression, &schema, &decode_ctx) }) @@ -219,14 +223,11 @@ impl PhysicalExtensionCodec for DistributedCodec { proto_converter, )?; - Ok(Arc::new( - new_network_broadcast_exec( - partitioning, - equivalence_properties, - parse_stage_proto(input_stage, inputs)?, - ) - .with_dynamic_filter_anchors(dynamic_filter_anchors), - )) + Ok(Arc::new(new_network_broadcast_exec( + partitioning, + equivalence_properties, + parse_stage_proto(input_stage, inputs, dynamic_filter_anchors)?, + ))) } DistributedExecNode::Broadcast(BroadcastExecProto { consumer_task_count, @@ -303,12 +304,22 @@ impl PhysicalExtensionCodec for DistributedCodec { buf: &mut Vec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { - fn encode_stage_proto(stage: &Stage) -> Result { + fn encode_stage_proto( + stage: &Stage, + codec: &DistributedCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result { + let dynamic_filter_anchors = stage + .dynamic_filter_anchors() + .iter() + .map(|expression| proto_converter.physical_expr_to_proto(expression, codec)) + .collect::>>()?; Ok(match stage { Stage::Local(local) => StageProto { query_id: serialize_uuid(&local.query_id).into(), num: local.num as u64, tasks: vec![ExecutionTaskProto::default(); local.tasks], + dynamic_filter_anchors, }, Stage::Remote(remote) => { let mut tasks = Vec::with_capacity(remote.workers.len()); @@ -321,6 +332,7 @@ impl PhysicalExtensionCodec for DistributedCodec { query_id: serialize_uuid(&remote.query_id).into(), num: remote.num as u64, tasks, + dynamic_filter_anchors, } } }) @@ -334,17 +346,16 @@ impl PhysicalExtensionCodec for DistributedCodec { self, proto_converter, )?), - input_stage: Some(encode_stage_proto(node.input_stage())?), + input_stage: Some(encode_stage_proto( + node.input_stage(), + self, + proto_converter, + )?), equivalence_classes: serialize_equivalence_group( node.properties().equivalence_properties(), self, proto_converter, )?, - dynamic_filter_anchors: node - .dynamic_filter_anchors() - .iter() - .map(|expression| proto_converter.physical_expr_to_proto(expression, self)) - .collect::>>()?, }; let wrapper = DistributedExecProto { @@ -360,17 +371,16 @@ impl PhysicalExtensionCodec for DistributedCodec { self, proto_converter, )?), - input_stage: Some(encode_stage_proto(node.input_stage())?), + input_stage: Some(encode_stage_proto( + node.input_stage(), + self, + proto_converter, + )?), equivalence_classes: serialize_equivalence_group( node.properties().equivalence_properties(), self, proto_converter, )?, - dynamic_filter_anchors: node - .dynamic_filter_anchors() - .iter() - .map(|expression| proto_converter.physical_expr_to_proto(expression, self)) - .collect::>>()?, }; let wrapper = DistributedExecProto { @@ -386,17 +396,16 @@ impl PhysicalExtensionCodec for DistributedCodec { self, proto_converter, )?), - input_stage: Some(encode_stage_proto(node.input_stage())?), + input_stage: Some(encode_stage_proto( + node.input_stage(), + self, + proto_converter, + )?), equivalence_classes: serialize_equivalence_group( node.properties().equivalence_properties(), self, proto_converter, )?, - dynamic_filter_anchors: node - .dynamic_filter_anchors() - .iter() - .map(|expression| proto_converter.physical_expr_to_proto(expression, self)) - .collect::>>()?, }; let wrapper = DistributedExecProto { @@ -515,6 +524,9 @@ pub struct StageProto { /// the plan #[prost(message, repeated, tag = "3")] pub tasks: Vec, + /// Dynamic-filter consumers retained after a remote stage's plan has moved to its workers. + #[prost(message, repeated, tag = "4")] + pub dynamic_filter_anchors: Vec, } #[derive(Clone, PartialEq, ::prost::Message)] @@ -561,8 +573,6 @@ pub struct NetworkShuffleExecProto { input_stage: Option, #[prost(message, repeated, tag = "4")] equivalence_classes: Vec, - #[prost(message, repeated, tag = "5")] - dynamic_filter_anchors: Vec, } #[derive(Clone, PartialEq, ::prost::Message)] @@ -621,7 +631,6 @@ fn new_network_hash_shuffle_exec( )), worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, - dynamic_filter_anchors: vec![], } } @@ -638,8 +647,6 @@ pub struct NetworkCoalesceExecProto { input_stage: Option, #[prost(message, repeated, tag = "4")] equivalence_classes: Vec, - #[prost(message, repeated, tag = "5")] - dynamic_filter_anchors: Vec, } fn new_network_coalesce_tasks_exec( @@ -656,7 +663,6 @@ fn new_network_coalesce_tasks_exec( )), worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, - dynamic_filter_anchors: vec![], } } @@ -670,8 +676,6 @@ pub struct NetworkBroadcastExecProto { input_stage: Option, #[prost(message, repeated, tag = "4")] equivalence_classes: Vec, - #[prost(message, repeated, tag = "5")] - dynamic_filter_anchors: Vec, } #[derive(Clone, PartialEq, ::prost::Message)] @@ -697,7 +701,6 @@ fn new_network_broadcast_exec( )), worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, - dynamic_filter_anchors: vec![], } } @@ -730,6 +733,7 @@ mod tests { num: 0, workers: vec![], runtime_stats: None, + dynamic_filter_anchors: vec![], }) } @@ -786,14 +790,35 @@ mod tests { lit(true), )) as Arc; let expected_id = dynamic_filter.expression_id(); - let network: Arc = Arc::new( - new_network_hash_shuffle_exec( - Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 4), - EquivalenceProperties::new(schema), - dummy_stage(), - ) - .with_dynamic_filter_anchors(vec![Arc::clone(&dynamic_filter)]), + let stage = Stage::Remote(RemoteStage { + query_id: Default::default(), + num: 0, + workers: vec![], + runtime_stats: None, + dynamic_filter_anchors: vec![Arc::clone(&dynamic_filter)], + }); + let network: Arc = Arc::new(new_network_hash_shuffle_exec( + Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 4), + EquivalenceProperties::new(schema), + stage, + )); + + let mut buf = vec![]; + DistributedCodec.try_encode(Arc::clone(&network), &mut buf, &default_proto_converter())?; + let encoded = DistributedExecProto::decode(buf.as_slice()) + .map_err(|error| proto_error(format!("{error}")))?; + let Some(DistributedExecNode::NetworkHashShuffle(encoded)) = encoded.node else { + panic!("expected a network shuffle") + }; + assert_eq!( + encoded + .input_stage + .expect("network shuffle should contain its input stage") + .dynamic_filter_anchors + .len(), + 1, ); + let plan: Arc = Arc::new(FilterExec::try_new(dynamic_filter, network)?); let decoded = roundtrip_pb(plan, &ctx)?; @@ -803,7 +828,7 @@ mod tests { .downcast_ref::() .unwrap(); let network = filter.input().downcast_ref::().unwrap(); - let anchor = network.dynamic_filter_anchors()[0] + let anchor = network.input_stage().dynamic_filter_anchors()[0] .downcast_ref::() .unwrap(); diff --git a/src/common/recursion.rs b/src/common/recursion.rs index b44df3462..b7f2ec6bd 100644 --- a/src/common/recursion.rs +++ b/src/common/recursion.rs @@ -1012,6 +1012,7 @@ mod tests { num: 0, workers: vec![], runtime_stats: None, + dynamic_filter_anchors: vec![], })) .unwrap() } diff --git a/src/coordinator/prepare_dynamic_plan.rs b/src/coordinator/prepare_dynamic_plan.rs index 964c01580..db318dac2 100644 --- a/src/coordinator/prepare_dynamic_plan.rs +++ b/src/coordinator/prepare_dynamic_plan.rs @@ -5,6 +5,7 @@ use crate::distributed_planner::{ InjectNetworkBoundaryContext, NetworkBoundaryBuilderResult, ProducerHead, calculate_cost, inject_network_boundaries, }; +use crate::dynamic_filtering::orphan_dynamic_filter_consumers; use crate::events::TaskCountAnnotation::{Desired, Maximum}; use crate::execution_plans::SamplerExec; use crate::stage::{LocalStage, RemoteStage}; @@ -79,6 +80,7 @@ pub(super) async fn prepare_dynamic_plan( // In order to infer the compute the cost of the stage above this one, here a sampler // is injected to gather runtime statistics. input_stage.plan = ProducerHead::insert_sampler(input_stage.plan)?; + let dynamic_filter_anchors = orphan_dynamic_filter_consumers(&input_stage.plan)?; let mut load_info_rxs = Vec::with_capacity(input_stage.tasks); @@ -128,6 +130,7 @@ pub(super) async fn prepare_dynamic_plan( num: input_stage.num, workers, runtime_stats: stats, + dynamic_filter_anchors, }), input_properties, }) diff --git a/src/coordinator/prepare_static_plan.rs b/src/coordinator/prepare_static_plan.rs index 5bb35b63e..f83838d5c 100644 --- a/src/coordinator/prepare_static_plan.rs +++ b/src/coordinator/prepare_static_plan.rs @@ -1,6 +1,7 @@ use crate::common::TreeNodeExt; use crate::coordinator::distributed::PreparedPlan; use crate::coordinator::query_coordinator::QueryCoordinator; +use crate::dynamic_filtering::orphan_dynamic_filter_consumers; use crate::stage::RemoteStage; use crate::{NetworkBoundaryExt, Stage}; use datafusion::common::tree_node::Transformed; @@ -31,6 +32,7 @@ pub(super) async fn prepare_static_plan( let Stage::Local(stage) = plan.input_stage() else { return exec_err!("Input stage from network boundary was not in Local state"); }; + let dynamic_filter_anchors = orphan_dynamic_filter_consumers(&stage.plan)?; let mut stage_coordinator = query_coordinator.stage_coordinator(stage); let mut futures = Vec::with_capacity(stage.tasks); @@ -52,6 +54,7 @@ pub(super) async fn prepare_static_plan( num: stage.num, workers, runtime_stats: None, + dynamic_filter_anchors, }, ))?)) }); diff --git a/src/distributed_planner/inject_network_boundaries.rs b/src/distributed_planner/inject_network_boundaries.rs index 9f1d61dd5..1452397bd 100644 --- a/src/distributed_planner/inject_network_boundaries.rs +++ b/src/distributed_planner/inject_network_boundaries.rs @@ -1,5 +1,4 @@ use crate::distributed_planner::insert_broadcast::is_left_broadcast_safe; -use crate::dynamic_filtering::orphan_dynamic_filter_consumers; use crate::events::TaskCountAnnotation::{Desired, Maximum}; use crate::events::{ DesiredTaskCountEvent, DesiredTaskCountHandlers, ScaleUpLeafNodeEvent, ScaleUpLeafNodeHandlers, @@ -332,15 +331,14 @@ async fn _inject_network_boundaries( tasks: task_count.as_usize(), metrics_set: Default::default(), }; - let dynamic_filter_anchors = orphan_dynamic_filter_consumers(&input_stage.plan)?; let result = nb_ctx .nb_builder .build(input_stage, TypeId::of::(), nb_ctx) .await?; - let nb = Arc::new( - NetworkShuffleExec::from_stage(result.input_stage, result.input_properties) - .with_dynamic_filter_anchors(dynamic_filter_anchors), - ); + let nb = Arc::new(NetworkShuffleExec::from_stage( + result.input_stage, + result.input_properties, + )); Ok(nb_ctx.plan_with_task_count(nb, result.consumer_task_count)) } // Upon reaching a broadcast, we need to introduce a network broadcast right above it. @@ -352,15 +350,14 @@ async fn _inject_network_boundaries( tasks: task_count.as_usize(), metrics_set: Default::default(), }; - let dynamic_filter_anchors = orphan_dynamic_filter_consumers(&input_stage.plan)?; let result = nb_ctx .nb_builder .build(input_stage, TypeId::of::(), nb_ctx) .await?; - let nb = Arc::new( - NetworkBroadcastExec::from_stage(result.input_stage, result.input_properties) - .with_dynamic_filter_anchors(dynamic_filter_anchors), - ); + let nb = Arc::new(NetworkBroadcastExec::from_stage( + result.input_stage, + result.input_properties, + )); Ok(nb_ctx.plan_with_task_count(nb, result.consumer_task_count)) } // If the parent of the current node is either a `CoalescePartitionsExec` or a @@ -375,7 +372,6 @@ async fn _inject_network_boundaries( tasks: task_count.as_usize(), metrics_set: Default::default(), }; - let dynamic_filter_anchors = orphan_dynamic_filter_consumers(&input_stage.plan)?; let result = nb_ctx .nb_builder .build(input_stage, TypeId::of::(), nb_ctx) @@ -388,10 +384,11 @@ async fn _inject_network_boundaries( // The parent that triggered this branch is a `CoalescePartitionsExec` or // `SortPreservingMergeExec`, both of which fold all partitions into one — so the // stage above this boundary must run in exactly one task. - let nb = Arc::new( - NetworkCoalesceExec::try_from_stage(result.input_stage, result.input_properties, 1)? - .with_dynamic_filter_anchors(dynamic_filter_anchors), - ); + let nb = Arc::new(NetworkCoalesceExec::try_from_stage( + result.input_stage, + result.input_properties, + 1, + )?); Ok(nb_ctx.plan_with_task_count(nb, result.consumer_task_count)) } else if parent.is_none() { // We've just finished walking the head stage's subplan. Run a final propagation so diff --git a/src/execution_plans/benchmarks/shuffle_bench.rs b/src/execution_plans/benchmarks/shuffle_bench.rs index 4a636b188..0656f79d9 100644 --- a/src/execution_plans/benchmarks/shuffle_bench.rs +++ b/src/execution_plans/benchmarks/shuffle_bench.rs @@ -217,6 +217,7 @@ impl ShuffleFixture { num: 0, workers: self.input_stage_workers.clone(), runtime_stats: None, + dynamic_filter_anchors: vec![], }); let mut join_set = JoinSet::default(); @@ -230,7 +231,6 @@ impl ShuffleFixture { )), input_stage: input_stage.clone(), worker_connections: WorkerConnectionPool::new(self.bench.producer_tasks), - dynamic_filter_anchors: vec![], }; let task_ctx = Arc::new(task_ctx_with_extension( &self.task_ctx, diff --git a/src/execution_plans/benchmarks/transport_bench.rs b/src/execution_plans/benchmarks/transport_bench.rs index 57c5ddae8..9cbaa15e4 100644 --- a/src/execution_plans/benchmarks/transport_bench.rs +++ b/src/execution_plans/benchmarks/transport_bench.rs @@ -265,6 +265,7 @@ impl TransportFixture { num: 0, workers: self.input_stage_tasks.clone(), runtime_stats: None, + dynamic_filter_anchors: vec![], }); let mut join_set = JoinSet::default(); @@ -280,7 +281,6 @@ impl TransportFixture { worker_connections: crate::worker::WorkerConnectionPool::new( self.bench.producer_tasks, ), - dynamic_filter_anchors: vec![], }; let task_ctx = Arc::new(task_ctx_with_extension( &self.task_ctx, diff --git a/src/execution_plans/network_broadcast.rs b/src/execution_plans/network_broadcast.rs index 0d3290d9e..10581506a 100644 --- a/src/execution_plans/network_broadcast.rs +++ b/src/execution_plans/network_broadcast.rs @@ -122,7 +122,6 @@ pub struct NetworkBroadcastExec { pub(crate) properties: Arc, pub(crate) input_stage: Stage, pub(crate) worker_connections: WorkerConnectionPool, - pub(crate) dynamic_filter_anchors: Vec>, } impl NetworkBroadcastExec { @@ -137,22 +136,9 @@ impl NetworkBroadcastExec { properties, worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, - dynamic_filter_anchors: vec![], } } - pub(crate) fn with_dynamic_filter_anchors( - mut self, - dynamic_filter_anchors: Vec>, - ) -> Self { - self.dynamic_filter_anchors = dynamic_filter_anchors; - self - } - - pub(crate) fn dynamic_filter_anchors(&self) -> &[Arc] { - &self.dynamic_filter_anchors - } - /// Creates a new [NetworkBroadcastExec] fed by the provided [BroadcastExec]. The input plan /// will be executed in a remote worker in `producer_tasks` number of tasks. pub fn try_new(input: Arc, producer_tasks: usize) -> Result { @@ -235,7 +221,7 @@ impl ExecutionPlan for NetworkBroadcastExec { &self, f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - apply_expression_roots(self.dynamic_filter_anchors.iter(), f) + apply_expression_roots(self.input_stage.dynamic_filter_anchors().iter(), f) } fn with_new_children( diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index a6768f60c..db9b964a8 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -81,7 +81,6 @@ pub struct NetworkCoalesceExec { pub(crate) properties: Arc, pub(crate) input_stage: Stage, pub(crate) worker_connections: WorkerConnectionPool, - pub(crate) dynamic_filter_anchors: Vec>, } impl NetworkCoalesceExec { @@ -100,22 +99,9 @@ impl NetworkCoalesceExec { properties: props, worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, - dynamic_filter_anchors: vec![], }) } - pub(crate) fn with_dynamic_filter_anchors( - mut self, - dynamic_filter_anchors: Vec>, - ) -> Self { - self.dynamic_filter_anchors = dynamic_filter_anchors; - self - } - - pub(crate) fn dynamic_filter_anchors(&self) -> &[Arc] { - &self.dynamic_filter_anchors - } - /// Creates a new [NetworkCoalesceExec] fed by the provided `input` plan. /// /// The `input` plan will be remotely executed in `producer_tasks` tasks, while the @@ -258,7 +244,7 @@ impl ExecutionPlan for NetworkCoalesceExec { &self, f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - apply_expression_roots(self.dynamic_filter_anchors.iter(), f) + apply_expression_roots(self.input_stage.dynamic_filter_anchors().iter(), f) } fn with_new_children( diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index 8172848d2..a32d41120 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -106,7 +106,6 @@ pub struct NetworkShuffleExec { pub(crate) properties: Arc, pub(crate) input_stage: Stage, pub(crate) worker_connections: WorkerConnectionPool, - pub(crate) dynamic_filter_anchors: Vec>, } impl NetworkShuffleExec { @@ -115,22 +114,9 @@ impl NetworkShuffleExec { properties: input_properties, worker_connections: WorkerConnectionPool::new(input_stage.task_count()), input_stage, - dynamic_filter_anchors: vec![], } } - pub(crate) fn with_dynamic_filter_anchors( - mut self, - dynamic_filter_anchors: Vec>, - ) -> Self { - self.dynamic_filter_anchors = dynamic_filter_anchors; - self - } - - pub(crate) fn dynamic_filter_anchors(&self) -> &[Arc] { - &self.dynamic_filter_anchors - } - /// Creates a new [NetworkShuffleExec] fed by the provided [RepartitionExec]. The input plan /// will be executed in a remote worker in `producer_tasks` number of tasks. pub fn try_new(input: Arc, producer_tasks: usize) -> Result { @@ -212,7 +198,7 @@ impl ExecutionPlan for NetworkShuffleExec { &self, f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - apply_expression_roots(self.dynamic_filter_anchors.iter(), f) + apply_expression_roots(self.input_stage.dynamic_filter_anchors().iter(), f) } fn with_new_children( diff --git a/src/stage.rs b/src/stage.rs index 6e1c520cc..edd7af909 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -5,6 +5,7 @@ use datafusion::common::{HashMap, Statistics, config_err}; use datafusion::common::{exec_err, plan_err}; use datafusion::error::Result; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::physical_plan::metrics::{Label, Metric, MetricsSet}; use datafusion::physical_plan::{ @@ -114,6 +115,8 @@ pub struct RemoteStage { pub workers: Vec, /// Statistics collected at runtime, if any. pub runtime_stats: Option>, + /// Dynamic-filter consumers retained after the stage's plan is moved to its workers. + pub dynamic_filter_anchors: Vec>, } impl Stage { @@ -145,6 +148,13 @@ impl Stage { } } + pub(crate) fn dynamic_filter_anchors(&self) -> &[Arc] { + match self { + Self::Local(_) => &[], + Self::Remote(remote) => &remote.dynamic_filter_anchors, + } + } + pub fn metrics(&self) -> MetricsSet { match &self { Self::Local(v) => v.metrics_set.clone(),