1use 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#[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 _ => codec.asymmetric_l2(query, encoded),
39 }
40}
41
42fn 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 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 let keep = rerank_k.min(scored.len());
91 scored.truncate(keep);
92
93 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 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
123fn 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
147fn 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
155fn 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 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 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 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 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 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 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 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 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 Ok((self.search(query, top_k, ef)?, false))
400 }
401 }
402 None => {
403 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 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 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 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 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 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 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 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 #[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 #[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 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}