Candidate-generation + rerank for a sealed segment that has a quantized codec attached. Generates a widened candidate pool via HNSW, re-scores candidates using the quantized codec (this is where SQ8/PQ actually pay off — the FP32 vectors need not be resident), and reranks the top `top_k` via exact FP32 distance from mmap or index storage.
(
seg: &SealedSegment,
query: &[f32],
top_k: usize,
ef: usize,
metric: DistanceMetric,
)
| 43 | /// off — the FP32 vectors need not be resident), and reranks the top |
| 44 | /// `top_k` via exact FP32 distance from mmap or index storage. |
| 45 | fn quantized_search( |
| 46 | seg: &SealedSegment, |
| 47 | query: &[f32], |
| 48 | top_k: usize, |
| 49 | ef: usize, |
| 50 | metric: DistanceMetric, |
| 51 | ) -> Result<Vec<SearchResult>, VectorError> { |
| 52 | let rerank_k = top_k.saturating_mul(3).max(20); |
| 53 | let hnsw_candidates = seg.index.search(query, rerank_k, ef); |
| 54 | |
| 55 | // Phase 1: rank candidates by quantized distance. |
| 56 | let mut scored: Vec<(u32, f32)> = if let Some((codec, codes)) = &seg.pq { |
| 57 | let table = codec.build_distance_table(query)?; |
| 58 | let m = codec.m; |
| 59 | hnsw_candidates |
| 60 | .into_iter() |
| 61 | .filter_map(|r| { |
| 62 | let start = (r.id as usize).checked_mul(m)?; |
| 63 | let end = start.checked_add(m)?; |
| 64 | let slice = codes.get(start..end)?; |
| 65 | Some((r.id, codec.asymmetric_distance(&table, slice))) |
| 66 | }) |
| 67 | .collect() |
| 68 | } else if let Some((codec, data)) = &seg.sq8 { |
| 69 | let dim = codec.dim(); |
| 70 | hnsw_candidates |
| 71 | .into_iter() |
| 72 | .filter_map(|r| { |
| 73 | let start = (r.id as usize).checked_mul(dim)?; |
| 74 | let end = start.checked_add(dim)?; |
| 75 | let slice = data.get(start..end)?; |
| 76 | Some((r.id, sq8_score(codec, query, slice, metric))) |
| 77 | }) |
| 78 | .collect() |
| 79 | } else { |
| 80 | hnsw_candidates |
| 81 | .into_iter() |
| 82 | .map(|r| (r.id, r.distance)) |
| 83 | .collect() |
| 84 | }; |
| 85 | scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); |
| 86 | |
| 87 | // Keep only the most promising candidates for FP32 rerank. |
| 88 | let keep = rerank_k.min(scored.len()); |
| 89 | scored.truncate(keep); |
| 90 | |
| 91 | // Prefetch FP32 vectors for reranking. |
| 92 | if let Some(mmap) = &seg.mmap_vectors { |
| 93 | let ids: Vec<u32> = scored.iter().map(|&(id, _)| id).collect(); |
| 94 | mmap.prefetch_batch(&ids); |
| 95 | } |
| 96 | |
| 97 | // Phase 2: rerank with exact FP32. |
| 98 | let mut reranked: Vec<SearchResult> = scored |
| 99 | .into_iter() |
| 100 | .filter_map(|(id, _)| { |
| 101 | let v = if let Some(mmap) = &seg.mmap_vectors { |
| 102 | mmap.get_vector(id)? |
no test coverage detected