Skip to content

PiPNN 1/6: extract shared RobustPrune - #1315

Open
weiyaoluo (SeliMeli) wants to merge 3 commits into
mainfrom
pipnn-stack/02-final-prune
Open

PiPNN 1/6: extract shared RobustPrune#1315
weiyaoluo (SeliMeli) wants to merge 3 commits into
mainfrom
pipnn-stack/02-final-prune

Conversation

@SeliMeli

@SeliMeli weiyaoluo (SeliMeli) commented Aug 6, 2026

Copy link
Copy Markdown

Purpose

This PR extracts the private Vamana RobustPrune selection loop. PiPNN can reuse this loop in a later PR.

The change adds no public API. It does not change Vamana pruning behavior.

Main changes

  • graph::internal::prune::robust_prune now owns the alpha-round selection loop.
  • The function accepts SortedNeighbors. Callers cannot pass an unsorted slice.
  • DiskANNIndex::occlude_list still prepares candidates, loads vectors, maps IDs, and applies saturation.
  • Existing comments, assertions, state updates, and unsafe checks remain with their original logic.

Preserved behavior

  • Candidate order remains unchanged.
  • Excluded candidates remain positional None entries.
  • Alpha progression and prune-kind behavior remain unchanged.
  • Saturation order and duplicate handling remain unchanged.
  • Provider errors remain in occlude_list.

Review order

  1. Review diskann/src/graph/internal/prune.rs.
  2. Compare the extracted loop with the old loop in index.rs.
  3. Review the call site in DiskANNIndex::occlude_list.
  4. Confirm that preparation, ID mapping, and saturation remain outside the extracted function.

Validation

  • All diskann library tests pass.
  • All-target Clippy passes with -Dwarnings.
  • Existing Vamana prune scenarios pass.

Stack

Stack 1/6. #1287 adds PiPNN numerical kernels.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR extracts Vamana’s reusable RobustPrune alpha-round occlusion state machine into a new allocation-free kernel (graph/internal/robust_prune.rs) while keeping Vamana-specific scratch space, provider error handling, and adapter logic in graph/internal/vamana_prune.rs. The DiskANNIndex::occlude_list path is updated to prepare “available-only” candidates for the shared kernel, then translate selected positions back to IDs and apply optional saturation.

Changes:

  • Introduces internal::robust_prune as a provider-independent RobustPrune kernel over prepared candidates + reusable per-candidate state.
  • Moves Vamana-owned scratch/context and ranked provider error types into internal::vamana_prune with co-located tests.
  • Refactors graph/index.rs::occlude_list to: validate bounds → prepare available candidates → call shared kernel → write adjacency → saturate from available-only candidates.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
diskann/src/graph/internal/vamana_prune.rs Adds Vamana-owned scratch/context + provider error ranking and Vamana integration tests.
diskann/src/graph/internal/robust_prune.rs Adds allocation-free RobustPrune kernel + pure state-machine tests.
diskann/src/graph/internal/prune.rs Removes the previous combined prune implementation/state.
diskann/src/graph/internal/mod.rs Rewires internal modules to expose robust_prune and vamana_prune.
diskann/src/graph/index.rs Switches pruning implementation to prepare candidates and call the shared robust_prune kernel; updates saturation behavior to be available-only.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread diskann/src/graph/internal/robust_prune.rs Outdated
@codecov-commenter

Codecov Comments Bot (codecov-commenter) commented Aug 6, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 98.75000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.32%. Comparing base (16c7451) to head (149417b).
⚠️ Report is 2 commits behind head on main.

Files with missing lines Patch % Lines
diskann/src/graph/internal/prune.rs 98.57% 1 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1315      +/-   ##
==========================================
+ Coverage   91.26%   91.32%   +0.05%     
==========================================
  Files         517      517              
  Lines       98511    98544      +33     
==========================================
+ Hits        89910    89992      +82     
+ Misses       8601     8552      -49     
Flag Coverage Δ
miri 91.32% <98.75%> (+0.05%) ⬆️
unittests 90.99% <98.75%> (+0.05%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
diskann/src/graph/index.rs 96.42% <100.00%> (-0.10%) ⬇️
diskann/src/graph/internal/prune.rs 85.45% <98.57%> (+22.95%) ⬆️

... and 32 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@SeliMeli
weiyaoluo (SeliMeli) force-pushed the pipnn-stack/02-final-prune branch from 25066a6 to 748d38f Compare August 6, 2026 10:43

@hildebrandmw Mark Hildebrand (hildebrandmw) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR would have been easy to review, but at nontrivially regresses internal documentation and comments. For such important internal algorithms, please avoid changing comments unnecessarily, particularly when unsafe is involved.

Please restore the deleted or altered comments, making only the necessary updates for the new location. Also, please avoid renaming variables when just moving code around. A simpler diff stands on its own.

Comment thread diskann/src/graph/robust_prune.rs Outdated
Comment thread diskann/src/graph/robust_prune.rs Outdated
Comment thread diskann/src/graph/robust_prune.rs Outdated
Comment thread diskann/src/graph/robust_prune.rs Outdated
Comment thread diskann/src/graph/mod.rs Outdated
Keep the prune module, comments, names, preparation, output, and saturation in place; expose only the selection loop and require the existing SortedNeighbors witness.
@SeliMeli
weiyaoluo (SeliMeli) force-pushed the pipnn-stack/02-final-prune branch from 748d38f to c2e6be6 Compare August 7, 2026 03:29
@SeliMeli

Copy link
Copy Markdown
Author

Addressed review in c2e6be62:

  • kept module and existing types in graph::internal::prune;
  • restored original loop comments, variable terminology, and checked-index unsafe justifications;
  • changed the helper boundary to &SortedNeighbors so source-distance ordering is enforced by the existing type;
  • added local pool/cache/state length checks before unchecked mutation;
  • retained Vamana preparation, ID mapping, saturation, errors, and alpha behavior unchanged.

Validation: 311 diskann library tests, all-target Clippy, LSP diagnostics, and git diff --check pass.

@SeliMeli

Copy link
Copy Markdown
Author

Local IAI-Callgrind check against main (16c7451f):

  • triangle 750→64: instructions +0.27259%, estimated cycles +0.29260%;
  • occluding+saturation 750→64: instructions -0.01482%, estimated cycles +0.01124%.

Measured under WSL2 with the archived local regression fixture; no submitted benchmark target. Full record: experiments/pipnn_local_regression/2026-08-07-wrapper-simplification-callgrind.md in the local analysis workspace.

@SeliMeli

Copy link
Copy Markdown
Author

All review threads now point to superseded lines. The final extraction keeps the helper in graph::internal::prune, accepts SortedNeighbors, restores the original loop/state/unsafe comments and existing debug_assert! sequence, and adds no assertions or pruning behavior changes. I am resolving the outdated threads.

@wuw92 Wei Wu (wuw92) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Since it is intended to be a focused extraction, I would keep the implementation as close to the original as possible and simplify the PR description.

Comment thread diskann/src/graph/internal/prune.rs Outdated
@SeliMeli
weiyaoluo (SeliMeli) force-pushed the pipnn-stack/02-final-prune branch from 4721ec5 to 149417b Compare August 11, 2026 11:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants