diff --git a/iris-mpc-cpu/src/execution/hawk_main/search.rs b/iris-mpc-cpu/src/execution/hawk_main/search.rs index 0a365b99b..33a1cb805 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/search.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/search.rs @@ -231,8 +231,9 @@ async fn classify_and_extend( ); metrics::counter!("supermatcher_extended_searches").increment(1); + let seeded_nbhd = SortedNeighborhood::from_ascending_vec(edges.to_vec()); let supermatch_neighbors = hnsw_supermatch - .search(aby3_store, graph_store, query, ef_supermatch) + .search_layer_0_seeded(aby3_store, graph_store, query, seeded_nbhd, ef_supermatch) .await?; let supermatch_classified = classify_edges( diff --git a/iris-mpc-cpu/src/hnsw/searcher.rs b/iris-mpc-cpu/src/hnsw/searcher.rs index abd89291f..d7a7c61f7 100644 --- a/iris-mpc-cpu/src/hnsw/searcher.rs +++ b/iris-mpc-cpu/src/hnsw/searcher.rs @@ -1382,6 +1382,44 @@ impl HnswSearcher { Ok(W) } + /// Extended layer-0-only search, seeded with the neighborhood produced by a + /// previous search. Used to confirm/deny a potential supermatcher without + /// repeating the upper-layer greedy descent. + /// + /// The edges of `seeded_nbhd` must be sorted ascending by distance. + #[instrument(level = "trace", target = "searcher::network", skip_all)] + pub async fn search_layer_0_seeded( + &self, + store: &mut V, + graph: &GraphMem, + query: &V::QueryRef, + seeded_nbhd: SortedNeighborhood, + ef: usize, + ) -> Result> { + let start = Instant::now(); + + if seeded_nbhd.is_empty() { + let W = self.search(store, graph, query, ef).await?; + return Ok(W); + } + + let mut W = seeded_nbhd; + Self::search_layer( + store, + graph, + query, + &mut W, + ef, + 0, + self.min_layer_search_batch_size, + ) + .await?; + + W.edges.truncate(ef); + metrics::histogram!("search_layer_0_seeded_duration").record(start.elapsed().as_secs_f64()); + Ok(W) + } + /// Insert `query` into the HNSW index represented by `store` and `graph`. /// Return a `VectorId` representing the inserted vector. #[instrument(level = "trace", skip_all, target = "searcher::network")] @@ -2042,4 +2080,87 @@ mod tests { assert!(ops.is_empty()); Ok(()) } + + /// Build a plaintext HNSW graph with `n` random elements, returning the + /// store, graph, the inserted queries, and their assigned `VectorId`s (in + /// insertion order). + async fn build_seeded_test_graph( + db: &HnswSearcher, + seed: u64, + n: usize, + ) -> Result<( + PlaintextStore, + GraphMem, + Vec>, + Vec, + )> { + let mut store = PlaintextStore::::new(); + let mut graph = GraphMem::new(); + let rng = &mut AesRng::seed_from_u64(seed); + + let queries = IrisDB::new_random_rng(n, rng) + .db + .into_iter() + .map(Arc::new) + .collect::>(); + + let mut ids = Vec::with_capacity(n); + for query in queries.iter() { + let insertion_layer = db.gen_layer_rng(rng)?; + let id = db + .insert(&mut store, &mut graph, query, insertion_layer) + .await?; + ids.push(id); + } + + Ok((store, graph, queries, ids)) + } + + /// For a node that lives in layer 0, a layer-0-only search seeded with + /// the result of a regular top-down search should reproduce that same + /// result. + #[tokio::test] + async fn test_search_layer_0_seeded_matches_full_search() -> Result<()> { + let db = HnswSearcher::new_with_test_parameters(); + let (mut store, graph, queries, ids) = build_seeded_test_graph(&db, 0, 256).await?; + + let query_idx = 0; + assert!(graph.layers[0].links.contains_key(&ids[query_idx])); + let query = &queries[query_idx]; + let ef = db.params.get_ef_search(1); + + // Regular top-down search. + let full = db.search(&mut store, &graph, query, ef).await?; + + // Seed a layer-0-only search with the full result; it must be a fixed point. + let seeded = db + .search_layer_0_seeded(&mut store, &graph, query, full.clone(), ef) + .await?; + + assert_eq!(full.edge_ids(), seeded.edge_ids()); + assert_eq!(seeded.edge_ids()[query_idx], ids[query_idx]); + Ok(()) + } + + /// An empty seeded neighborhood should fall back to a regular search and + /// therefore return exactly the same result as `search`. + #[tokio::test] + async fn test_search_layer_0_seeded_empty_falls_back_to_search() -> Result<()> { + let db = HnswSearcher::new_with_test_parameters(); + let (mut store, graph, queries, ids) = build_seeded_test_graph(&db, 0, 256).await?; + + let query_idx = 0; + assert!(graph.layers[0].links.contains_key(&ids[query_idx])); + let query = &queries[query_idx]; + let ef = db.params.get_ef_search(1); + + let full = db.search(&mut store, &graph, query, ef).await?; + let seeded = db + .search_layer_0_seeded(&mut store, &graph, query, SortedNeighborhood::new(), ef) + .await?; + + assert_eq!(full.edge_ids(), seeded.edge_ids()); + assert_eq!(seeded.edge_ids()[query_idx], ids[query_idx]); + Ok(()) + } }