feat: Add one config to limit max disk usage for spilling queries #14975
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
github-actions
bot
added
core
2010YOUY01
commented
Mar 3, 2025
| since = "46.0.0", | ||
| note = "This function is deprecated. Use `datafusion_execution::DiskManager::try_spill_record_batch_by_size` instead. Note this method is mainly used within DataFusion for spilling operators. If you only | ||
| want the functionality of writing `RecordBatch`es to disk, consider using `arrow::ipc::writer::StreamWriter` instead." | ||
| )] |
There was a problem hiding this comment.
This function doesn't have to be public, see spill_record_batches() is also non-public.
I believe this utility function is used mainly for tracking temp files and metrics for DataFusion executor's internal use, the major functionality for writing RecordBatches to disk is in arrow::ipc, so I guess deprecating it won't cause major issues.
| let disk_manager = DiskManager::try_new_without_arc(DiskManagerConfig::NewOs)? | ||
| .with_max_temp_directory_size(disk_limit)?; | ||
|
|
||
| let runtime = RuntimeEnvBuilder::new() |
There was a problem hiding this comment.
We can't use the builder to specify disk limit for now, because DiskManagerConfig is an enum instead of struct, so now the setup routine is a bit hacky.
Changing it I think will inevitably cause API change, so I prefer to leave it to a separate PR.
|
Nice PR thanks @2010YOUY01 I'm planning to review it today |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
For memory-limit queries, executors might write temporary results into the disk to reduce memory load. It's important to have a configuration option to limit the max disk usage, in case some query would bloat the disk and cause other issues.
DuckDB provides a similar configuration:
What changes are included in this PR?
This PR is mostly refactoring, the actual change is small. I made one big refactor into a separate commit, so you can only look at the second commit for the related changes.
Changes
1. refactor: Move
metricsmodule fromdatafusion_physical_plancrate todatafusion_executioncrateCurrent code organization:
datafusion_physical_plandepends ondatafusion_executiondatafusion_physical_planincludes physical executors likeSortExecdatafusion_executionincludes memory reservation and disk managerThis PR requires adding metrics tracking ability to
DiskManager, besides, I believe in the future we can also add moreMetricsintoMemoryPoolandDiskManager, for aggregated statistics across different executors, so this PR first made this refactoring change.2. refactor: Group all spilling-related metrics in executors into a single struct
SpillMetricsThis way we can ensure the metrics across different operators can be kept consistent.
3. refactor: Move
IPCStreamWriterfromdatafusion_physical_plantodatafusion_execution::DiskManager4. Implement the utility function to spill batches inside
DiskManagertry_spill_record_batchesis implemented to replace the old utility functionspill_record_batches, this way the total spilled file size can be tracked, also the interface for external operators to use is a little simpler.Are these changes tested?
Yes, integration test is included for queries exceed/not-exceed the disk limit.
Are there any user-facing changes?
No. However there is one deprecation, I'll explain separately.