diff --git a/laurus/src/vector/index/segment/fanout.rs b/laurus/src/vector/index/segment/fanout.rs index b686c98b..9e44a851 100644 --- a/laurus/src/vector/index/segment/fanout.rs +++ b/laurus/src/vector/index/segment/fanout.rs @@ -325,9 +325,111 @@ impl SegmentFanoutSearcher { .iter() .any(|r| r.contains_vector(doc_id, field)) } -} -impl SegmentFanoutSearcher { + /// Probe one segment: run the expanding-refill loop (#883) and drop + /// newest-wins-shadowed hits (#880). + /// + /// Over-fetch per segment: containment masking drops shadowed hits + /// AFTER the per-segment top-k, so stale copies would otherwise + /// consume result slots (#880). Start at 2x for the common case; if + /// masking dropped this segment below the requested limit AND the pass + /// was truncated at the budget (i.e. live hits may sit below the cut + /// behind a band of masked stale copies), EXPAND the budget + /// geometrically and re-query until enough live hits survive or the + /// segment is exhausted (#883). A fixed multiplier cannot bound recall + /// against an arbitrarily deep stale band — the true nearest neighbour + /// could sit below any constant cut — so the budget doubles each round. + /// Bounded: top_k reaches the segment size in O(log n) queries, and + /// the common (no-pollution) case stops after one pass. + /// + /// After the refill loop the surviving hits are rescored on the + /// shared dequantized-f32 basis (Issue #927, see + /// [`Self::rescore_on_shared_basis`]) unless the segment's scores + /// already sit on the exact f32 rerank basis, and `min_similarity` + /// is re-applied to the final scores. + /// + /// Entirely self-contained per segment (the concrete searcher is built + /// here, masking reads only immutable newer readers), so `search` can + /// run one probe per worker thread (#926). + /// + /// # Arguments + /// + /// * `idx` - Index of the segment's reader (newest-first ordering). + /// * `request` - The original query; its `top_k` is replaced by the + /// expanding budget internally. + /// * `limit` - The caller-requested `top_k`. + /// * `metric` - The index's distance metric. + /// * `prepared_query` - The raw query prepared once per search + /// (shared read-only across worker threads). + /// + /// # Returns + /// + /// The segment's surviving (unmasked, comparably-scored) hits and + /// the number of candidates its probes examined. + fn probe_segment( + &self, + idx: usize, + request: &VectorIndexQuery, + limit: usize, + metric: DistanceMetric, + prepared_query: &crate::vector::core::distance::PreparedQuery<'_>, + ) -> Result<( + Vec, + usize, + )> { + let searcher = (self.make_searcher)(self.readers[idx].clone())?; + + let mut probe = request.clone(); + probe.params.top_k = limit.saturating_mul(2); + let mut candidates_examined = 0usize; + // Whether this segment's scores are already on the exact f32 + // rerank basis (Issue #481 Stage 2) — comparable across + // segments and MORE precise than the dequantized rescore, so + // the #927 rescore below must not overwrite them. Stamped by + // the per-segment searcher; constant across refill rounds + // (same segment, same sidecar). + let mut exact_basis = false; + let mut kept = loop { + let results = searcher.search(&probe)?; + candidates_examined += results.candidates_examined; + exact_basis |= results + .query_metadata + .get(crate::vector::search::searcher::SCORE_BASIS_METADATA_KEY) + .is_some_and(|v| v == crate::vector::search::searcher::SCORE_BASIS_F32_RERANK); + let returned = results.results.len(); + let kept: Vec<_> = results + .results + .into_iter() + .filter(|hit| !self.shadowed(idx, hit.doc_id, &hit.field_name)) + .collect(); + + // Enough live hits survived masking, or the segment returned + // fewer than requested (nothing deeper to fetch) — done. + if kept.len() >= limit || returned < probe.params.top_k { + break kept; + } + let next = probe.params.top_k.saturating_mul(2); + if next == probe.params.top_k { + break kept; // budget saturated (overflow guard) + } + probe.params.top_k = next; + }; + + // Issue #927: overwrite the segment-local scores with the + // shared-basis ones, then re-apply `min_similarity` on the + // final scores — the per-segment filter ran on local scores, + // which clamping can only inflate (clamped distance ≤ true + // distance), so it never dropped a hit the shared basis would + // keep; the inverse (an inflated score sneaking past the + // threshold) is corrected here. Skipped when the segment's + // scores already sit on the exact f32 rerank basis — a MORE + // precise shared basis the dequantized rescore would degrade. + if !exact_basis { + self.rescore_on_shared_basis(&mut kept, idx, metric, prepared_query)?; + kept.retain(|hit| hit.similarity >= request.params.min_similarity); + } + Ok((kept, candidates_examined)) + } /// Rescore `hits` against the raw query in the shared /// dequantized-f32 space (Issue #927). /// @@ -391,75 +493,38 @@ impl VectorIndexSearcher for SegmentFanoutSearcher { } // Issue #927: shared scoring basis for the cross-segment merge — - // see `rescore_on_shared_basis`. + // see `rescore_on_shared_basis`. Prepared once and shared + // read-only across the per-segment probes. let metric = self.readers[0].distance_metric(); let prepared_query = metric.prepare_query(&request.query.data); - // Over-fetch per segment: containment masking drops shadowed hits - // AFTER the per-segment top-k, so stale copies would otherwise - // consume result slots (#880). Start at 2x for the common case; if - // masking dropped this segment below the requested limit AND the pass - // was truncated at the budget (i.e. live hits may sit below the cut - // behind a band of masked stale copies), EXPAND the budget - // geometrically and re-query until enough live hits survive or the - // segment is exhausted (#883). A fixed multiplier cannot bound recall - // against an arbitrarily deep stale band — the true nearest neighbour - // could sit below any constant cut — so the budget doubles each round. - // Bounded: top_k reaches the segment size in O(log n) queries, and - // the common (no-pollution) case stops after one pass. - for (idx, reader) in self.readers.iter().enumerate() { - let searcher = (self.make_searcher)(reader.clone())?; - - let mut probe = request.clone(); - probe.params.top_k = limit.saturating_mul(2); - let mut kept: Vec; - // Whether this segment's scores are already on the exact f32 - // rerank basis (Issue #481 Stage 2) — comparable across - // segments and MORE precise than the dequantized rescore, so - // the #927 rescore below must not overwrite them. Stamped by - // the per-segment searcher; constant across refill rounds - // (same segment, same sidecar). - let mut exact_basis = false; - loop { - let results = searcher.search(&probe)?; - merged.candidates_examined += results.candidates_examined; - exact_basis |= results - .query_metadata - .get(crate::vector::search::searcher::SCORE_BASIS_METADATA_KEY) - .is_some_and(|v| v == crate::vector::search::searcher::SCORE_BASIS_F32_RERANK); - let returned = results.results.len(); - kept = results - .results - .into_iter() - .filter(|hit| !self.shadowed(idx, hit.doc_id, &hit.field_name)) - .collect(); - - // Enough live hits survived masking, or the segment returned - // fewer than requested (nothing deeper to fetch) — done. - if kept.len() >= limit || returned < probe.params.top_k { - break; - } - let next = probe.params.top_k.saturating_mul(2); - if next == probe.params.top_k { - break; // budget saturated (overflow guard) - } - probe.params.top_k = next; - } - - // Issue #927: overwrite the segment-local scores with the - // shared-basis ones, then re-apply `min_similarity` on the - // final scores — the per-segment filter ran on local scores, - // which clamping can only inflate (clamped distance ≤ true - // distance), so it never dropped a hit the shared basis would - // keep; the inverse (an inflated score sneaking past the - // threshold) is corrected here. Skipped when the segment's - // scores already sit on the exact f32 rerank basis — a MORE - // precise shared basis the dequantized rescore would degrade. - if !exact_basis { - self.rescore_on_shared_basis(&mut kept, idx, metric, &prepared_query)?; - kept.retain(|hit| hit.similarity >= request.params.min_similarity); - } + // One self-contained probe per segment (see `probe_segment`). On + // native targets with more than one segment the probes run + // concurrently (#926) — the indexed collect keeps per-segment + // results in newest-first segment order, so the merged vector (and + // therefore the sort below) sees byte-identical input to the serial + // path and the results are bit-for-bit the same. wasm32 keeps the + // serial loop (no threads), matching the crate-wide convention. + let probe_all_serial = || -> Result> { + (0..self.readers.len()) + .map(|idx| self.probe_segment(idx, request, limit, metric, &prepared_query)) + .collect() + }; + #[cfg(not(target_arch = "wasm32"))] + let per_segment = if self.readers.len() > 1 { + use rayon::prelude::*; + (0..self.readers.len()) + .into_par_iter() + .map(|idx| self.probe_segment(idx, request, limit, metric, &prepared_query)) + .collect::>>()? + } else { + probe_all_serial()? + }; + #[cfg(target_arch = "wasm32")] + let per_segment = probe_all_serial()?; + for (mut kept, candidates_examined) in per_segment { + merged.candidates_examined += candidates_examined; merged.results.append(&mut kept); } diff --git a/laurus/tests/segment_fanout_parallel_test.rs b/laurus/tests/segment_fanout_parallel_test.rs new file mode 100644 index 00000000..eae0b187 --- /dev/null +++ b/laurus/tests/segment_fanout_parallel_test.rs @@ -0,0 +1,119 @@ +//! Determinism test for the parallel multi-segment fan-out (Issue #926). +//! +//! `SegmentFanoutSearcher::search` runs its per-segment probes concurrently +//! on native targets. This test pins that repeated identical searches over +//! a multi-segment fixture (stale duplicates + a deletion) return +//! bit-identical results — parallel probing must not leak scheduling +//! nondeterminism into the merged order. Cross-segment score correctness +//! itself is pinned by `segment_score_comparability_test.rs` (#927), which +//! exercises the same fan-out (parallel on native) against brute force. + +use std::sync::Arc; + +use laurus::storage::Storage; +use laurus::storage::memory::{MemoryStorage, MemoryStorageConfig}; +use laurus::vector::index::VectorIndex; +use laurus::vector::index::config::HnswIndexConfig; +use laurus::vector::index::hnsw::segmented::SegmentedHnswIndex; +use laurus::vector::search::searcher::{VectorIndexQuery, VectorIndexQueryParams}; +use laurus::vector::{DistanceMetric, Vector}; + +const DIM: usize = 8; +const TOP_K: usize = 5; + +fn config() -> HnswIndexConfig { + HnswIndexConfig { + dimension: DIM, + m: 8, + ef_construction: 64, + normalize_vectors: false, + distance_metric: DistanceMetric::Euclidean, + segmented: true, + ..Default::default() + } +} + +/// A vector whose components all equal `level * 10.0` — grid spacing keeps +/// every pairwise distance far outside quantization error. +fn grid(level: f32) -> Vector { + Vector::new(vec![level * 10.0; DIM]) +} + +fn query(vector: &Vector, ef: usize) -> VectorIndexQuery { + VectorIndexQuery { + query: vector.clone(), + params: VectorIndexQueryParams { + top_k: TOP_K, + ef_search: Some(ef), + ..Default::default() + }, + field_name: Some("v".to_string()), + filter: None, + } +} + +/// Build a 4-segment index with cross-segment stale duplicates and one +/// deletion (the same shape as #927's regression fixture). +fn build_fixture() -> SegmentedHnswIndex { + let storage = Arc::new(MemoryStorage::new(MemoryStorageConfig::default())); + let index = + SegmentedHnswIndex::open_or_create(storage as Arc, "vi", config()).unwrap(); + + // 4 commits = 4 sealed segments (no merge trigger in this test). + // Segment 1: docs 0..10 at levels 0..10. + // Segment 2: docs 10..20 at levels 10..20. + // Segment 3: docs 20..30 at levels 20..30, PLUS doc 3 re-added at + // level 25.5 (stale copy of doc 3 remains in segment 1). + // Segment 4: doc 7 re-added at level 12.5 (stale copy in segment 1). + let batches: Vec> = vec![ + (0..10u64) + .map(|i| (i, "v".to_string(), grid(i as f32))) + .collect(), + (10..20u64) + .map(|i| (i, "v".to_string(), grid(i as f32))) + .collect(), + (20..30u64) + .map(|i| (i, "v".to_string(), grid(i as f32))) + .chain(std::iter::once((3u64, "v".to_string(), grid(25.5)))) + .collect(), + vec![(7u64, "v".to_string(), grid(12.5))], + ]; + for batch in batches { + let mut w = index.writer().unwrap(); + w.add_vectors(batch).unwrap(); + w.commit().unwrap(); + } + + // Delete doc 15 (level 15) — must never appear in results. + index.soft_delete_document(15).unwrap(); + + index +} + +/// Repeated identical searches must return identical result lists — +/// guards against nondeterministic merge order under parallel probing. +#[test] +fn parallel_fanout_is_deterministic_across_runs() { + let index = build_fixture(); + let searcher = index.searcher().unwrap(); + + let reference: Vec<(u64, u32)> = searcher + .search(&query(&grid(13.7), 256)) + .unwrap() + .results + .iter() + .map(|r| (r.doc_id, r.similarity.to_bits())) + .collect(); + assert!(!reference.is_empty()); + + for run in 0..5 { + let again: Vec<(u64, u32)> = searcher + .search(&query(&grid(13.7), 256)) + .unwrap() + .results + .iter() + .map(|r| (r.doc_id, r.similarity.to_bits())) + .collect(); + assert_eq!(reference, again, "run {run}: results must be bit-identical"); + } +}