feat: reduce open file handles during IVF training by westonpace · Pull Request #6169 · lance-format/lance

westonpace · 2026-03-11T15:09:23Z

The previous shuffler used one open file per partition. At large scales this meant tedious re-adjusting of OS limits. The new shuffler uses two open files. We potentially introduce a bit more random access in the later read phase but the overall performance hasn't changed significantly.

github-actions · 2026-03-11T15:11:53Z

PR Review

Nice improvement — reducing open file descriptors from O(partitions) to O(1) is a meaningful usability win for large-scale IVF training. The two-file design with sorted data + cumulative offsets is clean and well-tested. A few items to consider:

P1: Potential u32 overflow in `partition_ranges`

In partition_ranges (shuffler.rs), the offset position is computed as:

let end_pos = (batch_idx as usize * self.num_partitions + partition_id) as u32;

This silently truncates to u32 if num_batches * num_partitions > u32::MAX. With the default 128MB batch size, processing ~10TB of data across 100K+ partitions would overflow. Since this PR specifically targets users with many partitions, consider using u64 positions (if ReadBatchParams supports it) or at minimum adding a checked cast / early validation that the offsets file won't exceed u32 rows.

Minor: `precomputed_shuffle_buffers` dead code

TwoFileShuffler declares the precomputed_shuffle_buffers field and with_precomputed_shuffle_buffers setter but never uses them (#[allow(dead_code)]). If this feature isn't planned for the new shuffler, removing the field/method avoids confusion. If it is planned, a tracking comment or TODO would help.

Tests

Good coverage of the core paths: round-trip, empty partitions, loss tracking, multi-batch spilling. The with_batch_size_bytes(16) test for forcing multiple batches is a nice touch.

🤖 Generated with Claude Code

codecov · 2026-03-11T15:44:05Z

Codecov Report

❌ Patch coverage is 91.79811% with 26 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
rust/lance-index/src/vector/v3/shuffler.rs	91.77%	14 Missing and 12 partials ⚠️

📢 Thoughts on this report? Let us know!

wkalt · 2026-03-11T18:11:07Z

+        let offsets = offsets_stream.try_collect::<Vec<_>>().await?;
+        let offsets = if offsets.len() == 0 {
+            // We should not hit this path if there is no batches
+            unreachable!()


should it panic with a message?

unreachable panics with the message "I didn't expect to get here" and the line number which is pretty much what I'd use if I did an explicit panic.

wkalt · 2026-03-11T18:15:53Z

+
+        let mut ranges = Vec::with_capacity(self.num_batches as usize);
+        for batch_idx in 0..self.num_batches {
+            if batch_idx == 0 && is_uneven {


is the is_uneven stuff equivalent to checking if partition_id = 0 and batch_idx = 0?

Yes, that's maybe simpler. Updated.

wkalt · 2026-03-11T18:24:45Z

here is the result on my test of 500M rows for posterity.

The previous behavior looked like this:

westonpace · 2026-03-11T18:30:25Z

P1: Potential u32 overflow in partition_ranges

This is a legitimate concern. Unfortunately, we don't actually support take with u64 offsets. This is because we're reusing some old paths from the v0.1 days where we used u32 offsets. It's all fixable but more follow-up territory.

By my math I think, even with 1536-dimension vectors, we are ok until we hit trillions of rows.

1T rows => sqrt(1T) * 96 * 1T bytes of PQ data => 96EB of data which, split into 128MB chunks, would give us ~1B chunks.

Either way, I turned it into a try_into so we get a nice error and can address in the future.

Xuanwo

Thank you for this PR! Only two small questions.

Xuanwo · 2026-03-12T10:17:18Z

+///
+/// This default is likely to be fine for most use cases.
+fn shuffle_batch_bytes() -> usize {
+    std::env::var("LANCE_SHUFFLE_BATCH_BYTES")


I think it's a good idea to add some protection here. This would prevent users from setting LANCE_SHUFFLE_BATCH_BYTES to 0, which could lead to using batch_size_bytes = 0 and producing incorrect results.

Good idea, now it will log a warning and use the default.

Xuanwo · 2026-03-12T10:23:23Z

+                let mut partition_counts = vec![0u64; np];
+                for i in 0..part_ids.len() {
+                    let pid = part_ids.value(i) as usize;
+                    if pid < np {


Is it possible for pid >= np? It looks like we will just write those data without offsets. global_row_count will always include them.

It shouldn't be possible (np is num_partitions) but I now log a warning instead of silently ignoring.

…num_partitions file handles Actually write the offsets to a file and don't accumulate Change progress reporting to report number of rows shuffled and not number of batches processed Address review suggestions More PR suggestions Remove dead code Address PR review

westonpace · 2026-03-16T16:30:42Z

CI failure seems unrelated. Will merge if remaining CI job passes.

github-actions Bot added the enhancement New feature or request label Mar 11, 2026

wkalt reviewed Mar 11, 2026

View reviewed changes

wkalt approved these changes Mar 11, 2026

View reviewed changes

Xuanwo reviewed Mar 12, 2026

View reviewed changes

westonpace force-pushed the feat/two-file-shuffler branch from a692013 to 203b465 Compare March 13, 2026 13:38

westonpace force-pushed the feat/two-file-shuffler branch from 203b465 to 08715e5 Compare March 16, 2026 14:05

Use windows-safe paths

7a56b42

westonpace merged commit bfe7f56 into lance-format:main Mar 16, 2026
27 of 28 checks passed

wjones127 mentioned this pull request Jun 16, 2026

Reduce TwoFileShuffler peak memory from 3-4x to 1x batch_size_bytes #7299

Closed

jackye1995 mentioned this pull request Sep 19, 2026

fix: prevent IVF-RQ merge stalls and speed up priority tracking #9431

Closed

Conversation

westonpace commented Mar 11, 2026

Uh oh!

github-actions Bot commented Mar 11, 2026

PR Review

P1: Potential u32 overflow in partition_ranges

Minor: precomputed_shuffle_buffers dead code

Tests

Uh oh!

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

Uh oh!

Codecov Report

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!

wkalt commented Mar 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

westonpace commented Mar 11, 2026

Uh oh!

Xuanwo 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!

westonpace commented Mar 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

P1: Potential u32 overflow in `partition_ranges`

Minor: `precomputed_shuffle_buffers` dead code

codecov Bot commented Mar 11, 2026 •

edited

Loading

wkalt commented Mar 11, 2026 •

edited

Loading