physical plan: move shared resources from exec nodes

To reuse a physical plan we need to move all shared resources from
them.

Some execs contain shared futures to initialize them lazy
once across all partitions. For example, `HashJoinExec` contain
a future that once collects all data from the build sie.

This patch moves such resources to the task context.
Now each plan can save to the task context its own state to reuse across partitions.

Author: askalt

datafusion/core/src/datasource/physical_plan/parquet/mod.rs

@@ -709,7 +709,6 @@ impl ExecutionPlan for ParquetExec {
             .clone()
             .unwrap_or_else(|| Arc::new(DefaultSchemaAdapterFactory::default()));
 
-        println!("parquet exec registered metrics...");
         let metrics = ctx.get_or_register_metric_set_with_default(self, || {
             ExecutionPlanMetricsSet::with_inner(self.base_metrics.clone_inner())
         });

datafusion/execution/src/lib.rs

@@ -39,4 +39,4 @@ pub use disk_manager::DiskManager;
 pub use metrics::Metric;
 pub use registry::FunctionRegistry;
 pub use stream::{RecordBatchStream, SendableRecordBatchStream};
-pub use task::TaskContext;
+pub use task::{PlanState, TaskContext};

datafusion/execution/src/task.rs

@@ -18,6 +18,7 @@
 use std::{
     any::Any,
     collections::{HashMap, HashSet},
+    fmt::Debug,
     sync::Arc,
 };
 
@@ -60,6 +61,9 @@ pub struct TaskContext {
     /// Metrics associated with a execution plan address.
     /// std mutex is used because too concurrent access is not assumed.
     metrics: std::sync::Mutex<HashMap<usize, ExecutionPlanMetricsSet>>,
+    /// Session plans state by an execution plan address.
+    /// Resources that shared across execution partitions.
+    plan_state: std::sync::Mutex<HashMap<usize, Arc<dyn PlanState>>>,
 }
 
 impl Default for TaskContext {
@@ -79,10 +83,16 @@ impl Default for TaskContext {
             runtime,
             param_values: None,
             metrics: Default::default(),
+            plan_state: Default::default(),
         }
     }
 }
 
+/// Generic plan state.
+pub trait PlanState: Debug + Any + Send + Sync {
+    fn as_any(&self) -> &dyn Any;
+}
+
 fn plan_addr(plan: &dyn Any) -> usize {
     plan as *const _ as *const () as usize
 }
@@ -112,6 +122,30 @@ impl TaskContext {
             runtime,
             param_values: None,
             metrics: Default::default(),
+            plan_state: Default::default(),
+        }
+    }
+
+    /// Fork a task context.
+    ///
+    /// Forked context contains the same
+    /// * session related attributes (id, udfs, etc),
+    /// * runtime environment
+    ///
+    /// But an empty metrics and plan state.
+    ///
+    pub fn fork(&self) -> Self {
+        Self {
+            task_id: self.task_id(),
+            session_id: self.session_id(),
+            session_config: self.session_config.clone(),
+            scalar_functions: self.scalar_functions.clone(),
+            aggregate_functions: self.aggregate_functions.clone(),
+            window_functions: self.window_functions.clone(),
+            runtime: Arc::clone(&self.runtime),
+            param_values: self.param_values.clone(),
+            metrics: Default::default(),
+            plan_state: Default::default(),
         }
     }
 
@@ -206,6 +240,26 @@ impl TaskContext {
             metric_set
         }
     }
+
+    /// Get state for specific plan or register a new state.
+    pub fn get_or_register_plan_state<F>(
+        &self,
+        plan: &dyn Any,
+        f: F,
+    ) -> Arc<dyn PlanState>
+    where
+        F: FnOnce() -> Arc<dyn PlanState>,
+    {
+        let addr = plan_addr(plan);
+        let mut plan_state = self.plan_state.lock().unwrap();
+        if let Some(state) = plan_state.get(&addr) {
+            Arc::clone(state)
+        } else {
+            let state = f();
+            plan_state.insert(addr, Arc::clone(&state));
+            state
+        }
+    }
 }
 
 impl FunctionRegistry for TaskContext {

datafusion/physical-plan/src/joins/cross_join.rs

@@ -38,7 +38,7 @@ use arrow_array::RecordBatchOptions;
 use datafusion_common::stats::Precision;
 use datafusion_common::{internal_err, JoinType, Result, ScalarValue};
 use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
-use datafusion_execution::TaskContext;
+use datafusion_execution::{PlanState, TaskContext};
 use datafusion_physical_expr::equivalence::join_equivalence_properties;
 
 use async_trait::async_trait;
@@ -57,9 +57,28 @@ pub struct CrossJoinExec {
     pub right: Arc<dyn ExecutionPlan>,
     /// The schema once the join is applied
     schema: SchemaRef,
+    cache: PlanProperties,
+}
+
+/// Exec state shared across partitions per one execution.
+#[derive(Debug)]
+struct CrossJoinExecState {
     /// Build-side data
     left_fut: OnceAsync<JoinLeftData>,
-    cache: PlanProperties,
+}
+
+impl CrossJoinExecState {
+    fn new() -> Self {
+        Self {
+            left_fut: Default::default(),
+        }
+    }
+}
+
+impl PlanState for CrossJoinExecState {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
 }
 
 impl CrossJoinExec {
@@ -87,7 +106,6 @@ impl CrossJoinExec {
             left,
             right,
             schema,
-            left_fut: Default::default(),
             cache,
         }
     }
@@ -239,14 +257,21 @@ impl ExecutionPlan for CrossJoinExec {
         let reservation =
             MemoryConsumer::new("CrossJoinExec").register(context.memory_pool());
 
-        let left_fut = self.left_fut.once(|| {
-            load_left_input(
-                Arc::clone(&self.left),
-                Arc::clone(&context),
-                join_metrics.clone(),
-                reservation,
-            )
-        });
+        let state = context
+            .get_or_register_plan_state(self, || Arc::new(CrossJoinExecState::new()));
+        let left_fut = state
+            .as_any()
+            .downcast_ref::<CrossJoinExecState>()
+            .unwrap()
+            .left_fut
+            .once(|| {
+                load_left_input(
+                    Arc::clone(&self.left),
+                    Arc::clone(&context),
+                    join_metrics.clone(),
+                    reservation,
+                )
+            });
 
         Ok(Box::pin(CrossJoinStream {
             schema: Arc::clone(&self.schema),

datafusion/physical-plan/src/joins/hash_join.rs

@@ -63,7 +63,7 @@ use datafusion_common::{
     JoinSide, JoinType, Result,
 };
 use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
-use datafusion_execution::TaskContext;
+use datafusion_execution::{PlanState, TaskContext};
 use datafusion_physical_expr::equivalence::{
     join_equivalence_properties, ProjectionMapping,
 };
@@ -306,8 +306,6 @@ pub struct HashJoinExec {
     /// The schema after join. Please be careful when using this schema,
     /// if there is a projection, the schema isn't the same as the output schema.
     join_schema: SchemaRef,
-    /// Future that consumes left input and builds the hash table
-    left_fut: OnceAsync<JoinLeftData>,
     /// Shared the `RandomState` for the hashing algorithm
     random_state: RandomState,
     /// Partitioning mode to use
@@ -325,6 +323,27 @@ pub struct HashJoinExec {
     cache: PlanProperties,
 }
 
+/// Exec state shared across partitions per one execution.
+#[derive(Debug)]
+struct HashJoinExecState {
+    /// Future that consumes left input and builds the hash table
+    left_fut: OnceAsync<JoinLeftData>,
+}
+
+impl HashJoinExecState {
+    fn new() -> Self {
+        Self {
+            left_fut: Default::default(),
+        }
+    }
+}
+
+impl PlanState for HashJoinExecState {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}
+
 impl HashJoinExec {
     /// Tries to create a new [HashJoinExec].
     ///
@@ -376,7 +395,6 @@ impl HashJoinExec {
             filter,
             join_type: *join_type,
             join_schema,
-            left_fut: Default::default(),
             random_state,
             mode: partition_mode,
             projection,
@@ -689,21 +707,32 @@ impl ExecutionPlan for HashJoinExec {
         let metrics = context.get_or_register_metric_set(self);
         let join_metrics = BuildProbeJoinMetrics::new(partition, &metrics);
         let left_fut = match self.mode {
-            PartitionMode::CollectLeft => self.left_fut.once(|| {
-                let reservation =
-                    MemoryConsumer::new("HashJoinInput").register(context.memory_pool());
-                collect_left_input(
-                    None,
-                    self.random_state.clone(),
-                    Arc::clone(&self.left),
-                    on_left.clone(),
-                    Arc::clone(&context),
-                    join_metrics.clone(),
-                    reservation,
-                    need_produce_result_in_final(self.join_type),
-                    self.right().output_partitioning().partition_count(),
-                )
-            }),
+            PartitionMode::CollectLeft => {
+                let state = context.get_or_register_plan_state(self, || {
+                    Arc::new(HashJoinExecState::new())
+                });
+
+                state
+                    .as_any()
+                    .downcast_ref::<HashJoinExecState>()
+                    .unwrap()
+                    .left_fut
+                    .once(|| {
+                        let reservation = MemoryConsumer::new("HashJoinInput")
+                            .register(context.memory_pool());
+                        collect_left_input(
+                            None,
+                            self.random_state.clone(),
+                            Arc::clone(&self.left),
+                            on_left.clone(),
+                            Arc::clone(&context),
+                            join_metrics.clone(),
+                            reservation,
+                            need_produce_result_in_final(self.join_type),
+                            self.right().output_partitioning().partition_count(),
+                        )
+                    })
+            }
             PartitionMode::Partitioned => {
                 let reservation =
                     MemoryConsumer::new(format!("HashJoinInput[{partition}]"))
@@ -3427,7 +3456,6 @@ mod tests {
     /// Test for parallelised HashJoinExec with PartitionMode::CollectLeft
     #[tokio::test]
     async fn test_collect_left_multiple_partitions_join() -> Result<()> {
-        let task_ctx = Arc::new(TaskContext::default());
         let left = build_table(
             ("a1", &vec![1, 2, 3]),
             ("b1", &vec![4, 5, 7]),
@@ -3522,6 +3550,7 @@ mod tests {
         ];
 
         for (join_type, expected) in test_cases {
+            let task_ctx = Arc::new(TaskContext::default());
             let (_, batches) = join_collect_with_partition_mode(
                 Arc::clone(&left),
                 Arc::clone(&right),

datafusion/physical-plan/src/joins/nested_loop_join.rs

@@ -47,7 +47,7 @@ use arrow::util::bit_util;
 use arrow_array::PrimitiveArray;
 use datafusion_common::{exec_datafusion_err, JoinSide, Result, Statistics};
 use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
-use datafusion_execution::TaskContext;
+use datafusion_execution::{PlanState, TaskContext};
 use datafusion_expr::JoinType;
 use datafusion_physical_expr::equivalence::join_equivalence_properties;
 
@@ -147,14 +147,33 @@ pub struct NestedLoopJoinExec {
     pub(crate) join_type: JoinType,
     /// The schema once the join is applied
     schema: SchemaRef,
-    /// Build-side data
-    inner_table: OnceAsync<JoinLeftData>,
     /// Information of index and left / right placement of columns
     column_indices: Vec<ColumnIndex>,
     /// Cache holding plan properties like equivalences, output partitioning etc.
     cache: PlanProperties,
 }
 
+/// Exec state shared across partitions per one execution.
+#[derive(Debug)]
+struct NestedLoopJoinExecState {
+    /// Build-side data.
+    inner_table: OnceAsync<JoinLeftData>,
+}
+
+impl NestedLoopJoinExecState {
+    fn new() -> Self {
+        Self {
+            inner_table: Default::default(),
+        }
+    }
+}
+
+impl PlanState for NestedLoopJoinExecState {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}
+
 impl NestedLoopJoinExec {
     /// Try to create a new [`NestedLoopJoinExec`]
     pub fn try_new(
@@ -178,7 +197,6 @@ impl NestedLoopJoinExec {
             filter,
             join_type: *join_type,
             schema,
-            inner_table: Default::default(),
             column_indices,
             cache,
         })
@@ -303,17 +321,26 @@ impl ExecutionPlan for NestedLoopJoinExec {
             MemoryConsumer::new(format!("NestedLoopJoinLoad[{partition}]"))
                 .register(context.memory_pool());
 
-        let inner_table = self.inner_table.once(|| {
-            collect_left_input(
-                Arc::clone(&self.left),
-                Arc::clone(&context),
-                join_metrics.clone(),
-                load_reservation,
-                need_produce_result_in_final(self.join_type),
-                self.right().output_partitioning().partition_count(),
-            )
+        let state = context.get_or_register_plan_state(self, || {
+            Arc::new(NestedLoopJoinExecState::new())
         });
 
+        let inner_table = state
+            .as_any()
+            .downcast_ref::<NestedLoopJoinExecState>()
+            .unwrap()
+            .inner_table
+            .once(|| {
+                collect_left_input(
+                    Arc::clone(&self.left),
+                    Arc::clone(&context),
+                    join_metrics.clone(),
+                    load_reservation,
+                    need_produce_result_in_final(self.join_type),
+                    self.right().output_partitioning().partition_count(),
+                )
+            });
+
         // Resolve placeholders in filter.
         let resolved_filter = if let Some(ref filter) = self.filter {
             Some(filter.resolve_placeholders(context.param_values())?)

datafusion/physical-plan/src/joins/symmetric_hash_join.rs

@@ -1642,6 +1642,7 @@ mod tests {
             Arc::clone(&task_ctx),
         )
         .await?;
+        let task_ctx = Arc::new(task_ctx.fork());
         let second_batches = partitioned_hash_join_with_filter(
             left, right, on, filter, &join_type, false, task_ctx,
         )