coordinator: display dynamic filters after execution

docs/upgrade/3.0.0.md

@@ -55,6 +55,14 @@ use datafusion_distributed::{WorkerChannel, grpc};
 let client: Box<dyn WorkerChannel> = grpc::create_worker_client(channel);
 ```
 
+Custom `WorkerChannel` transports must also carry the dynamic-filter display
+fields added to the coordinator protocol. Include `SetPlanRequest::dynamic_filter_ids`
+when constructing or encoding a plan request, and encode/decode the new
+`WorkerToCoordinatorMsg::TaskCompletedDynamicFilters` variant. The built-in gRPC
+transport handles both automatically. `TaskDynamicFilter::expression` contains a typed
+`datafusion_proto::protobuf::PhysicalExprNode`; custom transports should encode it only at their
+wire boundary.
+
 ## 2. Move `TaskEstimator` methods to event handlers
 
 `TaskEstimator`, `TaskEstimation`, `TaskRoutingContext`, and

src/common/dynamic_filtering.rs

@@ -0,0 +1,210 @@
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
+use datafusion::common::{HashMap, HashSet, Result};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr;
+use datafusion::physical_plan::ExecutionPlan;
+use std::sync::Arc;
+
+/// A dynamic-filter consumer discovered in an execution plan.
+#[derive(Clone)]
+pub(crate) struct DiscoveredDynamicFilter {
+    pub(crate) id: u64,
+    pub(crate) expression: Arc<dyn PhysicalExpr>,
+    pub(crate) input_schema: SchemaRef,
+}
+
+/// Finds dynamic-filter consumers in `plan`, optionally restricting the result to `allowed_ids`.
+///
+/// Producer and consumer occurrences intentionally share expression IDs. Producer occurrences are
+/// therefore removed only from the node that reports them through
+/// [`ExecutionPlan::dynamic_expressions_produced`], rather than subtracting producer IDs from the
+/// whole plan.
+pub(crate) fn discover_dynamic_filter_consumers(
+    plan: &Arc<dyn ExecutionPlan>,
+    allowed_ids: Option<&HashSet<u64>>,
+) -> Result<Vec<DiscoveredDynamicFilter>> {
+    let mut consumers = HashMap::new();
+
+    plan.apply(|node| {
+        let produced = node.dynamic_expressions_produced();
+        let input_schema = node
+            .children()
+            .first()
+            .map(|child| child.schema())
+            .unwrap_or_else(|| node.schema());
+
+        node.apply_expressions(&mut |root| {
+            root.apply(|expression| {
+                let Some(_) = expression.downcast_ref::<DynamicFilterPhysicalExpr>() else {
+                    return Ok(TreeNodeRecursion::Continue);
+                };
+
+                let id = expression
+                    .expression_id()
+                    .expect("DynamicFilterPhysicalExpr always has an expression ID");
+                let is_producer_occurrence = produced
+                    .iter()
+                    .any(|produced| Arc::ptr_eq(produced, expression));
+                let is_allowed = allowed_ids.is_none_or(|ids| ids.contains(&id));
+                if !is_producer_occurrence && is_allowed {
+                    consumers
+                        .entry(id)
+                        .or_insert_with(|| DiscoveredDynamicFilter {
+                            id,
+                            expression: Arc::clone(expression),
+                            input_schema: Arc::clone(&input_schema),
+                        });
+                }
+
+                Ok(TreeNodeRecursion::Continue)
+            })
+        })?;
+        Ok(TreeNodeRecursion::Continue)
+    })?;
+
+    let mut consumers: Vec<_> = consumers.into_values().collect();
+    consumers.sort_unstable_by_key(|consumer| consumer.id);
+    Ok(consumers)
+}
+
+pub(crate) fn dynamic_filter_consumer_ids(plan: &Arc<dyn ExecutionPlan>) -> Result<HashSet<u64>> {
+    Ok(discover_dynamic_filter_consumers(plan, None)?
+        .into_iter()
+        .map(|consumer| consumer.id)
+        .collect())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion::arrow::datatypes::{DataType, Field, Schema};
+    use datafusion::common::Result;
+    use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+    use datafusion::logical_expr::Operator;
+    use datafusion::physical_expr::expressions::{BinaryExpr, Column, lit};
+    use datafusion::physical_plan::empty::EmptyExec;
+    use datafusion::physical_plan::{
+        DisplayAs, DisplayFormatType, PlanProperties, apply_expression_roots,
+    };
+    use std::fmt::Formatter;
+
+    #[tokio::test]
+    async fn discovers_nested_consumer_but_not_its_producer_occurrence() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+        let input = Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc<dyn ExecutionPlan>;
+        let column = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
+        let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(
+            vec![Arc::clone(&column)],
+            lit(true),
+        )) as Arc<dyn PhysicalExpr>;
+        let nested = Arc::new(BinaryExpr::new(
+            Arc::clone(&dynamic_filter),
+            Operator::And,
+            lit(true),
+        )) as Arc<dyn PhysicalExpr>;
+
+        let consumer =
+            Arc::new(ExpressionExec::new(input, nested, false)) as Arc<dyn ExecutionPlan>;
+        let plan = Arc::new(ExpressionExec::new(
+            consumer,
+            Arc::clone(&dynamic_filter),
+            true,
+        )) as Arc<dyn ExecutionPlan>;
+
+        let discovered = discover_dynamic_filter_consumers(&plan, None)?;
+        assert_eq!(discovered.len(), 1);
+        assert_eq!(discovered[0].id, dynamic_filter.expression_id().unwrap());
+
+        dynamic_filter
+            .downcast_ref::<DynamicFilterPhysicalExpr>()
+            .unwrap()
+            .update(Arc::new(BinaryExpr::new(column, Operator::Gt, lit(10_i32))))?;
+        dynamic_filter
+            .downcast_ref::<DynamicFilterPhysicalExpr>()
+            .unwrap()
+            .mark_complete();
+
+        let current = discovered[0]
+            .expression
+            .downcast_ref::<DynamicFilterPhysicalExpr>()
+            .unwrap()
+            .current()?;
+        assert_eq!(current.to_string(), "a@0 > 10");
+        Ok(())
+    }
+
+    #[derive(Debug)]
+    struct ExpressionExec {
+        input: Arc<dyn ExecutionPlan>,
+        expression: Arc<dyn PhysicalExpr>,
+        produces_expression: bool,
+    }
+
+    impl ExpressionExec {
+        fn new(
+            input: Arc<dyn ExecutionPlan>,
+            expression: Arc<dyn PhysicalExpr>,
+            produces_expression: bool,
+        ) -> Self {
+            Self {
+                input,
+                expression,
+                produces_expression,
+            }
+        }
+    }
+
+    impl DisplayAs for ExpressionExec {
+        fn fmt_as(&self, _: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+            write!(f, "ExpressionExec")
+        }
+    }
+
+    impl ExecutionPlan for ExpressionExec {
+        fn name(&self) -> &str {
+            "ExpressionExec"
+        }
+
+        fn properties(&self) -> &Arc<PlanProperties> {
+            self.input.properties()
+        }
+
+        fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+            vec![&self.input]
+        }
+
+        fn dynamic_expressions_produced(&self) -> Vec<Arc<dyn PhysicalExpr>> {
+            self.produces_expression
+                .then(|| Arc::clone(&self.expression))
+                .into_iter()
+                .collect()
+        }
+
+        fn apply_expressions(
+            &self,
+            f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+        ) -> Result<TreeNodeRecursion> {
+            apply_expression_roots([&self.expression], f)
+        }
+
+        fn with_new_children(
+            self: Arc<Self>,
+            mut children: Vec<Arc<dyn ExecutionPlan>>,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            Ok(Arc::new(Self::new(
+                children.remove(0),
+                Arc::clone(&self.expression),
+                self.produces_expression,
+            )))
+        }
+
+        fn execute(
+            &self,
+            partition: usize,
+            context: Arc<TaskContext>,
+        ) -> Result<SendableRecordBatchStream> {
+            self.input.execute(partition, context)
+        }
+    }
+}

src/common/mod.rs

@@ -1,4 +1,5 @@
 mod children_helpers;
+mod dynamic_filtering;
 mod once_lock;
 mod recursion;
 mod task_context_helpers;
@@ -7,6 +8,9 @@ mod uuid;
 mod vec;
 
 pub(crate) use children_helpers::require_one_child;
+pub(crate) use dynamic_filtering::{
+    discover_dynamic_filter_consumers, dynamic_filter_consumer_ids,
+};
 pub(crate) use once_lock::OnceLockResult;
 pub(crate) use recursion::TreeNodeExt;
 pub(crate) use task_context_helpers::task_ctx_with_extension;

src/coordinator/distributed.rs

@@ -1,12 +1,12 @@
+use crate::DistributedConfig;
 use crate::common::require_one_child;
-use crate::coordinator::metrics_store::MetricsStore;
+use crate::coordinator::dynamic_filters::isolate_distributed_leaf_variants_for_display;
 use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan;
 use crate::coordinator::prepare_static_plan::prepare_static_plan;
 use crate::coordinator::query_coordinator::QueryCoordinator;
-use crate::distributed_planner::NetworkBoundaryExt;
-use crate::{DistributedConfig, TaskKey};
+use crate::coordinator::store::{MetricsStore, task_keys_for_plan};
 use datafusion::common::internal_datafusion_err;
-use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
+use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::common::{Result, exec_err};
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_expr::PhysicalExpr;
@@ -82,33 +82,13 @@ impl DistributedExec {
     ///
     /// [`rewrite_distributed_plan_with_metrics`]: crate::rewrite_distributed_plan_with_metrics
     pub async fn wait_for_metrics(&self) {
-        let mut expected_keys: Vec<TaskKey> = Vec::new();
         let Some(task_metrics) = &self.metrics_store else {
             return;
         };
         let Some(plan) = self.plan_for_viz.lock().unwrap().as_ref().cloned() else {
             return;
         };
-        let _ = plan.apply(|plan| {
-            if let Some(boundary) = plan.as_network_boundary() {
-                let stage = boundary.input_stage();
-                for i in 0..stage.task_count() {
-                    expected_keys.push(TaskKey {
-                        query_id: stage.query_id(),
-                        stage_id: stage.num(),
-                        task_number: i,
-                    });
-                }
-            }
-            Ok(TreeNodeRecursion::Continue)
-        });
-        if expected_keys.is_empty() {
-            return;
-        }
-        let mut rx = task_metrics.rx.clone();
-        let _ = rx
-            .wait_for(|map| expected_keys.iter().all(|key| map.contains_key(key)))
-            .await;
+        task_metrics.wait_for(&task_keys_for_plan(&plan)).await;
     }
 
     /// Returns the plan which is lazily prepared on `execute()` and actually gets executed.
@@ -124,6 +104,14 @@ impl DistributedExec {
             })
     }
 
+    pub(crate) fn plan_for_display(&self) -> Arc<dyn ExecutionPlan> {
+        self.plan_for_viz
+            .lock()
+            .ok()
+            .and_then(|plan| plan.clone())
+            .unwrap_or_else(|| Arc::clone(&self.base_plan))
+    }
+
     /// Returns the head stage that was actually executed. Unlike [`Self::plan_for_viz`] (which is
     /// reconstructed for visualization, with `Stage::Local` boundaries and rebuilt ancestor
     /// `Arc`s), this returns the original `Arc` instances whose metrics were populated during
@@ -223,10 +211,12 @@ impl ExecutionPlan for DistributedExec {
                 false => prepare_static_plan(&query_coordinator, &base_plan)?,
             };
 
+            let display_plan =
+                isolate_distributed_leaf_variants_for_display(result.plan_for_viz, &context)?;
             plan_for_viz
                 .lock()
                 .expect("poisoned lock")
-                .replace(result.plan_for_viz);
+                .replace(Arc::clone(&display_plan));
             head_stage
                 .lock()
                 .expect("poisoned lock")
@@ -237,8 +227,11 @@ impl ExecutionPlan for DistributedExec {
                     break; // channel closed
                 }
             }
-            drop(tx);
             drop(guard);
+            query_coordinator
+                .finish_dynamic_filter_display(&display_plan)
+                .await;
+            drop(tx);
             query_coordinator.drain_pending_tasks().await?;
             Ok(())
         });