Skip to content

Commit 01bf68c

Browse files
pantShreyalamb
andauthored
feat: introduce pluggable SpillFile trait and TempFileFactory for custom spill backends (apache#21882)
## Which issue does this PR close? - Closes apache#21215, depends on apache#22230 ## Rationale for this change DataFusion’s spill infrastructure is tightly coupled to OS-level files, with no extension points for alternative storage backends. `DiskManager` cannot be customized for file creation, and `IPCStreamWriter` depends on OS file paths. This prevents integration in environments where temporary storage must be managed by the host system. For example, Postgres extensions (e.g., ParadeDB) require spill files to go through `BufFile` APIs to respect `temp_tablespaces`, enforce `temp_file_limit`, and integrate with transaction-scoped cleanup. Since `BufFile` has no OS-visible path, it cannot work with the current design. A secondary motivation raised by @alamb is supporting object storage backends (S3, GCS) for spilling, which require async IO and cannot use `std::io::Write` or `std::io::Read`. ## What changes are included in this PR? - Introduced `SpillFile`, `SpillWriter`, and `TempFileFactory` traits to abstract spill file handling - Added `DiskManagerMode::Custom` to allow pluggable backends - Updated `DiskManager` to return `Arc<dyn SpillFile>` instead of OS-bound types - Refactored write path using `SpillWriteAdapter` to bridge sync Arrow writers with backend-agnostic writers - Refactored read path to use async streaming (`Stream<Item = Result<Bytes>>`) instead of blocking state machines - Updated spill-related components to operate on `Arc<dyn SpillFile>` - Migrated the Sort-Merge Join (SMJ) operator to use the async spill abstraction ## Are these changes tested? Yes. Existing spill tests cover the full read/write flow. - Fixed `test_disk_usage_decreases_as_files_consumed` by correcting a pre-existing off-by-one assumption in file rotation - Fixed `test_preserve_order_with_spilling` by just asserting spilling occurs (`spill_count>0`) and output batches are sorted ## Are there any user-facing changes? Yes this introduces API changes: - Spill-related APIs now use `Arc<dyn SpillFile>` instead of `RefCountedTempFile` - New public traits: `SpillFile`, `SpillWriter`, `TempFileFactory` - Added `DiskManagerMode::Custom` for custom backends Custom spill backends can now be implemented and plugged in via `DiskManager`. --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
1 parent 367f08e commit 01bf68c

18 files changed

Lines changed: 664 additions & 493 deletions

File tree

Cargo.lock

Lines changed: 4 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

datafusion/execution/Cargo.toml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ sql = []
5555
arrow = { workspace = true }
5656
arrow-buffer = { workspace = true }
5757
async-trait = { workspace = true }
58+
bytes = { workspace = true }
5859
dashmap = { workspace = true }
5960
datafusion-common = { workspace = true, default-features = false }
6061
datafusion-expr = { workspace = true, default-features = false }
@@ -66,7 +67,11 @@ parking_lot = { workspace = true }
6667
parquet = { workspace = true, optional = true }
6768
rand = { workspace = true }
6869
tempfile = { workspace = true }
70+
tokio = { workspace = true }
71+
tokio-util = { workspace = true, features = ["io"] }
6972
url = { workspace = true }
73+
[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
74+
tokio = { workspace = true, features = ["fs"] }
7075

7176
[dev-dependencies]
7277
chrono = { workspace = true }

0 commit comments

Comments
 (0)