perf(vector/search): parallelize SegmentFanoutSearcher's per-segment fan-out (#926)

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<crate::vector::search::searcher::VectorIndexQueryResult>,
+        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<crate::vector::search::searcher::VectorIndexQueryResult>;
-            // 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<Vec<_>> {
+            (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::<Result<Vec<_>>>()?
+        } 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);
         }
 

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<dyn Storage>, "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<(u64, String, Vector)>> = 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");
+    }
+}