adriangb opened a new pull request, #26270: URL: https://github.com/apache/datafusion/pull/26270
## Which issue does this PR close? - Closes apache/datafusion#26229. Stacked on apache/datafusion#25580. Only the last commit is new. ## Rationale for this change `MaterializedCteExec::reset_state` cleared its shared buffer in place and returned `self`. `ExecutionPlan::reset_state` must return a node with new state and leave the old node unchanged. Because the buffered output was in the plan, two runs of one physical plan also shared it. Example with one physical plan for `WITH c AS MATERIALIZED (SELECT random() AS x) SELECT * FROM c`: | | Before | After | |---|---|---| | Run 1 and run 2 of the same plan, each with its own `TaskContext` | Run 2 replays the buffer of run 1. | Each run computes the body. | | `reset_plan_states(plan)` while run 1 is open | The buffer that run 1 uses is replaced. | Run 1 is not changed. | No SQL query is known to give a wrong result on apache/datafusion#25580. ## What changes are included in this PR? - `TaskContext` has a new `RunState`. It holds state that operators keep for one run of a plan. - The output of a materialized CTE is now in the `RunState`, keyed by the CTE id. `MaterializedCteBuffer` holds plan-time data only. - `MaterializedCteExec` uses the default `reset_state`. - `RecursiveQueryExec` gives each iteration a child `RunState`, in which the CTEs that the recursive term declares are computed again. A CTE that is declared outside the recursive term keeps its output for the whole query. Alternative that was not used: rebind the scans to a new buffer in `reset_plan_states`. That function cannot reach a scan in a scalar subquery, and it does not separate two runs when nothing calls reset. Notes for reviewers: - A run is identified by its `TaskContext`. If each partition is executed with a different `TaskContext`, the body is computed one time for each context. - The output now lives until the `TaskContext` of the run and the replay streams are dropped. apache/datafusion#26233 releases it earlier. - **Open point:** the body future holds an `Arc<TaskContext>`, the `TaskContext` owns the `RunState`, and the `RunState` owns the future. While the body is not complete this is a reference cycle. A run that is abandoned before the body completes can keep its memory. This is not confirmed by a test and is not corrected yet. - `ScalarSubqueryExec::reset_state` has the same in-place reset and is not changed here. ## Are these changes tested? Yes. Unit tests: two overlapping runs of one plan have separate outputs, `reset_plan_states` leaves the old plan usable, a renewed context computes only the given CTEs again, and a test for `RunState`. The recursive cases in `cte_materialized.slt` pass with no change. There is no slt for nested recursion. ## Are there any user-facing changes? New public API: `datafusion_execution::RunState`, `TaskContext::run_state`, `TaskContext::with_run_state`, and `TaskContext: Clone`. No existing API changes. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
