Stateless execution plans for plan caching
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### Is your feature request related to a problem or challenge?
This issue is a reincarnation of https://github.com/apache/datafusion/issues/14342.
We are working on a DBMS that uses DataFusion as the SQL query engine. In general, there is a challenge: physical plan building bottlenecks. The impact can be reduced to some degree, but as the number of optimizations increases, query latency becomes uncontrollable.
It would be beneficial to reuse a single plan across many executions (like a prepared statement). The problem is not so much ideological as technical: the current implementation stores execution-stage state in the plan itself and generally does not assume that `execute(...)` can be called more than once or concurrently for the same `Arc`.
```rust
fn reset_state(self: Arc) -> Result> {
...
}
fn metrics(&self) -> Option {
...
}
fn with_new_state(
&self,
_state: Arc,
) -> Option> {
...
}
```
We implemented a reusable plans feature (based on release 42) in our [fork](https://github.com/tarantool/datafusion/tree/release-42.0.0). As it showed good results, we are currently working on porting it to the upstream.
### Describe the solution you'd like
Consider the following trait, which could be used to represent a plan state:
```rust
/// Generic plan state.
pub trait PlanState: std::fmt::Debug + std::any::Any + Send + Sync {
/// Represent this state as any.
fn as_any(&self) -> &dyn std::any::Any;
/// Return plan metrics if some.
fn metrics(&self) -> Option;
}
```
Suggest to store plan state in the existing `TaskContext` struct:
```rust
/// A [`TaskContext`] contains the state required during a single query's
/// execution. Please see the documentation on [`SessionContext`] for more
/// information.
```
`TaskContext` could store a mapping from the execution plan to its state. As `TaskContext` is only used per single query execution, so it can use a plan raw-pointer address:
```rust
struct TaskContext {
...
plan_state: std::sync::Mutex>>,
}
fn plan_addr(plan: &dyn Any) -> usize {
plan as *const _ as *const () as usize
}
impl TaskContext {
/// Get state of the passed plan or register a new one.
pub fn get_or_register_plan_state(
&self,
plan: &dyn Any,
f: F,
) -> Arc
where
F: FnOnce() -> Arc,
{
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
}
}
}
```
Note: if the plan is deallocated then its state is either stored in the streams or is not needed anymore.
Then, when a particular plan want to acquire its state, it calls `get_or_register_plan_state`, an example for `CrossJoinExec`:
```rust
fn execute(
&self,
partition: usize,
context: Arc,
) -> Result {
let state = context
.get_or_register_plan_state(self, || Arc::new(CrossJoinExecState::new()));
// Work with state: register metrics, spawn workers, etc.
let left_fut = state
.as_any()
.downcast_ref::()
.unwrap()
.left_fut
.once(|| {
load_left_input(
Arc::clone(&self.left),
Arc::clone(&context),
join_metrics.clone(),
reservation,
)
});
...
}
```
Note: this way a state is shared across partitions.
Note: in this case we should change public API as methods mentioned in the previous section (`reset_state`, `metrics`, `with_new_state`) should be removed from the `ExecutionPlan` trait itself.
### Describe alternatives you've considered
_No response_
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.