Skip to content

Comments

Refactor post-processing logic from the core search and implement rag_search using that pattern#794

Draft
narendatha wants to merge 6 commits intomainfrom
u/narendatha/post_process_refactor
Draft

Refactor post-processing logic from the core search and implement rag_search using that pattern#794
narendatha wants to merge 6 commits intomainfrom
u/narendatha/post_process_refactor

Conversation

@narendatha
Copy link
Contributor

Post-Process Refactor: Composable PostProcess Trait & RAG Search

Summary

Introduces a composable PostProcess<Processor, Provider, T, O> trait that decouples post-processing logic from the core search algorithm, and uses it to implement RAG-aware search as the first concrete consumer.

Changes

New: PostProcess trait (glue.rs)

  • Generic trait parameterized by a Processor marker type, enabling different post-processing strategies to be composed with any SearchStrategy without duplicating the search loop.
  • DefaultPostProcess marker type with a blanket impl that delegates to the strategy's existing SearchPostProcessor -- every SearchStrategy gets default post-processing for free.

New: rag_post_process.rs

  • Greedy orthogonalization algorithm for RAG result diversification.
  • Ridge-aware log-determinant reranking (eta > 0 path).
  • SIMD-accelerated dot_product via InnerProduct.
  • Comprehensive unit tests (3 tests).

New: rag_search.rs

  • RagSearchParams struct (rag_eta, rag_power) -- the processor marker for RAG post-processing.
  • RagSearch struct wrapping KnnWith<RagSearchParams>, implementing Search by delegating to the composed inner type.
  • Accessor methods: knn(), rag_eta(), rag_power(), k_value(), l_value().
  • Unit tests (3 tests).

Modified: knn_search.rs

  • KnnWith<PP> now accepts a generic Processor and requires S: PostProcess<PP, ...> instead of hardcoding a specific post-processor.
  • Updated doc examples to reference RagSearchParams.

Design

Search::search()
  -- KnnWith<PP>::search()
       -- search_core() -> greedy graph traversal
            -- PostProcess<PP>::post_process_with()
                 -- (e.g.) rag_post_process() for RagSearchParams
                 -- (e.g.) strategy.post_processor() for DefaultPostProcess

Any new post-processing strategy only needs to:

  1. Define a marker struct (like RagSearchParams).
  2. Implement PostProcess<MarkerStruct> on the relevant strategy.
  3. Wrap KnnWith<MarkerStruct> in a search type.

@codecov-commenter
Copy link

codecov-commenter commented Feb 23, 2026

Codecov Report

❌ Patch coverage is 43.35840% with 452 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.67%. Comparing base (6f8e8eb) to head (c06a1b1).

Files with missing lines Patch % Lines
diskann-disk/src/search/provider/disk_provider.rs 0.00% 134 Missing ⚠️
diskann-benchmark-core/src/search/graph/rag.rs 0.00% 103 Missing ⚠️
diskann-benchmark/src/backend/index/benchmarks.rs 11.94% 59 Missing ⚠️
diskann-benchmark/src/backend/index/search/rag.rs 0.00% 55 Missing ⚠️
diskann/src/graph/search/rag_search.rs 43.15% 54 Missing ⚠️
diskann/src/graph/search/knn_search.rs 82.65% 17 Missing ⚠️
diskann/src/graph/search/rag_post_process.rs 93.40% 13 Missing ⚠️
diskann-benchmark/src/inputs/async_.rs 0.00% 9 Missing ⚠️
diskann-benchmark/src/inputs/disk.rs 0.00% 8 Missing ⚠️

❌ Your patch status has failed because the patch coverage (43.35%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main     #794      +/-   ##
==========================================
- Coverage   89.16%   88.67%   -0.49%     
==========================================
  Files         431      435       +4     
  Lines       78895    79610     +715     
==========================================
+ Hits        70343    70597     +254     
- Misses       8552     9013     +461     
Flag Coverage Δ
miri 88.67% <43.35%> (-0.49%) ⬇️
unittests 88.67% <43.35%> (-0.49%) ⬇️

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

Files with missing lines Coverage Δ
diskann-benchmark-core/src/search/graph/mod.rs 100.00% <ø> (ø)
diskann-benchmark/src/backend/index/spherical.rs 100.00% <ø> (ø)
diskann/src/graph/glue.rs 96.11% <100.00%> (+0.18%) ⬆️
diskann/src/graph/index.rs 95.90% <100.00%> (-0.01%) ⬇️
diskann/src/graph/search/multihop_search.rs 98.02% <100.00%> (-0.03%) ⬇️
diskann/src/graph/search/range_search.rs 97.63% <100.00%> (-0.03%) ⬇️
diskann-benchmark/src/inputs/disk.rs 4.30% <0.00%> (-0.18%) ⬇️
diskann-benchmark/src/inputs/async_.rs 37.11% <0.00%> (-0.62%) ⬇️
diskann/src/graph/search/rag_post_process.rs 93.40% <93.40%> (ø)
diskann/src/graph/search/knn_search.rs 87.00% <82.65%> (-13.00%) ⬇️
... and 5 more

... and 1 file with indirect coverage changes

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

- Add RAG<DP,T,S> benchmark searcher in diskann-benchmark-core with
  Aggregator that reuses knn::Summary for recall/QPS reporting
- Add RagSearchSteps runner and Rag trait in diskann-benchmark
- Add TopkRag variant to SearchPhase enum with RagSearchPhase struct
  (rag_eta, rag_power fields)
- Add RunRagSearch trait with f32 specialization for type-safe dispatch
- Add rag-search.json example comparing Topk vs TopkRag on siftsmall
- Handle TopkRag in spherical.rs with unsupported bail
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.

2 participants