fix(index): keep HNSW IVF scratch dir alive during partition writing#6980
fix(index): keep HNSW IVF scratch dir alive during partition writing#6980geserdugarov wants to merge 2 commits into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
@Xuanwo, @wjones127, hi! |
geserdugarov
force-pushed
the
fix-hnsw-scratch-dir
branch
from
878bf4c to
72fb619
Compare
|
@Xuanwo, @wjones127, hi! |
| let tmp_part_dir = Path::from_filesystem_path(TempStdDir::default())?; | ||
| // Bind the guard to keep the scratch directory alive for the whole function; | ||
| // dropping it earlier would delete the partition files before they are read back. | ||
| let tmp_part_dir_guard = TempStdDir::default(); |
There was a problem hiding this comment.
This adds function-scoped scratch cleanup while the function can still return on the first partition-task error without draining the remaining spawned tasks. Those detached tasks can continue writing under a directory that is being removed, losing later build or I/O errors.
There was a problem hiding this comment.
Thanks for catching it! Need some time to work on this.
There was a problem hiding this comment.
I updated this to drain outstanding tasks on error before the temp guard can drop: remaining handles are kept as Option<JoinHandle<_>>, aborted, and awaited in drain_partition_tasks. I also made the scratch guard shared with partition tasks and removed the detached PQ storage spawn, so aux writes cannot outlive the partition task.
| // Bind the guard to keep the scratch directory alive for the whole function; | ||
| // dropping it earlier would delete the partition files before they are read back. | ||
| let tmp_part_dir_guard = TempStdDir::default(); | ||
| let tmp_part_dir = Path::from_filesystem_path(&tmp_part_dir_guard)?; |
There was a problem hiding this comment.
The cited existing test does not reach this legacy HNSW writer; HNSW_PQ/SQ creation goes through the V3 builder. This lifetime fix can regress without that test failing.
There was a problem hiding this comment.
Agreed. I added targeted coverage that calls the legacy build_ivf_hnsw_pq_index path in a child process with isolated TMPDIR, then asserts no scratch .tmp* dirs are leaked, plus small tests for guard/drain behavior.
geserdugarov
force-pushed
the
fix-hnsw-scratch-dir
branch
2 times, most recently
from
4bc592a to
4b9191a
Compare
geserdugarov
force-pushed
the
fix-hnsw-scratch-dir
branch
from
4b9191a to
f52f075
Compare
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughRefactors IVF_HNSW partition building in ChangesPartition build lifetime and cleanup
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant write_hnsw_quantization_index_partitions
participant TempStdDir
participant Partition task
participant drain_partition_tasks
write_hnsw_quantization_index_partitions->>TempStdDir: create shared scratch guard
write_hnsw_quantization_index_partitions->>Partition task: spawn with Arc clone
Partition task->>Partition task: build PQ storage + HNSW via try_join!
alt success
Partition task-->>write_hnsw_quantization_index_partitions: return partition result
write_hnsw_quantization_index_partitions->>write_hnsw_quantization_index_partitions: assemble staged files
else failure
write_hnsw_quantization_index_partitions->>drain_partition_tasks: abort and await remaining tasks
drain_partition_tasks-->>write_hnsw_quantization_index_partitions: collected errors
end
write_hnsw_quantization_index_partitions->>TempStdDir: drop guard and remove scratch dir
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance/src/index/vector/ivf/io.rs`:
- Around line 865-873: The test setup in the child-process path is hand-building
Arrow schema and batch objects, which should be replaced with the repository’s
test helper style. Update the code around the FixedSizeListArray setup in the
vector IVF I/O test to use record_batch!() from arrow_array (and any existing
test batch helpers) instead of manually constructing Schema, RecordBatch, and
RecordBatchIterator; keep the surrounding test logic and Dataset::write flow
unchanged.
- Line 795: Move the inline test-only imports out of the function body and into
the existing test module import block so the file follows the same convention as
the other `mod tests` imports. Update the test setup around `PathBuf` and the
other inline imports at the referenced locations in `ivf::io` to be grouped with
the rest of the test imports, keeping the import organization consistent and
avoiding local `use` statements inside functions.
- Line 482: The Vec allocations in the IVF I/O path should be preallocated
because their sizes are already bounded. Update the initialization in the
relevant vector-building code in io.rs, including the `errors` collection and
the other matching vector mentioned in the review, to use
`Vec::with_capacity(...)` with the known upper bounds (`tasks.len()` and
`NUM_SLOW + 1`) instead of empty Vec construction. Use the existing symbols
around the task-processing logic to locate both allocations and keep the
capacity estimate slightly generous if needed.
- Line 372: The semaphore acquisition in the IVF IO build path currently uses an
expect, which can panic inside library code and hide the real failure. Update
the acquire call in the relevant build task around sem.acquire to return the
error through the function’s Result instead of panicking, using proper error
propagation consistent with the surrounding lance::index::vector::ivf::io logic
and existing error types.
- Around line 538-548: The IVF_HNSW quantizer handling in the vector index IO
path still uses panic-prone fallbacks, specifically `aux_writer.unwrap()` and
`unreachable!()` inside the `match quantizer` block. Update this logic to return
a contextual `Error::index` instead of panicking when the auxiliary writer is
missing or the quantizer state is invalid, and thread the error through the
existing `Result` flow from `build_and_write_pq_storage` so the partition task
can fail normally.
- Around line 483-488: The drain loop in the IVF IO task cleanup is awaiting
each aborted handle immediately after aborting it, which delays cancellation of
later tasks; update the logic around the task iteration so all outstanding
handles from tasks are aborted first and only then awaited. Use the existing
task collection in the vector IVF IO path (the loop over tasks and the
handle.abort/handle.await sequence) to split cancellation and waiting into
separate phases.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 206641b6-417f-464a-ac84-c2ecba64e94b
📒 Files selected for processing (1)
rust/lance/src/index/vector/ivf/io.rs
3 participants
Summary
Fixes an HNSW-IVF index build bug analogous to #6957, but in the legacy HNSW partition writer path.
Changes
TempStdDirguard alive for the fullwrite_hnsw_quantization_index_partitionsfunction.hnsw_part_*files from disappearing before they are read back into the final index files.index::vector::ivf::tests::test_create_index_nullsalready coversIVF_HNSW_PQandIVF_HNSW_SQwithIndexFileVersion::Legacy.Testing
cargo test -p lance test_create_index_nulls -- --nocaptureIVF_HNSW_PQandIVF_HNSW_SQcases.Summary by CodeRabbit