coordinator: display consumer dynamic filters after execution

docs/source/user-guide/05-metrics.md

@@ -26,8 +26,11 @@ channel, so they are not lost even if the result stream is dropped early (for ex
 
 ## Rendering a plan with metrics
 
-Two functions, both exported from the crate root, do the work:
+These functions, all exported from the crate root, do the work:
 
+- `rewrite_distributed_plan_with_dynamic_filters(plan)` — folds the completed dynamic filters
+  reported by each worker task into an isolated copy of the plan. When displaying both dynamic
+  filters and metrics, apply the dynamic-filter rewrite first.
 - `rewrite_distributed_plan_with_metrics(plan, format)` — folds every task's metrics back into the
   coordinator's copy of the plan. It waits for all worker metrics to arrive, so the result is always
   complete. The `format` is a `DistributedMetricsFormat`:
@@ -54,11 +57,15 @@ execute_stream(plan.clone(), ctx.task_ctx())?
     .try_collect::<Vec<_>>()
     .await?;
 
-// 3. Fold the per-task metrics back into the plan...
+// 3. Fold the completed per-task dynamic filters back into the plan...
+let plan =
+    rewrite_distributed_plan_with_dynamic_filters(plan).await?;
+
+// 4. Fold the per-task metrics back into the plan...
 let plan =
     rewrite_distributed_plan_with_metrics(plan, DistributedMetricsFormat::Aggregated).await?;
 
-// 4. ...and render it.
+// 5. ...and render it.
 println!("{}", display_plan_ascii(plan.as_ref(), true));
 ```
 

src/codec/mod.rs

@@ -4,7 +4,8 @@ mod user_codec;
 
 pub use distributed_codec::DistributedCodec;
 pub(crate) use physical_plan::{
-    decode_execution_plan, decode_partitioning, encode_execution_plan, encode_partitioning,
+    decode_execution_plan, decode_partitioning, decode_physical_expr, encode_execution_plan,
+    encode_partitioning, encode_physical_expr,
 };
 pub(crate) use user_codec::{
     get_distributed_user_codecs, set_distributed_user_codec, set_distributed_user_codec_arc,

src/codec/physical_plan.rs

@@ -1,15 +1,17 @@
 use super::DistributedCodec;
-use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::arrow::datatypes::{Schema, SchemaRef};
 use datafusion::common::Result;
 use datafusion::execution::TaskContext;
-use datafusion::physical_expr::Partitioning;
+use datafusion::physical_expr::{Partitioning, PhysicalExpr};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion_proto::bytes::{
     physical_plan_from_bytes_with_proto_converter, physical_plan_to_bytes_with_proto_converter,
 };
 use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning;
 use datafusion_proto::physical_plan::to_proto::serialize_partitioning;
-use datafusion_proto::physical_plan::{DeduplicatingProtoConverter, PhysicalPlanDecodeContext};
+use datafusion_proto::physical_plan::{
+    DeduplicatingProtoConverter, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension,
+};
 use datafusion_proto::protobuf;
 use datafusion_proto::protobuf::proto_error;
 use prost::Message;
@@ -42,6 +44,26 @@ pub(crate) fn decode_execution_plan(
     physical_plan_from_bytes_with_proto_converter(encoded, task_ctx, &codec, &converter)
 }
 
+pub(crate) fn encode_physical_expr(
+    expression: &Arc<dyn PhysicalExpr>,
+    task_ctx: &TaskContext,
+) -> Result<protobuf::PhysicalExprNode> {
+    let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config());
+    let converter = new_proto_converter();
+    converter.physical_expr_to_proto(expression, &codec)
+}
+
+pub(crate) fn decode_physical_expr(
+    proto: &protobuf::PhysicalExprNode,
+    input_schema: &Schema,
+    task_ctx: &TaskContext,
+) -> Result<Arc<dyn PhysicalExpr>> {
+    let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config());
+    let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx, &codec);
+    let converter = new_proto_converter();
+    converter.proto_to_physical_expr(proto, input_schema, &decode_ctx)
+}
+
 pub(crate) fn encode_partitioning(
     partitioning: &Partitioning,
     task_ctx: &TaskContext,

src/common/dynamic_filtering.rs

@@ -0,0 +1,235 @@
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
+use datafusion::common::{HashMap, HashSet, Result, internal_err};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr;
+use datafusion::physical_plan::ExecutionPlan;
+use std::sync::Arc;
+
+/// A dynamic-filter consumer discovered in an execution plan along with the schema its evaluated
+/// against.
+#[derive(Clone)]
+pub(crate) struct DiscoveredDynamicFilter {
+    pub(crate) id: u64,
+    pub(crate) expression: Arc<dyn PhysicalExpr>,
+    pub(crate) input_schema: SchemaRef,
+}
+
+/// Finds dynamic-filter consumers in `plan`, deduplicated by expression ID.
+pub(crate) fn discover_dynamic_filter_consumers(
+    plan: &Arc<dyn ExecutionPlan>,
+) -> Result<Vec<DiscoveredDynamicFilter>> {
+    let mut consumers = HashMap::new();
+
+    plan.apply(|node| {
+        let produced_ids: HashSet<_> = node
+            .dynamic_expressions_produced()
+            .into_iter()
+            .map(|produced| {
+                let Some(id) = produced.expression_id() else {
+                    return internal_err!(
+                        "{}::dynamic_expressions_produced returned an expression without an expression ID",
+                        node.name()
+                    );
+                };
+                Ok(id)
+            })
+            .collect::<Result<_>>()?;
+        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 Some(id) = expression.expression_id() else {
+                    return internal_err!(
+                        "DynamicFilterPhysicalExpr did not have an expression ID"
+                    );
+                };
+                let is_producer_occurrence = produced_ids.contains(&id);
+                if !is_producer_occurrence {
+                    consumers
+                        .entry(id)
+                        .or_insert_with(|| DiscoveredDynamicFilter {
+                            id,
+                            expression: Arc::clone(expression),
+                            input_schema: Arc::clone(&input_schema),
+                        });
+                }
+
+                Ok(TreeNodeRecursion::Continue)
+            })
+        })?;
+        Ok(TreeNodeRecursion::Continue)
+    })?;
+
+    let mut consumers: Vec<_> = consumers.into_values().collect();
+    consumers.sort_unstable_by_key(|consumer| consumer.id);
+    Ok(consumers)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion::arrow::datatypes::{DataType, Field, Schema};
+    use datafusion::common::Result;
+    use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+    use datafusion::logical_expr::Operator;
+    use datafusion::physical_expr::expressions::{BinaryExpr, Column, lit};
+    use datafusion::physical_plan::empty::EmptyExec;
+    use datafusion::physical_plan::union::UnionExec;
+    use datafusion::physical_plan::{
+        DisplayAs, DisplayFormatType, PlanProperties, apply_expression_roots,
+    };
+    use std::fmt::Formatter;
+
+    #[tokio::test]
+    async fn discovers_nested_consumer_but_not_its_producer_occurrence() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+        let input = Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc<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)?;
+        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(())
+    }
+
+    #[test]
+    fn deduplicates_consumers_with_the_same_expression_id() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+        let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(
+            vec![Arc::new(Column::new("a", 0))],
+            lit(true),
+        )) as Arc<dyn PhysicalExpr>;
+        let consumers = (0..2)
+            .map(|_| {
+                Arc::new(ExpressionExec::new(
+                    Arc::new(EmptyExec::new(Arc::clone(&schema))),
+                    Arc::clone(&dynamic_filter),
+                    false,
+                )) as Arc<dyn ExecutionPlan>
+            })
+            .collect();
+        let plan = UnionExec::try_new(consumers)?;
+
+        let discovered = discover_dynamic_filter_consumers(&plan)?;
+
+        assert_eq!(discovered.len(), 1);
+        assert_eq!(discovered[0].id, dynamic_filter.expression_id().unwrap());
+        Ok(())
+    }
+
+    #[derive(Debug)]
+    struct ExpressionExec {
+        input: Arc<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 maybe_encoded;
 mod once_lock;
 mod recursion;
@@ -8,6 +9,7 @@ mod uuid;
 mod vec;
 
 pub(crate) use children_helpers::require_one_child;
+pub(crate) use dynamic_filtering::discover_dynamic_filter_consumers;
 pub use maybe_encoded::MaybeEncoded;
 pub(crate) use once_lock::OnceLockResult;
 pub(crate) use recursion::TreeNodeExt;