Skip to content

Commit 91d3f3d

Browse files
adriangbclaude
andcommitted
fix: keep the output of a materialized CTE per run, not in the plan
`MaterializedCteExec::reset_state` cleared the shared buffer in place and returned `self`. This did not follow the `ExecutionPlan::reset_state` contract, and two overlapping runs of one plan shared one buffer. The buffered output is now in a new `RunState` of the `TaskContext`, keyed by the CTE id. Each run of a plan has its own `TaskContext`, so each run has its own output. The plan nodes keep no state of a run, so the default `reset_state` is correct for them. `RecursiveQueryExec` gives each iteration a child `RunState` in which the CTEs that the recursive term declares have no output. A CTE that is declared outside the recursive term keeps its output for the whole query. Closes #26229. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 1a43b6c commit 91d3f3d

4 files changed

Lines changed: 411 additions & 36 deletions

File tree

‎datafusion/execution/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,4 +51,4 @@ pub use disk_manager::DiskManager;
5151
pub use registry::FunctionRegistry;
5252
pub use spill_file::{AsyncSpillWriter, SpillFile, SpillWriter, TempFileFactory};
5353
pub use stream::{RecordBatchStream, SendableRecordBatchStream};
54-
pub use task::{TaskContext, TaskContextProvider};
54+
pub use task::{RunState, TaskContext, TaskContextProvider};

‎datafusion/execution/src/task.rs‎

Lines changed: 104 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,10 @@ use crate::{
2222
use datafusion_common::{Result, internal_datafusion_err, plan_datafusion_err};
2323
use datafusion_expr::planner::ExprPlanner;
2424
use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF};
25+
use parking_lot::Mutex;
26+
use std::any::{Any, TypeId};
2527
use std::collections::HashSet;
28+
use std::fmt;
2629
use 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)]
5255
pub 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

71139
impl 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

Comments
 (0)