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: Memory limited merge join #5632

Merged
merged 4 commits into from
Mar 20, 2023
Merged

Conversation

korowa
Copy link
Contributor

@korowa korowa commented Mar 17, 2023

Which issue does this PR close?

Closes #5220.

Rationale for this change

Control over memory allocation in SortMergeJoinExec

What changes are included in this PR?

  • SMJStream.reservation -- memory reservation to control buffered data size -- it may require significant amount of memory, for example, in case of multiple batches with the same join key values
  • peak_memory_used -- during join execution, normally, both try_grow and shrink methods are called -- on adding new buffered batches and on flushing them if join key has changed, respectively. In this case it seems reasonable to track only peak memory for each partition -- after that it's summed up across all partitions which is not really precise, but still valuable in terms of showing "worst case" (in fact peak memory used, probably, can be lower)
  • memory limit tests now are able to run with configurable SessionConfig and perform multiple checks on returned error (used to ensure that memory allocation failed in specific operator/stream)

Are these changes tested?

  • existsing test cases (shrinks don't lead to negative size)
  • overallocation tests for all join types (if accepted by merge join operator)
  • SQL-level test for overallocation

Are there any user-facing changes?

SortMergeJoinExec should now respect runtime memory limitations.

@github-actions github-actions bot added the core Core DataFusion crate label Mar 17, 2023
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.

Thank you @korowa -- this looks great. I reviewed the code and made sure the reservation was updated when the buffered batch was pushed/poped.

cc @richox who I believe contributed the original implementation in #2242


assert_contains!(
err.to_string(),
"Resources exhausted: Failed to allocate additional"
Copy link
Contributor

Choose a reason for hiding this comment

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

Thank you for the unit test

.map(|arr| arr.get_array_memory_size())
.sum::<usize>()
+ batch.num_rows().next_power_of_two() * 8
+ 24;
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
+ 24;
+ sizeof::<Range>()
+ sizeof::<usize>();

Copy link
Contributor Author

Choose a reason for hiding this comment

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

True, it may vary. I've fixed it in separate commit, thank you!

@alamb alamb merged commit 6d6079b into apache:main Mar 20, 2023
@andygrove andygrove added the enhancement New feature or request label Apr 7, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
core Core DataFusion crate enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Use MemoryManager in Join operators, and return errors if the memory budget is exceeded
4 participants