From 520e39b6b402705943d6ca486f109ed8a379cd5f Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 11 Aug 2026 21:40:50 +0000 Subject: [PATCH 01/11] feat: apply merged dynamic filter updates on worker consumers Apply successive full dynamic filter snapshots in producer coordinates and mark completion only when reported. Reuse deduplicated consumer discovery, preserve alias remapping and task-local relationships, and verify remote join, TopK, and aggregate filters with SQL integration tests and enabled/disabled result comparisons. --- src/codec/mod.rs | 4 +- src/codec/physical_plan.rs | 61 ++++- src/coordinator/dynamic_filter_registry.rs | 19 +- src/dynamic_filtering/display.rs | 110 +++++--- src/worker/impl_coordinator_channel.rs | 90 +++++- tests/dynamic_filtering/aggregates.rs | 15 +- tests/dynamic_filtering/collect_left_join.rs | 85 +++++- tests/dynamic_filtering/common.rs | 251 ++++++++++++++--- tests/dynamic_filtering/config.rs | 1 - tests/dynamic_filtering/partitioned_join.rs | 274 +++++++++++++++++-- tests/dynamic_filtering/sorts.rs | 27 +- 11 files changed, 787 insertions(+), 150 deletions(-) diff --git a/src/codec/mod.rs b/src/codec/mod.rs index fc03f09fb..59ed4730e 100644 --- a/src/codec/mod.rs +++ b/src/codec/mod.rs @@ -4,8 +4,8 @@ mod user_codec; pub use distributed_codec::DistributedCodec; pub(crate) use physical_plan::{ - decode_execution_plan, decode_partitioning, decode_physical_expr, encode_execution_plan, - encode_partitioning, encode_physical_expr, roundtrip_pb, + decode_execution_plan, decode_partitioning, decode_physical_expr, dynamic_filter_update_target, + encode_execution_plan, encode_partitioning, encode_physical_expr, roundtrip_pb, }; pub(crate) use user_codec::{ get_distributed_user_codecs, set_distributed_user_codec, set_distributed_user_codec_arc, diff --git a/src/codec/physical_plan.rs b/src/codec/physical_plan.rs index f1a970ace..b7aeef6b0 100644 --- a/src/codec/physical_plan.rs +++ b/src/codec/physical_plan.rs @@ -1,7 +1,8 @@ use super::DistributedCodec; use datafusion::arrow::datatypes::{Schema, SchemaRef}; -use datafusion::common::Result; +use datafusion::common::{Result, internal_err}; use datafusion::execution::TaskContext; +use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_expr::{Partitioning, PhysicalExpr}; use datafusion::physical_plan::ExecutionPlan; use datafusion_proto::bytes::{ @@ -13,6 +14,7 @@ use datafusion_proto::physical_plan::{ DeduplicatingProtoConverter, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; use datafusion_proto::protobuf; +use datafusion_proto::protobuf::physical_expr_node::ExprType; use datafusion_proto::protobuf::proto_error; use prost::Message; use std::sync::Arc; @@ -78,6 +80,63 @@ pub(crate) fn decode_physical_expr( converter.proto_to_physical_expr(proto, input_schema, &decode_ctx) } +/// Creates a sythetic "producer" [`DynamicFilterPhysicalExpr`] which can be used to +/// update consumer dynamic filters via a shared state. Rather than updating +/// the provicder consumer expression directly, the returned expression should +/// be updated instead. +/// +/// DataFusion does a particular song and dance to update dynamic filters. This function +/// helps us do the same. +/// +/// ```text +/// HashJoinExec: on=[build.key = probe.key] +/// │ └── producer: original=[key], children=[key] ──────────────┐ +/// ├── build │ +/// └── UnionExec: probe │ +/// ├── DataSourceExec: phone_number AS key │ shared inner +/// │ └── consumer 1: original=[key], children=[phone_number] ─┤ +/// └── DataSourceExec: telephone AS key │ +/// └── consumer 2: original=[key], children=[telephone] ───┘ +/// ``` +/// +/// Behavior: +/// 1. [`DynamicFilterPhysicalExpr::update()`] remaps any occurrences of `original` to `children` +/// and stores the result in the shared state. +/// +/// 2. [`DynamicFilterPhysicalExpr::current()`] reads the remapped expression and remaps it again, +/// mapping `original` to `children` and returns it without storing. +/// +/// In the above plan, the producer calls update() so an expression like `key > 123` is mapped to +/// `key > 123` and this is stored. Then, the consumers call current(), reading from the same state, +/// to get `phone_number > 123` and `telephone > 123` respectively. +/// +/// You cannot update() consumers directly. In the above example, updating consumer 1 would remap +/// `key > 123` to `phone_number > 123` and store this in the shared state. Consumer 2 would be +/// unable to apply this filter now. +pub(crate) fn dynamic_filter_update_target( + consumer: &Arc, + input_schema: &Schema, + task_ctx: &TaskContext, +) -> Result> { + // Since consumer.children() returns the remapped children, we use the proto as a workaround to get the + // original children from the producer. + let consumer: Arc = consumer.clone(); + let proto = encode_physical_expr(&consumer, task_ctx)?; + let Some(ExprType::DynamicFilter(dynamic_filter)) = proto.expr_type else { + return internal_err!("expected a dynamic filter expression"); + }; + let original_children = dynamic_filter + .children + .iter() + .map(|child| decode_physical_expr(child, input_schema, task_ctx)) + .collect::>>()?; + let update_target = Arc::clone(&consumer).with_new_children(original_children)?; + let Ok(update_target) = Arc::downcast::(update_target) else { + return internal_err!("expected a dynamic filter update target"); + }; + Ok(update_target) +} + pub(crate) fn encode_partitioning( partitioning: &Partitioning, task_ctx: &TaskContext, diff --git a/src/coordinator/dynamic_filter_registry.rs b/src/coordinator/dynamic_filter_registry.rs index 4905f5c9c..c5f35c91f 100644 --- a/src/coordinator/dynamic_filter_registry.rs +++ b/src/coordinator/dynamic_filter_registry.rs @@ -29,8 +29,7 @@ pub(super) enum DynamicFilterMergeMode { /// Wait for any producer to report a complete dynamic filter and forward it. /// Used for collect left joins. FirstProducerComplete, - /// Forward any update received from any producer, merging updates from - /// different producers. Used for TopK dynamic filters in aggregates and sorts. + /// Forward any update received from any producer. Incremental, } @@ -288,6 +287,7 @@ impl DynamicFilterRegistry { .iter() .filter_map(|(_, report)| report.inner_expr.as_deref().cloned()) .collect(), + mode, ) .map(Box::new); let is_complete = mode == DynamicFilterMergeMode::FirstProducerComplete || all_complete; @@ -356,8 +356,12 @@ impl DynamicFilterRegistry { } } -/// Merges [`PhysicalExprNode`] together by ORing them. -fn merge_predicates(mut predicates: Vec) -> Option { +/// Merges [`PhysicalExprNode`] together. Merge behavior (ex. AND vs OR) is determined +/// by the provided [`DynamicFilterMergeMode`]. +fn merge_predicates( + mut predicates: Vec, + mode: DynamicFilterMergeMode, +) -> Option { match predicates.len() { 0 => None, 1 => predicates.pop(), @@ -366,7 +370,12 @@ fn merge_predicates(mut predicates: Vec) -> Option, + stage: &LocalStage, reports: &HashMap, task_ctx: &Arc, -) { - let _ = plan.apply(|node| { - let Some(leaf) = node.downcast_ref::() else { - return Ok(TreeNodeRecursion::Continue); +) -> Result<()> { + for task_number in 0..stage.tasks { + let task_key = TaskKey { + query_id: stage.query_id, + stage_id: stage.num, + task_number, }; - - for (task_key, report) in reports { - let Some(variant) = leaf.variants().get(task_key.task_number) else { - continue; + let Some(report) = reports.get(&task_key) else { + continue; + }; + let updates: HashMap<_, _> = report + .filters + .iter() + .map(|filter| (filter.expression_id, &filter.expression)) + .collect(); + let d_ctx = DistributedTaskContext { + task_index: task_number, + task_count: stage.tasks, + }; + // Union children can use different task indices from their enclosing stage. + stage.plan.apply_with_dt_ctx(d_ctx, |node, d_ctx| { + let Some(leaf) = node.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); }; - let updates: HashMap<_, _> = report - .filters - .iter() - .map(|filter| (filter.expression_id, &filter.expression)) - .collect(); - let Ok(discovered) = discover_dynamic_filter_consumers(variant) else { - continue; + let Some(variant) = leaf.variants().get(d_ctx.task_index) else { + return internal_err!( + "missing leaf variant {} for task {} in stage {}", + d_ctx.task_index, + task_number, + stage.num + ); }; + let discovered = discover_dynamic_filter_consumers(variant)?; for consumer in discovered.consumers { let Some(expression) = updates.get(&consumer.id).copied() else { continue; }; - let Ok(proto) = expression.to_proto(task_ctx) else { - continue; - }; + let proto = expression.to_proto(task_ctx)?; let Some(ExprType::DynamicFilter(dynamic_filter_proto)) = proto.expr_type.as_ref() else { - continue; - }; - let Ok(reported_expression) = - decode_physical_expr(&proto, consumer.input_schema.as_ref(), task_ctx) - else { - continue; - }; - let Some(reported_dynamic_filter) = - reported_expression.downcast_ref::() - else { - continue; + return internal_err!("expected a dynamic filter in the completed task report"); }; - let Ok(expression) = reported_dynamic_filter.current() else { + if dynamic_filter_proto.generation <= 1 { continue; - }; - if dynamic_filter_proto.generation > 1 { - let _ = consumer.expression.update(expression); } + let Some(predicate) = dynamic_filter_proto.inner_expr.as_deref() else { + return internal_err!("reported dynamic filter has no predicate"); + }; + let predicate = + decode_physical_expr(predicate, consumer.input_schema.as_ref(), task_ctx)?; + let dynamic_filter = dynamic_filter_update_target( + &consumer.expression, + consumer.input_schema.as_ref(), + task_ctx, + )?; + dynamic_filter.update(predicate)?; } - } - - Ok(TreeNodeRecursion::Continue) - }); + Ok(TreeNodeRecursion::Continue) + })?; + } + Ok(()) } #[cfg(test)] @@ -264,7 +281,14 @@ mod tests { report, )]); - apply_reports_to_distributed_leaves(&isolated, &reports, &task_ctx); + let stage = LocalStage { + query_id: Uuid::nil(), + num: 1, + plan: Arc::clone(&isolated), + tasks: 2, + metrics_set: Default::default(), + }; + apply_reports_to_distributed_leaves(&stage, &reports, &task_ctx)?; let leaf = isolated.downcast_ref::().unwrap(); let task_0 = displayable(leaf.variants()[0].as_ref()) .one_line() diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index e1dcd1305..6bf0a8f6c 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,6 +1,7 @@ +use crate::codec::{decode_physical_expr, dynamic_filter_update_target}; use crate::common::TreeNodeExt; use crate::dynamic_filtering::{ - discover_dynamic_filter_consumers, discover_dynamic_filter_producers, + DiscoveredDynamicFilter, discover_dynamic_filter_consumers, discover_dynamic_filter_producers, }; use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers}; use crate::execution_plans::SamplerExec; @@ -10,25 +11,31 @@ use crate::protocol::grpc::on_drop_stream; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::task_data::TaskDataMetrics; use crate::{ - CoordinatorToWorkerMsg, DistributedConfig, DistributedExt, DistributedTaskContext, - MaybeEncoded, ProducedDynamicFilter, SetPlanRequest, TaskCompletedDynamicFilters, TaskData, - TaskDynamicFilter, TaskMetrics, Worker, WorkerQueryContext, WorkerToCoordinatorMsg, + ApplyDynamicFilter, CoordinatorToWorkerMsg, DistributedConfig, DistributedExt, + DistributedTaskContext, MaybeEncoded, ProducedDynamicFilter, SetPlanRequest, + TaskCompletedDynamicFilters, TaskData, TaskDynamicFilter, TaskMetrics, Worker, + WorkerQueryContext, WorkerToCoordinatorMsg, }; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{DataFusionError, HashSet, Result, exec_datafusion_err}; +use datafusion::common::{ + DataFusionError, HashMap, HashSet, Result, exec_datafusion_err, internal_err, +}; use datafusion::execution::{SessionStateBuilder, TaskContext}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; +use datafusion_proto::protobuf::physical_expr_node::ExprType; use futures::stream::{BoxStream, FuturesUnordered, select_all}; use futures::{FutureExt, StreamExt, TryStreamExt}; use http::HeaderMap; #[cfg(feature = "integration")] use std::sync::atomic::Ordering; use std::sync::{Arc, OnceLock}; +use tokio::sync::mpsc::channel; use tokio::sync::oneshot::Sender; use tokio::sync::{oneshot, watch}; +use tokio_stream::wrappers::ReceiverStream; /// Return value of the [Worker::coordinator_channel] method. pub struct CoordinatorChannelResult { @@ -61,6 +68,7 @@ impl Worker { let (metrics_tx, metrics_rx) = oneshot::channel(); let (dynamic_filters_tx, dynamic_filters_rx) = oneshot::channel(); + let (dynamic_filters_error_tx, dynamic_filters_error_rx) = channel(1); let task_data = || async { let mut cfg = SessionConfig::default() @@ -141,11 +149,17 @@ impl Worker { let producer_filters = discover_dynamic_filter_producers(&task_data.base_plan)? .into_iter() .filter(|producer| dynamic_filter_remote_producer_ids.contains(&producer.id)); + let dynamic_filter_consumers = discover_dynamic_filter_consumers(&task_data.base_plan)? + .consumers + .into_iter() + .map(|consumer| (consumer.id, consumer)) + .collect(); let (producer_cancel_tx, producer_cancel_rx) = watch::channel(false); // Continue reading remaining messages (work unit feed data) in the background. let mut work_unit_senders = Some(remote_work_unit_feed_registry.senders); let task_data_entries = Arc::clone(&self.task_data_entries); + let dynamic_filter_task_ctx = Arc::clone(&task_data.task_ctx); // This tokio task takes ownership of the final-report senders that keep the // worker->coordinator stream alive. As soon as this task ends, the runtime metrics and @@ -195,9 +209,15 @@ impl Worker { CoordinatorToWorkerMsg::KickOffSampling => { sampler_gate.kick_off(); } - CoordinatorToWorkerMsg::ApplyDynamicFilter(_) => { - // Runtime application is introduced independently from the routing - // protocol. Until then, accepting the message is intentionally a no-op. + CoordinatorToWorkerMsg::ApplyDynamicFilter(filter) => { + if let Err(error) = apply_merged_dynamic_filter( + *filter, + &dynamic_filter_consumers, + &dynamic_filter_task_ctx, + ) { + let _ = dynamic_filters_error_tx.try_send(error); + break; + } } } } @@ -273,12 +293,14 @@ impl Worker { })); let stream = select_all([ - produced_dynamic_filters_stream.boxed(), - load_info_stream.boxed(), - metrics_stream.boxed(), - dynamic_filters_stream.boxed(), + produced_dynamic_filters_stream.map(Ok).boxed(), + load_info_stream.map(Ok).boxed(), + metrics_stream.map(Ok).boxed(), + dynamic_filters_stream.map(Ok).boxed(), + ReceiverStream::new(dynamic_filters_error_rx) + .map(Err) + .boxed(), ]) - .map(Ok) .boxed(); #[cfg(feature = "integration")] @@ -337,6 +359,48 @@ fn produced_dynamic_filter_stream( .boxed() } +fn apply_merged_dynamic_filter( + filter: ApplyDynamicFilter, + consumers: &HashMap, + task_ctx: &Arc, +) -> Result<()> { + let Some(consumer) = consumers.get(&filter.expression_id) else { + return internal_err!( + "received dynamic filter update for unknown consumer {}", + filter.expression_id + ); + }; + + // Bytes to PhysicalExprNode + let expression = filter.expression.to_proto(task_ctx)?; + + if expression.expr_id != Some(filter.expression_id) { + return internal_err!("dynamic filter update has a mismatched expression ID"); + } + + let Some(ExprType::DynamicFilter(update)) = expression.expr_type else { + return internal_err!("expected a dynamic filter update"); + }; + + let Some(predicate) = update.inner_expr else { + return internal_err!("dynamic filter update has no predicate"); + }; + + // PhysicalExprNode to PhysicalExpr + let predicate = decode_physical_expr(&predicate, consumer.input_schema.as_ref(), task_ctx)?; + + let dynamic_filter = dynamic_filter_update_target( + &consumer.expression, + consumer.input_schema.as_ref(), + task_ctx, + )?; + dynamic_filter.update(predicate)?; + if update.is_complete { + dynamic_filter.mark_complete(); + } + Ok(()) +} + /// Finds all consumed dynamic filters for the completed task report. /// /// Note that it's possible that a dynamic filter is consumed by the leaf, updated by diff --git a/tests/dynamic_filtering/aggregates.rs b/tests/dynamic_filtering/aggregates.rs index c650dbcd8..e6b6464fb 100644 --- a/tests/dynamic_filtering/aggregates.rs +++ b/tests/dynamic_filtering/aggregates.rs @@ -16,6 +16,8 @@ mod tests { ) .execute() .await?; + + // Each task in stage 1 produces a distinct filter to push down. assert_snapshot!(display, @" ┌───── DistributedExec │ AggregateExec: mode=Final, gby=[], aggr=[min(weather.MinTemp)] @@ -26,14 +28,14 @@ mod tests { │ AggregateExec: mode=Partial, gby=[], aggr=[min(weather.MinTemp)] │ FilterExec: RainToday@1 = Yes, projection=[MinTemp@0] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp, RainToday], file_type=parquet, predicate=RainToday@19 = Yes AND DynamicFilter [ MinTemp@0 < 4.3 ], dynamic_rg_pruning=eligible, pruning_predicate=RainToday_null_count@2 != row_count@3 AND RainToday_min@0 <= Yes AND Yes <= RainToday_max@1 AND MinTemp_null_count@5 != row_count@3 AND MinTemp_min@4 < 4.3, required_guarantees=[RainToday in (Yes)] - │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp, RainToday], file_type=parquet, predicate=RainToday@19 = Yes AND DynamicFilter [ MinTemp@0 < -1.6 ], dynamic_rg_pruning=eligible, pruning_predicate=RainToday_null_count@2 != row_count@3 AND RainToday_min@0 <= Yes AND Yes <= RainToday_max@1 AND MinTemp_null_count@5 != row_count@3 AND MinTemp_min@4 < -1.6, required_guarantees=[RainToday in (Yes)] + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp, RainToday], file_type=parquet, predicate=RainToday@19 = Yes AND DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp, RainToday], file_type=parquet, predicate=RainToday@19 = Yes AND DynamicFilter [ expression_id_0_hash_1 ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── "); Ok(()) } - /// A partial aggregate does not yet update a data source across a shuffle. + /// A partial aggregate updates a data source across a shuffle. #[tokio::test] async fn remote_dynamic_filters() -> Result<()> { let display = TestQuery::new( @@ -45,9 +47,10 @@ mod tests { ) "#, ) - .expect_dynamic_filter_updates() .execute() .await?; + + // The per-task partial aggregate filters in stage 2 are merged into one and passed to both consumers. assert_snapshot!(display, @" ┌───── DistributedExec │ AggregateExec: mode=Final, gby=[], aggr=[min(key)] @@ -63,8 +66,8 @@ mod tests { │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp@0 as key], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible - │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp@0 as key], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp@0 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp@0 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── "); Ok(()) diff --git a/tests/dynamic_filtering/collect_left_join.rs b/tests/dynamic_filtering/collect_left_join.rs index 49cd54e70..55ff2e18e 100644 --- a/tests/dynamic_filtering/collect_left_join.rs +++ b/tests/dynamic_filtering/collect_left_join.rs @@ -1,6 +1,6 @@ #[cfg(test)] mod tests { - use crate::common::TestQuery; + use crate::common::{LOCAL_AND_REMOTE_UNION_QUERY, TestQuery}; use datafusion::common::Result; use datafusion_distributed::assert_snapshot; @@ -20,6 +20,9 @@ mod tests { .with_broadcast_joins() .execute() .await?; + + // Each per-task join in stage 3 produces an identical filter (due to the common build side) and passes + // it down to the local consumer. assert_snapshot!(display, @r" ┌───── DistributedExec │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] @@ -33,8 +36,8 @@ mod tests { │ CoalescePartitionsExec │ [Stage 2] => NetworkBroadcastExec: partitions_per_consumer=3, stage_partitions=6, input_tasks=2 │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ RainToday@19 >= No AND RainToday@19 <= Yes AND RainToday@19 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=RainToday_null_count@1 != row_count@2 AND RainToday_max@0 >= No AND RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND (RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND Yes <= RainToday_max@0 OR RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= No AND No <= RainToday_max@0), required_guarantees=[RainToday in (No, Yes)] - │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ RainToday@19 >= No AND RainToday@19 <= Yes AND RainToday@19 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=RainToday_null_count@1 != row_count@2 AND RainToday_max@0 >= No AND RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND (RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= Yes AND Yes <= RainToday_max@0 OR RainToday_null_count@1 != row_count@2 AND RainToday_min@3 <= No AND No <= RainToday_max@0), required_guarantees=[RainToday in (No, Yes)] + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── ┌───── Stage 2 ── tasks=2, partitions=12 │ BroadcastExec: input_partitions=3, consumer_tasks=2, output_partitions=6 @@ -52,7 +55,7 @@ mod tests { Ok(()) } - /// A CollectLeft HashJoinExec does not propagate dynamic filters to a remote consumer. + /// A CollectLeft HashJoinExec propagates its first complete filter to remote consumers. #[tokio::test] async fn remote_dynamic_filters() -> Result<()> { let display = TestQuery::new( @@ -69,10 +72,12 @@ mod tests { "#, ) .with_broadcast_joins() - .expect_dynamic_filter_updates() .execute() .await?; - assert_snapshot!(display, @r" + + // The dynamic filter is propagated from the per-task joins in stage 4 + // to the consumers. All consumers see the same filter due to a common build side. + assert_snapshot!(display, @" ┌───── DistributedExec │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] @@ -103,13 +108,73 @@ mod tests { │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible - │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── "); Ok(()) } + /// Task 1's filter reports must reach variant 0 of its isolated union child. + #[tokio::test] + async fn union_with_local_and_remote_probe_children() -> Result<()> { + let display = TestQuery::new(LOCAL_AND_REMOTE_UNION_QUERY) + .with_broadcast_joins() + .execute() + .await?; + // The consumer in c0 of t0 in stage 5 sees the dynamic filter from the topmost join. + // The consumer in c1 of t1 in stage 5 sees distinct local dynamic filters from both joins. + // DataFusion ANDs distinct dynamic filter predicates. + assert_snapshot!(display, @" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 5] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 5 ── tasks=2, partitions=6 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(key@0, key@0)], projection=[] + │ CoalescePartitionsExec + │ [Stage 2] => NetworkBroadcastExec: partitions_per_consumer=3, stage_partitions=6, input_tasks=2 + │ DistributedUnionExec: t0:[c0] t1:[c1] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir@5 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ ProjectionExec: expr=[WindGustDir@0 as key] + │ HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(key@0, RainToday@1)], projection=[WindGustDir@0] + │ CoalescePartitionsExec + │ [Stage 4] => NetworkBroadcastExec: partitions_per_consumer=3, stage_partitions=3, input_tasks=2 + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir, RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_1_hash_1 ] AND DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=12 + │ BroadcastExec: input_partitions=3, consumer_tasks=2, output_partitions=6 + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir@5 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir@5 as key], file_type=parquet + └────────────────────────────────────────────────── + ┌───── Stage 4 ── tasks=2, partitions=6 + │ BroadcastExec: input_partitions=3, consumer_tasks=1, output_partitions=3 + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 3] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 3 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + └────────────────────────────────────────────────── + "); + Ok(()) + } + #[tokio::test] async fn union_probe_deduplicates_dynamic_filters() -> Result<()> { let display = TestQuery::new( @@ -154,11 +219,11 @@ mod tests { │ [Stage 1] => NetworkBroadcastExec: partitions_per_consumer=3, stage_partitions=6, input_tasks=1 │ DistributedUnionExec: t0:[c0, c2, c4] t1:[c1, c3] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp], file_type=parquet, predicate=DynamicFilter [ MinTemp@0 >= -5.3 AND MinTemp@0 <= 20.9 AND true ], dynamic_rg_pruning=eligible, pruning_predicate=MinTemp_null_count@1 != row_count@2 AND MinTemp_max@0 >= -5.3 AND MinTemp_null_count@1 != row_count@2 AND MinTemp_min@3 <= 20.9, required_guarantees=[] + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible │ ProjectionExec: expr=[CAST(-1000 AS Float64) as MinTemp] │ PlaceholderRowExec │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp], file_type=parquet, predicate=DynamicFilter [ MinTemp@0 >= -5.3 AND MinTemp@0 <= 20.9 AND true ], dynamic_rg_pruning=eligible, pruning_predicate=MinTemp_null_count@1 != row_count@2 AND MinTemp_max@0 >= -5.3 AND MinTemp_null_count@1 != row_count@2 AND MinTemp_min@3 <= 20.9, required_guarantees=[] + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible │ ProjectionExec: expr=[CAST(-1001 AS Float64) as MinTemp] │ PlaceholderRowExec │ ProjectionExec: expr=[CAST(-1002 AS Float64) as MinTemp] diff --git a/tests/dynamic_filtering/common.rs b/tests/dynamic_filtering/common.rs index 104a76eb2..58a668f2d 100644 --- a/tests/dynamic_filtering/common.rs +++ b/tests/dynamic_filtering/common.rs @@ -1,11 +1,15 @@ use datafusion::arrow::datatypes::DataType; -use datafusion::common::{Result, ScalarValue, SplitPoint}; +use datafusion::common::test_util::batches_to_sort_string; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion::common::{HashMap, Result, ScalarValue, SplitPoint, internal_err}; use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; use datafusion::logical_expr::{Partitioning, RangePartitioning}; -use datafusion::physical_plan::collect; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, UnKnownColumn}; +use datafusion::physical_plan::{ExecutionPlan, collect}; use datafusion::prelude::{SessionContext, col}; use datafusion_distributed::test_utils::localhost::start_localhost_context; use datafusion_distributed::test_utils::parquet::register_parquet_tables; @@ -13,19 +17,42 @@ use datafusion_distributed::test_utils::routing::{ ColocateAllTasksHandler, UrlEmitterRouteTaskHandler, }; use datafusion_distributed::{ - DefaultSessionBuilder, DistributedExt, display_plan_ascii, - rewrite_distributed_plan_with_dynamic_filters, + DefaultSessionBuilder, DistributedConfig, DistributedExt, DistributedLeafExec, + display_plan_ascii, rewrite_distributed_plan_with_dynamic_filters, }; use std::sync::Arc; +pub(crate) const LOCAL_AND_REMOTE_UNION_QUERY: &str = r#" + WITH remote_probe AS ( + SELECT probe."WindGustDir" AS key + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) nested_build + JOIN weather probe + ON nested_build.key = probe."RainToday" + ) + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "WindGustDir" AS key + FROM weather + ) build + JOIN ( + SELECT "WindGustDir" AS key FROM weather + UNION ALL + SELECT key FROM remote_probe + ) probe ON build.key = probe.key +"#; + pub(crate) struct TestQuery<'a> { sql: &'a str, expected_rows: usize, broadcast_joins: bool, one_task_per_leaf: bool, + dynamic_task_count: bool, collect_dynamic_filters: bool, remote_dynamic_filters: bool, - expect_dynamic_filter_updates: Option, + normalized_filter_hashes: bool, } impl<'a> TestQuery<'a> { @@ -35,9 +62,10 @@ impl<'a> TestQuery<'a> { expected_rows: 1, broadcast_joins: false, one_task_per_leaf: false, + dynamic_task_count: false, collect_dynamic_filters: true, remote_dynamic_filters: true, - expect_dynamic_filter_updates: None, + normalized_filter_hashes: false, } } @@ -59,27 +87,29 @@ impl<'a> TestQuery<'a> { self } + /// Enables the dynamic task-count planner. + pub(crate) fn with_dynamic_task_count(mut self) -> Self { + self.dynamic_task_count = true; + self + } + /// Disables dynamic filter collection. pub(crate) fn without_dynamic_filter_collection(mut self) -> Self { self.collect_dynamic_filters = false; self } - /// Disables distributing dynamic filters across network boundaries. + /// Disables distributing dynamic filters across network boundaries and asserts that + /// the coordinator received no dynamic filter updates. pub(crate) fn without_remote_dynamic_filters(mut self) -> Self { self.remote_dynamic_filters = false; self } - /// Asserts that the coordinator received at least one dynamic filter update. - pub(crate) fn expect_dynamic_filter_updates(mut self) -> Self { - self.expect_dynamic_filter_updates = Some(true); - self - } - - /// Asserts that the coordinator received no dynamic filter updates. - pub(crate) fn expect_no_dynamic_filter_updates(mut self) -> Self { - self.expect_dynamic_filter_updates = Some(false); + /// Adds a normalized hash to the dynamic filter display formatter. See + /// [`DynamicFilterLabels`] below. + pub(crate) fn with_normalized_filter_hashes(mut self) -> Self { + self.normalized_filter_hashes = true; self } @@ -89,6 +119,7 @@ impl<'a> TestQuery<'a> { .with_distributed_broadcast_joins(self.broadcast_joins)? .with_distributed_dynamic_filter_collection(self.collect_dynamic_filters)? .with_distributed_remote_dynamic_filters(self.remote_dynamic_filters)?; + ctx.set_distributed_dynamic_task_count(self.dynamic_task_count)?; if self.one_task_per_leaf { ctx = ctx.with_distributed_desired_task_count_handler(1usize); } @@ -108,7 +139,7 @@ impl<'a> TestQuery<'a> { self.sql, self.expected_rows, self.collect_dynamic_filters, - self.expect_dynamic_filter_updates, + self.normalized_filter_hashes, ) .await } @@ -140,7 +171,7 @@ pub(crate) async fn execute_range_partitioned_query( register_range_partitioned_table(&ctx, "dim", "testdata/join/parquet/dim", "d_dkey").await?; register_range_partitioned_table(&ctx, "fact", "testdata/join/parquet/fact", "f_dkey").await?; - execute_query_and_display(&ctx, sql, expected_rows, true, None).await + execute_query_and_display(&ctx, sql, expected_rows, true, false).await } async fn register_range_partitioned_table( @@ -172,17 +203,31 @@ async fn execute_query_and_display( sql: &str, expected_rows: usize, collect_dynamic_filters: bool, - expect_dynamic_filter_updates: Option, + normalized_filter_hashes: bool, ) -> Result { + set_dynamic_filter_pushdown(ctx, true)?; let plan = ctx.sql(sql).await?.create_physical_plan().await?; let task_ctx = ctx.task_ctx(); - let results = collect(Arc::clone(&plan), Arc::clone(&task_ctx)).await?; + let results_with_dynamic_filters = collect(Arc::clone(&plan), Arc::clone(&task_ctx)).await?; assert_eq!( - results.iter().map(|batch| batch.num_rows()).sum::(), + results_with_dynamic_filters + .iter() + .map(|batch| batch.num_rows()) + .sum::(), expected_rows ); + set_dynamic_filter_pushdown(ctx, false)?; + let plan_without_dynamic_filters = ctx.sql(sql).await?.create_physical_plan().await?; + let results_without_dynamic_filters = + collect(plan_without_dynamic_filters, ctx.task_ctx()).await?; + assert_eq!( + batches_to_sort_string(&results_with_dynamic_filters), + batches_to_sort_string(&results_without_dynamic_filters), + "query results changed when dynamic filtering was enabled", + ); + let original_display = display_plan_ascii(plan.as_ref(), false); let plan_with_dynamic_filters = rewrite_distributed_plan_with_dynamic_filters(Arc::clone(&plan), &task_ctx).await?; @@ -192,27 +237,161 @@ async fn execute_query_and_display( ); assert_eq!(display_plan_ascii(plan.as_ref(), false), original_display); - if let Some(expect_updates) = expect_dynamic_filter_updates { + if !DistributedConfig::from_task_context(&task_ctx)?.remote_dynamic_filters { let updates = plan .metrics() .expect("DistributedExec has metrics") .sum(|metric| metric.value().name() == "dynamic_filter_updates_received") .map_or(0, |metric| metric.as_usize()); - if expect_updates { - assert!( - updates > 0, - "expected dynamic_filter_updates_received > 0, got {updates}" - ); - } else { - assert_eq!( - updates, 0, - "expected dynamic_filter_updates_received == 0, got {updates}" - ); + assert_eq!( + updates, 0, + "expected dynamic_filter_updates_received == 0, got {updates}" + ); + } + DynamicFilterLabels { + normalized_predicates: normalized_filter_hashes.then(Vec::new), + ..Default::default() + } + .normalize(plan_with_dynamic_filters) +} + +/// Encodes dynamic filter expressions in the plan deterministically for consistent snapshots. +/// +/// They are encoded as `expression_id_{}_hash_{}` where expression_id is the expression id +/// of the filter (which generally denotes the producer it came from) and hash of the expression. +/// +/// Both of these are normalized to be monotonic integers starting from `0` for readability, +/// assigned during a pre-order traversal of the plan. +/// +/// If `normalized_predicates` is non empty, then we append the suffix `_normalized_hash_{}` +/// containing the hash of the expression without column names. +#[derive(Default)] +struct DynamicFilterLabels { + expression_ids: HashMap, + // InListExpr ignores element order for equality, but not hashing. + predicates: Vec>, + normalized_predicates: Option>>, +} + +impl DynamicFilterLabels { + fn normalize(&mut self, plan: Arc) -> Result { + self.label_plan(&plan)?; + Ok(remove_runtime_pruning_details(display_plan_ascii( + plan.as_ref(), + false, + ))) + } + + fn label_plan(&mut self, plan: &Arc) -> Result<()> { + let mut updates = vec![]; + plan.apply(|node| { + if let Some(leaf) = node.downcast_ref::() { + for variant in leaf.variants() { + self.label_variant(variant, &mut updates)?; + } + } + Ok(TreeNodeRecursion::Continue) + })?; + // Read every predicate before changing any potentially shared filter state. + for (dynamic_filter, label) in updates { + dynamic_filter.update(Arc::new(UnKnownColumn::new(&label)))?; } + Ok(()) } - Ok(display_plan_ascii( - plan_with_dynamic_filters.as_ref(), - false, - )) + fn label_variant( + &mut self, + variant: &Arc, + updates: &mut Vec<(Arc, String)>, + ) -> Result<()> { + variant.apply(|node| { + node.apply_expressions(&mut |root| { + root.apply(|expression| { + let Ok(dynamic_filter) = + Arc::downcast::(expression.clone()) + else { + return Ok(TreeNodeRecursion::Continue); + }; + if expression.snapshot_generation() == 1 { + return Ok(TreeNodeRecursion::Continue); + } + let Some(expression_id) = expression.expression_id() else { + return internal_err!("dynamic filter did not have an expression ID"); + }; + let next_expression_id = self.expression_ids.len(); + let expression_id = *self + .expression_ids + .entry(expression_id) + .or_insert(next_expression_id); + let predicate = dynamic_filter.current()?; + let predicate_id = + intern_predicate(&mut self.predicates, Arc::clone(&predicate)); + let mut label = format!("expression_id_{expression_id}_hash_{predicate_id}"); + if let Some(predicates) = &mut self.normalized_predicates { + let normalized_id = + intern_predicate(predicates, normalize_columns(predicate)?); + label.push_str(&format!("_normalized_hash_{normalized_id}")); + } + updates.push((dynamic_filter, label)); + Ok(TreeNodeRecursion::Continue) + }) + })?; + Ok(TreeNodeRecursion::Continue) + })?; + Ok(()) + } +} + +fn intern_predicate( + predicates: &mut Vec>, + predicate: Arc, +) -> usize { + predicates + .iter() + .position(|existing| existing.eq(&predicate)) + .unwrap_or_else(|| { + let id = predicates.len(); + predicates.push(predicate); + id + }) +} + +fn normalize_columns(predicate: Arc) -> Result> { + let mut columns = HashMap::new(); + Ok(predicate + .transform_down(|expression| { + let Some(column) = expression.downcast_ref::() else { + return Ok(Transformed::no(expression)); + }; + let next_id = columns.len(); + let id = *columns.entry(column.clone()).or_insert(next_id); + Ok(Transformed::yes( + Arc::new(Column::new(&format!("column_{id}"), id)) as Arc, + )) + })? + .data) +} + +fn remove_runtime_pruning_details(display: String) -> String { + const START: &str = "DynamicFilter [ expression_id_"; + const END: &str = "dynamic_rg_pruning=eligible"; + + display + .lines() + .map(|line| { + let (Some(start), Some(end)) = (line.find(START), line.find(END)) else { + return line.to_owned(); + }; + format!("{}{}", &line[..start], &line[start..end + END.len()]) + }) + .collect::>() + .join("\n") +} + +fn set_dynamic_filter_pushdown(ctx: &SessionContext, enabled: bool) -> Result<()> { + let state = ctx.state_ref(); + state.write().config_mut().options_mut().set( + "datafusion.optimizer.enable_dynamic_filter_pushdown", + &enabled.to_string(), + ) } diff --git a/tests/dynamic_filtering/config.rs b/tests/dynamic_filtering/config.rs index 101f609c0..098bab5d1 100644 --- a/tests/dynamic_filtering/config.rs +++ b/tests/dynamic_filtering/config.rs @@ -39,7 +39,6 @@ mod tests { "#, ) .without_remote_dynamic_filters() - .expect_no_dynamic_filter_updates() .execute() .await?; assert_snapshot!(display, @" diff --git a/tests/dynamic_filtering/partitioned_join.rs b/tests/dynamic_filtering/partitioned_join.rs index 1743c5da3..e86ccfb2a 100644 --- a/tests/dynamic_filtering/partitioned_join.rs +++ b/tests/dynamic_filtering/partitioned_join.rs @@ -1,9 +1,10 @@ #[cfg(test)] mod tests { - use crate::common::{TestQuery, execute_range_partitioned_query}; + use crate::common::{LOCAL_AND_REMOTE_UNION_QUERY, TestQuery, execute_range_partitioned_query}; use datafusion::common::Result; use datafusion_distributed::assert_snapshot; use datafusion_distributed::test_utils::insta::insta::allow_duplicates; + use test_case::test_case; /// A Partitioned HashJoinExec propagates dynamic filters to local consumers. #[tokio::test] @@ -30,6 +31,8 @@ mod tests { "#; fn assert_local_dynamic_filters_plan(display: String) { + // Each task in stage 1 produces a distinct dynamic filter and pushes it down + // to the local consumers. allow_duplicates! { assert_snapshot!(display, @r" ┌───── DistributedExec @@ -50,29 +53,115 @@ mod tests { │ t0: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/dim/d_dkey=A/data0.parquet], [/testdata/join/parquet/dim/d_dkey=C/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] │ t1: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/dim/d_dkey=B/data0.parquet], [/testdata/join/parquet/dim/d_dkey=D/data0.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Range([d_dkey@2 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=DynamicFilter [ f_dkey@2 >= A AND f_dkey@2 <= A AND f_dkey@2 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=f_dkey_null_count@1 != row_count@2 AND f_dkey_max@0 >= A AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= A AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= A AND A <= f_dkey_max@0, required_guarantees=[f_dkey in (A)] - │ t1: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/fact/f_dkey=B/data0.parquet], [/testdata/join/parquet/fact/f_dkey=D/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=DynamicFilter [ f_dkey@2 >= B AND f_dkey@2 <= B AND f_dkey@2 IN (SET) ([]) ], dynamic_rg_pruning=eligible, pruning_predicate=f_dkey_null_count@1 != row_count@2 AND f_dkey_max@0 >= B AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= B AND f_dkey_null_count@1 != row_count@2 AND f_dkey_min@3 <= B AND B <= f_dkey_max@0, required_guarantees=[f_dkey in (B)] + │ t0: DataSourceExec: file_groups={2 groups: [[/testdata/join/parquet/fact/f_dkey=A/data0.parquet], [/testdata/join/parquet/fact/f_dkey=C/data0.parquet]]}, projection=[f_dkey], output_partitioning=Range([f_dkey@0 ASC NULLS LAST], [(C)], 2), file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ 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 [ expression_id_0_hash_1 ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── "); } } - /// A Partitioned HashJoinExec does not propagate dynamic filters to a remote consumer. + /// A Partitioned HashJoinExec propagates merged dynamic filters to a remote consumer. + #[test_case(false; "static_planner")] + #[test_case(true; "adaptive_planner")] #[tokio::test] - async fn remote_dynamic_filters() -> Result<()> { + async fn remote_dynamic_filters(adaptive: bool) -> Result<()> { + let display = remote_dynamic_filters_plan(adaptive).await?; + + // Each join in stage 3 merges its filters at the coordinator which pushes + // down one expression to the remote consumers in stage 2. + if !adaptive { + assert_snapshot!(display, @" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 3] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 3 ── tasks=2, partitions=3 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(key@0, RainToday@0)], projection=[] + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + │ [Stage 2] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([RainToday@0], 6), input_partitions=3 + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + └────────────────────────────────────────────────── + "); + } else { + assert_snapshot!(display, @" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 3] => NetworkCoalesceExec: output_partitions=3, input_tasks=1 + └────────────────────────────────────────────────── + ┌───── Stage 3 ── tasks=1, partitions=3 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(key@0, RainToday@0)], projection=[] + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + │ [Stage 2] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=3 + │ RepartitionExec: partitioning=Hash([key@0], 3), input_partitions=3 + │ SamplerExec: partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=3 + │ RepartitionExec: partitioning=Hash([RainToday@0], 3), input_partitions=3 + │ SamplerExec: partitions=3 + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + └────────────────────────────────────────────────── + "); + } + Ok(()) + } + + /// Multiple consumers of one remote filter can remap it to different source columns. + #[tokio::test] + async fn remote_union_remaps_each_consumer() -> Result<()> { let display = TestQuery::new( r#" - SELECT COUNT(*) - FROM ( - SELECT DISTINCT "RainToday" AS key - FROM weather - ) build - JOIN weather probe ON build.key = probe."RainToday" - "#, + WITH probe AS ( + SELECT "RainToday" AS key FROM weather + UNION ALL + SELECT "WindGustDir" AS key FROM weather + UNION ALL + SELECT "WindDir9am" AS key FROM weather + ) + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + JOIN probe ON build.key = probe.key + "#, ) - .expect_dynamic_filter_updates() + .with_one_task_per_leaf() + .with_normalized_filter_hashes() .execute() .await?; + + // The each hash join in stage 3 propagates its dynamic filter to the coordindator where + // it's merged and forwarded to the consumers in stage 2. Each consumer applies the dynamic + // filter to a different column (RainToday, WindGustDir, WindDir9am), so the hashes are different but + // the hashes with normalized column names match. assert_snapshot!(display, @" ┌───── DistributedExec │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] @@ -82,10 +171,71 @@ mod tests { └────────────────────────────────────────────────── ┌───── Stage 3 ── tasks=2, partitions=3 │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(key@0, key@0)], projection=[] + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=1 + │ [Stage 2] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=1, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=6 + │ DistributedUnionExec: t0:[c0, c2] t1:[c1] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0_normalized_hash_0 ], dynamic_rg_pruning=eligible + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir@5 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_1_normalized_hash_0 ], dynamic_rg_pruning=eligible + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindDir9am@7 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_2_normalized_hash_0 ], dynamic_rg_pruning=eligible + └────────────────────────────────────────────────── + "); + Ok(()) + } + + #[tokio::test] + async fn remote_dynamic_filters_propagate_through_multiple_stages() -> Result<()> { + let display = TestQuery::new( + r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + JOIN ( + SELECT probe."RainToday" + FROM ( + SELECT "WindGustDir" AS key + FROM weather + ) nested_build + RIGHT JOIN weather probe + ON nested_build.key = probe."WindGustDir" + ) probe ON build.key = probe."RainToday" + "#, + ) + .execute() + .await?; + + // The filters produced in stage 5 are merged into one filter and forwarded + // to the consumers in stage 3. + // + // Note that the right join in stage 4 does not produce dynamic filters. + assert_snapshot!(display, @" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 5] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 5 ── tasks=2, partitions=3 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] │ HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(key@0, RainToday@0)], projection=[] │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] │ [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=2 - │ [Stage 2] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + │ [Stage 4] => NetworkShuffleExec: output_partitions=3, input_tasks=2 └────────────────────────────────────────────────── ┌───── Stage 1 ── tasks=2, partitions=6 │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 @@ -94,13 +244,101 @@ mod tests { │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet └────────────────────────────────────────────────── - ┌───── Stage 2 ── tasks=2, partitions=6 + ┌───── Stage 4 ── tasks=2, partitions=6 │ RepartitionExec: partitioning=Hash([RainToday@0], 6), input_partitions=3 - │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible - │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + │ HashJoinExec: mode=Partitioned, join_type=Right, on=[(key@0, WindGustDir@0)], projection=[RainToday@2] + │ [Stage 2] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + │ [Stage 3] => NetworkShuffleExec: output_partitions=3, input_tasks=2 └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir@5 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir@5 as key], file_type=parquet + └────────────────────────────────────────────────── + ┌───── Stage 3 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([WindGustDir@0], 6), input_partitions=3 + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir, RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir, RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + └────────────────────────────────────────────────── "); Ok(()) } + + #[tokio::test] + async fn union_with_local_and_remote_probe_children() -> Result<()> { + let display = TestQuery::new(LOCAL_AND_REMOTE_UNION_QUERY) + .execute() + .await?; + + // Complex test. + // Dynamic filters from the stage 5 tasks are merged and forwarded to both c0 of t0 in stage + // 4 and the consumers in stage 3. Dynamic filters from the join in stage 4 are also merged + // and forwarded to the consumers in stage 3 where they are AND'ed with the other filters. + assert_snapshot!(display, @" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 5] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 5 ── tasks=2, partitions=3 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(key@0, key@0)], projection=[] + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + │ [Stage 4] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir@5 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir@5 as key], file_type=parquet + └────────────────────────────────────────────────── + ┌───── Stage 4 ── tasks=2, partitions=6 + │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 + │ DistributedUnionExec: t0:[c0] t1:[c1] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir@5 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ ProjectionExec: expr=[WindGustDir@0 as key] + │ HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(key@0, RainToday@1)], projection=[WindGustDir@0] + │ AggregateExec: mode=FinalPartitioned, gby=[key@0 as key], aggr=[] + │ [Stage 2] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + │ [Stage 3] => NetworkShuffleExec: output_partitions=3, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=3 + │ RepartitionExec: partitioning=Hash([key@0], 3), input_partitions=3 + │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[] + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[RainToday@19 as key], file_type=parquet + └────────────────────────────────────────────────── + ┌───── Stage 3 ── tasks=2, partitions=3 + │ RepartitionExec: partitioning=Hash([RainToday@1], 3), input_partitions=3 + │ DistributedLeafExec: + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir, RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_1_hash_1 ] AND DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[WindGustDir, RainToday], file_type=parquet, predicate=DynamicFilter [ expression_id_1_hash_1 ] AND DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + └────────────────────────────────────────────────── + "); + Ok(()) + } + + const REMOTE_DYNAMIC_FILTER_QUERY: &str = r#" + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) build + JOIN weather probe ON build.key = probe."RainToday" + "#; + + async fn remote_dynamic_filters_plan(dynamic_planning: bool) -> Result { + let mut query = TestQuery::new(REMOTE_DYNAMIC_FILTER_QUERY); + if dynamic_planning { + query = query.with_dynamic_task_count(); + } + query.execute().await + } } diff --git a/tests/dynamic_filtering/sorts.rs b/tests/dynamic_filtering/sorts.rs index 3522b3444..e3e2254c5 100644 --- a/tests/dynamic_filtering/sorts.rs +++ b/tests/dynamic_filtering/sorts.rs @@ -2,7 +2,7 @@ mod tests { use crate::common::TestQuery; use datafusion::common::Result; - use datafusion_distributed::{assert_snapshot, test_utils::insta::settings}; + use datafusion_distributed::assert_snapshot; /// A TopK SortExec applies dynamic filters to local data sources. #[tokio::test] @@ -18,13 +18,9 @@ mod tests { .with_expected_rows(10) .execute() .await?; - let mut settings = settings(); - settings.add_filter( - r"(DynamicFilter \[[^\]\n]*? > )-?\d+(?:\.\d+)?( \])", - "${1}${2}", - ); - settings.add_filter(r"(_max@\d+ > )-?\d+(?:\.\d+)?", "${1}"); - settings.bind(|| assert_snapshot!(display, @" + + // Each task in stage 1 produces a distinct filter to push down. + assert_snapshot!(display, @" ┌───── DistributedExec │ SortPreservingMergeExec: [MinTemp@0 DESC], fetch=10 │ [Stage 1] => NetworkCoalesceExec: output_partitions=6, input_tasks=2 @@ -32,14 +28,14 @@ mod tests { ┌───── Stage 1 ── tasks=2, partitions=6 │ SortExec: TopK(fetch=10), expr=[MinTemp@0 DESC], preserve_partitioning=[true] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp], file_type=parquet, predicate=DynamicFilter [ MinTemp@0 IS NULL OR MinTemp@0 > ], dynamic_rg_pruning=eligible, pruning_predicate=MinTemp_null_count@0 > 0 OR MinTemp_null_count@0 != row_count@2 AND MinTemp_max@1 > , required_guarantees=[] - │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp], file_type=parquet, predicate=DynamicFilter [ MinTemp@0 IS NULL OR MinTemp@0 > ], dynamic_rg_pruning=eligible, pruning_predicate=MinTemp_null_count@0 > 0 OR MinTemp_null_count@0 != row_count@2 AND MinTemp_max@1 > , required_guarantees=[] + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_1 ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── - ")); + "); Ok(()) } - /// A TopK sort does not yet update dynamic filters to remote consumers. + /// A TopK sort updates dynamic filters on remote consumers across a shuffle. #[tokio::test] async fn remote_dynamic_filters() -> Result<()> { let display = TestQuery::new( @@ -54,9 +50,10 @@ mod tests { "#, ) .with_expected_rows(10) - .expect_dynamic_filter_updates() .execute() .await?; + + // The per-task TopK filters in stage 2 are merged into one and passed to both consumers. assert_snapshot!(display, @" ┌───── DistributedExec │ SortPreservingMergeExec: [key@0 DESC], fetch=10 @@ -71,8 +68,8 @@ mod tests { │ RepartitionExec: partitioning=Hash([key@0], 6), input_partitions=3 │ AggregateExec: mode=Partial, gby=[key@0 as key], aggr=[], lim=[10] │ DistributedLeafExec: - │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp@0 as key], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible - │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp@0 as key], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + │ t0: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000000.parquet:.., /testdata/weather/result-000001.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp@0 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible + │ t1: DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet:..], [/testdata/weather/result-000001.parquet:.., /testdata/weather/result-000002.parquet:..], [/testdata/weather/result-000002.parquet:..]]}, projection=[MinTemp@0 as key], file_type=parquet, predicate=DynamicFilter [ expression_id_0_hash_0 ], dynamic_rg_pruning=eligible └────────────────────────────────────────────────── "); Ok(()) From 47b5f82066a11cbd793bd7c42043d5bc26e0ff28 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Wed, 23 Sep 2026 01:21:22 +0000 Subject: [PATCH 02/11] fix: send producer schema with dynamic filters When decoding a physical expression, we have to decode it against a schema. Previously, we were decoding producer dynamic filters against the consumer schema, which was causing errors. Example: A producer may send a filter on `my_col@2` but the consumer data source may have this column at a different index ex. `my_col@8`. To fix this, we now store the producer schema in the registry and send it along with dynamic filter updates. --- src/codec/mod.rs | 2 +- src/codec/physical_plan.rs | 23 ++++++----- src/coordinator/dynamic_filter_registry.rs | 38 +++++++++++++++-- src/dynamic_filtering/discovery.rs | 39 +++++++++++------- src/dynamic_filtering/display.rs | 47 +++++++++++++++++----- src/protocol/grpc/generated/worker.rs | 3 ++ src/protocol/grpc/worker.proto | 2 + src/protocol/grpc/worker_client.rs | 3 ++ src/protocol/grpc/worker_service.rs | 9 +++++ src/protocol/worker_channel.rs | 3 ++ src/worker/impl_coordinator_channel.rs | 13 +++--- 11 files changed, 134 insertions(+), 48 deletions(-) diff --git a/src/codec/mod.rs b/src/codec/mod.rs index 59ed4730e..f18db9478 100644 --- a/src/codec/mod.rs +++ b/src/codec/mod.rs @@ -4,7 +4,7 @@ mod user_codec; pub use distributed_codec::DistributedCodec; pub(crate) use physical_plan::{ - decode_execution_plan, decode_partitioning, decode_physical_expr, dynamic_filter_update_target, + apply_dynamic_filter_update, decode_execution_plan, decode_partitioning, decode_physical_expr, encode_execution_plan, encode_partitioning, encode_physical_expr, roundtrip_pb, }; pub(crate) use user_codec::{ diff --git a/src/codec/physical_plan.rs b/src/codec/physical_plan.rs index b7aeef6b0..c0a9fae59 100644 --- a/src/codec/physical_plan.rs +++ b/src/codec/physical_plan.rs @@ -80,13 +80,12 @@ pub(crate) fn decode_physical_expr( converter.proto_to_physical_expr(proto, input_schema, &decode_ctx) } -/// Creates a sythetic "producer" [`DynamicFilterPhysicalExpr`] which can be used to -/// update consumer dynamic filters via a shared state. Rather than updating -/// the provicder consumer expression directly, the returned expression should -/// be updated instead. +/// Applies a predicate to a consumer dynamic filter expression. /// -/// DataFusion does a particular song and dance to update dynamic filters. This function -/// helps us do the same. +/// To update a consumer, this function creates a synthetic "producer" +/// [`DynamicFilterPhysicalExpr`] which can be used to update consumer +/// dynamic filters via a shared state. The reason for this is to match +/// the same behavior as DataFusion. /// /// ```text /// HashJoinExec: on=[build.key = probe.key] @@ -113,11 +112,13 @@ pub(crate) fn decode_physical_expr( /// You cannot update() consumers directly. In the above example, updating consumer 1 would remap /// `key > 123` to `phone_number > 123` and store this in the shared state. Consumer 2 would be /// unable to apply this filter now. -pub(crate) fn dynamic_filter_update_target( +pub(crate) fn apply_dynamic_filter_update( consumer: &Arc, - input_schema: &Schema, + predicate: &protobuf::PhysicalExprNode, + producer_schema: &Schema, task_ctx: &TaskContext, -) -> Result> { +) -> Result<()> { + let predicate = decode_physical_expr(predicate, producer_schema, task_ctx)?; // Since consumer.children() returns the remapped children, we use the proto as a workaround to get the // original children from the producer. let consumer: Arc = consumer.clone(); @@ -128,13 +129,13 @@ pub(crate) fn dynamic_filter_update_target( let original_children = dynamic_filter .children .iter() - .map(|child| decode_physical_expr(child, input_schema, task_ctx)) + .map(|child| decode_physical_expr(child, producer_schema, task_ctx)) .collect::>>()?; let update_target = Arc::clone(&consumer).with_new_children(original_children)?; let Ok(update_target) = Arc::downcast::(update_target) else { return internal_err!("expected a dynamic filter update target"); }; - Ok(update_target) + update_target.update(predicate) } pub(crate) fn encode_partitioning( diff --git a/src/coordinator/dynamic_filter_registry.rs b/src/coordinator/dynamic_filter_registry.rs index c5f35c91f..34b499b21 100644 --- a/src/coordinator/dynamic_filter_registry.rs +++ b/src/coordinator/dynamic_filter_registry.rs @@ -1,7 +1,8 @@ -use crate::dynamic_filtering::discover_dynamic_filter_consumers; +use crate::dynamic_filtering::{discover_dynamic_filter_consumers, dynamic_filter_producer_schema}; use crate::{ ApplyDynamicFilter, CoordinatorToWorkerMsg, MaybeEncoded, ProducedDynamicFilter, TaskKey, }; +use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::common::{HashMap, HashSet, Result, exec_err, internal_err}; use datafusion::execution::TaskContext; @@ -35,15 +36,25 @@ pub(super) enum DynamicFilterMergeMode { #[derive(Default)] pub(super) struct PlannedDynamicFilter { - pub(super) merge_mode: Option, // 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. + // /// Registered producers and their latest accepted snapshots, if any. pub(super) producers: HashMap>, pub(super) consumer_tasks: HashSet, + + /// Populated upon task registration when a producer is found. + /// + /// The behavior to take when an update arrives from a producer. + pub(super) merge_mode: Option, + /// Schema of the original producer arguments, shared by all producer tasks. + pub(super) producer_schema: Option, + + /// Populated when an update is ready to be sent to consumers. + /// /// Full dynamic filter containing the merged predicate and its completion state. pub(super) merged: Option, /// Immutable snapshot shared by local and remote delivery, encoded once per change. @@ -130,7 +141,14 @@ impl DynamicFilterRegistry { } }) .collect::>()?; - producers.extend(produced_ids.iter().map(|id| (*id, merge_mode))); + if !produced_ids.is_empty() { + let schema = dynamic_filter_producer_schema(node.as_ref())?; + producers.extend( + produced_ids + .iter() + .map(|id| (*id, merge_mode, Arc::clone(&schema))), + ); + } Ok(TreeNodeRecursion::Continue) })?; // We can safely ignore anchors because they are not evaluated by network boundaries. This @@ -138,8 +156,14 @@ impl DynamicFilterRegistry { let consumers = discover_dynamic_filter_consumers(task_specialized_plan)?.consumers; let mut state = self.state.lock().expect("dynamic filter registry poisoned"); - for (id, merge_mode) in producers { + for (id, merge_mode, producer_schema) in producers { let filter = state.filters.entry(id).or_default(); + if let Some(existing) = &filter.producer_schema + && existing != &producer_schema + { + return internal_err!("Dynamic filter {id} has conflicting producer schemas"); + } + filter.producer_schema.get_or_insert(producer_schema); filter.merge_mode = Some(match filter.merge_mode { Some(existing) if existing != merge_mode => { return internal_err!( @@ -328,6 +352,11 @@ impl DynamicFilterRegistry { let Some(expression) = &filter.merged_bytes else { return Ok(()); }; + let Some(producer_schema) = &filter.producer_schema else { + return internal_err!( + "Dynamic filter {id} has a merged predicate but no producer schema" + ); + }; for &task_key in &filter.consumer_tasks { // A task-local consumer is already updated directly by its producer. if filter.producers.contains_key(&task_key) || state.delivered.contains(&(id, task_key)) @@ -340,6 +369,7 @@ impl DynamicFilterRegistry { let update = CoordinatorToWorkerMsg::ApplyDynamicFilter(Box::new(ApplyDynamicFilter { expression_id: id, expression: MaybeEncoded::Encoded(expression.clone()), + producer_schema: Arc::clone(producer_schema), })); if sender.send(update).is_err() { // Closing the channel is expected only once query shutdown has started. diff --git a/src/dynamic_filtering/discovery.rs b/src/dynamic_filtering/discovery.rs index b1a7f05a6..0bf1ce2f7 100644 --- a/src/dynamic_filtering/discovery.rs +++ b/src/dynamic_filtering/discovery.rs @@ -5,6 +5,7 @@ use datafusion::common::{HashMap, HashSet, Result, internal_err}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::joins::HashJoinExec; use std::sync::Arc; /// A dynamic filter produced by an [`ExecutionPlan`]. @@ -12,15 +13,14 @@ use std::sync::Arc; pub(crate) struct DiscoveredDynamicFilterProducer { pub(crate) id: u64, pub(crate) expression: Arc, + pub(crate) input_schema: SchemaRef, } -/// A dynamic-filter consumer discovered in an execution plan along with the schema it is evaluated -/// against. +/// A dynamic-filter consumer discovered in an execution plan. #[derive(Clone)] pub(crate) struct DiscoveredDynamicFilter { pub(crate) id: u64, pub(crate) expression: Arc, - pub(crate) input_schema: SchemaRef, } /// An anchor is an artificial dynamic filter consumer injected into network boundaries @@ -62,11 +62,6 @@ pub(crate) fn discover_dynamic_filter_consumers( Ok(id) }) .collect::>()?; - let input_schema = node - .children() - .first() - .map(|child| child.schema()) - .unwrap_or_else(|| node.schema()); let is_network_boundary = node.is_network_boundary(); node.apply_expressions(&mut |root| { @@ -93,11 +88,7 @@ pub(crate) fn discover_dynamic_filter_consumers( } else if !produced_ids.contains(&id) { consumers .entry(id) - .or_insert_with(|| DiscoveredDynamicFilter { - id, - expression, - input_schema: Arc::clone(&input_schema), - }); + .or_insert_with(|| DiscoveredDynamicFilter { id, expression }); } Ok(TreeNodeRecursion::Continue) @@ -129,9 +120,14 @@ pub(crate) fn discover_dynamic_filter_producers( let Some(id) = expression.expression_id() else { return internal_err!("DynamicFilterPhysicalExpr did not have an expression ID"); }; + let input_schema = dynamic_filter_producer_schema(node.as_ref())?; producers .entry(id) - .or_insert_with(|| DiscoveredDynamicFilterProducer { id, expression }); + .or_insert(DiscoveredDynamicFilterProducer { + id, + expression, + input_schema, + }); } Ok(TreeNodeRecursion::Continue) })?; @@ -141,6 +137,21 @@ pub(crate) fn discover_dynamic_filter_producers( Ok(producers) } +/// Schema referenced by a producer's predicates, before any consumer remapping. +pub(crate) fn dynamic_filter_producer_schema(plan: &dyn ExecutionPlan) -> Result { + if let Some(join) = plan.downcast_ref::() { + return Ok(join.right().schema()); + } + match plan.children().as_slice() { + [input] => Ok(input.schema()), + [] => Ok(plan.schema()), + _ => internal_err!( + "cannot determine dynamic filter producer schema for {}", + plan.name() + ), + } +} + /// Returns producer IDs with at least one remote consumer. /// /// If a producer ID is present in the dynamic-filter anchors of any [`NetworkBoundary`], the plan diff --git a/src/dynamic_filtering/display.rs b/src/dynamic_filtering/display.rs index 0f19eb7ea..dc5943381 100644 --- a/src/dynamic_filtering/display.rs +++ b/src/dynamic_filtering/display.rs @@ -1,10 +1,13 @@ -use crate::codec::{decode_physical_expr, dynamic_filter_update_target, roundtrip_pb}; +use crate::codec::{apply_dynamic_filter_update, roundtrip_pb}; use crate::common::TreeNodeExt; use crate::coordinator::DistributedExec; -use crate::dynamic_filtering::discover_dynamic_filter_consumers; +use crate::dynamic_filtering::{ + discover_dynamic_filter_consumers, discover_dynamic_filter_producers, +}; use crate::execution_plans::DistributedLeafExec; use crate::stage::{LocalStage, Stage, find_all_stages}; use crate::{DistributedTaskContext, TaskCompletedDynamicFilters, TaskKey}; +use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::common::{HashMap, Result, internal_err}; use datafusion::execution::TaskContext; @@ -38,6 +41,24 @@ pub async fn rewrite_distributed_plan_with_dynamic_filters( let Some(reports) = distributed_exec.wait_for_dynamic_filters().await else { return internal_err!("dynamic filters were enabled but the execution was not prepared"); }; + + // Collect producer schemas so we can decode collected dynamic filters. + // + // Note that we display dynamic filters on consumers but decode them with respect + // to the producer schemas. When updating the consumers, they automatically remap + // the expressions to their own schema. + let mut producers = discover_dynamic_filter_producers(&plan_for_viz)?; + plan_for_viz.apply(|node| { + if let Some(leaf) = node.downcast_ref::() { + producers.extend(discover_dynamic_filter_producers(leaf.original())?); + } + Ok(TreeNodeRecursion::Continue) + })?; + let producer_schemas = producers + .into_iter() + .map(|producer| (producer.id, producer.input_schema)) + .collect(); + // Avoids mutating the `plan_for_viz` of the incoming DistributedExec. let plan_for_viz = sever_dynamic_filter_relationships_in_plan_for_display(plan_for_viz, task_ctx)?; @@ -45,7 +66,7 @@ pub async fn rewrite_distributed_plan_with_dynamic_filters( let Stage::Local(stage) = stage else { return internal_err!("expected local stages in the visualization plan"); }; - apply_reports_to_distributed_leaves(stage, &reports, task_ctx)?; + apply_reports_to_distributed_leaves(stage, &reports, &producer_schemas, task_ctx)?; } let plan = distributed_exec.with_plan_for_viz(Arc::clone(&plan_for_viz))?; plan.replace_children( @@ -153,6 +174,7 @@ fn isolate_sort_dynamic_filters_for_display( fn apply_reports_to_distributed_leaves( stage: &LocalStage, reports: &HashMap, + producer_schemas: &HashMap, task_ctx: &Arc, ) -> Result<()> { for task_number in 0..stage.tasks { @@ -202,14 +224,18 @@ fn apply_reports_to_distributed_leaves( let Some(predicate) = dynamic_filter_proto.inner_expr.as_deref() else { return internal_err!("reported dynamic filter has no predicate"); }; - let predicate = - decode_physical_expr(predicate, consumer.input_schema.as_ref(), task_ctx)?; - let dynamic_filter = dynamic_filter_update_target( + let Some(producer_schema) = producer_schemas.get(&consumer.id) else { + return internal_err!( + "missing producer schema for dynamic filter {}", + consumer.id + ); + }; + apply_dynamic_filter_update( &consumer.expression, - consumer.input_schema.as_ref(), + predicate, + producer_schema.as_ref(), task_ctx, )?; - dynamic_filter.update(predicate)?; } Ok(TreeNodeRecursion::Continue) })?; @@ -244,7 +270,7 @@ mod tests { vec![Arc::clone(&column)], lit(true), )) as Arc; - let input = Arc::new(EmptyExec::new(schema)) as Arc; + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc; let variant = Arc::new(FilterExec::try_new( Arc::clone(&dynamic_filter), Arc::clone(&input), @@ -288,7 +314,8 @@ mod tests { tasks: 2, metrics_set: Default::default(), }; - apply_reports_to_distributed_leaves(&stage, &reports, &task_ctx)?; + let producer_schemas = HashMap::from([(dynamic_filter.expression_id().unwrap(), schema)]); + apply_reports_to_distributed_leaves(&stage, &reports, &producer_schemas, &task_ctx)?; let leaf = isolated.downcast_ref::().unwrap(); let task_0 = displayable(leaf.variants()[0].as_ref()) .one_line() diff --git a/src/protocol/grpc/generated/worker.rs b/src/protocol/grpc/generated/worker.rs index 91bcc9d2c..8939e0064 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/grpc/generated/worker.rs @@ -37,6 +37,9 @@ pub struct ApplyDynamicFilter { /// Serialized datafusion.proto.PhysicalExprNode containing a full DynamicFilter expression. #[prost(bytes = "vec", tag = "2")] pub expression_proto: ::prost::alloc::vec::Vec, + /// Serialized datafusion.proto.Schema referenced by the producer predicate. + #[prost(bytes = "vec", tag = "3")] + pub producer_schema: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct WorkerToCoordinatorMsg { diff --git a/src/protocol/grpc/worker.proto b/src/protocol/grpc/worker.proto index 74ccab378..d9a621a85 100644 --- a/src/protocol/grpc/worker.proto +++ b/src/protocol/grpc/worker.proto @@ -38,6 +38,8 @@ message ApplyDynamicFilter { uint64 expression_id = 1; // Serialized datafusion.proto.PhysicalExprNode containing a full DynamicFilter expression. bytes expression_proto = 2; + // Serialized datafusion.proto.Schema referenced by the producer predicate. + bytes producer_schema = 3; } message WorkerToCoordinatorMsg { diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index ea7355db7..f9b3bdee0 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -26,6 +26,7 @@ use datafusion::execution::TaskContext; use datafusion::execution::memory_pool::MemoryConsumer; use datafusion::physical_expr_common::metrics::{Count, Label, MetricBuilder, MetricValue, Time}; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, Gauge}; +use datafusion_proto::protobuf; use futures::future::{Either, pending, ready}; use futures::stream::{BoxStream, select}; use futures::{FutureExt, Stream, StreamExt, TryStreamExt}; @@ -498,6 +499,8 @@ fn encode_coordinator_to_worker_msg( pb::coordinator_to_worker_msg::Inner::ApplyDynamicFilter(pb::ApplyDynamicFilter { expression_id: filter.expression_id, expression_proto: filter.expression.encode(task_ctx)?, + producer_schema: protobuf::Schema::try_from(filter.producer_schema)? + .encode_to_vec(), }) } }), diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index 7a71db92a..7ee621787 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -19,10 +19,12 @@ use arrow_flight::error::FlightError; use arrow_select::dictionary::garbage_collect_any_dictionary; use async_trait::async_trait; use datafusion::arrow::array::{Array, AsArray, RecordBatch, RecordBatchOptions}; +use datafusion::arrow::datatypes::Schema; use datafusion::arrow::ipc::CompressionType; use datafusion::arrow::ipc::writer::IpcWriteOptions; use datafusion::common::DataFusionError; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_proto::protobuf; use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; use prost::Message; @@ -220,9 +222,16 @@ fn decode_coordinator_to_worker_msg( CoordinatorToWorkerMsg::KickOffSampling } pb::coordinator_to_worker_msg::Inner::ApplyDynamicFilter(filter) => { + let producer_schema = + protobuf::Schema::decode(filter.producer_schema.as_slice()) + .map_err(|error| Status::invalid_argument(error.to_string()))?; + let producer_schema = Schema::try_from(&producer_schema).map_err(|error| { + datafusion_error_to_tonic_status(DataFusionError::from(error)) + })?; CoordinatorToWorkerMsg::ApplyDynamicFilter(Box::new(ApplyDynamicFilter { expression_id: filter.expression_id, expression: MaybeEncoded::Encoded(filter.expression_proto), + producer_schema: Arc::new(producer_schema), })) } }, diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index e548f69a4..77772600b 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -1,5 +1,6 @@ use crate::{MaybeEncoded, ProducerHead, WorkUnit}; use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::Result; use datafusion::execution::TaskContext; @@ -64,6 +65,8 @@ pub struct ApplyDynamicFilter { pub expression_id: u64, /// A full dynamic filter containing the merged predicate. pub expression: MaybeEncoded>, + /// Schema referenced by the producer predicate, before consumer-specific remapping. + pub producer_schema: SchemaRef, } #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index 6bf0a8f6c..6c9db46a9 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,4 +1,4 @@ -use crate::codec::{decode_physical_expr, dynamic_filter_update_target}; +use crate::codec::apply_dynamic_filter_update; use crate::common::TreeNodeExt; use crate::dynamic_filtering::{ DiscoveredDynamicFilter, discover_dynamic_filter_consumers, discover_dynamic_filter_producers, @@ -386,17 +386,14 @@ fn apply_merged_dynamic_filter( return internal_err!("dynamic filter update has no predicate"); }; - // PhysicalExprNode to PhysicalExpr - let predicate = decode_physical_expr(&predicate, consumer.input_schema.as_ref(), task_ctx)?; - - let dynamic_filter = dynamic_filter_update_target( + apply_dynamic_filter_update( &consumer.expression, - consumer.input_schema.as_ref(), + &predicate, + filter.producer_schema.as_ref(), task_ctx, )?; - dynamic_filter.update(predicate)?; if update.is_complete { - dynamic_filter.mark_complete(); + consumer.expression.mark_complete(); } Ok(()) } From 3f4c23f05c93a839b3bbf9be2c355e6e9dfb793b Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Wed, 7 Oct 2026 14:23:54 +0000 Subject: [PATCH 03/11] pr comments round 1 - stream style - avoid sending task complete dynamic filters with no updates - handle error in worker coordinator channel - DynamicFilterMergeMode --- src/codec/mod.rs | 4 +- src/codec/physical_plan.rs | 62 +------------------- src/coordinator/dynamic_filter_registry.rs | 36 ++++++------ src/dynamic_filtering/display.rs | 8 +-- src/dynamic_filtering/mod.rs | 67 +++++++++++++++++++++- src/worker/impl_coordinator_channel.rs | 18 ++++-- 6 files changed, 102 insertions(+), 93 deletions(-) diff --git a/src/codec/mod.rs b/src/codec/mod.rs index f18db9478..fc03f09fb 100644 --- a/src/codec/mod.rs +++ b/src/codec/mod.rs @@ -4,8 +4,8 @@ mod user_codec; pub use distributed_codec::DistributedCodec; pub(crate) use physical_plan::{ - apply_dynamic_filter_update, decode_execution_plan, decode_partitioning, decode_physical_expr, - encode_execution_plan, encode_partitioning, encode_physical_expr, roundtrip_pb, + decode_execution_plan, decode_partitioning, decode_physical_expr, encode_execution_plan, + encode_partitioning, encode_physical_expr, roundtrip_pb, }; pub(crate) use user_codec::{ get_distributed_user_codecs, set_distributed_user_codec, set_distributed_user_codec_arc, diff --git a/src/codec/physical_plan.rs b/src/codec/physical_plan.rs index c0a9fae59..f1a970ace 100644 --- a/src/codec/physical_plan.rs +++ b/src/codec/physical_plan.rs @@ -1,8 +1,7 @@ use super::DistributedCodec; use datafusion::arrow::datatypes::{Schema, SchemaRef}; -use datafusion::common::{Result, internal_err}; +use datafusion::common::Result; use datafusion::execution::TaskContext; -use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_expr::{Partitioning, PhysicalExpr}; use datafusion::physical_plan::ExecutionPlan; use datafusion_proto::bytes::{ @@ -14,7 +13,6 @@ use datafusion_proto::physical_plan::{ DeduplicatingProtoConverter, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; use datafusion_proto::protobuf; -use datafusion_proto::protobuf::physical_expr_node::ExprType; use datafusion_proto::protobuf::proto_error; use prost::Message; use std::sync::Arc; @@ -80,64 +78,6 @@ pub(crate) fn decode_physical_expr( converter.proto_to_physical_expr(proto, input_schema, &decode_ctx) } -/// Applies a predicate to a consumer dynamic filter expression. -/// -/// To update a consumer, this function creates a synthetic "producer" -/// [`DynamicFilterPhysicalExpr`] which can be used to update consumer -/// dynamic filters via a shared state. The reason for this is to match -/// the same behavior as DataFusion. -/// -/// ```text -/// HashJoinExec: on=[build.key = probe.key] -/// │ └── producer: original=[key], children=[key] ──────────────┐ -/// ├── build │ -/// └── UnionExec: probe │ -/// ├── DataSourceExec: phone_number AS key │ shared inner -/// │ └── consumer 1: original=[key], children=[phone_number] ─┤ -/// └── DataSourceExec: telephone AS key │ -/// └── consumer 2: original=[key], children=[telephone] ───┘ -/// ``` -/// -/// Behavior: -/// 1. [`DynamicFilterPhysicalExpr::update()`] remaps any occurrences of `original` to `children` -/// and stores the result in the shared state. -/// -/// 2. [`DynamicFilterPhysicalExpr::current()`] reads the remapped expression and remaps it again, -/// mapping `original` to `children` and returns it without storing. -/// -/// In the above plan, the producer calls update() so an expression like `key > 123` is mapped to -/// `key > 123` and this is stored. Then, the consumers call current(), reading from the same state, -/// to get `phone_number > 123` and `telephone > 123` respectively. -/// -/// You cannot update() consumers directly. In the above example, updating consumer 1 would remap -/// `key > 123` to `phone_number > 123` and store this in the shared state. Consumer 2 would be -/// unable to apply this filter now. -pub(crate) fn apply_dynamic_filter_update( - consumer: &Arc, - predicate: &protobuf::PhysicalExprNode, - producer_schema: &Schema, - task_ctx: &TaskContext, -) -> Result<()> { - let predicate = decode_physical_expr(predicate, producer_schema, task_ctx)?; - // Since consumer.children() returns the remapped children, we use the proto as a workaround to get the - // original children from the producer. - let consumer: Arc = consumer.clone(); - let proto = encode_physical_expr(&consumer, task_ctx)?; - let Some(ExprType::DynamicFilter(dynamic_filter)) = proto.expr_type else { - return internal_err!("expected a dynamic filter expression"); - }; - let original_children = dynamic_filter - .children - .iter() - .map(|child| decode_physical_expr(child, producer_schema, task_ctx)) - .collect::>>()?; - let update_target = Arc::clone(&consumer).with_new_children(original_children)?; - let Ok(update_target) = Arc::downcast::(update_target) else { - return internal_err!("expected a dynamic filter update target"); - }; - update_target.update(predicate) -} - pub(crate) fn encode_partitioning( partitioning: &Partitioning, task_ctx: &TaskContext, diff --git a/src/coordinator/dynamic_filter_registry.rs b/src/coordinator/dynamic_filter_registry.rs index 34b499b21..ae041c31a 100644 --- a/src/coordinator/dynamic_filter_registry.rs +++ b/src/coordinator/dynamic_filter_registry.rs @@ -23,15 +23,17 @@ use tokio::sync::mpsc::UnboundedSender; use tokio_util::sync::CancellationToken; #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(super) enum DynamicFilterMergeMode { +enum DynamicFilterMergePolicy { /// Wait for every planned producer task to report a complete dynamic filter, then - /// merge and forward the filter. Used for partitioned joins. - AllProducersComplete, + /// union and forward the filter. Used for partitioned joins. + UnionWhenAllComplete, /// Wait for any producer to report a complete dynamic filter and forward it. /// Used for collect left joins. - FirstProducerComplete, - /// Forward any update received from any producer. - Incremental, + ForwardFirstComplete, + /// Intersect all updates from producers and immediately forward them. Used + /// for TopK and Min/Max aggregate dynamic filters where intersecting updates + /// is safe. + IntersectOnUpdate, } #[derive(Default)] @@ -49,7 +51,7 @@ pub(super) struct PlannedDynamicFilter { /// Populated upon task registration when a producer is found. /// /// The behavior to take when an update arrives from a producer. - pub(super) merge_mode: Option, + pub(super) merge_mode: Option, /// Schema of the original producer arguments, shared by all producer tasks. pub(super) producer_schema: Option, @@ -120,11 +122,11 @@ impl DynamicFilterRegistry { .downcast_ref::() .is_some_and(|join| matches!(join.partition_mode(), PartitionMode::CollectLeft)) { - DynamicFilterMergeMode::FirstProducerComplete + DynamicFilterMergePolicy::ForwardFirstComplete } else if node.is::() || node.is::() { - DynamicFilterMergeMode::Incremental + DynamicFilterMergePolicy::IntersectOnUpdate } else { - DynamicFilterMergeMode::AllProducersComplete + DynamicFilterMergePolicy::UnionWhenAllComplete }; let produced_ids: HashSet<_> = node .dynamic_expressions_produced() @@ -257,7 +259,7 @@ impl DynamicFilterRegistry { report.expression_id ); }; - if (filter.merge_mode != Some(DynamicFilterMergeMode::Incremental) + if (filter.merge_mode != Some(DynamicFilterMergePolicy::IntersectOnUpdate) && !dynamic_filter.is_complete) || previous .as_ref() @@ -290,7 +292,7 @@ impl DynamicFilterRegistry { state.sealed_stages.contains(&task.stage_id) && report.as_ref().is_some_and(|report| report.is_complete) }); - if mode == DynamicFilterMergeMode::AllProducersComplete && !all_complete { + if mode == DynamicFilterMergePolicy::UnionWhenAllComplete && !all_complete { return false; } @@ -300,7 +302,7 @@ impl DynamicFilterRegistry { .filter_map(|(task, report)| report.as_ref().map(|report| (task, report))) .collect(); reports.sort_unstable_by_key(|(key, _)| (key.stage_id, key.task_number)); - if mode == DynamicFilterMergeMode::FirstProducerComplete { + if mode == DynamicFilterMergePolicy::ForwardFirstComplete { reports.truncate(1); } let Some((_, template)) = reports.first() else { @@ -314,7 +316,7 @@ impl DynamicFilterRegistry { mode, ) .map(Box::new); - let is_complete = mode == DynamicFilterMergeMode::FirstProducerComplete || all_complete; + let is_complete = mode == DynamicFilterMergePolicy::ForwardFirstComplete || all_complete; if previous.is_some_and(|previous| { previous.inner_expr == inner_expr && previous.is_complete == is_complete }) { @@ -387,10 +389,10 @@ impl DynamicFilterRegistry { } /// Merges [`PhysicalExprNode`] together. Merge behavior (ex. AND vs OR) is determined -/// by the provided [`DynamicFilterMergeMode`]. +/// by the provided [`DynamicFilterMergePolicy`]. fn merge_predicates( mut predicates: Vec, - mode: DynamicFilterMergeMode, + mode: DynamicFilterMergePolicy, ) -> Option { match predicates.len() { 0 => None, @@ -400,7 +402,7 @@ fn merge_predicates( expr_type: Some(ExprType::BinaryExpr(Box::new(PhysicalBinaryExprNode { l: None, r: None, - op: if mode == DynamicFilterMergeMode::Incremental { + op: if mode == DynamicFilterMergeMode::IntersectOnUpdate { "And" } else { "Or" diff --git a/src/dynamic_filtering/display.rs b/src/dynamic_filtering/display.rs index dc5943381..87aba894c 100644 --- a/src/dynamic_filtering/display.rs +++ b/src/dynamic_filtering/display.rs @@ -1,8 +1,9 @@ -use crate::codec::{apply_dynamic_filter_update, roundtrip_pb}; +use crate::codec::roundtrip_pb; use crate::common::TreeNodeExt; use crate::coordinator::DistributedExec; use crate::dynamic_filtering::{ - discover_dynamic_filter_consumers, discover_dynamic_filter_producers, + apply_dynamic_filter_update, discover_dynamic_filter_consumers, + discover_dynamic_filter_producers, }; use crate::execution_plans::DistributedLeafExec; use crate::stage::{LocalStage, Stage, find_all_stages}; @@ -218,9 +219,6 @@ fn apply_reports_to_distributed_leaves( else { return internal_err!("expected a dynamic filter in the completed task report"); }; - if dynamic_filter_proto.generation <= 1 { - continue; - } let Some(predicate) = dynamic_filter_proto.inner_expr.as_deref() else { return internal_err!("reported dynamic filter has no predicate"); }; diff --git a/src/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs index 3eadda4d9..f0a638858 100644 --- a/src/dynamic_filtering/mod.rs +++ b/src/dynamic_filtering/mod.rs @@ -2,11 +2,16 @@ mod discovery; mod display; use crate::DistributedConfig; -use crate::codec::roundtrip_pb; -use datafusion::common::Result; +use crate::codec::{decode_physical_expr, encode_physical_expr, roundtrip_pb}; +use datafusion::arrow::datatypes::Schema; +use datafusion::common::{Result, internal_err}; use datafusion::execution::TaskContext; use datafusion::execution::config::SessionConfig; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::protobuf; +use datafusion_proto::protobuf::physical_expr_node::ExprType; use std::sync::Arc; pub(crate) use discovery::*; @@ -187,3 +192,61 @@ pub(crate) fn maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationshi Ok(plan) } } + +/// Applies a predicate to a consumer dynamic filter expression. +/// +/// To update a consumer, this function creates a synthetic "producer" +/// [`DynamicFilterPhysicalExpr`] which can be used to update consumer +/// dynamic filters via a shared state. The reason for this is to match +/// the same behavior as DataFusion. +/// +/// ```text +/// HashJoinExec: on=[build.key = probe.key] +/// │ └── producer: original=[key], children=[key] ──────────────┐ +/// ├── build │ +/// └── UnionExec: probe │ +/// ├── DataSourceExec: phone_number AS key │ shared inner +/// │ └── consumer 1: original=[key], children=[phone_number] ─┤ +/// └── DataSourceExec: telephone AS key │ +/// └── consumer 2: original=[key], children=[telephone] ───┘ +/// ``` +/// +/// Behavior: +/// 1. [`DynamicFilterPhysicalExpr::update()`] remaps any occurrences of `original` to `children` +/// and stores the result in the shared state. +/// +/// 2. [`DynamicFilterPhysicalExpr::current()`] reads the remapped expression and remaps it again, +/// mapping `original` to `children` and returns it without storing. +/// +/// In the above plan, the producer calls update() so an expression like `key > 123` is mapped to +/// `key > 123` and this is stored. Then, the consumers call current(), reading from the same state, +/// to get `phone_number > 123` and `telephone > 123` respectively. +/// +/// You cannot update() consumers directly. In the above example, updating consumer 1 would remap +/// `key > 123` to `phone_number > 123` and store this in the shared state. Consumer 2 would be +/// unable to apply this filter now. +pub(crate) fn apply_dynamic_filter_update( + consumer: &Arc, + predicate: &protobuf::PhysicalExprNode, + producer_schema: &Schema, + task_ctx: &TaskContext, +) -> Result<()> { + let predicate = decode_physical_expr(predicate, producer_schema, task_ctx)?; + // Since consumer.children() returns the remapped children, we use the proto as a workaround to get the + // original children from the producer. + let consumer: Arc = consumer.clone(); + let proto = encode_physical_expr(&consumer, task_ctx)?; + let Some(ExprType::DynamicFilter(dynamic_filter)) = proto.expr_type else { + return internal_err!("expected a dynamic filter expression"); + }; + let original_children = dynamic_filter + .children + .iter() + .map(|child| decode_physical_expr(child, producer_schema, task_ctx)) + .collect::>>()?; + let update_target = Arc::clone(&consumer).with_new_children(original_children)?; + let Ok(update_target) = Arc::downcast::(update_target) else { + return internal_err!("expected a dynamic filter update target"); + }; + update_target.update(predicate) +} diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index 6c9db46a9..185f255c5 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,7 +1,7 @@ -use crate::codec::apply_dynamic_filter_update; use crate::common::TreeNodeExt; use crate::dynamic_filtering::{ - DiscoveredDynamicFilter, discover_dynamic_filter_consumers, discover_dynamic_filter_producers, + DiscoveredDynamicFilter, apply_dynamic_filter_update, discover_dynamic_filter_consumers, + discover_dynamic_filter_producers, }; use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers}; use crate::execution_plans::SamplerExec; @@ -216,7 +216,9 @@ impl Worker { &dynamic_filter_task_ctx, ) { let _ = dynamic_filters_error_tx.try_send(error); - break; + // Continue to read messages. Let the coordinator handle the error + // sent on the channel and gracefully terminate the worker.impl + continue; } } } @@ -292,14 +294,14 @@ impl Worker { ) })); + let dynamic_filters_error_stream = ReceiverStream::new(dynamic_filters_error_rx); + let stream = select_all([ produced_dynamic_filters_stream.map(Ok).boxed(), load_info_stream.map(Ok).boxed(), metrics_stream.map(Ok).boxed(), dynamic_filters_stream.map(Ok).boxed(), - ReceiverStream::new(dynamic_filters_error_rx) - .map(Err) - .boxed(), + dynamic_filters_error_stream.map(Err).boxed(), ]) .boxed(); @@ -411,6 +413,10 @@ fn build_task_completed_dynamic_filters( ) -> Result { let mut filters = vec![]; for consumer in discover_dynamic_filter_consumers(plan)?.consumers { + // Generation numbers start at 1. Avoid sending any empty filters. + if consumer.expression.snapshot_generation() <= 1 { + continue; + } filters.push(TaskDynamicFilter { expression_id: consumer.id, expression: MaybeEncoded::Decoded(consumer.expression), From 2e70c56a37b1d304448316dd611debf183caa292 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Wed, 7 Oct 2026 14:46:52 +0000 Subject: [PATCH 04/11] pr comments - use insta to normalize filters in plan snapshots - inline repeat query --- src/coordinator/dynamic_filter_registry.rs | 2 +- src/test_utils/insta.rs | 5 +++ tests/dynamic_filtering/collect_left_join.rs | 32 ++++++++++++--- tests/dynamic_filtering/common.rs | 43 +------------------- tests/dynamic_filtering/partitioned_join.rs | 30 ++++++++++++-- 5 files changed, 60 insertions(+), 52 deletions(-) diff --git a/src/coordinator/dynamic_filter_registry.rs b/src/coordinator/dynamic_filter_registry.rs index ae041c31a..493204088 100644 --- a/src/coordinator/dynamic_filter_registry.rs +++ b/src/coordinator/dynamic_filter_registry.rs @@ -402,7 +402,7 @@ fn merge_predicates( expr_type: Some(ExprType::BinaryExpr(Box::new(PhysicalBinaryExprNode { l: None, r: None, - op: if mode == DynamicFilterMergeMode::IntersectOnUpdate { + op: if mode == DynamicFilterMergePolicy::IntersectOnUpdate { "And" } else { "Or" diff --git a/src/test_utils/insta.rs b/src/test_utils/insta.rs index d1ce9d438..1e280b76e 100644 --- a/src/test_utils/insta.rs +++ b/src/test_utils/insta.rs @@ -30,6 +30,11 @@ pub fn settings() -> insta::Settings { r"(DynamicFilter \[[^\[\]]*IN \(SET\) \(\[)[^\]]*(\]\))", "${1}${2}", ); + // Omit runtime pruning details when dynamic filters have normalized labels. + settings.add_filter( + r"(DynamicFilter \[ expression_id_[^\n]*?dynamic_rg_pruning=eligible)[^\n]*", + "${1}", + ); // When a scan has a limit, DataFusion stops scanning files as soon as the row count limit // is satisfied. Because file metadata is fetched asynchronously (buffer_unordered) from // the directory listing, whichever file finishes reading first is selected as the sole diff --git a/tests/dynamic_filtering/collect_left_join.rs b/tests/dynamic_filtering/collect_left_join.rs index 55ff2e18e..d5a90b24c 100644 --- a/tests/dynamic_filtering/collect_left_join.rs +++ b/tests/dynamic_filtering/collect_left_join.rs @@ -1,6 +1,6 @@ #[cfg(test)] mod tests { - use crate::common::{LOCAL_AND_REMOTE_UNION_QUERY, TestQuery}; + use crate::common::TestQuery; use datafusion::common::Result; use datafusion_distributed::assert_snapshot; @@ -118,10 +118,32 @@ mod tests { /// Task 1's filter reports must reach variant 0 of its isolated union child. #[tokio::test] async fn union_with_local_and_remote_probe_children() -> Result<()> { - let display = TestQuery::new(LOCAL_AND_REMOTE_UNION_QUERY) - .with_broadcast_joins() - .execute() - .await?; + let display = TestQuery::new( + r#" + WITH remote_probe AS ( + SELECT probe."WindGustDir" AS key + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) nested_build + JOIN weather probe + ON nested_build.key = probe."RainToday" + ) + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "WindGustDir" AS key + FROM weather + ) build + JOIN ( + SELECT "WindGustDir" AS key FROM weather + UNION ALL + SELECT key FROM remote_probe + ) probe ON build.key = probe.key + "#, + ) + .with_broadcast_joins() + .execute() + .await?; // The consumer in c0 of t0 in stage 5 sees the dynamic filter from the topmost join. // The consumer in c1 of t1 in stage 5 sees distinct local dynamic filters from both joins. // DataFusion ANDs distinct dynamic filter predicates. diff --git a/tests/dynamic_filtering/common.rs b/tests/dynamic_filtering/common.rs index 58a668f2d..37330a747 100644 --- a/tests/dynamic_filtering/common.rs +++ b/tests/dynamic_filtering/common.rs @@ -22,28 +22,6 @@ use datafusion_distributed::{ }; use std::sync::Arc; -pub(crate) const LOCAL_AND_REMOTE_UNION_QUERY: &str = r#" - WITH remote_probe AS ( - SELECT probe."WindGustDir" AS key - FROM ( - SELECT DISTINCT "RainToday" AS key - FROM weather - ) nested_build - JOIN weather probe - ON nested_build.key = probe."RainToday" - ) - SELECT COUNT(*) - FROM ( - SELECT DISTINCT "WindGustDir" AS key - FROM weather - ) build - JOIN ( - SELECT "WindGustDir" AS key FROM weather - UNION ALL - SELECT key FROM remote_probe - ) probe ON build.key = probe.key -"#; - pub(crate) struct TestQuery<'a> { sql: &'a str, expected_rows: usize, @@ -276,10 +254,7 @@ struct DynamicFilterLabels { impl DynamicFilterLabels { fn normalize(&mut self, plan: Arc) -> Result { self.label_plan(&plan)?; - Ok(remove_runtime_pruning_details(display_plan_ascii( - plan.as_ref(), - false, - ))) + Ok(display_plan_ascii(plan.as_ref(), false)) } fn label_plan(&mut self, plan: &Arc) -> Result<()> { @@ -372,22 +347,6 @@ fn normalize_columns(predicate: Arc) -> Result String { - const START: &str = "DynamicFilter [ expression_id_"; - const END: &str = "dynamic_rg_pruning=eligible"; - - display - .lines() - .map(|line| { - let (Some(start), Some(end)) = (line.find(START), line.find(END)) else { - return line.to_owned(); - }; - format!("{}{}", &line[..start], &line[start..end + END.len()]) - }) - .collect::>() - .join("\n") -} - fn set_dynamic_filter_pushdown(ctx: &SessionContext, enabled: bool) -> Result<()> { let state = ctx.state_ref(); state.write().config_mut().options_mut().set( diff --git a/tests/dynamic_filtering/partitioned_join.rs b/tests/dynamic_filtering/partitioned_join.rs index e86ccfb2a..e8eb3ced2 100644 --- a/tests/dynamic_filtering/partitioned_join.rs +++ b/tests/dynamic_filtering/partitioned_join.rs @@ -1,6 +1,6 @@ #[cfg(test)] mod tests { - use crate::common::{LOCAL_AND_REMOTE_UNION_QUERY, TestQuery, execute_range_partitioned_query}; + 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; @@ -268,9 +268,31 @@ mod tests { #[tokio::test] async fn union_with_local_and_remote_probe_children() -> Result<()> { - let display = TestQuery::new(LOCAL_AND_REMOTE_UNION_QUERY) - .execute() - .await?; + let display = TestQuery::new( + r#" + WITH remote_probe AS ( + SELECT probe."WindGustDir" AS key + FROM ( + SELECT DISTINCT "RainToday" AS key + FROM weather + ) nested_build + JOIN weather probe + ON nested_build.key = probe."RainToday" + ) + SELECT COUNT(*) + FROM ( + SELECT DISTINCT "WindGustDir" AS key + FROM weather + ) build + JOIN ( + SELECT "WindGustDir" AS key FROM weather + UNION ALL + SELECT key FROM remote_probe + ) probe ON build.key = probe.key + "#, + ) + .execute() + .await?; // Complex test. // Dynamic filters from the stage 5 tasks are merged and forwarded to both c0 of t0 in stage From b090e56fc6686ef2748aaaef564299699907a7e8 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Wed, 7 Oct 2026 16:15:57 +0000 Subject: [PATCH 05/11] reduce dynamic filter display size --- tests/dynamic_filtering/common.rs | 168 +++++++++++++----------------- 1 file changed, 70 insertions(+), 98 deletions(-) diff --git a/tests/dynamic_filtering/common.rs b/tests/dynamic_filtering/common.rs index 37330a747..434e9248d 100644 --- a/tests/dynamic_filtering/common.rs +++ b/tests/dynamic_filtering/common.rs @@ -1,7 +1,7 @@ use datafusion::arrow::datatypes::DataType; use datafusion::common::test_util::batches_to_sort_string; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; -use datafusion::common::{HashMap, Result, ScalarValue, SplitPoint, internal_err}; +use datafusion::common::{Result, ScalarValue, SplitPoint, internal_err}; use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, @@ -84,8 +84,7 @@ impl<'a> TestQuery<'a> { self } - /// Adds a normalized hash to the dynamic filter display formatter. See - /// [`DynamicFilterLabels`] below. + /// Adds column-independent predicate labels for comparing remapped consumers. pub(crate) fn with_normalized_filter_hashes(mut self) -> Self { self.normalized_filter_hashes = true; self @@ -226,11 +225,15 @@ async fn execute_query_and_display( "expected dynamic_filter_updates_received == 0, got {updates}" ); } - DynamicFilterLabels { - normalized_predicates: normalized_filter_hashes.then(Vec::new), - ..Default::default() - } - .normalize(plan_with_dynamic_filters) + display_with_dynamic_filter_labels(&plan_with_dynamic_filters, normalized_filter_hashes) +} + +fn set_dynamic_filter_pushdown(ctx: &SessionContext, enabled: bool) -> Result<()> { + let state = ctx.state_ref(); + state.write().config_mut().options_mut().set( + "datafusion.optimizer.enable_dynamic_filter_pushdown", + &enabled.to_string(), + ) } /// Encodes dynamic filter expressions in the plan deterministically for consistent snapshots. @@ -241,116 +244,85 @@ async fn execute_query_and_display( /// Both of these are normalized to be monotonic integers starting from `0` for readability, /// assigned during a pre-order traversal of the plan. /// -/// If `normalized_predicates` is non empty, then we append the suffix `_normalized_hash_{}` +/// If `normalized_filter_hashes` is non empty, then we append the suffix `_normalized_hash_{}` /// containing the hash of the expression without column names. -#[derive(Default)] -struct DynamicFilterLabels { - expression_ids: HashMap, - // InListExpr ignores element order for equality, but not hashing. - predicates: Vec>, - normalized_predicates: Option>>, -} - -impl DynamicFilterLabels { - fn normalize(&mut self, plan: Arc) -> Result { - self.label_plan(&plan)?; - Ok(display_plan_ascii(plan.as_ref(), false)) - } - - fn label_plan(&mut self, plan: &Arc) -> Result<()> { - let mut updates = vec![]; - plan.apply(|node| { - if let Some(leaf) = node.downcast_ref::() { - for variant in leaf.variants() { - self.label_variant(variant, &mut updates)?; - } +fn display_with_dynamic_filter_labels( + plan: &Arc, + normalized_filter_hashes: bool, +) -> Result { + let mut nodes = Vec::new(); + plan.apply(|node| { + if let Some(leaf) = node.downcast_ref::() { + for variant in leaf.variants() { + variant.apply(|node| { + nodes.push(Arc::clone(node)); + Ok(TreeNodeRecursion::Continue) + })?; } - Ok(TreeNodeRecursion::Continue) - })?; - // Read every predicate before changing any potentially shared filter state. - for (dynamic_filter, label) in updates { - dynamic_filter.update(Arc::new(UnKnownColumn::new(&label)))?; } - Ok(()) - } + Ok(TreeNodeRecursion::Continue) + })?; - fn label_variant( - &mut self, - variant: &Arc, - updates: &mut Vec<(Arc, String)>, - ) -> Result<()> { - variant.apply(|node| { - node.apply_expressions(&mut |root| { - root.apply(|expression| { - let Ok(dynamic_filter) = - Arc::downcast::(expression.clone()) - else { - return Ok(TreeNodeRecursion::Continue); - }; - if expression.snapshot_generation() == 1 { - return Ok(TreeNodeRecursion::Continue); - } - let Some(expression_id) = expression.expression_id() else { - return internal_err!("dynamic filter did not have an expression ID"); - }; - let next_expression_id = self.expression_ids.len(); - let expression_id = *self - .expression_ids - .entry(expression_id) - .or_insert(next_expression_id); - let predicate = dynamic_filter.current()?; - let predicate_id = - intern_predicate(&mut self.predicates, Arc::clone(&predicate)); - let mut label = format!("expression_id_{expression_id}_hash_{predicate_id}"); - if let Some(predicates) = &mut self.normalized_predicates { - let normalized_id = - intern_predicate(predicates, normalize_columns(predicate)?); - label.push_str(&format!("_normalized_hash_{normalized_id}")); - } - updates.push((dynamic_filter, label)); - Ok(TreeNodeRecursion::Continue) - }) - })?; - Ok(TreeNodeRecursion::Continue) + let mut expression_ids = Vec::new(); + let mut predicates = Vec::new(); + let mut normalized_predicates = Vec::new(); + let mut updates = Vec::new(); + for node in nodes { + node.apply_expressions(&mut |root| { + root.apply(|expression| { + let Ok(dynamic_filter) = + Arc::downcast::(expression.clone()) + else { + return Ok(TreeNodeRecursion::Continue); + }; + if expression.snapshot_generation() == 1 { + return Ok(TreeNodeRecursion::Continue); + } + let Some(expression_id) = expression.expression_id() else { + return internal_err!("dynamic filter did not have an expression ID"); + }; + let expression_id = intern(&mut expression_ids, expression_id); + let predicate = dynamic_filter.current()?; + let predicate_id = intern(&mut predicates, Arc::clone(&predicate)); + let mut label = format!("expression_id_{expression_id}_hash_{predicate_id}"); + if normalized_filter_hashes { + let normalized_id = + intern(&mut normalized_predicates, normalize_columns(predicate)?); + label.push_str(&format!("_normalized_hash_{normalized_id}")); + } + updates.push((dynamic_filter, label)); + Ok(TreeNodeRecursion::Continue) + }) })?; - Ok(()) } + // Read every predicate before changing any potentially shared filter state. + for (dynamic_filter, label) in updates { + dynamic_filter.update(Arc::new(UnKnownColumn::new(&label)))?; + } + Ok(display_plan_ascii(plan.as_ref(), false)) } -fn intern_predicate( - predicates: &mut Vec>, - predicate: Arc, -) -> usize { - predicates +fn intern(values: &mut Vec, value: T) -> usize { + values .iter() - .position(|existing| existing.eq(&predicate)) + .position(|existing| existing.eq(&value)) .unwrap_or_else(|| { - let id = predicates.len(); - predicates.push(predicate); + let id = values.len(); + values.push(value); id }) } +/// Replaces every column with `column_0@0`. fn normalize_columns(predicate: Arc) -> Result> { - let mut columns = HashMap::new(); Ok(predicate .transform_down(|expression| { - let Some(column) = expression.downcast_ref::() else { + if expression.downcast_ref::().is_none() { return Ok(Transformed::no(expression)); - }; - let next_id = columns.len(); - let id = *columns.entry(column.clone()).or_insert(next_id); + } Ok(Transformed::yes( - Arc::new(Column::new(&format!("column_{id}"), id)) as Arc, + Arc::new(Column::new("column_0", 0)) as Arc )) })? .data) } - -fn set_dynamic_filter_pushdown(ctx: &SessionContext, enabled: bool) -> Result<()> { - let state = ctx.state_ref(); - state.write().config_mut().options_mut().set( - "datafusion.optimizer.enable_dynamic_filter_pushdown", - &enabled.to_string(), - ) -} From 0b3a283be04dc657d3053323f7ccca350bd41523 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Wed, 7 Oct 2026 18:39:19 +0000 Subject: [PATCH 06/11] cleanup proto encoding --- src/dynamic_filtering/display.rs | 11 +--- src/dynamic_filtering/mod.rs | 80 ++++++++++++++++---------- src/worker/impl_coordinator_channel.rs | 24 +------- 3 files changed, 54 insertions(+), 61 deletions(-) diff --git a/src/dynamic_filtering/display.rs b/src/dynamic_filtering/display.rs index 87aba894c..000254359 100644 --- a/src/dynamic_filtering/display.rs +++ b/src/dynamic_filtering/display.rs @@ -17,7 +17,6 @@ use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::{ ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, ReplaceChildrenOptions, }; -use datafusion_proto::protobuf::physical_expr_node::ExprType; use std::sync::Arc; /// Rewrites an executed distributed plan with the dynamic filters reported by its completed @@ -214,14 +213,6 @@ fn apply_reports_to_distributed_leaves( let Some(expression) = updates.get(&consumer.id).copied() else { continue; }; - let proto = expression.to_proto(task_ctx)?; - let Some(ExprType::DynamicFilter(dynamic_filter_proto)) = proto.expr_type.as_ref() - else { - return internal_err!("expected a dynamic filter in the completed task report"); - }; - let Some(predicate) = dynamic_filter_proto.inner_expr.as_deref() else { - return internal_err!("reported dynamic filter has no predicate"); - }; let Some(producer_schema) = producer_schemas.get(&consumer.id) else { return internal_err!( "missing producer schema for dynamic filter {}", @@ -230,7 +221,7 @@ fn apply_reports_to_distributed_leaves( }; apply_dynamic_filter_update( &consumer.expression, - predicate, + expression, producer_schema.as_ref(), task_ctx, )?; diff --git a/src/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs index f0a638858..6c30ccabf 100644 --- a/src/dynamic_filtering/mod.rs +++ b/src/dynamic_filtering/mod.rs @@ -1,8 +1,8 @@ mod discovery; mod display; -use crate::DistributedConfig; -use crate::codec::{decode_physical_expr, encode_physical_expr, roundtrip_pb}; +use crate::codec::{decode_physical_expr, roundtrip_pb}; +use crate::{DistributedConfig, MaybeEncoded}; use datafusion::arrow::datatypes::Schema; use datafusion::common::{Result, internal_err}; use datafusion::execution::TaskContext; @@ -10,7 +10,6 @@ use datafusion::execution::config::SessionConfig; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_plan::ExecutionPlan; -use datafusion_proto::protobuf; use datafusion_proto::protobuf::physical_expr_node::ExprType; use std::sync::Arc; @@ -193,12 +192,12 @@ pub(crate) fn maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationshi } } -/// Applies a predicate to a consumer dynamic filter expression. +/// Applies a `source` dynamic filter expression to a `target` dynamic filter expression. /// -/// To update a consumer, this function creates a synthetic "producer" -/// [`DynamicFilterPhysicalExpr`] which can be used to update consumer -/// dynamic filters via a shared state. The reason for this is to match -/// the same behavior as DataFusion. +/// This function always uses the original children of the dynamic filter expression +/// when applying an update. See below. +/// +/// ## Normal Behavior /// /// ```text /// HashJoinExec: on=[build.key = probe.key] @@ -211,42 +210,65 @@ pub(crate) fn maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationshi /// └── consumer 2: original=[key], children=[telephone] ───┘ /// ``` /// -/// Behavior: /// 1. [`DynamicFilterPhysicalExpr::update()`] remaps any occurrences of `original` to `children` /// and stores the result in the shared state. /// /// 2. [`DynamicFilterPhysicalExpr::current()`] reads the remapped expression and remaps it again, /// mapping `original` to `children` and returns it without storing. /// -/// In the above plan, the producer calls update() so an expression like `key > 123` is mapped to -/// `key > 123` and this is stored. Then, the consumers call current(), reading from the same state, -/// to get `phone_number > 123` and `telephone > 123` respectively. +/// Example: the producer calls update(`key > 123`) which trivially maps to `key > 123` +/// and is stored. Then, the consumers call current(), reading from the shared inner state, +/// to get the remapped `phone_number > 123` and `telephone > 123` expressions respectively. +/// +/// ## Failiure Mode +/// +/// The source expression may be a consumer and cause remapping to fail. /// -/// You cannot update() consumers directly. In the above example, updating consumer 1 would remap -/// `key > 123` to `phone_number > 123` and store this in the shared state. Consumer 2 would be -/// unable to apply this filter now. +/// Say that the source and target are both `consumer 1: original=[key], children=[phone_number]`. +/// The source is a populated filter from a worker `phone_number > 123` and the target is empty, +/// used for display. +/// +/// In this situation, we should use the inner expression `key > 123` from the source to call +/// update(`key > 123`). We should not call `update(source.current())` or +/// `update(phone_number > 123)`, or else the update fails. pub(crate) fn apply_dynamic_filter_update( - consumer: &Arc, - predicate: &protobuf::PhysicalExprNode, - producer_schema: &Schema, + target: &Arc, + source: &MaybeEncoded>, + source_schema: &Schema, task_ctx: &TaskContext, ) -> Result<()> { - let predicate = decode_physical_expr(predicate, producer_schema, task_ctx)?; - // Since consumer.children() returns the remapped children, we use the proto as a workaround to get the - // original children from the producer. - let consumer: Arc = consumer.clone(); - let proto = encode_physical_expr(&consumer, task_ctx)?; - let Some(ExprType::DynamicFilter(dynamic_filter)) = proto.expr_type else { - return internal_err!("expected a dynamic filter expression"); + // We need two private fields from the source, so we convert to the proto + // representation as a workaround. + // 1. `is_complete` (which is public in df56) + // 2. The inner expression which has unremapped children (this is only + // required when the source is another consumer) + // + // TODO: Avoid needing to encode to proto. + let source_proto = source.to_proto(task_ctx)?; + if source_proto.expr_id != target.expression_id() { + return internal_err!("dynamic filter update has a mismatched expression ID"); + } + let Some(ExprType::DynamicFilter(source_dynamic_filter)) = source_proto.expr_type else { + return internal_err!("expected a dynamic filter update"); }; - let original_children = dynamic_filter + let Some(source_predicate) = source_dynamic_filter.inner_expr.as_deref() else { + return internal_err!("dynamic filter update has no predicate"); + }; + let predicate = decode_physical_expr(source_predicate, source_schema, task_ctx)?; + + let source_original_children = source_dynamic_filter .children .iter() - .map(|child| decode_physical_expr(child, producer_schema, task_ctx)) + .map(|child| decode_physical_expr(child, source_schema, task_ctx)) .collect::>>()?; - let update_target = Arc::clone(&consumer).with_new_children(original_children)?; + + let update_target = Arc::clone(target).with_new_children(source_original_children)?; let Ok(update_target) = Arc::downcast::(update_target) else { return internal_err!("expected a dynamic filter update target"); }; - update_target.update(predicate) + update_target.update(predicate)?; + if source_dynamic_filter.is_complete { + update_target.mark_complete(); + } + Ok(()) } diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index 185f255c5..335adb9a2 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -25,7 +25,6 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; -use datafusion_proto::protobuf::physical_expr_node::ExprType; use futures::stream::{BoxStream, FuturesUnordered, select_all}; use futures::{FutureExt, StreamExt, TryStreamExt}; use http::HeaderMap; @@ -373,31 +372,12 @@ fn apply_merged_dynamic_filter( ); }; - // Bytes to PhysicalExprNode - let expression = filter.expression.to_proto(task_ctx)?; - - if expression.expr_id != Some(filter.expression_id) { - return internal_err!("dynamic filter update has a mismatched expression ID"); - } - - let Some(ExprType::DynamicFilter(update)) = expression.expr_type else { - return internal_err!("expected a dynamic filter update"); - }; - - let Some(predicate) = update.inner_expr else { - return internal_err!("dynamic filter update has no predicate"); - }; - apply_dynamic_filter_update( &consumer.expression, - &predicate, + &filter.expression, filter.producer_schema.as_ref(), task_ctx, - )?; - if update.is_complete { - consumer.expression.mark_complete(); - } - Ok(()) + ) } /// Finds all consumed dynamic filters for the completed task report. From 5dca6917281f747b5b411d234e7fcd4758f18da7 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Wed, 7 Oct 2026 18:42:00 +0000 Subject: [PATCH 07/11] fmt and clippy --- src/coordinator/dynamic_filter_registry.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coordinator/dynamic_filter_registry.rs b/src/coordinator/dynamic_filter_registry.rs index 493204088..661535dca 100644 --- a/src/coordinator/dynamic_filter_registry.rs +++ b/src/coordinator/dynamic_filter_registry.rs @@ -51,7 +51,7 @@ pub(super) struct PlannedDynamicFilter { /// Populated upon task registration when a producer is found. /// /// The behavior to take when an update arrives from a producer. - pub(super) merge_mode: Option, + merge_mode: Option, /// Schema of the original producer arguments, shared by all producer tasks. pub(super) producer_schema: Option, From 1059c038fa2d594b1af849e826ab2b181943d1b9 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 8 Oct 2026 10:22:12 -0400 Subject: [PATCH 08/11] Update src/dynamic_filtering/mod.rs Co-authored-by: Gabriel <45515538+gabotechs@users.noreply.github.com> --- src/dynamic_filtering/mod.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs index 6c30ccabf..f61c7869b 100644 --- a/src/dynamic_filtering/mod.rs +++ b/src/dynamic_filtering/mod.rs @@ -243,7 +243,8 @@ pub(crate) fn apply_dynamic_filter_update( // 2. The inner expression which has unremapped children (this is only // required when the source is another consumer) // - // TODO: Avoid needing to encode to proto. + // TODO: Avoid needing to encode to proto: + // https://github.com/datafusion-contrib/datafusion-distributed/issues/768 let source_proto = source.to_proto(task_ctx)?; if source_proto.expr_id != target.expression_id() { return internal_err!("dynamic filter update has a mismatched expression ID"); From 9bed440597a388f88fceb3975ff55ce234d4cc0c Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 8 Oct 2026 10:22:24 -0400 Subject: [PATCH 09/11] Update src/dynamic_filtering/mod.rs Co-authored-by: Gabriel <45515538+gabotechs@users.noreply.github.com> --- src/dynamic_filtering/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs index f61c7869b..2325e876b 100644 --- a/src/dynamic_filtering/mod.rs +++ b/src/dynamic_filtering/mod.rs @@ -220,7 +220,7 @@ pub(crate) fn maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationshi /// and is stored. Then, the consumers call current(), reading from the shared inner state, /// to get the remapped `phone_number > 123` and `telephone > 123` expressions respectively. /// -/// ## Failiure Mode +/// ## Failure Mode /// /// The source expression may be a consumer and cause remapping to fail. /// From 15cb572e178c0a4a6b73f5eb28c3bd0578915de7 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 8 Oct 2026 14:29:47 +0000 Subject: [PATCH 10/11] move code around --- src/dynamic_filtering/apply.rs | 90 +++++++++++++++++++++++++ src/dynamic_filtering/mod.rs | 93 ++------------------------ src/worker/impl_coordinator_channel.rs | 64 ++++++++---------- 3 files changed, 123 insertions(+), 124 deletions(-) create mode 100644 src/dynamic_filtering/apply.rs diff --git a/src/dynamic_filtering/apply.rs b/src/dynamic_filtering/apply.rs new file mode 100644 index 000000000..2a04acedc --- /dev/null +++ b/src/dynamic_filtering/apply.rs @@ -0,0 +1,90 @@ +use crate::MaybeEncoded; +use crate::codec::decode_physical_expr; +use datafusion::arrow::datatypes::Schema; +use datafusion::common::{Result, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; +use datafusion_proto::protobuf::physical_expr_node::ExprType; +use std::sync::Arc; + +/// Applies a `source` dynamic filter expression to a `target` dynamic filter expression. +/// +/// This function always uses the original children of the dynamic filter expression +/// when applying an update. See below. +/// +/// ## Normal Behavior +/// +/// ```text +/// HashJoinExec: on=[build.key = probe.key] +/// │ └── producer: original=[key], children=[key] ──────────────┐ +/// ├── build │ +/// └── UnionExec: probe │ +/// ├── DataSourceExec: phone_number AS key │ shared inner +/// │ └── consumer 1: original=[key], children=[phone_number] ─┤ +/// └── DataSourceExec: telephone AS key │ +/// └── consumer 2: original=[key], children=[telephone] ───┘ +/// ``` +/// +/// 1. [`DynamicFilterPhysicalExpr::update()`] remaps any occurrences of `original` to `children` +/// and stores the result in the shared state. +/// +/// 2. [`DynamicFilterPhysicalExpr::current()`] reads the remapped expression and remaps it again, +/// mapping `original` to `children` and returns it without storing. +/// +/// Example: the producer calls update(`key > 123`) which trivially maps to `key > 123` +/// and is stored. Then, the consumers call current(), reading from the shared inner state, +/// to get the remapped `phone_number > 123` and `telephone > 123` expressions respectively. +/// +/// ## Failiure Mode +/// +/// The source expression may be a consumer and cause remapping to fail. +/// +/// Say that the source and target are both `consumer 1: original=[key], children=[phone_number]`. +/// The source is a populated filter from a worker `phone_number > 123` and the target is empty, +/// used for display. +/// +/// In this situation, we should use the inner expression `key > 123` from the source to call +/// update(`key > 123`). We should not call `update(source.current())` or +/// `update(phone_number > 123)`, or else the update fails. +pub(crate) fn apply_dynamic_filter_update( + target: &Arc, + source: &MaybeEncoded>, + source_schema: &Schema, + task_ctx: &TaskContext, +) -> Result<()> { + // We need two private fields from the source, so we convert to the proto + // representation as a workaround. + // 1. `is_complete` (which is public in df56) + // 2. The inner expression which has unremapped children (this is only + // required when the source is another consumer) + // + // TODO: Avoid needing to encode to proto. + let source_proto = source.to_proto(task_ctx)?; + if source_proto.expr_id != target.expression_id() { + return internal_err!("dynamic filter update has a mismatched expression ID"); + } + let Some(ExprType::DynamicFilter(source_dynamic_filter)) = source_proto.expr_type else { + return internal_err!("expected a dynamic filter update"); + }; + let Some(source_predicate) = source_dynamic_filter.inner_expr.as_deref() else { + return internal_err!("dynamic filter update has no predicate"); + }; + let predicate = decode_physical_expr(source_predicate, source_schema, task_ctx)?; + + let source_original_children = source_dynamic_filter + .children + .iter() + .map(|child| decode_physical_expr(child, source_schema, task_ctx)) + .collect::>>()?; + + let update_target = Arc::clone(target).with_new_children(source_original_children)?; + let Ok(update_target) = Arc::downcast::(update_target) else { + return internal_err!("expected a dynamic filter update target"); + }; + update_target.update(predicate)?; + if source_dynamic_filter.is_complete { + update_target.mark_complete(); + } + Ok(()) +} diff --git a/src/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs index 6c30ccabf..668744216 100644 --- a/src/dynamic_filtering/mod.rs +++ b/src/dynamic_filtering/mod.rs @@ -1,18 +1,16 @@ +mod apply; mod discovery; mod display; -use crate::codec::{decode_physical_expr, roundtrip_pb}; -use crate::{DistributedConfig, MaybeEncoded}; -use datafusion::arrow::datatypes::Schema; -use datafusion::common::{Result, internal_err}; +use crate::DistributedConfig; +use crate::codec::roundtrip_pb; +use datafusion::common::Result; use datafusion::execution::TaskContext; use datafusion::execution::config::SessionConfig; -use datafusion::physical_expr::PhysicalExpr; -use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion::physical_plan::ExecutionPlan; -use datafusion_proto::protobuf::physical_expr_node::ExprType; use std::sync::Arc; +pub(crate) use apply::apply_dynamic_filter_update; 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; @@ -191,84 +189,3 @@ pub(crate) fn maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationshi Ok(plan) } } - -/// Applies a `source` dynamic filter expression to a `target` dynamic filter expression. -/// -/// This function always uses the original children of the dynamic filter expression -/// when applying an update. See below. -/// -/// ## Normal Behavior -/// -/// ```text -/// HashJoinExec: on=[build.key = probe.key] -/// │ └── producer: original=[key], children=[key] ──────────────┐ -/// ├── build │ -/// └── UnionExec: probe │ -/// ├── DataSourceExec: phone_number AS key │ shared inner -/// │ └── consumer 1: original=[key], children=[phone_number] ─┤ -/// └── DataSourceExec: telephone AS key │ -/// └── consumer 2: original=[key], children=[telephone] ───┘ -/// ``` -/// -/// 1. [`DynamicFilterPhysicalExpr::update()`] remaps any occurrences of `original` to `children` -/// and stores the result in the shared state. -/// -/// 2. [`DynamicFilterPhysicalExpr::current()`] reads the remapped expression and remaps it again, -/// mapping `original` to `children` and returns it without storing. -/// -/// Example: the producer calls update(`key > 123`) which trivially maps to `key > 123` -/// and is stored. Then, the consumers call current(), reading from the shared inner state, -/// to get the remapped `phone_number > 123` and `telephone > 123` expressions respectively. -/// -/// ## Failiure Mode -/// -/// The source expression may be a consumer and cause remapping to fail. -/// -/// Say that the source and target are both `consumer 1: original=[key], children=[phone_number]`. -/// The source is a populated filter from a worker `phone_number > 123` and the target is empty, -/// used for display. -/// -/// In this situation, we should use the inner expression `key > 123` from the source to call -/// update(`key > 123`). We should not call `update(source.current())` or -/// `update(phone_number > 123)`, or else the update fails. -pub(crate) fn apply_dynamic_filter_update( - target: &Arc, - source: &MaybeEncoded>, - source_schema: &Schema, - task_ctx: &TaskContext, -) -> Result<()> { - // We need two private fields from the source, so we convert to the proto - // representation as a workaround. - // 1. `is_complete` (which is public in df56) - // 2. The inner expression which has unremapped children (this is only - // required when the source is another consumer) - // - // TODO: Avoid needing to encode to proto. - let source_proto = source.to_proto(task_ctx)?; - if source_proto.expr_id != target.expression_id() { - return internal_err!("dynamic filter update has a mismatched expression ID"); - } - let Some(ExprType::DynamicFilter(source_dynamic_filter)) = source_proto.expr_type else { - return internal_err!("expected a dynamic filter update"); - }; - let Some(source_predicate) = source_dynamic_filter.inner_expr.as_deref() else { - return internal_err!("dynamic filter update has no predicate"); - }; - let predicate = decode_physical_expr(source_predicate, source_schema, task_ctx)?; - - let source_original_children = source_dynamic_filter - .children - .iter() - .map(|child| decode_physical_expr(child, source_schema, task_ctx)) - .collect::>>()?; - - let update_target = Arc::clone(target).with_new_children(source_original_children)?; - let Ok(update_target) = Arc::downcast::(update_target) else { - return internal_err!("expected a dynamic filter update target"); - }; - update_target.update(predicate)?; - if source_dynamic_filter.is_complete { - update_target.mark_complete(); - } - Ok(()) -} diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index 335adb9a2..57ae19640 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,6 +1,6 @@ use crate::common::TreeNodeExt; use crate::dynamic_filtering::{ - DiscoveredDynamicFilter, apply_dynamic_filter_update, discover_dynamic_filter_consumers, + apply_dynamic_filter_update, discover_dynamic_filter_consumers, discover_dynamic_filter_producers, }; use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers}; @@ -11,14 +11,13 @@ use crate::protocol::grpc::on_drop_stream; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::task_data::TaskDataMetrics; use crate::{ - ApplyDynamicFilter, CoordinatorToWorkerMsg, DistributedConfig, DistributedExt, - DistributedTaskContext, MaybeEncoded, ProducedDynamicFilter, SetPlanRequest, - TaskCompletedDynamicFilters, TaskData, TaskDynamicFilter, TaskMetrics, Worker, - WorkerQueryContext, WorkerToCoordinatorMsg, + CoordinatorToWorkerMsg, DistributedConfig, DistributedExt, DistributedTaskContext, + MaybeEncoded, ProducedDynamicFilter, SetPlanRequest, TaskCompletedDynamicFilters, TaskData, + TaskDynamicFilter, TaskMetrics, Worker, WorkerQueryContext, WorkerToCoordinatorMsg, }; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{ - DataFusionError, HashMap, HashSet, Result, exec_datafusion_err, internal_err, + DataFusionError, HashMap, HashSet, Result, exec_datafusion_err, internal_datafusion_err, }; use datafusion::execution::{SessionStateBuilder, TaskContext}; use datafusion::physical_expr::PhysicalExpr; @@ -148,11 +147,12 @@ impl Worker { let producer_filters = discover_dynamic_filter_producers(&task_data.base_plan)? .into_iter() .filter(|producer| dynamic_filter_remote_producer_ids.contains(&producer.id)); - let dynamic_filter_consumers = discover_dynamic_filter_consumers(&task_data.base_plan)? - .consumers - .into_iter() - .map(|consumer| (consumer.id, consumer)) - .collect(); + let dynamic_filter_consumers: HashMap<_, _> = + discover_dynamic_filter_consumers(&task_data.base_plan)? + .consumers + .into_iter() + .map(|consumer| (consumer.id, consumer)) + .collect(); let (producer_cancel_tx, producer_cancel_rx) = watch::channel(false); // Continue reading remaining messages (work unit feed data) in the background. @@ -209,11 +209,23 @@ impl Worker { sampler_gate.kick_off(); } CoordinatorToWorkerMsg::ApplyDynamicFilter(filter) => { - if let Err(error) = apply_merged_dynamic_filter( - *filter, - &dynamic_filter_consumers, - &dynamic_filter_task_ctx, - ) { + let result = dynamic_filter_consumers + .get(&filter.expression_id) + .ok_or_else(|| { + internal_datafusion_err!( + "received dynamic filter update for unknown consumer {}", + filter.expression_id + ) + }) + .and_then(|consumer| { + apply_dynamic_filter_update( + &consumer.expression, + &filter.expression, + filter.producer_schema.as_ref(), + &dynamic_filter_task_ctx, + ) + }); + if let Err(error) = result { let _ = dynamic_filters_error_tx.try_send(error); // Continue to read messages. Let the coordinator handle the error // sent on the channel and gracefully terminate the worker.impl @@ -360,26 +372,6 @@ fn produced_dynamic_filter_stream( .boxed() } -fn apply_merged_dynamic_filter( - filter: ApplyDynamicFilter, - consumers: &HashMap, - task_ctx: &Arc, -) -> Result<()> { - let Some(consumer) = consumers.get(&filter.expression_id) else { - return internal_err!( - "received dynamic filter update for unknown consumer {}", - filter.expression_id - ); - }; - - apply_dynamic_filter_update( - &consumer.expression, - &filter.expression, - filter.producer_schema.as_ref(), - task_ctx, - ) -} - /// Finds all consumed dynamic filters for the completed task report. /// /// Note that it's possible that a dynamic filter is consumed by the leaf, updated by From 6d03f14a22a2b3f22a2355a897c02ea6ce5e19ec Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 8 Oct 2026 15:04:16 +0000 Subject: [PATCH 11/11] fmt --- src/dynamic_filtering/mod.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/dynamic_filtering/mod.rs b/src/dynamic_filtering/mod.rs index 2d2e77b32..668744216 100644 --- a/src/dynamic_filtering/mod.rs +++ b/src/dynamic_filtering/mod.rs @@ -189,4 +189,3 @@ pub(crate) fn maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationshi Ok(plan) } } -