@@ -22,7 +22,10 @@ use crate::{
2222use datafusion_common:: { Result , internal_datafusion_err, plan_datafusion_err} ;
2323use datafusion_expr:: planner:: ExprPlanner ;
2424use datafusion_expr:: { AggregateUDF , HigherOrderUDF , ScalarUDF , WindowUDF } ;
25+ use parking_lot:: Mutex ;
26+ use std:: any:: { Any , TypeId } ;
2527use std:: collections:: HashSet ;
28+ use std:: fmt;
2629use std:: { collections:: HashMap , sync:: Arc } ;
2730
2831/// Task Execution Context
@@ -48,7 +51,7 @@ use std::{collections::HashMap, sync::Arc};
4851///
4952/// [`SessionContext`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html
5053/// [`ExecutionProps`]: datafusion_expr::execution_props::ExecutionProps
51- #[ derive( Debug ) ]
54+ #[ derive( Debug , Clone ) ]
5255pub struct TaskContext {
5356 /// Session Id
5457 session_id : String ,
@@ -66,6 +69,71 @@ pub struct TaskContext {
6669 window_functions : HashMap < String , Arc < WindowUDF > > ,
6770 /// Runtime environment associated with this task context
6871 runtime : Arc < RuntimeEnv > ,
72+ /// State that operators keep for one run of a plan
73+ run_state : Arc < RunState > ,
74+ }
75+
76+ /// State that the operators of a plan keep for one run of that plan.
77+ ///
78+ /// An operator must not keep the result of a run in its plan node when another
79+ /// node reads that result, because the plan can be executed again, or two
80+ /// times at the same time. Such an operator keeps the result here. All
81+ /// `execute` calls of one run get the same [`TaskContext`], so they share one
82+ /// [`RunState`]. A different [`TaskContext`] is a different run.
83+ ///
84+ /// A [`RunState`] holds at most one value of each type. An operator defines
85+ /// its own type, for example a map from an operator id to a result.
86+ ///
87+ /// An operator that executes a part of the plan more than one time in a run
88+ /// (such as the recursive term of a recursive query) can give each execution a
89+ /// child state. See [`RunState::new_child`].
90+ #[ derive( Default ) ]
91+ pub struct RunState {
92+ /// Types that are not in `values` are read from the parent.
93+ parent : Option < Arc < RunState > > ,
94+ values : Mutex < HashMap < TypeId , Arc < dyn Any + Send + Sync > > > ,
95+ }
96+
97+ impl fmt:: Debug for RunState {
98+ fn fmt ( & self , f : & mut fmt:: Formatter < ' _ > ) -> fmt:: Result {
99+ f. debug_struct ( "RunState" )
100+ . field ( "parent" , & self . parent )
101+ . finish_non_exhaustive ( )
102+ }
103+ }
104+
105+ impl RunState {
106+ /// Create an empty state.
107+ pub fn new ( ) -> Self {
108+ Self :: default ( )
109+ }
110+
111+ /// Create a state in which `value` replaces the value of type `T` of
112+ /// `parent`. The values of all other types are those of `parent`.
113+ pub fn new_child < T : Send + Sync + ' static > (
114+ parent : Arc < RunState > ,
115+ value : Arc < T > ,
116+ ) -> Self {
117+ let value = value as Arc < dyn Any + Send + Sync > ;
118+ Self {
119+ parent : Some ( parent) ,
120+ values : Mutex :: new ( HashMap :: from ( [ ( TypeId :: of :: < T > ( ) , value) ] ) ) ,
121+ }
122+ }
123+
124+ /// Return the value of type `T`. The first call creates it.
125+ pub fn get_or_init < T : Default + Send + Sync + ' static > ( & self ) -> Arc < T > {
126+ let mut values = self . values . lock ( ) ;
127+ let id = TypeId :: of :: < T > ( ) ;
128+ let value = match ( values. get ( & id) , & self . parent ) {
129+ ( Some ( value) , _) => Arc :: clone ( value) ,
130+ ( None , Some ( parent) ) => return parent. get_or_init ( ) ,
131+ ( None , None ) => {
132+ Arc :: clone ( values. entry ( id) . or_insert_with ( || Arc :: new ( T :: default ( ) ) ) )
133+ }
134+ } ;
135+ value. downcast ( ) . expect ( "the value has the type of its key" )
136+ }
69137}
70138
71139impl Default for TaskContext {
@@ -82,6 +150,7 @@ impl Default for TaskContext {
82150 aggregate_functions : HashMap :: new ( ) ,
83151 window_functions : HashMap :: new ( ) ,
84152 runtime,
153+ run_state : Arc :: default ( ) ,
85154 }
86155 }
87156}
@@ -112,6 +181,7 @@ impl TaskContext {
112181 aggregate_functions,
113182 window_functions,
114183 runtime,
184+ run_state : Arc :: default ( ) ,
115185 }
116186 }
117187
@@ -168,6 +238,17 @@ impl TaskContext {
168238 self
169239 }
170240
241+ /// Return the [`RunState`] of the run that this [`TaskContext`] is for
242+ pub fn run_state ( & self ) -> & Arc < RunState > {
243+ & self . run_state
244+ }
245+
246+ /// Update the [`RunState`]
247+ pub fn with_run_state ( mut self , run_state : Arc < RunState > ) -> Self {
248+ self . run_state = run_state;
249+ self
250+ }
251+
171252 /// Update the `task_id`
172253 pub fn with_task_id ( mut self , task_id : String ) -> Self {
173254 self . task_id = Some ( task_id) ;
@@ -294,6 +375,28 @@ mod tests {
294375 const PREFIX : & ' static str = "test" ;
295376 }
296377
378+ #[ test]
379+ fn run_state_child_replaces_one_type ( ) {
380+ let context = TaskContext :: default ( ) ;
381+ let state = context. run_state ( ) ;
382+ let number: Arc < Mutex < usize > > = state. get_or_init ( ) ;
383+ * number. lock ( ) = 1 ;
384+ // A copy of the context is the same run.
385+ let copy = context. clone ( ) ;
386+ assert_eq ! ( * copy. run_state( ) . get_or_init:: <Mutex <usize >>( ) . lock( ) , 1 ) ;
387+
388+ let child = RunState :: new_child ( Arc :: clone ( state) , Arc :: new ( Mutex :: new ( 2_usize ) ) ) ;
389+ assert_eq ! ( * child. get_or_init:: <Mutex <usize >>( ) . lock( ) , 2 ) ;
390+ assert_eq ! ( * state. get_or_init:: <Mutex <usize >>( ) . lock( ) , 1 ) ;
391+ // The child reads all other types from the parent.
392+ * child. get_or_init :: < Mutex < String > > ( ) . lock ( ) = "a" . to_string ( ) ;
393+ assert_eq ! ( * state. get_or_init:: <Mutex <String >>( ) . lock( ) , "a" ) ;
394+
395+ // A new context is a new run.
396+ let other = TaskContext :: default ( ) ;
397+ assert_eq ! ( * other. run_state( ) . get_or_init:: <Mutex <usize >>( ) . lock( ) , 0 ) ;
398+ }
399+
297400 #[ test]
298401 fn task_context_extensions ( ) -> Result < ( ) > {
299402 let runtime = Arc :: new ( RuntimeEnv :: default ( ) ) ;
0 commit comments