feat: generalize dynamic state injection via with_new_state on ExecutionPlan

Introduce a generic `with_new_state` method to the `ExecutionPlan` trait,
superseding the specialized `with_work_table`.  This provides a uniform
mechanism for execution nodes to accept late-bound, run-time state of any
type, improving extensibility for custom operators.

Motivation:
`RecursiveQueryExec` and other potential operators need to retrofit
execution plans with run-time state (e.g. the working table for recursive
CTEs).  The previous `with_work_table` API was hard-coded to one concrete
type, preventing wrapper or third-party nodes from re-using the same
mechanism for different kinds of state.

Changes:
- Replace `with_work_table` with `with_new_state(&self, state: Arc<dyn Any + Send + Sync>)`
  on `ExecutionPlan`; default implementation returns `None`.
- Update rustdoc on the trait to describe the generic contract and thread-
  safety (`Send + Sync`) requirements.
- Implement `with_new_state` in `WorkTableExec`, down-casting the supplied
  state to `Arc<WorkTable>` and returning an updated plan when applicable.
- Refactor `RecursiveQueryExec::assign_work_table` to use the new trait
  method, passing the work table as generic state.
- Remove the now-unused re-export of `WorkTable` from `execution_plan.rs`.
- Refresh documentation throughout to emphasise that the mechanism is
  generic, using `WorkTableExec` solely as an illustrative example.

This refactor future-proofs DataFusion’s execution API, enabling any
custom node to inject and propagate arbitrary shared state without
introducing new trait methods each time.

Author: geoffreyclaude

datafusion/physical-plan/src/execution_plan.rs

@@ -23,7 +23,6 @@ use crate::filter_pushdown::{
 pub use crate::metrics::Metric;
 pub use crate::ordering::InputOrderMode;
 pub use crate::stream::EmptyRecordBatchStream;
-pub use crate::work_table::WorkTable;
 
 pub use datafusion_common::hash_utils;
 pub use datafusion_common::utils::project_schema;
@@ -572,20 +571,25 @@ pub trait ExecutionPlan: Debug + DisplayAs + Send + Sync {
         ))
     }
 
-    /// Returns a new execution plan that uses the provided work table, if supported.
-    /// This enables recursive query execution by allowing work table injection.
+    /// Injects arbitrary run-time state into this execution plan, returning a new plan
+    /// instance that incorporates that state *if* it is relevant to the concrete
+    /// node implementation.
     ///
-    /// Primarily implemented by [`WorkTableExec`], but custom execution nodes that wrap
-    /// or contain `WorkTableExec` instances should also implement this to propagate
-    /// work table injection to their inner components.
+    /// This is a generic entry point: the `state` can be any type wrapped in
+    /// `Arc<dyn Any + Send + Sync>`.  A node that cares about the state should
+    /// down-cast it to the concrete type it expects and, if successful, return a
+    /// modified copy of itself that captures the provided value.  If the state is
+    /// not applicable, the default behaviour is to return `None` so that parent
+    /// nodes can continue propagating the attempt further down the plan tree.
     ///
-    /// See [`WorkTableExec::with_work_table`] for the reference implementation.
-    ///
-    /// [`WorkTableExec`]: crate::work_table::WorkTableExec
-    /// [`WorkTableExec::with_work_table`]: crate::work_table::WorkTableExec::with_work_table
-    fn with_work_table(
+    /// For example, [`WorkTableExec`](crate::work_table::WorkTableExec)
+    /// down-casts the supplied state to an `Arc<WorkTable>`
+    /// in order to wire up the working table used during recursive-CTE execution.
+    /// Similar patterns can be followed by custom nodes that need late-bound
+    /// dependencies or shared state.
+    fn with_new_state(
         &self,
-        _work_table: Arc<WorkTable>,
+        _state: Arc<dyn Any + Send + Sync>,
     ) -> Option<Arc<dyn ExecutionPlan>> {
         None
     }

datafusion/physical-plan/src/lib.rs

@@ -50,7 +50,6 @@ pub use crate::ordering::InputOrderMode;
 pub use crate::stream::EmptyRecordBatchStream;
 pub use crate::topk::TopK;
 pub use crate::visitor::{accept, visit_execution_plan, ExecutionPlanVisitor};
-pub use crate::work_table::WorkTable;
 pub use spill::spill_manager::SpillManager;
 
 mod ordering;

datafusion/physical-plan/src/recursive_query.rs

@@ -351,7 +351,9 @@ fn assign_work_table(
 ) -> Result<Arc<dyn ExecutionPlan>> {
     let mut work_table_refs = 0;
     plan.transform_down(|plan| {
-        if let Some(new_plan) = plan.with_work_table(Arc::clone(&work_table)) {
+        if let Some(new_plan) =
+            plan.with_new_state(Arc::clone(&work_table) as Arc<dyn Any + Send + Sync>)
+        {
             if work_table_refs > 0 {
                 not_impl_err!(
                     "Multiple recursive references to the same CTE are not supported"

datafusion/physical-plan/src/work_table.rs

@@ -57,7 +57,7 @@ impl ReservedBatches {
 /// See <https://wiki.postgresql.org/wiki/CTEReadme#How_Recursion_Works>
 /// This table serves as a mirror or buffer between each iteration of a recursive query.
 #[derive(Debug)]
-pub struct WorkTable {
+pub(super) struct WorkTable {
     batches: Mutex<Option<ReservedBatches>>,
 }
 
@@ -225,20 +225,28 @@ impl ExecutionPlan for WorkTableExec {
         Ok(Statistics::new_unknown(&self.schema()))
     }
 
-    /// Creates a new `WorkTableExec` with the provided work table for recursive query execution.
-    /// During query planning, `WorkTableExec` nodes are created as placeholders; this method
-    /// "wires up" the actual work table that coordinates data between recursive iterations.
-    fn with_work_table(
+    /// Injects run-time state into this `WorkTableExec`.
+    ///
+    /// The only state this node currently understands is an [`Arc<WorkTable>`].
+    /// If `state` can be down-cast to that type, a new `WorkTableExec` backed
+    /// by the provided work table is returned.  Otherwise `None` is returned
+    /// so that callers can attempt to propagate the state further down the
+    /// execution plan tree.
+    fn with_new_state(
         &self,
-        work_table: Arc<WorkTable>,
+        state: Arc<dyn Any + Send + Sync>,
     ) -> Option<Arc<dyn ExecutionPlan>> {
-        Some(Arc::new(Self {
-            name: self.name.clone(),
-            schema: Arc::clone(&self.schema),
-            metrics: ExecutionPlanMetricsSet::new(),
-            work_table,
-            cache: self.cache.clone(),
-        }))
+        // Attempt to down-cast the supplied state to a WorkTable
+        match state.downcast::<WorkTable>() {
+            Ok(work_table) => Some(Arc::new(Self {
+                name: self.name.clone(),
+                schema: Arc::clone(&self.schema),
+                metrics: ExecutionPlanMetricsSet::new(),
+                work_table,
+                cache: self.cache.clone(),
+            })),
+            Err(_) => None,
+        }
     }
 }