Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: track memory usage for recursive CTE, enable recursive CTEs by default #9619

Merged
merged 8 commits into from
Mar 16, 2024

Conversation

jonahgao
Copy link
Member

@jonahgao jonahgao commented Mar 15, 2024

Which issue does this PR close?

Closes #9554.

Rationale for this change

  1. Add a MemoryReservation to track the memory usage of the RecursiveQuery's buffer.
  2. When a round of iteration ends, the buffer and its MemoryReservation are moved together to the WorkTable, and then passed on to the MemoryStream in the next iteration.
  3. The MemoryReservation will finally be freed after the MemoryStream is dropped.

What changes are included in this PR?

  • Track memory usage for recursive CTEs.
  • Enable recursive CTE support by default.

Are these changes tested?

Yes

Are there any user-facing changes?

No. Enable a new feature by default.

@github-actions github-actions bot added the core Core DataFusion crate label Mar 15, 2024
@jonahgao jonahgao marked this pull request as draft March 15, 2024 09:37
@github-actions github-actions bot added the sql SQL Planner label Mar 15, 2024
@github-actions github-actions bot added the sqllogictest SQL Logic Tests (.slt) label Mar 15, 2024
@jonahgao jonahgao marked this pull request as ready for review March 15, 2024 14:01
@alamb alamb changed the title feat: track memory usage for recursive CTE feat: track memory usage for recursive CTE, enable recursive CTEs by default Mar 15, 2024
Copy link
Contributor

@alamb alamb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very nice -- thank you @jonahgao

cc @matthewgapp

)
SELECT * FROM nodes;",
)
.with_expected_errors(vec![
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

@@ -236,6 +237,8 @@ impl MemoryExec {
pub struct MemoryStream {
/// Vector of record batches
data: Vec<RecordBatch>,
/// Optional memory reservation bound to the data, freed on drop
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was a little worried at first that this optional API makes it easy to forget to provide reservation. However I see now that the reservation is used only with the recursive CTE case

(not for this PR) In general I wondered if we should always have a memory reservation to MemoryStream 🤔 I think that would double count batches from a MemTable however, so it isn't an obviously good improvement

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This concern is quite reasonable; previously there was no need to add a MemoryReservation to the MemoryStream, and the with_reservation function is a bit confused.

I plan to no longer use MemoryStream for future recursive CTEs and instead use a new stream class that can read the WorkTable multiple times, in order to support
https://github.com/apache/arrow-datafusion/blob/81b0a011705b17a09f494f550a5382b0c3414597/datafusion/physical-plan/src/recursive_query.rs#L316

So that the MemoryStream::with_reservation API can be removed.

@matthewgapp
Copy link
Contributor

matthewgapp commented Mar 15, 2024

Looks great, @jonahgao—thanks for taking this on. Would it be sensible to have the worktable's update method take a ReservedBatches as an argument, instead of constructing it internally? My thinking is that, given worktable's sole purpose is to facilitate the recursive query, tightening the integration with ReservedBatches could better illustrate to future contributors the coupled ownership relationship between a RecordBatch's buffer's and its corresponding memory reservation. It's more of a nit than anything so feel free to ignore :)

And apologies that I've been a bit radio silent on this front @alamb! Been heads down on other stuff :/

@alamb
Copy link
Contributor

alamb commented Mar 15, 2024

And apologies that I've been a bit radio silent on this front @alamb! Been heads down on other stuff :/

No worries @matthewgapp -- I totally understand. Good luck and thank you for contributing the original PR 🙏

@jonahgao
Copy link
Member Author

Looks great, @jonahgao—thanks for taking this on. Would it be sensible to have the worktable's update method take a ReservedBatches as an argument, instead of constructing it internally? My thinking is that, given worktable's sole purpose is to facilitate the recursive query, tightening the integration with ReservedBatches could better illustrate to future contributors the coupled ownership relationship between a RecordBatch's buffer's and its corresponding memory reservation. It's more of a nit than anything so feel free to ignore :)

And apologies that I've been a bit radio silent on this front @alamb! Been heads down on other stuff :/

Make sense to me! Updated and thank you @matthewgapp

@alamb alamb merged commit 80d223f into apache:main Mar 16, 2024
24 checks passed
@alamb
Copy link
Contributor

alamb commented Mar 16, 2024

Thanks @jonahgao and @matthewgapp 🙏

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
core Core DataFusion crate sql SQL Planner sqllogictest SQL Logic Tests (.slt)
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Enable recursive CTE support by default
3 participants