feat: add SpillStore trait with local-disk implementation by wjones127 · Pull Request #7311 · lance-format/lance

wjones127 · 2026-06-17T03:08:58Z

Summary

Adds a generic SpillStore — reclaimable RAII scratch storage for intermediate state that overflows memory (e.g. index-build posting lists, shuffle runs, BTree pages). Mechanism only; consumer migration (IVF shuffler) is a follow-up.

SpillStore / SpillFile (lance-io): create_spill_file() vends a write-once RAII handle whose writer() / reader() hand back Box<dyn Writer> / Box<dyn Reader>, so callers feed spill files straight into FileWriter::try_new and a v2 FileReader without leaking an ObjectStore + path. Dropping the handle deletes the file and releases its bytes back to the store's budget.
LocalSpillStore: writes to an OS temp directory; with_cap enforces an optional byte budget shared across all handles, returning a typed Error::DiskCapExceeded instead of silently filling the disk. Enforcement lives entirely in the spill store — the spill file decorates the writer with a QuotaWriter (reserve-on-write, release-on-drop-by-stat) rather than threading a field through ObjectStore and every provider, so it works for any backend the store opens.
From<io::Error> recovers a wrapped lance Error, so typed errors such as DiskCapExceeded survive the AsyncWrite boundary.
ScanScheduler::open_reader builds a FileScheduler over an already-open Reader (no path/size lookup), bridging a bare spill reader into the v2 reader path.
Session gains a spill_store field (default: uncapped LocalSpillStore), a with_spill_store() builder for injection, and a spill_store() accessor.

Closes #7300

🤖 Generated with Claude Code

Adds a `SpillStore` trait on `Session` providing uniform, reclaimable scratch space for intermediate state that overflows memory (e.g. index build posting lists, shuffle runs, BTree pages). - `SpillStore` / `SpillFile` (lance-io): `create_spill_file()` vends a RAII handle; `writer()` / `reader()` hand back `Box<dyn Writer>` / `Box<dyn Reader>` so callers feed spill files directly into `FileWriter::try_new` and a v2 `FileReader` without leaking an `ObjectStore` + path. The file is deleted on drop and its bytes are released back to the store's usage counter. - `LocalSpillStore`: writes to an OS temp directory; optionally enforces a byte cap. Enforcement lives entirely in the spill store: the spill file decorates the writer with a quota-enforcing `QuotaWriter` (reserve-on-write, release-on-drop-by-stat) rather than threading a field through `ObjectStore` and every provider, so it works for any backend the store opens. - `From<io::Error>` recovers a wrapped lance `Error`, so typed errors such as `DiskCapExceeded` survive the `AsyncWrite` boundary. - `ScanScheduler::open_reader` builds a `FileScheduler` over an already-open `Reader` (no path/size lookup). - `Session` gains a `spill_store` field (defaults to uncapped `LocalSpillStore`), a `with_spill_store()` builder, and a `spill_store()` accessor so callers and tests can inject alternatives. Mechanism only; consumer migration (IVF shuffler) is a follow-up. Closes lance-format#7300 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

codecov · 2026-06-17T22:36:16Z

Codecov Report

❌ Patch coverage is 90.96774% with 28 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
rust/lance-io/src/spill.rs	88.78%	19 Missing and 6 partials ⚠️
rust/lance-core/src/error.rs	93.10%	1 Missing and 1 partial ⚠️
rust/lance/src/session.rs	96.77%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

westonpace

Just some nits

westonpace · 2026-06-18T21:37:04Z

+//! called only once, so the bytes reserved while writing match the file size
+//! released on drop. Two minor inexactnesses are not engineered around: a write
+//! aborted at the cap leaks its reservation until the store is dropped, and a
+//! file whose size cannot be stat-ed on drop is not released.


What does "a file whose size cannot be stat-ed on drop" mean? I guess it's not super important.

westonpace · 2026-06-18T21:40:37Z

+pub trait SpillFile: Send + Sync {
+    /// Open a writer over this spill file.
+    ///
+    /// For a capped store, writes that would exceed the cap fail with
+    /// [`lance_core::Error::DiskCapExceeded`]. Spill files are write-once:
+    /// implementations may reject a second call (the [`LocalSpillStore`] impl
+    /// returns [`lance_core::Error::invalid_input`]).
+    async fn writer(&self) -> Result<Box<dyn Writer>>;
+
+    /// Open a reader over this spill file.
+    ///
+    /// The data must have been fully written (the writer shut down) first.
+    async fn reader(&self) -> Result<Box<dyn Reader>>;
+}


It's not super important but it seems there is an implicit lifecycle here that could be captured at compile time. The file must be fully written before it can be opened for reads and the writer can only be obtained once.

That's a good point. I might refactor based on that.

westonpace · 2026-06-18T21:40:53Z

+/// [`SpillStore::create_spill_file`], allowing implementations not backed by a
+/// local file (e.g. in-memory buffers, remote object stores).
+#[async_trait]
+pub trait SpillFile: Send + Sync {


Why a trait and not a struct? Are you expecting multiple impls?

Yeah, potentially multiple custom impls. They might custom how and which disks are used. If we didn't do a trait, I would probably re-use DataFusions DiskManager rather than create out own.

westonpace · 2026-06-18T21:42:06Z

+        loop {
+            let current = self.used.load(Ordering::Relaxed);
+            let next = current.saturating_add(n);
+            if next > self.cap_bytes {
+                return Err(Error::disk_cap_exceeded(self.cap_bytes, current));
+            }
+            if self
+                .used
+                .compare_exchange(current, next, Ordering::Relaxed, Ordering::Relaxed)
+                .is_ok()
+            {
+                return Ok(());
+            }
+            // Another thread won the CAS — retry.
+        }


Personal nit but anytime I start writing CAS loops I give up and just use a mutex 😆

The CAS loop guarded a single counter on a cold path (gated by disk I/O), so lock-freedom bought nothing. A plain `Mutex<u64>` is shorter and removes the atomic-ordering reasoning entirely. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

Replaces the `SpillFile { writer(), reader() }` handle with `SpillStore::new_spill() -> (Box<dyn Writer>, Box<dyn Spill>)`. Handing the writer out at allocation makes the lifecycle constraints structural instead of guarded: - Write-once is free — there is no `writer()` method to call twice, so the `AtomicBool` guard is gone. - `Spill::reader()` fails until the writer is shut down (the writer flips a shared flag on shutdown, which `FileWriter::finish` triggers), so partially written bytes can't be read back. The `Spill` handle owns the file's RAII cleanup and disk-budget release. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

- `test_reader_ready_after_async_shutdown`: shuts the writer down through the `AsyncWrite` surface, covering the `poll_shutdown` flag-set path that the `Writer::shutdown` tests don't reach. - Extends the quota test to pin the exact-cap boundary (reserve == cap succeeds, one byte past fails). - `test_empty_spill`: a zero-byte spill round-trips empty under a cap. Also documents why the `finished` flag is set in both shutdown paths and why its `Relaxed` load in `reader()` is sufficient. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

github-actions Bot added the enhancement New feature or request label Jun 17, 2026

wjones127 mentioned this pull request Jun 17, 2026

Add SpillStore trait with local-disk implementation #7300

Closed

github-actions Bot added the A-encoding Encoding, IO, file reader/writer label Jun 17, 2026

wjones127 force-pushed the worktree-piped-wiggling-metcalfe branch from 2a5a2b4 to ad19afe Compare June 17, 2026 20:56

wjones127 force-pushed the worktree-piped-wiggling-metcalfe branch from ad19afe to b2c0c08 Compare June 17, 2026 21:10

wjones127 marked this pull request as ready for review June 17, 2026 22:01

wjones127 requested a review from westonpace June 17, 2026 22:01

westonpace approved these changes Jun 18, 2026

View reviewed changes

wjones127 and others added 3 commits June 23, 2026 11:52

wjones127 merged commit bae3c6a into lance-format:main Jun 23, 2026
31 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: add SpillStore trait with local-disk implementation#7311

feat: add SpillStore trait with local-disk implementation#7311
wjones127 merged 4 commits into
lance-format:mainfrom
wjones127:worktree-piped-wiggling-metcalfe

wjones127 commented Jun 17, 2026 •

edited

Loading

Uh oh!

codecov Bot commented Jun 17, 2026 •

edited

Loading

Uh oh!

westonpace left a comment

Uh oh!

westonpace Jun 18, 2026

Uh oh!

westonpace Jun 18, 2026

Uh oh!

wjones127 Jun 18, 2026

Uh oh!

westonpace Jun 18, 2026

Uh oh!

wjones127 Jun 23, 2026

Uh oh!

westonpace Jun 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

wjones127 commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Uh oh!

codecov Bot commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

westonpace left a comment

Choose a reason for hiding this comment

Uh oh!

westonpace Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

westonpace Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

wjones127 Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

westonpace Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

wjones127 Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

westonpace Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

wjones127 commented Jun 17, 2026 •

edited

Loading

codecov Bot commented Jun 17, 2026 •

edited

Loading