worker: apply merged dynamic filters during execution

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,
+    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::{
     get_distributed_user_codecs, set_distributed_user_codec, set_distributed_user_codec_arc,

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,62 @@ pub(crate) fn decode_physical_expr(
     converter.proto_to_physical_expr(proto, input_schema, &decode_ctx)
 }
 
+/// Applies a producer-space predicate through a synthetic producer sharing the consumer's state.
+///
+/// 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 occurences 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<DynamicFilterPhysicalExpr>,
+    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<dyn PhysicalExpr> = 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::<Result<Vec<_>>>()?;
+    let update_target = Arc::clone(&consumer).with_new_children(original_children)?;
+    let Ok(update_target) = Arc::downcast::<DynamicFilterPhysicalExpr>(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,

src/coordinator/distributed.rs

@@ -222,27 +222,34 @@ impl ExecutionPlan for DistributedExec {
         let prepared_plan = Arc::clone(&self.prepared_plan);
         let collect_dynamic_filters = self.completed_dynamic_filter_store.is_some();
 
-        let query_coordinator = Arc::new(QueryCoordinator::new(
+        let (query_coordinator, mut errors) = QueryCoordinator::new(
             Arc::clone(&context),
             &self.metrics,
             self.metrics_store.clone(),
             self.completed_dynamic_filter_store.clone(),
-        ));
+        );
+        let query_coordinator = Arc::new(query_coordinator);
 
         let mut builder = RecordBatchReceiverStreamBuilder::new(self.schema(), 1);
         let tx = builder.tx();
 
+        // Handle coordinator errors on the execution stream.
+        builder.spawn(async move {
+            match errors.recv().await {
+                Some(error) => Err(error),
+                None => Ok(()),
+            }
+        });
+
         builder.spawn(async move {
             // Dropping this `guard` is what signals the coordinator->worker channel to be dropped,
             // which triggers a chain reaction that ends up also gracefully closing the
             // worker->coordinator channel. The flow looks like this:
             // 1. The query ends normally, as all Arrow RecordBatches are already streamed.
-            // 2. The `guard` here is dropped.
-            // 3. In StageCoordinator::send_plan_task(), `end_stream_notifier` fires and the
-            //    coordinator->worker channel is gracefully ended.
-            // 4. The coordinator->worker channel EOS is received in `impl_coordinator_channel.rs`.
-            // 5. The metrics are send back in the worker->coordinator channel, and then that
-            //    channel is closed.
+            // 2. The `guard` here is dropped, ending the coordinator->worker stream.
+            // 3. The worker observes end-of-stream in `impl_coordinator_channel.rs`.
+            // 4. The worker sends final metrics and completed dynamic filters, if enabled.
+            // 5. The the worker->coordinator response stream ends.
             let guard = query_coordinator.end_query_guard();
 
             let d_cfg = DistributedConfig::from_config_options(context.session_config().options())?;