adriangb opened a new pull request, #26272: URL: https://github.com/apache/datafusion/pull/26272
## Which issue does this PR close? - Closes apache/datafusion#26233. Stacked on apache/datafusion#25580 and on #26270 (apache/datafusion#26229). Only the last commit is new. ## Rationale for this change The buffered output of a materialized CTE (in-memory batches, memory reservations, spill files) was held until the end of the run. A CTE is often read early, and the rest of the query then runs while the buffer holds memory and disk space that nothing reads again. ```sql WITH c AS MATERIALIZED (SELECT value AS v FROM generate_series(1, 20000)) SELECT count(*) FROM c a JOIN c b ON a.v = b.v WHERE a.v <= (SELECT max(v) FROM c); ``` Before: after the output stream ends, the memory pool still has 196896 bytes reserved for `MaterializedCte[c]` until the `TaskContext` is dropped. After: the reservation is 0 when the last of the three scans finishes. ## What changes are included in this PR? - The output of a CTE in a run counts its readers. It is released when each partition of each bound scan was executed in the run, and each of these executions has finished or was dropped. - The count is of the partitions that can still be executed, not of the open streams. Thus the output is not released between two scans that are executed one after the other. - A scan stream that is dropped before its end (for example below a `LIMIT`) counts as finished. If all scan streams are dropped before the body completes, the body is cancelled. - A scan partition that is never executed keeps the output until the run ends, as before. - A scan partition that is executed after the release computes the body again. It never reads an empty output. - `RecursiveQueryExec` holds the outputs of the CTEs that its recursive term reads and does not declare, so these are computed one time for the whole query. Not included: the release of each body partition when all scans have read it. Known limits: - If an operator other than `RecursiveQueryExec` executes a scan partition again after the release, the body is computed again. For a volatile body, the second read then has different rows. No operator in the repository does this. - The hold of `RecursiveQueryExec` ends when its stream is dropped, not when the stream ends. ## Are these changes tested? Yes. Unit tests check that the pool reservation is 0 and the spill files are deleted after the last scan partition finishes while the plan and the `TaskContext` are alive, and cover streams dropped early, body cancellation, a partition that is never executed, a scan after the release, and the hold. A test in `datafusion/core/tests/sql/select.rs` checks the release for a planned SQL query. Four existing unit tests asserted the old lifetime and now assert the new one. ## Are there any user-facing changes? No API change. Memory and spill files of a materialized CTE are released earlier. A plan that is executed two times with the same `TaskContext` computes the body again in the second execution, if the first has finished. 🤖 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]
