feat: Add config `max_temp_directory_size` to limit max disk usage for spilling queries #14975

2010YOUY01 · 2025-03-03T09:16:53Z

Which issue does this PR close?

Closes #.

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:

max_temp_directory_size | The maximum amount of data stored inside the 'temp_directory' (when set) (e.g., 1GB)

This PR adds a configuration option max_temp_directory_size to disk manager (default to 100GB), if the limit is reached for all spilled files, an error will be returned. Currently, it can be only configured by constructing DiskManager with an additional argument.

See example in datafusion-cli (Manually modified datafusion-cli's startup code to change the disk limit to 1GB)

yongting@Mac ~/C/d/datafusion-cli (spill-disk-limit *)> cargo run --profile release-nonlto -- --mem-pool-type fair -m 100M
DataFusion CLI v46.0.0
> select * from generate_series(1, 1000000000000) as t1(v1) order by v1;
Resources exhausted: The used disk space during the spilling process has exceeded the allowable limit of 1024.0 MB. Try increasing the `max_temp_directory_size` in the disk manager 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 `metrics` module from `datafusion_physical_plan` crate to `datafusion_execution` crate

Current code organization:

datafusion_physical_plan depends on datafusion_execution
datafusion_physical_plan includes physical executors like SortExec
datafusion_execution includes memory reservation and disk manager

This PR requires adding metrics tracking ability to DiskManager, besides, I believe in the future we can also add more Metrics into MemoryPool and DiskManager, 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 `SpillMetrics`

This way we can ensure the metrics across different operators can be kept consistent.

3. refactor: Move `IPCStreamWriter` from `datafusion_physical_plan` to `datafusion_execution::DiskManager`

4. Implement the utility function to spill batches inside `DiskManager`

try_spill_record_batches is implemented to replace the old utility function spill_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.

…fusion_execution`

2010YOUY01 · 2025-03-03T09:22:07Z

datafusion/physical-plan/src/spill.rs

+    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."
+)]


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.

2010YOUY01 · 2025-03-03T09:25:09Z

datafusion/core/tests/memory_limit/mod.rs

+    let disk_manager = DiskManager::try_new_without_arc(DiskManagerConfig::NewOs)?
+        .with_max_temp_directory_size(disk_limit)?;
+
+    let runtime = RuntimeEnvBuilder::new()


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.

I think you coudl do this iinstead of adding a new function:

let disk_manager = DiskManager::try_new(DiskManagerConfig::NewOs)?; let disk_manager = Arc::try_unwrap(disk_manager) .expect("DiskManager should be unique") .with_max_temp_directory_size(disk_limit)?;

I filed Make DiskManagerBuilder to construct DiskManagers #15319 to track

comphead · 2025-03-03T15:57:12Z

Nice PR thanks @2010YOUY01 I'm planning to review it today

adriangb · 2025-03-17T11:51:42Z

This makes total sense to me, I’ve certainly wanted this feature!

comphead · 2025-03-17T17:49:46Z

Sorry, it totally slipped my mind

alamb

Thanks @2010YOUY01 -- I really like where this is heading. Sorry for the very late review

I think we should try and simplify the configuration system a bit (I left some comment) as well as perhaps avoid adding new methods to the DiskManager).

I feel like some good follow on work would be (I can file tickets):

Connect this setting to datafusion-cli so we can test it more easily
Make the DiskManagerConfig more a builder style

Something like

let manager = DiskManagerBuilder::new()
  .with_max_temp_directory_size(100*1024*1024)
  .build_arc();

alamb · 2025-03-19T20:13:29Z

datafusion/execution/src/disk_manager.rs

        })
    }
+
+    /// Write record batches to a temporary file, and return the spill file handle.


I am not sure about putting this method on the DiskManager itself -- I think managing the contents of the spill files is something that will likely get more sophisticated (e.g. if we want to explore using mmap and other techniques to speed up reading)

I have been meaning to file tickets about this and I will do so shortly.

What would you think about introducing a new struct that would be responsible for managing the spill files for a particular operation

Something like this perhaps:

struct SpillFiles { env: Arc<RuntimeEnv> files: Vec<RefCountedTempFile> ... } ... impl SpillFiles { pub fn try_spill_record_batches( &self, batches: &[RecordBatch], request_description: &str, caller_spill_metrics: &mut SpillMetrics, ) -> Result<()> {..} }

alamb · 2025-03-19T20:21:09Z

datafusion/execution/src/disk_manager.rs

        }
    }

+    pub fn try_new_without_arc(config: DiskManagerConfig) -> Result<Self> {


I don't understand the need for this function and find the name confusing. It seems like the only difference is that it will error if it doesn't have exclusive access to the disk manager. I have an alternate suggestion above

alamb · 2025-03-19T20:21:47Z

datafusion/execution/src/disk_manager.rs

+    }
+
+    /// Set the maximum amount of data (in bytes) stored inside the temporary directories.
+    pub fn with_max_temp_directory_size(


I would expect that this function should be on the DiskManagerConfig rather than DiskManager 🤔

Then you won't have to add the new try_uwnrap either

alamb · 2025-03-19T20:22:23Z

datafusion/core/tests/memory_limit/mod.rs

+    let disk_manager = DiskManager::try_new_without_arc(DiskManagerConfig::NewOs)?
+        .with_max_temp_directory_size(disk_limit)?;
+
+    let runtime = RuntimeEnvBuilder::new()


I think you coudl do this iinstead of adding a new function:

let disk_manager = DiskManager::try_new(DiskManagerConfig::NewOs)?; let disk_manager = Arc::try_unwrap(disk_manager) .expect("DiskManager should be unique") .with_max_temp_directory_size(disk_limit)?;

2010YOUY01 · 2025-03-20T09:37:54Z

@alamb Thanks for the review, I agree there are two things we can improve:

Create a new struct for spilling related utilities, instead of putting it inside DiskManager
Easier-to-use builder for DiskManager for configuring max_temp_dicrectory_size (@Standing-Man already start helping in Make DiskManagerBuilder to construct DiskManagers #15319, thank you 🙏🏼 )

I also found another limitation in this spill interface, for certain cases we want:

let in_progress_temp_file = spiller.create_in_progress_temp_file();
in_progress_temp_file.append(batch);
...
in_progress_temp_file.append(batch);
let temp_file = in_progress_temp_file.finish();

Instead of always writing batches to a spill file at once.

I plan to address point 1 and the above-mentioned limitation in another PR, and after that rework this PR to add disk limit feature. So closed this PR for now.

alamb · 2025-03-20T09:45:38Z

I plan to address point 1 and the above-mentioned limitation in another PR, and after that rework this PR to add disk limit feature. So closed this PR for now.

THank you -- I think consolidating the spilling related utilities and making the interface clear will make it much easier as we go forward to add new features.

Maybe we can take inspiration from comet's spill / shuffle writer code: https://github.com/apache/datafusion-comet/blob/382ac938f2d10666eefc08ec5c1c82025ddf3726/native/core/src/execution/shuffle/shuffle_writer.rs

2010YOUY01 added 2 commits March 3, 2025 12:35

refactor: move metrics mod from datafusion_physical_plan to `data…

c663e88

…fusion_execution`

Implement disk limit tracking

df410e6

github-actions bot added core Core DataFusion crate execution Related to the execution crate labels Mar 3, 2025

2010YOUY01 commented Mar 3, 2025

View reviewed changes

2010YOUY01 changed the title ~~feat: Add one config to limit max disk usage for spilling queries~~ feat: Add config max_temp_directory_size to limit max disk usage for spilling queries Mar 4, 2025

This was referenced Mar 5, 2025

Weekly Plan (Andrew Lamb) March 3, 2025 #14978

Closed

Weekly Plan (Andrew Lamb) March 10, 2025 #15121

Closed

This was referenced Mar 17, 2025

[EPIC] A collection of tickets for improving sorting larger than memory datasets / spilling sorts #15271

Closed

March 17, 2025: This week(s) in DataFusion #15269

Closed

alamb reviewed Mar 19, 2025

View reviewed changes

alamb mentioned this pull request Mar 19, 2025

Make DiskManagerBuilder to construct DiskManagers #15319

Closed

2010YOUY01 closed this Mar 20, 2025

2010YOUY01 mentioned this pull request Mar 22, 2025

fix: Redundant files spilled during external sort + introduce SpillManager #15355

Merged

alamb mentioned this pull request Mar 22, 2025

limit max disk usage for spilling queries #15358

Closed

2010YOUY01 mentioned this pull request Apr 1, 2025

feat: Add config max_temp_directory_size to limit max disk usage for spilling queries #15520

Merged

feat: Add config max_temp_directory_size to limit max disk usage for spilling queries #14975

feat: Add config max_temp_directory_size to limit max disk usage for spilling queries #14975

Uh oh!

Conversation

2010YOUY01 commented Mar 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Changes

1. refactor: Move metrics module from datafusion_physical_plan crate to datafusion_execution crate

2. refactor: Group all spilling-related metrics in executors into a single struct SpillMetrics

3. refactor: Move IPCStreamWriter from datafusion_physical_plan to datafusion_execution::DiskManager

4. Implement the utility function to spill batches inside DiskManager

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

comphead commented Mar 3, 2025

Uh oh!

adriangb commented Mar 17, 2025

Uh oh!

comphead commented Mar 17, 2025

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

2010YOUY01 commented Mar 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

alamb commented Mar 20, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

feat: Add config `max_temp_directory_size` to limit max disk usage for spilling queries #14975

feat: Add config `max_temp_directory_size` to limit max disk usage for spilling queries #14975

2010YOUY01 commented Mar 3, 2025 •

edited

Loading

1. refactor: Move `metrics` module from `datafusion_physical_plan` crate to `datafusion_execution` crate

2. refactor: Group all spilling-related metrics in executors into a single struct `SpillMetrics`

3. refactor: Move `IPCStreamWriter` from `datafusion_physical_plan` to `datafusion_execution::DiskManager`

4. Implement the utility function to spill batches inside `DiskManager`

2010YOUY01 commented Mar 20, 2025 •

edited

Loading