Simplify the DataProvider contract for graph search#1067
Draft
hildebrandmw wants to merge 20 commits into
Draft
Conversation
added 2 commits
May 13, 2026 17:49
Contributor
|
get_element is basically unused in garnet, except for implementing VEMB (return an vector) and other debugging-style functions. This sounds like a very promising direction, and I am very supportive. |
DataProvider contract for graph search
hildebrandmw
added a commit
that referenced
this pull request
May 16, 2026
Remove `flat_search` from `DiskANNIndex` and the `IdIterator` trait from `diskann`. Since the only caller was from `diskann-disk`, add a new `flat_search` inherent method to `DiskIndexSearcher`. The flat search method is not compatible with the experimental direction in #1067 and with #983 on the horizon, this is safe to move for now.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
hildebrandmw
added a commit
that referenced
this pull request
May 21, 2026
Paged search has been causing all kinds of issues for our code base and is actively getting in the way of simplifications in #1067 due to interactions with the `PagedSearchState`. The TLDR of the issue is that `PagedSearchState` requires types to be `'static` and introduces the need to "pause" and "resume" search state in a way that is complex to describe in trait bounds. Since our code is already async, we can lean into that and use the usual Rust machinery to embed non-`'static` paged searcher inside an otherwise `'static` future. The recommended way to now interact with paged search is via [channels](https://docs.rs/tokio/latest/tokio/sync/mpsc/fn.channel.html). [Rendered RFC](https://github.com/microsoft/DiskANN/blob/b63d009893f362e07ace518314f7c85733cba212/rfcs/01078-paged-search.md) ## API Migration Guide | Old pattern | New pattern | |---|---| | `index.start_paged_search(s, ctx, q, l).await` | `index.paged_search(s, ctx, q, l).await` | | `index.next_search_results(ctx, &mut state, k, &mut buf).await` | `search.next_page(k).await` | | `SearchState<Id, (S, C)>` | `PagedSearch<'a, DP, S, T>` | | `PagedSearchState<DP, S, C>` | `PagedSearch<'a, DP, S, T>` | | Check return count for exhaustion | Check `page.is_empty()` | If existing code embedded the `SearchState` in some `'static` container, that is no longer viable because of the borrow. Instead, channels can be used for this communication: ```rust // Types are illustrative — adapt names to your crate. type PageResult = ANNResult<Vec<Neighbor<ExternalId>>>; /// Spawn a paged search session. The index is held by Arc so the task is 'static. /// /// Returns a request channel and a result channel. The caller sends the desired /// page size (`k`) and awaits the corresponding result on the other end. fn spawn_paged_session( index: Arc<DiskANNIndex<DP>>, context: Arc<DP::Context>, query: T, l: usize, ) -> (mpsc::Sender<usize>, mpsc::Receiver<PageResult>) { let (req_tx, mut req_rx) = mpsc::channel::<usize>(1); let (res_tx, res_rx) = mpsc::channel::<PageResult>(1); tokio::spawn(async move { // Borrow from the Arc — these references are scoped to the task. let mut search = index.paged_search(strategy, &*context, query, l).await.unwrap(); while let Some(k) = req_rx.recv().await { let page = search.next_page(k).await; if res_tx.send(page).await.is_err() { break; // caller dropped the result receiver } } // Request channel closed -> caller dropped sender -> clean shutdown. }); (req_tx, res_rx) } ``` If code was already explicitly using a `.await` loop with `SearchState`, then minimal changes should be needed. ## For Users of Paged Search via `wrapped_async` Users of paged search via `wrapped_async::DiskANNIndex` that know their inner futures will never suspend can use the new `wrapped_async::DiskANNIndex::paged_search_no_await`. This will use the new API transparently via `wrapped_async::noawait::PagedSearch` and efficiently run paged searches with minimal synchronization overhead. This should **only** be used if the implementation of `Accessor`, `BuildQueryComputer`, `SearchExt`, `DataProvider`, and `ExpandBeam` are known to never yield and always complete with `Poll::Ready`. --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
hildebrandmw
added a commit
that referenced
this pull request
May 21, 2026
In preparation for #1067, switch the implementation of `DynIndex::search_element` from using the `Accessor` trait to an inherent method on the underlying provider.
hildebrandmw
added a commit
that referenced
this pull request
May 21, 2026
Refactor `DiskIndexSearcher::flat_search` to use the bulk `pq_distances` method instead of one-by-one computation. This does two things: 1. The bulk method is presumably a little more efficient. 2. This moves the implementation away from `Accessor`/`BuildQueryComputer` to help with #1067. The small tricky bit is deriving the maximum batch size from the size of scratch. To that end, I added a private accessor for `PQScratch` to return this bound. --------- Co-authored-by: Mark Hildebrand <mhildebrand@microsoft.com>
arkrishn94
added a commit
that referenced
this pull request
May 22, 2026
This PR introduces a trait interface and a light index to support brute-force search for providers that can be used as/are a flat-index. There is an associated RFC that walks through the interface and associated implementation in `diskann` as a new `flat` module. Rendered RFC [link](https://github.com/microsoft/DiskANN/blob/u/adkrishnan/flat-index/rfcs/00983-flat-search.md). ## Motivation The repo has no first-class surface for brute-force search. This PR adds a small trait hierarchy that gives flat search the same provider-agnostic shape that graph search has, so any backend (in-memory, quantized, disk, remote) can plug in once and reuse a shared algorithm. ## Traits (`flat/strategy.rs`) **`DistancesUnordered<C>`** — the single trait a backend must implement. Fuses iteration and scoring into one method: the implementation drives a full scan, scoring each element with a precomputed query computer `C`, and invokes a callback with `(id, distance)` pairs. Key associated types: - `ElementRef<'a>` -- the reference shape `C` scores against. - `Id` -- the id type yielded to the callback (decoupled from `HasId` so visitors can yield any id shape). - `C : for<'a> PreprocessedDistanceFunction<Self::ElementRef<'a>, f32>` -- the precomputer query computer. ```rust pub trait DistancesUnordered<C>: Send + Sync where C: for<'a> PreprocessedDistanceFunction<Self::ElementRef<'a>, f32>, { type ElementRef<'a>; type Id; type Error: ToRanked + Debug + Send + Sync + 'static; fn distances_unordered<F>( &mut self, computer: &C, f: F, ) -> impl SendFuture<Result<(), Self::Error>> where F: Send + FnMut(Self::Id, f32); } ``` **`SearchStrategy<P, T>`** — factory that creates a `DistancesUnordered` visitor from a provider + context, and builds the per-query computer. Mirrors the graph-side strategy pattern. Two fallible methods: - `create_visitor` — borrows provider + context, returns a `Visitor` - `build_query_computer` — preprocesses the query `T` into a `QueryComputer` ```rust pub trait SearchStrategy<P, T>: Send + Sync where P: DataProvider, { type ElementRef<'a>; type Id; type QueryComputer: for<'a> PreprocessedDistanceFunction<Self::ElementRef<'a>, f32>; type QueryComputerError: StandardError; type Visitor<'a>: for<'b> DistancesUnordered< Self::QueryComputer, ElementRef<'b> = Self::ElementRef<'b>, Id = Self::Id, >; type Error: StandardError; fn create_visitor<'a>(&'a self, provider: &'a P, context: &'a P::Context) -> Result<Self::Visitor<'a>, Self::Error>; fn build_query_computer(&self, query: T) -> Result<Self::QueryComputer, Self::QueryComputerError>; } ``` ## Index (`flat/index.rs`) **`FlatIndex<P>`** — thin `'static` wrapper around a `DataProvider`. Currently we have implemented the naive kNN search algorithm for the flat index. `knn_search` asks the strategy for a visitor, builds the query computer, drives `distances_unordered` through a priority queue, and writes results via `SearchPostProcess`. ## Test infrastructure (`flat/test/`) A self-contained test provider with dimension-validated `Strategy`, transient-error injection, and a `KnnOracleRun` harness that compares `knn_search` results against a brute-force reference with baseline caching for regression detection. ## Future work - **Post-processing** — `knn_search` currently uses the graph-side `SearchPostProcess` trait to write results into the output buffer. #1067 will introduce a flat-specific post-processing step that decouples flat search from the graph module's output machinery. - **Vector ids** - The vector id over which `DistancesUnordered` acts is an associated type, instead of the `InternalId` of the provider. This is due to overly restrictive trait bounds on `VectorId` trait. We plan on relaxing this allowing for more generic id types. --------- Co-authored-by: Alex Razumov (from Dev Box) <alrazu@microsoft.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Executive Summary: What happens if we start simplifying our trait hierarchy and rely on
ExpandBeamto do the heavy lifting?This is controversial and I have no intention (yet) of turning this into a full PR, but it demonstrates the feasibility of something I've come to realize recently. The whole
Accessor->BuildQueryComputer->ExpandBeamoriginated from earlier patterns in the library of building up indexing algorithms from tiny pieces, enabling a user to implement a small set of methods and get the whole indexing algorithm for free.This breaks down, though, when we consider performance: the more information we provide a backend database about the operation we want to perform, the better the database can optimize. This lead to
on_elements_unordered(bulk element retrieval, allowing for example prefetching),distances_unordered(bulk retrieval with a concrete distance computations), and finallyExpandBeam. Of these, the latter has proved by far the most useful, being used bydiskann-diskanddiskann-garnetto specialize the algorithm to their preferred data access pattern.I've been exploring how to make filtered search more of a first-class citizen to allow us to batch together predicate queries and potentially fuse predicate evaluation with data access. I was halfway through writing a parallel Accessor/Computer/ExpandBeam/Strategy hierarchy for predicate aware search (and struggling with expressing the more complicated type relationships required), when I realized that all we really needed were a few slightly different flavors of
ExpandBeam. Instead of building a complex nest of interrelated traits to express these subtly differentExpandBeammethods, why not just require users to implement the necessaryExpandBeam? The required semantics are not that complicated and it gives the implementor full control of the data access.This PR is an experiment in removing all but the
ElementRefGAT of theAccessortrait and instead makingExpandBeamthe most relevant method. There's more simplification that can be done: fusingSearchExtwithExpandBeam, getting rid of the peskyAsNeighborbound fromExpandBeam, simplifyingBuildQueryComputerto maybe not even need thePreprocessedDistanceFunctionbound. However, this is already a large reduction in code and has a couple extra benefits:ExpandBeamimplementations for inmem are actually more efficient than the previous provided implementations because we can more aggressively evaluate the visited predicate (and merge the neighbor buffer in the prefetchid_bufferto save an allocation on eachExpandBeamcall. This is not a significant increase in code because we simply move the oldon_elements_unordered.diskanntesting the provided implementations that are now no longer needed.ExpandBeamwithSearchExt, then all the relevant bits for search are grouped in a single place instead of spread out for easier discovery.diskann, with potentially more simplification if we continue down this path.ExpandBeamis perhaps a little more prone to error. This is a distinct negative, but if everything is collected into a single place, it's easier to audit.Now for some tradeoffs:
flat_searchfunction. I kind of view this as a bonus: the existing flat search relies on an ID iterator (not easy for most backends) and uses an inefficient loop. The call indiskann-diskwould likely be much more efficient if it was just done directly on the PQ dataset.diskann-label-filter. Coarsening the granularity of data access methods means the wrapping done by the label provider is less effective. Again, though, this is probably a feature: since real implementations are moving towardsExpandBeamfor performance, the label provider may not be a good approach anyways. We certainly had this experience in Remove the Caching Provider #1052.Accessor's get-element. This can be grouped into the largerdiskann-providerscleanup.Fillfordiskann::graph::workingset::Mapis no longer applicable. Users have to write their own relatively easyFillmethod. I don't think this is as big of a problem as it sounds because I don't think any serious implementation uses the provided implementation: it's too slow.Why is this relevant now? Aside from the filtered search work, #983 is working towards adding flat search and introduces an even finer granularity of traits in the hierarchy, when my gut tells me we should be moving the other direction towards slightly coarser traits. Also, if we have an opportunity to simplify our codebase, I think we should take it.
Why is this possible now? #859 moved one of the last big
get_elementrelated operations behind a bulk operation. Surprisingly, that removed all calls to the rawget_elementfrom insert or inplace delete.