Skip to main content

nodedb_vector/collection/
search.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! VectorCollection search: multi-segment merging with SQ8 reranking. A
4//! trained IVF-PQ index answers beside the segments under global ids.
5//!
6//! The `search_with_payload_filter` method wires payload bitmap pre-filtering
7//! into the search path. When all referenced fields in the predicate are
8//! indexed, the bitmap is built and passed to `search_with_bitmap_bytes`.
9//! When any field is un-indexed, the search falls back to the full unfiltered
10//! path and lets the caller apply post-filtering — the un-indexed predicate is
11//! never silently dropped.
12
13use crate::distance::{DistanceMetric, distance};
14use crate::error::{VectorError, check_dim};
15use crate::hnsw::SearchResult;
16use crate::hnsw::search::decode_filter_bitmap;
17
18use super::lifecycle::VectorCollection;
19use super::payload_index::FilterPredicate;
20use super::segment::SealedSegment;
21
22/// Score a single candidate via the SQ8 codec, using the metric-appropriate
23/// asymmetric distance.
24#[inline]
25fn sq8_score(
26    codec: &crate::quantize::sq8::Sq8Codec,
27    query: &[f32],
28    encoded: &[u8],
29    metric: DistanceMetric,
30) -> f32 {
31    match metric {
32        DistanceMetric::Cosine => codec.asymmetric_cosine(query, encoded),
33        DistanceMetric::InnerProduct => codec.asymmetric_ip(query, encoded),
34        // L2 (and all other metrics that don't have a specialized asymmetric
35        // form yet) fall back to squared L2 — correct for ordering when the
36        // metric is L2 and a reasonable proxy otherwise since we rerank with
37        // exact FP32 below.
38        _ => codec.asymmetric_l2(query, encoded),
39    }
40}
41
42/// Candidate-generation + rerank for a sealed segment that has a quantized
43/// codec attached. Generates a widened candidate pool via HNSW, re-scores
44/// candidates using the quantized codec (this is where SQ8/PQ actually pay
45/// off — the FP32 vectors need not be resident), and reranks the top
46/// `top_k` via exact FP32 distance from mmap or index storage.
47fn quantized_search(
48    seg: &SealedSegment,
49    query: &[f32],
50    top_k: usize,
51    ef: usize,
52    metric: DistanceMetric,
53) -> Result<Vec<SearchResult>, VectorError> {
54    let rerank_k = top_k.saturating_mul(3).max(20);
55    let hnsw_candidates = seg.index.search(query, rerank_k, ef)?;
56
57    // Phase 1: rank candidates by quantized distance.
58    let mut scored: Vec<(u32, f32)> = if let Some((codec, codes)) = &seg.pq {
59        let table = codec.build_distance_table(query)?;
60        let m = codec.m;
61        hnsw_candidates
62            .into_iter()
63            .filter_map(|r| {
64                let start = (r.id as usize).checked_mul(m)?;
65                let end = start.checked_add(m)?;
66                let slice = codes.get(start..end)?;
67                Some((r.id, codec.asymmetric_distance(&table, slice)))
68            })
69            .collect()
70    } else if let Some((codec, data)) = &seg.sq8 {
71        let dim = codec.dim();
72        hnsw_candidates
73            .into_iter()
74            .filter_map(|r| {
75                let start = (r.id as usize).checked_mul(dim)?;
76                let end = start.checked_add(dim)?;
77                let slice = data.get(start..end)?;
78                Some((r.id, sq8_score(codec, query, slice, metric)))
79            })
80            .collect()
81    } else {
82        hnsw_candidates
83            .into_iter()
84            .map(|r| (r.id, r.distance))
85            .collect()
86    };
87    scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
88
89    // Keep only the most promising candidates for FP32 rerank.
90    let keep = rerank_k.min(scored.len());
91    scored.truncate(keep);
92
93    // Prefetch FP32 vectors for reranking.
94    if let Some(mmap) = &seg.mmap_vectors {
95        let ids: Vec<u32> = scored.iter().map(|&(id, _)| id).collect();
96        mmap.prefetch_batch(&ids);
97    }
98
99    // Phase 2: rerank with exact FP32.
100    let mut reranked: Vec<SearchResult> = scored
101        .into_iter()
102        .filter_map(|(id, _)| {
103            let v = if let Some(mmap) = &seg.mmap_vectors {
104                mmap.get_vector(id)?
105            } else {
106                seg.index.get_vector(id)?
107            };
108            Some(SearchResult {
109                id,
110                distance: distance(query, v, metric),
111            })
112        })
113        .collect();
114    reranked.sort_by(|a, b| {
115        a.distance
116            .partial_cmp(&b.distance)
117            .unwrap_or(std::cmp::Ordering::Equal)
118    });
119    reranked.truncate(top_k);
120    Ok(reranked)
121}
122
123/// Search one sealed segment: through its quantized codec when it has one,
124/// else its HNSW graph. A codec pass that exceeds the memory budget falls
125/// back to the HNSW graph, which answers the same query from FP32 vectors.
126/// Every other error fails the search.
127fn search_sealed(
128    seg: &SealedSegment,
129    query: &[f32],
130    top_k: usize,
131    ef: usize,
132    metric: DistanceMetric,
133) -> Result<Vec<SearchResult>, VectorError> {
134    if seg.pq.is_none() && seg.sq8.is_none() {
135        return seg.index.search(query, top_k, ef);
136    }
137    match quantized_search(seg, query, top_k, ef, metric) {
138        Ok(results) => Ok(results),
139        Err(VectorError::BudgetExhausted(e)) => {
140            tracing::warn!(error = %e, "quantized search over budget; searching the HNSW graph");
141            seg.index.search(query, top_k, ef)
142        }
143        Err(e) => Err(e),
144    }
145}
146
147/// Shift segment-local result ids to global ids and append them.
148fn push_shifted(all: &mut Vec<SearchResult>, results: Vec<SearchResult>, base_id: u32) {
149    all.extend(results.into_iter().map(|mut r| {
150        r.id += base_id;
151        r
152    }));
153}
154
155/// Order merged results by distance and keep the `top_k` nearest.
156fn finish(mut all: Vec<SearchResult>, top_k: usize) -> Vec<SearchResult> {
157    all.sort_by(|a, b| {
158        a.distance
159            .partial_cmp(&b.distance)
160            .unwrap_or(std::cmp::Ordering::Equal)
161    });
162    all.truncate(top_k);
163    all
164}
165
166impl VectorCollection {
167    /// Search across all segments, merging results by distance.
168    ///
169    /// A query without the collection dimension fails with
170    /// [`VectorError::DimensionMismatch`] before any segment is read.
171    pub fn search(
172        &self,
173        query: &[f32],
174        top_k: usize,
175        ef: usize,
176    ) -> Result<Vec<SearchResult>, VectorError> {
177        check_dim(self.dim, query.len())?;
178        let mut all: Vec<SearchResult> = Vec::new();
179
180        // Codec-dispatch fast path: a collection-level HnswCodecIndex (RaBitQ
181        // or BBQ) answers for the sealed segments; the growing and building
182        // segments are read by brute force beside it.
183        if let Some(ref dispatch) = self.codec_dispatch {
184            all.extend(
185                dispatch
186                    .search(query, top_k, ef)?
187                    .into_iter()
188                    .map(|r| SearchResult {
189                        id: r.id,
190                        distance: r.distance,
191                    }),
192            );
193        } else {
194            for seg in &self.sealed {
195                let results = search_sealed(seg, query, top_k, ef, self.params.metric)?;
196                push_shifted(&mut all, results, seg.base_id);
197            }
198        }
199        if let Some(ivf) = &self.ivf {
200            all.extend(ivf.search(query, top_k)?);
201        }
202
203        push_shifted(
204            &mut all,
205            self.growing.search(query, top_k)?,
206            self.growing_base_id,
207        );
208        for seg in &self.building {
209            push_shifted(&mut all, seg.flat.search(query, top_k)?, seg.base_id);
210        }
211        Ok(finish(all, top_k))
212    }
213
214    /// Search across all segments using an explicit metric override.
215    ///
216    /// For sealed segments with quantized codecs, the metric override is applied
217    /// during candidate reranking. Growing and building segments apply it exactly
218    /// via brute-force. The HNSW graph structure was built with the collection
219    /// metric; using a different metric affects the scoring but not graph traversal.
220    /// A codec-dispatch index scores with the collection metric.
221    pub fn search_with_metric(
222        &self,
223        query: &[f32],
224        top_k: usize,
225        ef: usize,
226        metric: DistanceMetric,
227    ) -> Result<Vec<SearchResult>, VectorError> {
228        check_dim(self.dim, query.len())?;
229        let mut all: Vec<SearchResult> = Vec::new();
230
231        if let Some(ref dispatch) = self.codec_dispatch {
232            all.extend(
233                dispatch
234                    .search(query, top_k, ef)?
235                    .into_iter()
236                    .map(|r| SearchResult {
237                        id: r.id,
238                        distance: r.distance,
239                    }),
240            );
241        } else {
242            for seg in &self.sealed {
243                let results = search_sealed(seg, query, top_k, ef, metric)?;
244                push_shifted(&mut all, results, seg.base_id);
245            }
246        }
247        if let Some(ivf) = &self.ivf {
248            all.extend(ivf.search_with(query, top_k, metric, None)?);
249        }
250
251        push_shifted(
252            &mut all,
253            self.growing.search_with_metric(query, top_k, metric)?,
254            self.growing_base_id,
255        );
256        for seg in &self.building {
257            push_shifted(
258                &mut all,
259                seg.flat.search_with_metric(query, top_k, metric)?,
260                seg.base_id,
261            );
262        }
263        Ok(finish(all, top_k))
264    }
265
266    /// Search with a pre-filter bitmap (byte-array format) and explicit metric override.
267    ///
268    /// Bitmap bytes that do not decode fail with
269    /// [`VectorError::InvalidFilterBitmap`].
270    pub fn search_with_bitmap_bytes_and_metric(
271        &self,
272        query: &[f32],
273        top_k: usize,
274        ef: usize,
275        bitmap: &[u8],
276        metric: DistanceMetric,
277    ) -> Result<Vec<SearchResult>, VectorError> {
278        check_dim(self.dim, query.len())?;
279        let mut all: Vec<SearchResult> = Vec::new();
280
281        if let Some(ivf) = &self.ivf {
282            let filter = decode_filter_bitmap(bitmap)?;
283            all.extend(ivf.search_with(query, top_k, metric, Some(&filter))?);
284        }
285        push_shifted(
286            &mut all,
287            self.growing.search_filtered_offset_with_metric(
288                query,
289                top_k,
290                bitmap,
291                self.growing_base_id,
292                metric,
293            )?,
294            self.growing_base_id,
295        );
296
297        for seg in &self.sealed {
298            let mut results =
299                seg.index
300                    .search_with_bitmap_bytes_offset(query, top_k, ef, bitmap, seg.base_id)?;
301            // Rerank with the requested metric using the stored FP32 vector.
302            for r in &mut results {
303                if let Some(v) = seg.index.get_vector(r.id) {
304                    r.distance = distance(query, v, metric);
305                }
306            }
307            push_shifted(&mut all, results, seg.base_id);
308        }
309
310        for seg in &self.building {
311            push_shifted(
312                &mut all,
313                seg.flat.search_filtered_offset_with_metric(
314                    query,
315                    top_k,
316                    bitmap,
317                    seg.base_id,
318                    metric,
319                )?,
320                seg.base_id,
321            );
322        }
323        Ok(finish(all, top_k))
324    }
325
326    /// Search with a pre-filter bitmap (byte-array format).
327    ///
328    /// Bitmap bytes that do not decode fail with
329    /// [`VectorError::InvalidFilterBitmap`].
330    pub fn search_with_bitmap_bytes(
331        &self,
332        query: &[f32],
333        top_k: usize,
334        ef: usize,
335        bitmap: &[u8],
336    ) -> Result<Vec<SearchResult>, VectorError> {
337        check_dim(self.dim, query.len())?;
338        let mut all: Vec<SearchResult> = Vec::new();
339
340        if let Some(ivf) = &self.ivf {
341            let filter = decode_filter_bitmap(bitmap)?;
342            all.extend(ivf.search_with(query, top_k, self.params.metric, Some(&filter))?);
343        }
344        push_shifted(
345            &mut all,
346            self.growing
347                .search_filtered_offset(query, top_k, bitmap, self.growing_base_id)?,
348            self.growing_base_id,
349        );
350        for seg in &self.sealed {
351            push_shifted(
352                &mut all,
353                seg.index
354                    .search_with_bitmap_bytes_offset(query, top_k, ef, bitmap, seg.base_id)?,
355                seg.base_id,
356            );
357        }
358        for seg in &self.building {
359            push_shifted(
360                &mut all,
361                seg.flat
362                    .search_filtered_offset(query, top_k, bitmap, seg.base_id)?,
363                seg.base_id,
364            );
365        }
366        Ok(finish(all, top_k))
367    }
368
369    /// Search with a structured payload predicate.
370    ///
371    /// If `predicate` is fully covered by indexed fields (all leaf fields have
372    /// a bitmap index), the bitmap is built and HNSW traversal uses it as a
373    /// pre-filter.
374    ///
375    /// If any field in `predicate` is un-indexed, the method returns
376    /// `(results, false)` where `false` signals that the predicate was NOT
377    /// applied and the caller must apply it as a post-filter. This guarantees
378    /// the un-indexed predicate is never silently dropped.
379    ///
380    /// Returns `(results, filter_was_applied)`.
381    pub fn search_with_payload_filter(
382        &self,
383        query: &[f32],
384        top_k: usize,
385        ef: usize,
386        predicate: &FilterPredicate,
387    ) -> Result<(Vec<SearchResult>, bool), VectorError> {
388        match self.payload.pre_filter(predicate) {
389            Some(bm) => {
390                // Serialize the bitmap to the byte format expected by
391                // `search_with_bitmap_bytes`.
392                let mut bm_bytes = Vec::new();
393                if bm.serialize_into(&mut bm_bytes).is_ok() {
394                    let results = self.search_with_bitmap_bytes(query, top_k, ef, &bm_bytes)?;
395                    Ok((results, true))
396                } else {
397                    // Serialization failure: unfiltered search, and the
398                    // caller applies the predicate as a post-filter.
399                    Ok((self.search(query, top_k, ef)?, false))
400                }
401            }
402            None => {
403                // Un-indexed field present: full scan, caller must post-filter.
404                Ok((self.search(query, top_k, ef)?, false))
405            }
406        }
407    }
408}
409
410#[cfg(test)]
411mod tests {
412    use crate::collection::lifecycle::VectorCollection;
413    use crate::collection::segment::DEFAULT_SEAL_THRESHOLD;
414    use crate::distance::DistanceMetric;
415    use crate::hnsw::{HnswIndex, HnswParams};
416    use crate::test_support::test_memory;
417
418    fn make_collection() -> VectorCollection {
419        VectorCollection::new(
420            3,
421            HnswParams {
422                metric: DistanceMetric::L2,
423                ..HnswParams::default()
424            },
425        )
426    }
427
428    #[test]
429    fn insert_and_search() {
430        let mut coll = make_collection();
431        for i in 0..100u32 {
432            coll.insert(vec![i as f32, 0.0, 0.0]).unwrap();
433        }
434        assert_eq!(coll.len(), 100);
435        let results = coll.search(&[50.0, 0.0, 0.0], 3, 64).unwrap();
436        assert_eq!(results.len(), 3);
437        assert_eq!(results[0].id, 50);
438    }
439
440    #[test]
441    fn seal_moves_to_building() {
442        let mut coll = VectorCollection::new(2, HnswParams::default());
443        for i in 0..DEFAULT_SEAL_THRESHOLD {
444            coll.insert(vec![i as f32, 0.0]).unwrap();
445        }
446        assert!(coll.needs_seal());
447
448        let req = coll.seal("test_key").unwrap();
449        assert_eq!(req.vectors.len(), DEFAULT_SEAL_THRESHOLD);
450        assert_eq!(coll.building.len(), 1);
451        assert_eq!(coll.growing.len(), 0);
452
453        let results = coll.search(&[100.0, 0.0], 1, 64).unwrap();
454        assert!(!results.is_empty());
455    }
456
457    #[test]
458    fn complete_build_promotes_to_sealed() {
459        let mut coll = VectorCollection::new(2, HnswParams::default());
460        for i in 0..100 {
461            coll.insert(vec![i as f32, 0.0]).unwrap();
462        }
463        let req = coll.seal("test").unwrap();
464
465        let mut index = HnswIndex::new(req.dim, req.params);
466        for v in &req.vectors {
467            index.insert(v.clone()).unwrap();
468        }
469        coll.complete_build(req.segment_id, index, test_memory());
470
471        assert_eq!(coll.building.len(), 0);
472        assert_eq!(coll.sealed.len(), 1);
473
474        let results = coll.search(&[50.0, 0.0], 3, 64).unwrap();
475        assert!(!results.is_empty());
476    }
477
478    #[test]
479    fn multi_segment_search_merges() {
480        let mut coll = VectorCollection::new(
481            2,
482            HnswParams {
483                metric: DistanceMetric::L2,
484                ..HnswParams::default()
485            },
486        );
487
488        for i in 0..100 {
489            coll.insert(vec![i as f32, 0.0]).unwrap();
490        }
491        let req = coll.seal("test").unwrap();
492        let mut idx = HnswIndex::new(2, req.params);
493        for v in &req.vectors {
494            idx.insert(v.clone()).unwrap();
495        }
496        coll.complete_build(req.segment_id, idx, test_memory());
497
498        for i in 100..200 {
499            coll.insert(vec![i as f32, 0.0]).unwrap();
500        }
501
502        let results = coll.search(&[150.0, 0.0], 3, 64).unwrap();
503        assert_eq!(results.len(), 3);
504        assert_eq!(results[0].id, 150);
505    }
506
507    #[test]
508    fn delete_across_segments() {
509        let mut coll = VectorCollection::new(2, HnswParams::default());
510        for i in 0..10 {
511            coll.insert(vec![i as f32, 0.0]).unwrap();
512        }
513        assert!(coll.delete(5));
514        assert_eq!(coll.live_count(), 9);
515
516        let results = coll.search(&[5.0, 0.0], 10, 64).unwrap();
517        assert!(results.iter().all(|r| r.id != 5));
518    }
519
520    /// Build a sealed HNSW segment from `n` vectors of `dim=2`, where vector `i`
521    /// is `[i as f32, 0.0]`. Returns the collection with one sealed segment.
522    fn make_sealed_collection(n: usize) -> VectorCollection {
523        let mut coll = VectorCollection::new(
524            2,
525            HnswParams {
526                metric: DistanceMetric::L2,
527                ..HnswParams::default()
528            },
529        );
530        for i in 0..n {
531            coll.insert(vec![i as f32, 0.0]).unwrap();
532        }
533        let req = coll.seal("seg").unwrap();
534        let mut idx = HnswIndex::new(req.dim, req.params);
535        for v in &req.vectors {
536            idx.insert(v.clone()).unwrap();
537        }
538        coll.complete_build(req.segment_id, idx, test_memory());
539        coll
540    }
541
542    /// Attach SQ8 quantization to the first sealed segment of `coll`.
543    fn attach_sq8(coll: &mut VectorCollection) {
544        use crate::quantize::sq8::Sq8Codec;
545
546        let sealed = &mut coll.sealed[0];
547        let dim = sealed.index.dim();
548        let n = sealed.index.len();
549        let vecs: Vec<Vec<f32>> = (0..n)
550            .filter_map(|i| sealed.index.get_vector(i as u32).map(|v| v.to_vec()))
551            .collect();
552        let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect();
553        let codec = Sq8Codec::calibrate(&refs, dim).unwrap();
554        let sq8_data: Vec<u8> = vecs.iter().flat_map(|v| codec.quantize(v)).collect();
555        sealed.sq8 = Some((codec, sq8_data));
556    }
557
558    #[test]
559    fn sq8_search_returns_correct_nearest_neighbor() {
560        let mut coll = make_sealed_collection(200);
561        attach_sq8(&mut coll);
562
563        let results = coll.search(&[100.0, 0.0], 5, 64).unwrap();
564        assert!(!results.is_empty(), "expected non-empty results");
565        assert_eq!(
566            results[0].id, 100,
567            "nearest neighbor of [100,0] should be id=100, got id={}",
568            results[0].id
569        );
570    }
571
572    #[test]
573    fn sq8_search_recall_matches_hnsw() {
574        // Build two identical collections — one without SQ8, one with.
575        let coll_plain = make_sealed_collection(500);
576        let mut coll_sq8 = make_sealed_collection(500);
577        attach_sq8(&mut coll_sq8);
578
579        let query = [250.0f32, 0.0];
580        let top_k = 5;
581
582        let plain_results = coll_plain.search(&query, top_k, 64).unwrap();
583        let sq8_results = coll_sq8.search(&query, top_k, 64).unwrap();
584
585        let plain_ids: std::collections::HashSet<u32> =
586            plain_results.iter().map(|r| r.id).collect();
587        let sq8_ids: std::collections::HashSet<u32> = sq8_results.iter().map(|r| r.id).collect();
588
589        let overlap = plain_ids.intersection(&sq8_ids).count();
590        assert!(
591            overlap >= 4,
592            "SQ8 recall too low: {overlap}/5 results matched plain HNSW (need >=4)"
593        );
594    }
595
596    /// A collection whose first 50 vectors (`[i, 0, 0, 0]`) sit in one
597    /// sealed segment, with an empty growing segment.
598    fn sealed_dim4_collection() -> VectorCollection {
599        let mut coll = VectorCollection::new(
600            4,
601            HnswParams {
602                metric: DistanceMetric::L2,
603                m: 8,
604                ef_construction: 50,
605                ..HnswParams::default()
606            },
607        );
608        for i in 0u32..50 {
609            coll.insert(vec![i as f32, 0.0, 0.0, 0.0]).unwrap();
610        }
611        let req = coll.seal("codec").unwrap();
612        let mut idx = HnswIndex::new(req.dim, req.params);
613        for v in &req.vectors {
614            idx.insert(v.clone()).unwrap();
615        }
616        coll.complete_build(req.segment_id, idx, test_memory());
617        coll
618    }
619
620    #[test]
621    fn codec_dispatch_bbq_search_returns_results_and_stats_report_bbq() {
622        let mut coll = sealed_dim4_collection();
623
624        // Build the collection-level BBQ dispatch index over the sealed vectors.
625        let dispatch = coll.build_codec_dispatch("bbq").unwrap();
626        assert!(
627            dispatch.is_some(),
628            "build_codec_dispatch(bbq) should return Some"
629        );
630
631        // Query near id=25.
632        let query = [25.0f32, 0.0, 0.0, 0.0];
633        let results = coll.search(&query, 5, 32).unwrap();
634        assert!(
635            !results.is_empty(),
636            "BBQ codec-dispatch search should return results"
637        );
638
639        // Stats should report Bbq quantization.
640        let stats = coll.stats();
641        assert_eq!(
642            stats.quantization,
643            nodedb_types::VectorIndexQuantization::Bbq,
644            "stats quantization should be Bbq after build_codec_dispatch(bbq)"
645        );
646    }
647
648    #[test]
649    fn codec_dispatch_rabitq_search_non_empty() {
650        let mut coll = sealed_dim4_collection();
651        coll.build_codec_dispatch("rabitq").unwrap().unwrap();
652
653        let results = coll.search(&[10.0, 0.0, 0.0, 0.0], 3, 32).unwrap();
654        assert!(
655            !results.is_empty(),
656            "RaBitQ dispatch search should return results"
657        );
658
659        let stats = coll.stats();
660        assert_eq!(
661            stats.quantization,
662            nodedb_types::VectorIndexQuantization::RaBitQ
663        );
664    }
665
666    /// The codec index covers the sealed segments under their global ids;
667    /// the growing segment is read beside it. A growing vector is found under
668    /// its own id, once.
669    #[test]
670    fn codec_dispatch_keeps_global_ids_and_reads_growing_once() {
671        let mut coll = sealed_dim4_collection();
672        coll.build_codec_dispatch("bbq").unwrap().unwrap();
673        let far = coll.insert(vec![1000.0, 0.0, 0.0, 0.0]).unwrap();
674        assert_eq!(far, 50, "the first growing vector takes the next global id");
675
676        let results = coll.search(&[1000.0, 0.0, 0.0, 0.0], 3, 32).unwrap();
677        assert_eq!(results[0].id, far);
678        assert_eq!(
679            results.iter().filter(|r| r.id == far).count(),
680            1,
681            "{results:?}"
682        );
683        assert!(results.iter().all(|r| r.id <= far), "{results:?}");
684    }
685
686    /// Every collection search entry point refuses a query of the wrong
687    /// dimension, on each segment kind and on the codec-dispatch path.
688    #[test]
689    fn wrong_dimension_query_is_a_typed_error() {
690        use crate::error::VectorError;
691        let mut coll = sealed_dim4_collection();
692        coll.insert(vec![1.0, 0.0, 0.0, 0.0]).unwrap();
693        let bytes = {
694            let bm: roaring::RoaringBitmap = (0..51u32).collect();
695            let mut out = Vec::new();
696            bm.serialize_into(&mut out).unwrap();
697            out
698        };
699        let short = [1.0_f32, 0.0];
700        let check = |coll: &VectorCollection| {
701            for result in [
702                coll.search(&short, 3, 32),
703                coll.search_with_metric(&short, 3, 32, DistanceMetric::Cosine),
704                coll.search_with_bitmap_bytes(&short, 3, 32, &bytes),
705                coll.search_with_bitmap_bytes_and_metric(&short, 3, 32, &bytes, DistanceMetric::L2),
706            ] {
707                assert!(
708                    matches!(
709                        result,
710                        Err(VectorError::DimensionMismatch {
711                            expected: 4,
712                            got: 2
713                        })
714                    ),
715                    "{result:?}"
716                );
717            }
718        };
719        check(&coll);
720        coll.build_codec_dispatch("rabitq").unwrap().unwrap();
721        check(&coll);
722        assert!(matches!(
723            coll.insert(vec![1.0; 3]),
724            Err(VectorError::DimensionMismatch {
725                expected: 4,
726                got: 3
727            })
728        ));
729    }
730
731    #[test]
732    fn sq8_search_does_not_scan_all_vectors() {
733        // This test validates correctness of the SQ8 search path for a large
734        // segment. The bug being guarded against is an O(N) linear scan instead
735        // of graph-guided traversal: the fix must use HNSW with SQ8 as the
736        // distance function. Correctness (correct nearest neighbor) is the
737        // invariant that must be preserved when the implementation changes.
738        let mut coll = make_sealed_collection(2000);
739        attach_sq8(&mut coll);
740
741        let results = coll.search(&[1000.0, 0.0], 5, 64).unwrap();
742        assert!(!results.is_empty(), "expected non-empty results");
743        assert_eq!(
744            results[0].id, 1000,
745            "nearest neighbor of [1000,0] should be id=1000, got id={}",
746            results[0].id
747        );
748    }
749}