From 353a1f17e5b85c87ba2417ee9026ba2c84829a7e Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Mon, 27 Jul 2026 03:30:04 -0700 Subject: [PATCH 1/7] Optimize filtered IVF-PQ batch scanning --- core/src/ivfpq.rs | 468 ++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 412 insertions(+), 56 deletions(-) diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index 7237d93..3975dc3 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -374,6 +374,25 @@ impl IVFPQIndex { let use_precomputed = !self.precomputed_table.is_empty(); let use_fastscan = !self.fastscan_codes.is_empty() && self.pq.nbits == 4; + let matching_positions_by_list = filter.map(|filter| { + let mut probed_lists = vec![false; self.nlist]; + for probe_indices in &all_probe_indices { + for &list_id in probe_indices { + probed_lists[list_id] = true; + } + } + self.ids + .iter() + .zip(probed_lists) + .map(|(ids, probed)| { + if probed { + matching_positions(ids, Some(filter)).unwrap() + } else { + Vec::new() + } + }) + .collect::>() + }); let results: Vec> = (0..nq) .into_par_iter() @@ -398,6 +417,9 @@ impl IVFPQIndex { if count == 0 { continue; } + let matching_positions = matching_positions_by_list + .as_ref() + .map(|positions| positions[list_id].as_slice()); // Precomputed sim_table omits ||q-c||²; add it as dis0. // Non-precomputed path computes from residual_query, already full distance. @@ -419,7 +441,11 @@ impl IVFPQIndex { self.compute_list_table(query, list_id, &mut sim_table); } - if use_fastscan { + if use_fastscan + && !matching_positions + .map(|positions| should_scan_matching_positions(count, positions)) + .unwrap_or(false) + { let mut dists = vec![0.0f32; count]; crate::fastscan::fastscan_4bit( &sim_table, @@ -428,13 +454,14 @@ impl IVFPQIndex { m, &mut dists, ); - for i in 0..count { - if let Some(f) = filter { - if !f.contains(self.ids[list_id][i]) { - continue; - } + if let Some(positions) = matching_positions { + for &position in positions { + heap.push(dis0 + dists[position], self.ids[list_id][position]); + } + } else { + for i in 0..count { + heap.push(dis0 + dists[i], self.ids[list_id][i]); } - heap.push(dis0 + dists[i], self.ids[list_id][i]); } } else if self.pq.nbits == 4 { scan_codes_4bit( @@ -445,7 +472,7 @@ impl IVFPQIndex { m, ksub, dis0, - filter, + matching_positions, &mut heap, ); } else { @@ -457,7 +484,7 @@ impl IVFPQIndex { m, ksub, dis0, - filter, + matching_positions, &mut heap, ); } @@ -747,6 +774,19 @@ fn invalid_merge_input(message: impl Into) -> io::Error { io::Error::new(io::ErrorKind::InvalidInput, message.into()) } +fn matching_positions(ids: &[i64], filter: Option<&dyn RowIdFilter>) -> Option> { + filter.map(|filter| { + ids.iter() + .enumerate() + .filter_map(|(position, &id)| filter.contains(id).then_some(position)) + .collect() + }) +} + +fn should_scan_matching_positions(count: usize, matching_positions: &[usize]) -> bool { + matching_positions.len() <= count / 2 +} + /// Scan 4-bit packed codes using u8-domain accumulation. fn scan_codes_4bit( sim_table: &[f32], @@ -756,19 +796,37 @@ fn scan_codes_4bit( m: usize, _ksub: usize, dis0: f32, - filter: Option<&dyn RowIdFilter>, + matching_positions: Option<&[usize]>, heap: &mut TopKHeap, ) { + if let Some(positions) = + matching_positions.filter(|positions| should_scan_matching_positions(count, positions)) + { + let code_size = m / 2; + for &position in positions { + let mut distance = dis0; + let code_offset = position * code_size; + for pair in 0..code_size { + let code = codes[code_offset + pair]; + distance += sim_table[(pair * 2) * 16 + (code & 0x0f) as usize]; + distance += sim_table[(pair * 2 + 1) * 16 + (code >> 4) as usize]; + } + heap.push(distance, ids[position]); + } + return; + } + let mut dists = vec![0.0f32; count]; crate::distance::scan_4bit_simd(sim_table, codes, count, m, &mut dists); - for i in 0..count { - if let Some(f) = filter { - if !f.contains(ids[i]) { - continue; - } + if let Some(positions) = matching_positions { + for &position in positions { + heap.push(dis0 + dists[position], ids[position]); + } + } else { + for i in 0..count { + heap.push(dis0 + dists[i], ids[i]); } - heap.push(dis0 + dists[i], ids[i]); } } @@ -781,7 +839,7 @@ fn scan_codes_4bit_transposed( count: usize, m: usize, dis0: f32, - filter: Option<&dyn RowIdFilter>, + matching_positions: Option<&[usize]>, heap: &mut TopKHeap, ) { let cs = m / 2; @@ -803,6 +861,44 @@ fn scan_codes_4bit_transposed( dists[i] = d; } + if let Some(positions) = + matching_positions.filter(|positions| should_scan_matching_positions(count, positions)) + { + if count <= FLAT_NUM { + for &position in positions { + heap.push(dis0 + dists[position], ids[position]); + } + return; + } + + let qmin = sim_table.iter().cloned().fold(f32::INFINITY, f32::min); + let qmax = dists[..flat_end].iter().cloned().fold(f32::MIN, f32::max); + let range = (qmax - qmin).max(1e-10); + let factor = 255.0 / range; + let qtable: Vec = sim_table + .iter() + .map(|&d| ((d - qmin) * factor).clamp(0.0, 255.0) as u8) + .collect(); + let inv_factor = range / 255.0; + let base_dist = qmin * m as f32; + + for &position in positions { + let distance = if position < flat_end { + dists[position] + } else { + let mut quantized_distance = 0u16; + for pair in 0..cs { + let code = codes[pair * count + position]; + quantized_distance += qtable[(pair * 2) * 16 + (code & 0x0f) as usize] as u16; + quantized_distance += qtable[(pair * 2 + 1) * 16 + (code >> 4) as usize] as u16; + } + quantized_distance as f32 * inv_factor + base_dist + }; + heap.push(dis0 + distance, ids[position]); + } + return; + } + if count > FLAT_NUM { let qmin = sim_table.iter().cloned().fold(f32::INFINITY, f32::min); let qmax = dists[..flat_end].iter().cloned().fold(f32::MIN, f32::max); @@ -835,13 +931,14 @@ fn scan_codes_4bit_transposed( } } - for i in 0..count { - if let Some(f) = filter { - if !f.contains(ids[i]) { - continue; - } + if let Some(positions) = matching_positions { + for &position in positions { + heap.push(dis0 + dists[position], ids[position]); + } + } else { + for i in 0..count { + heap.push(dis0 + dists[i], ids[i]); } - heap.push(dis0 + dists[i], ids[i]); } } @@ -856,11 +953,24 @@ fn scan_codes_transposed_with_scratch( m: usize, ksub: usize, dis0: f32, - filter: Option<&dyn RowIdFilter>, + matching_positions: Option<&[usize]>, heap: &mut TopKHeap, dists: &mut Vec, ) { debug_assert!(m > 0); + if let Some(positions) = + matching_positions.filter(|positions| should_scan_matching_positions(count, positions)) + { + for &row in positions { + let mut distance = dis0; + for sub in 0..m { + distance += sim_table[sub * ksub + codes[sub * count + row] as usize]; + } + heap.push(distance, ids[row]); + } + return; + } + dists.resize(count, 0.0); let table = &sim_table[..ksub]; let column = &codes[..count]; @@ -902,13 +1012,14 @@ fn scan_codes_transposed_with_scratch( } } - for i in 0..count { - if let Some(f) = filter { - if !f.contains(ids[i]) { - continue; - } + if let Some(positions) = matching_positions { + for &position in positions { + heap.push(dists[position], ids[position]); + } + } else { + for i in 0..count { + heap.push(dists[i], ids[i]); } - heap.push(dists[i], ids[i]); } } @@ -921,9 +1032,54 @@ fn scan_codes_batched( m: usize, ksub: usize, dis0: f32, - filter: Option<&dyn RowIdFilter>, + matching_positions: Option<&[usize]>, heap: &mut TopKHeap, ) { + if let Some(positions) = + matching_positions.filter(|positions| should_scan_matching_positions(count, positions)) + { + for &position in positions { + let code = &codes[position * m..(position + 1) * m]; + let distance = dis0 + pq_distance_from_table(sim_table, code, m, ksub); + heap.push(distance, ids[position]); + } + return; + } + + if let Some(positions) = matching_positions { + let mut distances = vec![0.0f32; count]; + let mut position = 0; + while position + 4 <= count { + let batch = pq_distance_four_codes( + sim_table, + codes, + m, + ksub, + [ + position * m, + (position + 1) * m, + (position + 2) * m, + (position + 3) * m, + ], + ); + distances[position..position + 4].copy_from_slice(&batch); + position += 4; + } + while position < count { + distances[position] = pq_distance_from_table( + sim_table, + &codes[position * m..(position + 1) * m], + m, + ksub, + ); + position += 1; + } + for &position in positions { + heap.push(dis0 + distances[position], ids[position]); + } + return; + } + let mut i = 0; while i + 4 <= count { @@ -938,11 +1094,6 @@ fn scan_codes_batched( for j in 0..4 { let idx = i + j; let id = ids[idx]; - if let Some(f) = filter { - if !f.contains(id) { - continue; - } - } heap.push(dis0 + dists[j], id); } i += 4; @@ -952,12 +1103,6 @@ fn scan_codes_batched( let code = &codes[i * m..(i + 1) * m]; let dist = dis0 + pq_distance_from_table(sim_table, code, m, ksub); let id = ids[i]; - if let Some(f) = filter { - if !f.contains(id) { - i += 1; - continue; - } - } heap.push(dist, id); i += 1; } @@ -967,7 +1112,6 @@ struct ReaderSearchContext<'a> { q: &'a [f32], ip_table: &'a [f32], use_precomputed: bool, - filter: Option<&'a dyn RowIdFilter>, d: usize, m: usize, ksub: usize, @@ -1088,6 +1232,7 @@ pub fn search_with_reader_filter( let transposed_codes = reader.transposed_codes; let mut scratch = ReaderScanScratch::default(); reader.for_each_streamed_list_chunk(first_list, |ids, codes| { + let positions = matching_positions(ids, filter); scan_reader_codes( &sim_table, codes, @@ -1097,7 +1242,7 @@ pub fn search_with_reader_filter( pq_nbits, transposed_codes, dis0, - filter, + positions.as_deref(), &mut scratch.distances, &mut heap, ); @@ -1132,7 +1277,6 @@ pub fn search_with_reader_filter( q: &q, ip_table: &ip_table, use_precomputed, - filter, d, m, ksub, @@ -1147,7 +1291,15 @@ pub fn search_with_reader_filter( .par_iter() .map_init(ReaderScanScratch::default, |scratch, (entry, dis0)| { let mut local_heap = TopKHeap::new(k); - scan_reader_list(entry, *dis0, &ctx, scratch, &mut local_heap); + let positions = matching_positions(&entry.ids, filter); + scan_reader_list( + entry, + *dis0, + &ctx, + positions.as_deref(), + scratch, + &mut local_heap, + ); local_heap.into_sorted() }) .collect::>(); @@ -1183,9 +1335,13 @@ fn scan_reader_list( entry: &InvertedListPayload, dis0: f32, ctx: &ReaderSearchContext<'_>, + matching_positions: Option<&[usize]>, scratch: &mut ReaderScanScratch, heap: &mut TopKHeap, ) { + if matching_positions.is_some_and(|positions| positions.is_empty()) { + return; + } fill_reader_sim_table(entry.list_id, ctx, &mut scratch.sim_table); scan_reader_codes( &scratch.sim_table, @@ -1196,7 +1352,7 @@ fn scan_reader_list( ctx.pq.nbits, ctx.transposed_codes, dis0, - ctx.filter, + matching_positions, &mut scratch.distances, heap, ); @@ -1241,7 +1397,6 @@ fn reader_sim_table( q: query, ip_table, use_precomputed, - filter: None, d: reader.d, m: reader.m, ksub: reader.ksub, @@ -1267,22 +1422,63 @@ fn scan_reader_codes( pq_nbits: usize, transposed_codes: bool, dis0: f32, - filter: Option<&dyn RowIdFilter>, + matching_positions: Option<&[usize]>, distances: &mut Vec, heap: &mut TopKHeap, ) { + if matching_positions.is_some_and(|positions| positions.is_empty()) { + return; + } let is_4bit = pq_nbits == 4; let count = ids.len(); if is_4bit && transposed_codes { - scan_codes_4bit_transposed(sim_table, codes, ids, count, m, dis0, filter, heap); + scan_codes_4bit_transposed( + sim_table, + codes, + ids, + count, + m, + dis0, + matching_positions, + heap, + ); } else if is_4bit { - scan_codes_4bit(sim_table, codes, ids, count, m, ksub, dis0, filter, heap); + scan_codes_4bit( + sim_table, + codes, + ids, + count, + m, + ksub, + dis0, + matching_positions, + heap, + ); } else if transposed_codes { scan_codes_transposed_with_scratch( - sim_table, codes, ids, count, m, ksub, dis0, filter, heap, distances, + sim_table, + codes, + ids, + count, + m, + ksub, + dis0, + matching_positions, + heap, + distances, ); } else { - scan_codes_batched(sim_table, codes, ids, count, m, ksub, dis0, filter, heap); + scan_codes_batched( + sim_table, + codes, + ids, + count, + m, + ksub, + dis0, + matching_positions, + heap, + ); } } @@ -1442,6 +1638,7 @@ pub fn search_batch_reader_filter( // distance buffer instead of retaining one per query. let mut distances = Vec::new(); reader.for_each_streamed_list_chunk(first_list, |ids, codes| { + let positions = matching_positions(ids, filter); for (query_index, dis0, sim_table) in &query_tables { scan_reader_codes( sim_table, @@ -1452,7 +1649,7 @@ pub fn search_batch_reader_filter( pq_nbits, transposed_codes, *dis0, - filter, + positions.as_deref(), &mut distances, &mut heaps[*query_index], ); @@ -1469,6 +1666,10 @@ pub fn search_batch_reader_filter( for (position, list) in loaded_lists.iter().enumerate() { list_positions[list.list_id] = position; } + let matching_positions_by_list = loaded_lists + .iter() + .map(|list| matching_positions(&list.ids, filter)) + .collect::>(); let rows = (0..nq) .into_par_iter() @@ -1482,7 +1683,6 @@ pub fn search_batch_reader_filter( &[] }, use_precomputed, - filter, d, m, ksub, @@ -1505,7 +1705,14 @@ pub fn search_batch_reader_filter( } else { 0.0 }; - scan_reader_list(&loaded_lists[position], dis0, &ctx, &mut scratch, &mut heap); + scan_reader_list( + &loaded_lists[position], + dis0, + &ctx, + matching_positions_by_list[position].as_deref(), + &mut scratch, + &mut heap, + ); } heap.into_sorted() }) @@ -1646,8 +1853,28 @@ mod tests { use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use std::io::Cursor; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; + struct CountingFilter { + contains_calls: AtomicUsize, + } + + impl CountingFilter { + fn new() -> Self { + Self { + contains_calls: AtomicUsize::new(0), + } + } + } + + impl RowIdFilter for CountingFilter { + fn contains(&self, id: i64) -> bool { + self.contains_calls.fetch_add(1, Ordering::Relaxed); + id % 7 == 0 + } + } + #[derive(Default)] struct ReaderStats { pread_calls: usize, @@ -1804,6 +2031,72 @@ mod tests { } } + #[test] + fn in_memory_batch_filter_only_evaluates_probed_lists() { + let d = 16; + let nlist = 8; + let m = 4; + let n = 800; + let nq = 4; + let k = 5; + let nprobe = 2; + let data = generate_clustered_data(n, d, nlist, 51); + let ids = (0..n as i64).collect::>(); + let mut index = IVFPQIndex::new(d, nlist, m, MetricType::L2, false); + index.train(&data, n); + index.add(&data, &ids, n); + + let queries = &data[..nq * d]; + let processed_queries = index.preprocess_queries(queries, nq); + let (probe_indices, _) = kmeans::find_topk_batch( + &processed_queries, + nq, + &index.quantizer_centroids, + nlist, + d, + nprobe, + ); + let mut probed_lists = vec![false; nlist]; + for query_probes in probe_indices { + for list_id in query_probes { + probed_lists[list_id] = true; + } + } + let expected_calls = index + .ids + .iter() + .zip(probed_lists) + .filter_map(|(ids, probed)| probed.then_some(ids.len())) + .sum::(); + + let filter = CountingFilter::new(); + let mut distances = vec![0.0f32; nq * k]; + let mut labels = vec![-1i64; nq * k]; + index.search_with_filter( + queries, + nq, + k, + nprobe, + Some(&filter), + &mut distances, + &mut labels, + ); + + assert!( + labels.iter().filter(|&&id| id >= 0).all(|id| id % 7 == 0), + "filtered search returned a disallowed row ID" + ); + assert_eq!( + filter.contains_calls.load(Ordering::Relaxed), + expected_calls, + "the filter should only evaluate rows from lists probed by this batch" + ); + assert!( + expected_calls < n, + "the test must leave at least one inverted list unprobed" + ); + } + #[test] fn test_batch_search() { let d = 16; @@ -2127,6 +2420,26 @@ mod tests { .collect::>(); expected.sort_by(|left, right| left.0.total_cmp(&right.0)); assert_eq!(heap.into_sorted(), expected); + + let matching_positions = (0..count).step_by(5).collect::>(); + let mut filtered_heap = TopKHeap::new(matching_positions.len()); + scan_codes_transposed_with_scratch( + &table, + &codes, + &ids, + count, + m, + ksub, + dis0, + Some(&matching_positions), + &mut filtered_heap, + &mut scratch, + ); + let filtered_expected = expected + .into_iter() + .filter(|(_, id)| (id - ids[0]) % 5 == 0) + .collect::>(); + assert_eq!(filtered_heap.into_sorted(), filtered_expected); } #[test] @@ -2663,6 +2976,49 @@ mod tests { } } + #[test] + fn test_batch_reader_evaluates_filter_once_per_loaded_row() { + use crate::io::{write_index, IVFPQIndexReader, PosWriter}; + + let d = 16; + let nlist = 1; + let m = 4; + let n = 500; + let k = 5; + let nq = 4; + let nprobe = 1; + + let data = generate_clustered_data(n, d, 1, 42); + let ids: Vec = (0..n as i64).collect(); + + let mut index = IVFPQIndex::new(d, nlist, m, MetricType::L2, false); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut buf = Vec::new(); + let mut writer = PosWriter::new(&mut buf); + write_index(&index, &mut writer).unwrap(); + + let filter = CountingFilter::new(); + let queries = &data[..nq * d]; + let mut reader = IVFPQIndexReader::open(Cursor::new(buf)).unwrap(); + let (result_ids, _) = + search_batch_reader_filter(&mut reader, queries, nq, k, nprobe, Some(&filter)).unwrap(); + + assert!( + result_ids + .iter() + .filter(|&&id| id >= 0) + .all(|id| id % 7 == 0), + "filtered batch search returned a disallowed row ID" + ); + assert_eq!( + filter.contains_calls.load(Ordering::Relaxed), + n, + "the shared filter should be evaluated once per loaded row, not once per query" + ); + } + #[test] fn test_batch_reader_empty_roaring_filter_returns_empty_results() { use crate::io::{write_index, IVFPQIndexReader, PosWriter}; From 831bc0c082ccdaf5c61f58d261c11dc442e3411c Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Wed, 29 Jul 2026 00:25:58 -0700 Subject: [PATCH 2/7] Fix filtered IVF-PQ scan regressions --- core/src/ivfpq.rs | 495 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 367 insertions(+), 128 deletions(-) diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index 3975dc3..25789f8 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -374,7 +374,7 @@ impl IVFPQIndex { let use_precomputed = !self.precomputed_table.is_empty(); let use_fastscan = !self.fastscan_codes.is_empty() && self.pq.nbits == 4; - let matching_positions_by_list = filter.map(|filter| { + let matching_rows_by_list = filter.map(|filter| { let mut probed_lists = vec![false; self.nlist]; for probe_indices in &all_probe_indices { for &list_id in probe_indices { @@ -386,9 +386,9 @@ impl IVFPQIndex { .zip(probed_lists) .map(|(ids, probed)| { if probed { - matching_positions(ids, Some(filter)).unwrap() + matching_rows(ids, Some(filter)).unwrap() } else { - Vec::new() + MatchingRows::Sparse(Vec::new()) } }) .collect::>() @@ -417,9 +417,10 @@ impl IVFPQIndex { if count == 0 { continue; } - let matching_positions = matching_positions_by_list - .as_ref() - .map(|positions| positions[list_id].as_slice()); + let matching_rows = matching_rows_by_list.as_ref().map(|rows| &rows[list_id]); + if matching_rows.is_some_and(MatchingRows::is_empty) { + continue; + } // Precomputed sim_table omits ||q-c||²; add it as dis0. // Non-precomputed path computes from residual_query, already full distance. @@ -441,11 +442,7 @@ impl IVFPQIndex { self.compute_list_table(query, list_id, &mut sim_table); } - if use_fastscan - && !matching_positions - .map(|positions| should_scan_matching_positions(count, positions)) - .unwrap_or(false) - { + if use_fastscan { let mut dists = vec![0.0f32; count]; crate::fastscan::fastscan_4bit( &sim_table, @@ -454,8 +451,8 @@ impl IVFPQIndex { m, &mut dists, ); - if let Some(positions) = matching_positions { - for &position in positions { + if let Some(rows) = matching_rows { + for position in rows.positions() { heap.push(dis0 + dists[position], self.ids[list_id][position]); } } else { @@ -472,7 +469,7 @@ impl IVFPQIndex { m, ksub, dis0, - matching_positions, + matching_rows, &mut heap, ); } else { @@ -484,7 +481,7 @@ impl IVFPQIndex { m, ksub, dis0, - matching_positions, + matching_rows, &mut heap, ); } @@ -774,17 +771,147 @@ fn invalid_merge_input(message: impl Into) -> io::Error { io::Error::new(io::ErrorKind::InvalidInput, message.into()) } -fn matching_positions(ids: &[i64], filter: Option<&dyn RowIdFilter>) -> Option> { +enum MatchingRows { + Sparse(Vec), + Bitmap { words: Vec, len: usize }, +} + +impl MatchingRows { + fn len(&self) -> usize { + match self { + Self::Sparse(positions) => positions.len(), + Self::Bitmap { len, .. } => *len, + } + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn contains(&self, position: usize) -> bool { + match self { + Self::Sparse(positions) => positions.binary_search(&position).is_ok(), + Self::Bitmap { words, .. } => words + .get(position / 64) + .is_some_and(|word| word & (1u64 << (position % 64)) != 0), + } + } + + fn positions(&self) -> MatchingRowIter<'_> { + match self { + Self::Sparse(positions) => MatchingRowIter::Sparse { + positions, + index: 0, + }, + Self::Bitmap { words, .. } => MatchingRowIter::Bitmap { + words, + word_index: 0, + word: 0, + word_base: 0, + }, + } + } + + #[cfg(test)] + fn storage_bytes(&self) -> usize { + match self { + Self::Sparse(positions) => positions.capacity() * std::mem::size_of::(), + Self::Bitmap { words, .. } => words.capacity() * std::mem::size_of::(), + } + } +} + +enum MatchingRowIter<'a> { + Sparse { + positions: &'a [usize], + index: usize, + }, + Bitmap { + words: &'a [u64], + word_index: usize, + word: u64, + word_base: usize, + }, +} + +impl Iterator for MatchingRowIter<'_> { + type Item = usize; + + fn next(&mut self) -> Option { + match self { + Self::Sparse { positions, index } => { + let position = positions.get(*index).copied(); + *index += usize::from(position.is_some()); + position + } + Self::Bitmap { + words, + word_index, + word, + word_base, + } => loop { + if *word != 0 { + let bit = word.trailing_zeros() as usize; + *word &= *word - 1; + return Some(*word_base + bit); + } + let next_word = words.get(*word_index).copied()?; + *word = next_word; + *word_base = *word_index * 64; + *word_index += 1; + }, + } + } +} + +fn matching_rows(ids: &[i64], filter: Option<&dyn RowIdFilter>) -> Option { filter.map(|filter| { - ids.iter() - .enumerate() - .filter_map(|(position, &id)| filter.contains(id).then_some(position)) - .collect() + let bitmap_words = ids.len().div_ceil(64); + let sparse_limit = + bitmap_words.saturating_mul(std::mem::size_of::()) / std::mem::size_of::(); + let mut positions = Vec::new(); + let mut bitmap = None::>; + let mut matching_count = 0usize; + + for (position, &id) in ids.iter().enumerate() { + if !filter.contains(id) { + continue; + } + matching_count += 1; + if let Some(words) = bitmap.as_mut() { + words[position / 64] |= 1u64 << (position % 64); + } else if positions.len() < sparse_limit { + positions.push(position); + } else { + let mut words = vec![0u64; bitmap_words]; + for previous in positions.drain(..) { + words[previous / 64] |= 1u64 << (previous % 64); + } + words[position / 64] |= 1u64 << (position % 64); + bitmap = Some(words); + } + } + + match bitmap { + Some(words) => MatchingRows::Bitmap { + words, + len: matching_count, + }, + None => MatchingRows::Sparse(positions), + } }) } -fn should_scan_matching_positions(count: usize, matching_positions: &[usize]) -> bool { - matching_positions.len() <= count / 2 +// Sparse row-major scans give up four-code ILP, while sparse transposed scans +// replace sequential column reads with random row lookups. Keep separate, +// conservative crossover points instead of applying one threshold to every +// kernel. Packed 4-bit and FastScan paths retain their normal distance kernel +// to preserve score semantics. +const ROW_MAJOR_SPARSE_SCAN_DIVISOR: usize = 4; +const TRANSPOSED_SPARSE_SCAN_DIVISOR: usize = 8; + +fn should_scan_sparse(count: usize, matching_rows: &MatchingRows, divisor: usize) -> bool { + matching_rows.len().saturating_mul(divisor) <= count } /// Scan 4-bit packed codes using u8-domain accumulation. @@ -796,31 +923,14 @@ fn scan_codes_4bit( m: usize, _ksub: usize, dis0: f32, - matching_positions: Option<&[usize]>, + matching_rows: Option<&MatchingRows>, heap: &mut TopKHeap, ) { - if let Some(positions) = - matching_positions.filter(|positions| should_scan_matching_positions(count, positions)) - { - let code_size = m / 2; - for &position in positions { - let mut distance = dis0; - let code_offset = position * code_size; - for pair in 0..code_size { - let code = codes[code_offset + pair]; - distance += sim_table[(pair * 2) * 16 + (code & 0x0f) as usize]; - distance += sim_table[(pair * 2 + 1) * 16 + (code >> 4) as usize]; - } - heap.push(distance, ids[position]); - } - return; - } - let mut dists = vec![0.0f32; count]; crate::distance::scan_4bit_simd(sim_table, codes, count, m, &mut dists); - if let Some(positions) = matching_positions { - for &position in positions { + if let Some(rows) = matching_rows { + for position in rows.positions() { heap.push(dis0 + dists[position], ids[position]); } } else { @@ -839,7 +949,7 @@ fn scan_codes_4bit_transposed( count: usize, m: usize, dis0: f32, - matching_positions: Option<&[usize]>, + matching_rows: Option<&MatchingRows>, heap: &mut TopKHeap, ) { let cs = m / 2; @@ -861,11 +971,11 @@ fn scan_codes_4bit_transposed( dists[i] = d; } - if let Some(positions) = - matching_positions.filter(|positions| should_scan_matching_positions(count, positions)) + if let Some(rows) = + matching_rows.filter(|rows| should_scan_sparse(count, rows, TRANSPOSED_SPARSE_SCAN_DIVISOR)) { if count <= FLAT_NUM { - for &position in positions { + for position in rows.positions() { heap.push(dis0 + dists[position], ids[position]); } return; @@ -882,7 +992,7 @@ fn scan_codes_4bit_transposed( let inv_factor = range / 255.0; let base_dist = qmin * m as f32; - for &position in positions { + for position in rows.positions() { let distance = if position < flat_end { dists[position] } else { @@ -931,8 +1041,8 @@ fn scan_codes_4bit_transposed( } } - if let Some(positions) = matching_positions { - for &position in positions { + if let Some(rows) = matching_rows { + for position in rows.positions() { heap.push(dis0 + dists[position], ids[position]); } } else { @@ -953,15 +1063,15 @@ fn scan_codes_transposed_with_scratch( m: usize, ksub: usize, dis0: f32, - matching_positions: Option<&[usize]>, + matching_rows: Option<&MatchingRows>, heap: &mut TopKHeap, dists: &mut Vec, ) { debug_assert!(m > 0); - if let Some(positions) = - matching_positions.filter(|positions| should_scan_matching_positions(count, positions)) + if let Some(rows) = + matching_rows.filter(|rows| should_scan_sparse(count, rows, TRANSPOSED_SPARSE_SCAN_DIVISOR)) { - for &row in positions { + for row in rows.positions() { let mut distance = dis0; for sub in 0..m { distance += sim_table[sub * ksub + codes[sub * count + row] as usize]; @@ -1012,8 +1122,8 @@ fn scan_codes_transposed_with_scratch( } } - if let Some(positions) = matching_positions { - for &position in positions { + if let Some(rows) = matching_rows { + for position in rows.positions() { heap.push(dists[position], ids[position]); } } else { @@ -1032,13 +1142,13 @@ fn scan_codes_batched( m: usize, ksub: usize, dis0: f32, - matching_positions: Option<&[usize]>, + matching_rows: Option<&MatchingRows>, heap: &mut TopKHeap, ) { - if let Some(positions) = - matching_positions.filter(|positions| should_scan_matching_positions(count, positions)) + if let Some(rows) = + matching_rows.filter(|rows| should_scan_sparse(count, rows, ROW_MAJOR_SPARSE_SCAN_DIVISOR)) { - for &position in positions { + for position in rows.positions() { let code = &codes[position * m..(position + 1) * m]; let distance = dis0 + pq_distance_from_table(sim_table, code, m, ksub); heap.push(distance, ids[position]); @@ -1046,40 +1156,6 @@ fn scan_codes_batched( return; } - if let Some(positions) = matching_positions { - let mut distances = vec![0.0f32; count]; - let mut position = 0; - while position + 4 <= count { - let batch = pq_distance_four_codes( - sim_table, - codes, - m, - ksub, - [ - position * m, - (position + 1) * m, - (position + 2) * m, - (position + 3) * m, - ], - ); - distances[position..position + 4].copy_from_slice(&batch); - position += 4; - } - while position < count { - distances[position] = pq_distance_from_table( - sim_table, - &codes[position * m..(position + 1) * m], - m, - ksub, - ); - position += 1; - } - for &position in positions { - heap.push(dis0 + distances[position], ids[position]); - } - return; - } - let mut i = 0; while i + 4 <= count { @@ -1093,17 +1169,19 @@ fn scan_codes_batched( for j in 0..4 { let idx = i + j; - let id = ids[idx]; - heap.push(dis0 + dists[j], id); + if matching_rows.is_none_or(|rows| rows.contains(idx)) { + heap.push(dis0 + dists[j], ids[idx]); + } } i += 4; } while i < count { - let code = &codes[i * m..(i + 1) * m]; - let dist = dis0 + pq_distance_from_table(sim_table, code, m, ksub); - let id = ids[i]; - heap.push(dist, id); + if matching_rows.is_none_or(|rows| rows.contains(i)) { + let code = &codes[i * m..(i + 1) * m]; + let dist = dis0 + pq_distance_from_table(sim_table, code, m, ksub); + heap.push(dist, ids[i]); + } i += 1; } } @@ -1232,7 +1310,7 @@ pub fn search_with_reader_filter( let transposed_codes = reader.transposed_codes; let mut scratch = ReaderScanScratch::default(); reader.for_each_streamed_list_chunk(first_list, |ids, codes| { - let positions = matching_positions(ids, filter); + let positions = matching_rows(ids, filter); scan_reader_codes( &sim_table, codes, @@ -1242,7 +1320,7 @@ pub fn search_with_reader_filter( pq_nbits, transposed_codes, dis0, - positions.as_deref(), + positions.as_ref(), &mut scratch.distances, &mut heap, ); @@ -1291,12 +1369,12 @@ pub fn search_with_reader_filter( .par_iter() .map_init(ReaderScanScratch::default, |scratch, (entry, dis0)| { let mut local_heap = TopKHeap::new(k); - let positions = matching_positions(&entry.ids, filter); + let positions = matching_rows(&entry.ids, filter); scan_reader_list( entry, *dis0, &ctx, - positions.as_deref(), + positions.as_ref(), scratch, &mut local_heap, ); @@ -1335,11 +1413,11 @@ fn scan_reader_list( entry: &InvertedListPayload, dis0: f32, ctx: &ReaderSearchContext<'_>, - matching_positions: Option<&[usize]>, + matching_rows: Option<&MatchingRows>, scratch: &mut ReaderScanScratch, heap: &mut TopKHeap, ) { - if matching_positions.is_some_and(|positions| positions.is_empty()) { + if matching_rows.is_some_and(MatchingRows::is_empty) { return; } fill_reader_sim_table(entry.list_id, ctx, &mut scratch.sim_table); @@ -1352,7 +1430,7 @@ fn scan_reader_list( ctx.pq.nbits, ctx.transposed_codes, dis0, - matching_positions, + matching_rows, &mut scratch.distances, heap, ); @@ -1422,26 +1500,17 @@ fn scan_reader_codes( pq_nbits: usize, transposed_codes: bool, dis0: f32, - matching_positions: Option<&[usize]>, + matching_rows: Option<&MatchingRows>, distances: &mut Vec, heap: &mut TopKHeap, ) { - if matching_positions.is_some_and(|positions| positions.is_empty()) { + if matching_rows.is_some_and(MatchingRows::is_empty) { return; } let is_4bit = pq_nbits == 4; let count = ids.len(); if is_4bit && transposed_codes { - scan_codes_4bit_transposed( - sim_table, - codes, - ids, - count, - m, - dis0, - matching_positions, - heap, - ); + scan_codes_4bit_transposed(sim_table, codes, ids, count, m, dis0, matching_rows, heap); } else if is_4bit { scan_codes_4bit( sim_table, @@ -1451,7 +1520,7 @@ fn scan_reader_codes( m, ksub, dis0, - matching_positions, + matching_rows, heap, ); } else if transposed_codes { @@ -1463,7 +1532,7 @@ fn scan_reader_codes( m, ksub, dis0, - matching_positions, + matching_rows, heap, distances, ); @@ -1476,7 +1545,7 @@ fn scan_reader_codes( m, ksub, dis0, - matching_positions, + matching_rows, heap, ); } @@ -1638,7 +1707,7 @@ pub fn search_batch_reader_filter( // distance buffer instead of retaining one per query. let mut distances = Vec::new(); reader.for_each_streamed_list_chunk(first_list, |ids, codes| { - let positions = matching_positions(ids, filter); + let positions = matching_rows(ids, filter); for (query_index, dis0, sim_table) in &query_tables { scan_reader_codes( sim_table, @@ -1649,7 +1718,7 @@ pub fn search_batch_reader_filter( pq_nbits, transposed_codes, *dis0, - positions.as_deref(), + positions.as_ref(), &mut distances, &mut heaps[*query_index], ); @@ -1666,9 +1735,9 @@ pub fn search_batch_reader_filter( for (position, list) in loaded_lists.iter().enumerate() { list_positions[list.list_id] = position; } - let matching_positions_by_list = loaded_lists + let matching_rows_by_list = loaded_lists .iter() - .map(|list| matching_positions(&list.ids, filter)) + .map(|list| matching_rows(&list.ids, filter)) .collect::>(); let rows = (0..nq) @@ -1709,7 +1778,7 @@ pub fn search_batch_reader_filter( &loaded_lists[position], dis0, &ctx, - matching_positions_by_list[position].as_deref(), + matching_rows_by_list[position].as_ref(), &mut scratch, &mut heap, ); @@ -2422,7 +2491,8 @@ mod tests { assert_eq!(heap.into_sorted(), expected); let matching_positions = (0..count).step_by(5).collect::>(); - let mut filtered_heap = TopKHeap::new(matching_positions.len()); + let matching_rows = MatchingRows::Sparse(matching_positions); + let mut filtered_heap = TopKHeap::new(matching_rows.len()); scan_codes_transposed_with_scratch( &table, &codes, @@ -2431,7 +2501,7 @@ mod tests { m, ksub, dis0, - Some(&matching_positions), + Some(&matching_rows), &mut filtered_heap, &mut scratch, ); @@ -2442,6 +2512,175 @@ mod tests { assert_eq!(filtered_heap.into_sorted(), filtered_expected); } + #[test] + fn row_major_4bit_filtered_scan_matches_unfiltered_kernel() { + let count = 400; + let m = 8; + let ksub = 16; + let dis0 = 1.25; + let ids = (10_000..10_000 + count as i64).collect::>(); + let codes = (0..count * (m / 2)) + .map(|index| ((index * 37 + 11) & 0xff) as u8) + .collect::>(); + let table = (0..m * ksub) + .map(|index| ((index * 29 + 7) % 113) as f32 * 0.03125) + .collect::>(); + let matching_positions = (0..count).step_by(3).collect::>(); + + let mut expected_distances = vec![0.0f32; count]; + crate::distance::scan_4bit_simd(&table, &codes, count, m, &mut expected_distances); + let mut expected = matching_positions + .iter() + .map(|&position| (dis0 + expected_distances[position], ids[position])) + .collect::>(); + expected.sort_unstable_by_key(|&(_, id)| id); + + let mut filtered_heap = TopKHeap::new(matching_positions.len()); + let matching_rows = MatchingRows::Sparse(matching_positions); + scan_codes_4bit( + &table, + &codes, + &ids, + count, + m, + ksub, + dis0, + Some(&matching_rows), + &mut filtered_heap, + ); + + let mut actual = filtered_heap.into_sorted(); + actual.sort_unstable_by_key(|&(_, id)| id); + assert_eq!(actual, expected); + } + + #[test] + fn matching_rows_adapts_sparse_positions_to_bounded_bitmap() { + let ids = (0..1024i64).collect::>(); + let sparse_filter = [3i64, 511, 900].into_iter().collect::>(); + let sparse = matching_rows(&ids, Some(&sparse_filter)).unwrap(); + assert!(matches!(sparse, MatchingRows::Sparse(_))); + assert_eq!(sparse.positions().collect::>(), vec![3, 511, 900]); + + let dense_filter = ids.iter().copied().collect::>(); + let dense = matching_rows(&ids, Some(&dense_filter)).unwrap(); + assert!(matches!(dense, MatchingRows::Bitmap { .. })); + assert_eq!(dense.len(), ids.len()); + assert_eq!( + dense.positions().collect::>(), + (0..ids.len()).collect::>() + ); + assert!( + dense.storage_bytes() <= ids.len().div_ceil(64) * size_of::(), + "dense match storage must be bounded to one bit per row" + ); + } + + #[test] + fn row_major_dense_bitmap_scan_matches_exact_distances() { + let count = 1024; + let m = 8; + let ksub = 256; + let dis0 = 2.5; + let ids = (20_000..20_000 + count as i64).collect::>(); + let codes = (0..count * m) + .map(|index| ((index * 73 + 19) % ksub) as u8) + .collect::>(); + let table = (0..m * ksub) + .map(|index| ((index * 31 + 5) % 127) as f32 * 0.015625) + .collect::>(); + let filter = ids + .iter() + .copied() + .filter(|id| id % 4 != 0) + .collect::>(); + let matching_rows = matching_rows(&ids, Some(&filter)).unwrap(); + assert!(matches!(matching_rows, MatchingRows::Bitmap { .. })); + + let mut expected = matching_rows + .positions() + .map(|position| { + let code = &codes[position * m..(position + 1) * m]; + ( + dis0 + pq_distance_from_table(&table, code, m, ksub), + ids[position], + ) + }) + .collect::>(); + expected.sort_unstable_by_key(|&(_, id)| id); + + let mut heap = TopKHeap::new(matching_rows.len()); + scan_codes_batched( + &table, + &codes, + &ids, + count, + m, + ksub, + dis0, + Some(&matching_rows), + &mut heap, + ); + let mut actual = heap.into_sorted(); + actual.sort_unstable_by_key(|&(_, id)| id); + assert_eq!(actual, expected); + } + + #[test] + fn fastscan_filtered_search_keeps_fastscan_distance_semantics() { + let d = 32; + let nlist = 1; + let m = 8; + let n = 4000; + let k = 50; + let data = generate_clustered_data(n, d, 1, 777); + let ids = (0..n as i64).collect::>(); + let mut index = IVFPQIndex::with_nbits(d, nlist, m, 4, MetricType::L2, false); + index.train(&data, n); + index.add(&data, &ids, n); + index.build_search_structures(); + + let query = &data[..d]; + let mut sim_table = vec![0.0f32; m * index.pq.ksub]; + index.compute_list_table(query, 0, &mut sim_table); + let mut all_distances = vec![0.0f32; n]; + crate::fastscan::fastscan_4bit( + &sim_table, + &index.fastscan_codes[0], + n, + m, + &mut all_distances, + ); + let filter = ids + .iter() + .copied() + .filter(|id| id % 3 == 0) + .collect::>(); + let mut expected_heap = TopKHeap::new(k); + for position in (0..n).filter(|&position| filter.contains(&ids[position])) { + expected_heap.push(all_distances[position], ids[position]); + } + let expected = expected_heap.into_sorted(); + + let mut actual_distances = vec![0.0f32; k]; + let mut actual_labels = vec![-1i64; k]; + index.search_with_filter( + query, + 1, + k, + 1, + Some(&filter), + &mut actual_distances, + &mut actual_labels, + ); + let actual = actual_distances + .into_iter() + .zip(actual_labels) + .collect::>(); + + assert_eq!(actual, expected); + } + #[test] fn test_precomputed_table_matches_normal_search() { let d = 16; From 6ad8b49be76e33ee5caf7c612f4017d17a1e0233 Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Wed, 29 Jul 2026 01:13:12 -0700 Subject: [PATCH 3/7] Scope filter pushdown to 8-bit IVF-PQ --- core/src/ivfpq.rs | 114 +++++----------------------------------------- 1 file changed, 12 insertions(+), 102 deletions(-) diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index 25789f8..0f25512 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -971,44 +971,6 @@ fn scan_codes_4bit_transposed( dists[i] = d; } - if let Some(rows) = - matching_rows.filter(|rows| should_scan_sparse(count, rows, TRANSPOSED_SPARSE_SCAN_DIVISOR)) - { - if count <= FLAT_NUM { - for position in rows.positions() { - heap.push(dis0 + dists[position], ids[position]); - } - return; - } - - let qmin = sim_table.iter().cloned().fold(f32::INFINITY, f32::min); - let qmax = dists[..flat_end].iter().cloned().fold(f32::MIN, f32::max); - let range = (qmax - qmin).max(1e-10); - let factor = 255.0 / range; - let qtable: Vec = sim_table - .iter() - .map(|&d| ((d - qmin) * factor).clamp(0.0, 255.0) as u8) - .collect(); - let inv_factor = range / 255.0; - let base_dist = qmin * m as f32; - - for position in rows.positions() { - let distance = if position < flat_end { - dists[position] - } else { - let mut quantized_distance = 0u16; - for pair in 0..cs { - let code = codes[pair * count + position]; - quantized_distance += qtable[(pair * 2) * 16 + (code & 0x0f) as usize] as u16; - quantized_distance += qtable[(pair * 2 + 1) * 16 + (code >> 4) as usize] as u16; - } - quantized_distance as f32 * inv_factor + base_dist - }; - heap.push(dis0 + distance, ids[position]); - } - return; - } - if count > FLAT_NUM { let qmin = sim_table.iter().cloned().fold(f32::INFINITY, f32::min); let qmax = dists[..flat_end].iter().cloned().fold(f32::MIN, f32::max); @@ -2513,31 +2475,34 @@ mod tests { } #[test] - fn row_major_4bit_filtered_scan_matches_unfiltered_kernel() { + fn row_major_sparse_scan_matches_exact_distances() { let count = 400; let m = 8; - let ksub = 16; + let ksub = 256; let dis0 = 1.25; let ids = (10_000..10_000 + count as i64).collect::>(); - let codes = (0..count * (m / 2)) - .map(|index| ((index * 37 + 11) & 0xff) as u8) + let codes = (0..count * m) + .map(|index| ((index * 37 + 11) % ksub) as u8) .collect::>(); let table = (0..m * ksub) .map(|index| ((index * 29 + 7) % 113) as f32 * 0.03125) .collect::>(); let matching_positions = (0..count).step_by(3).collect::>(); - - let mut expected_distances = vec![0.0f32; count]; - crate::distance::scan_4bit_simd(&table, &codes, count, m, &mut expected_distances); let mut expected = matching_positions .iter() - .map(|&position| (dis0 + expected_distances[position], ids[position])) + .map(|&position| { + let code = &codes[position * m..(position + 1) * m]; + ( + dis0 + pq_distance_from_table(&table, code, m, ksub), + ids[position], + ) + }) .collect::>(); expected.sort_unstable_by_key(|&(_, id)| id); let mut filtered_heap = TopKHeap::new(matching_positions.len()); let matching_rows = MatchingRows::Sparse(matching_positions); - scan_codes_4bit( + scan_codes_batched( &table, &codes, &ids, @@ -2626,61 +2591,6 @@ mod tests { assert_eq!(actual, expected); } - #[test] - fn fastscan_filtered_search_keeps_fastscan_distance_semantics() { - let d = 32; - let nlist = 1; - let m = 8; - let n = 4000; - let k = 50; - let data = generate_clustered_data(n, d, 1, 777); - let ids = (0..n as i64).collect::>(); - let mut index = IVFPQIndex::with_nbits(d, nlist, m, 4, MetricType::L2, false); - index.train(&data, n); - index.add(&data, &ids, n); - index.build_search_structures(); - - let query = &data[..d]; - let mut sim_table = vec![0.0f32; m * index.pq.ksub]; - index.compute_list_table(query, 0, &mut sim_table); - let mut all_distances = vec![0.0f32; n]; - crate::fastscan::fastscan_4bit( - &sim_table, - &index.fastscan_codes[0], - n, - m, - &mut all_distances, - ); - let filter = ids - .iter() - .copied() - .filter(|id| id % 3 == 0) - .collect::>(); - let mut expected_heap = TopKHeap::new(k); - for position in (0..n).filter(|&position| filter.contains(&ids[position])) { - expected_heap.push(all_distances[position], ids[position]); - } - let expected = expected_heap.into_sorted(); - - let mut actual_distances = vec![0.0f32; k]; - let mut actual_labels = vec![-1i64; k]; - index.search_with_filter( - query, - 1, - k, - 1, - Some(&filter), - &mut actual_distances, - &mut actual_labels, - ); - let actual = actual_distances - .into_iter() - .zip(actual_labels) - .collect::>(); - - assert_eq!(actual, expected); - } - #[test] fn test_precomputed_table_matches_normal_search() { let d = 16; From 27c9bbfe620cfa12e8f12af263ad21efe6de4856 Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Tue, 28 Jul 2026 05:39:46 -0700 Subject: [PATCH 4/7] Reuse IVF-PQ list tables in batch search --- core/src/ivfpq.rs | 446 ++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 435 insertions(+), 11 deletions(-) diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index 0f25512..c15f7cd 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -16,8 +16,8 @@ // under the License. use crate::distance::{ - fvec_inner_product, fvec_madd, fvec_normalize, pq_distance_four_codes, pq_distance_from_table, - MetricType, + fvec_inner_product, fvec_l2sqr, fvec_madd, fvec_normalize, pq_distance_four_codes, + pq_distance_from_table, MetricType, }; use crate::index_io_util::ivf_payload_is_oversized; use crate::io::{IVFPQIndexReader, InvertedListPayload, SeekRead}; @@ -914,6 +914,53 @@ fn should_scan_sparse(count: usize, matching_rows: &MatchingRows, divisor: usize matching_rows.len().saturating_mul(divisor) <= count } +fn has_matching_rows(matching_rows: Option<&MatchingRows>) -> bool { + match matching_rows { + Some(rows) => !rows.is_empty(), + None => true, + } +} + +// Below this size, table construction and Rayon scheduling dominate the saved +// per-query/list distance-table work. +const MIN_EPHEMERAL_PRECOMPUTE_QUERIES: usize = 64; + +fn should_use_ephemeral_precomputation( + matching_list_count: usize, + active_query_count: usize, + probe_count: usize, +) -> bool { + let setup_tables = matching_list_count.saturating_add(active_query_count); + // Require at least 2x reuse over the list-table and query-table setup work. + setup_tables > 0 && probe_count >= setup_tables.saturating_mul(2) +} + +fn fill_list_precomputed_table( + coarse_centroid: &[f32], + pq: &ProductQuantizer, + pq_norms: &[f32], + table: &mut Vec, +) { + debug_assert_eq!(coarse_centroid.len(), pq.d); + debug_assert_eq!(pq_norms.len(), pq.m * pq.ksub); + table.resize(pq.m * pq.ksub, 0.0); + for sub in 0..pq.m { + let range = pq.chunk_range(sub); + let chunk_dim = range.len(); + let pq_base = range.start * pq.ksub; + for code in 0..pq.ksub { + let pq_offset = pq_base + code * chunk_dim; + let mut inner_product = 0.0f32; + for dimension in 0..chunk_dim { + inner_product += + coarse_centroid[range.start + dimension] * pq.centroids[pq_offset + dimension]; + } + let table_offset = sub * pq.ksub + code; + table[table_offset] = pq_norms[table_offset] + 2.0 * inner_product; + } + } +} + /// Scan 4-bit packed codes using u8-domain accumulation. fn scan_codes_4bit( sim_table: &[f32], @@ -1166,6 +1213,7 @@ struct ReaderSearchContext<'a> { #[derive(Default)] struct ReaderScanScratch { sim_table: Vec, + ip_table: Vec, distances: Vec, } @@ -1533,6 +1581,18 @@ pub fn search_batch_reader_filter( k: usize, nprobe: usize, filter: Option<&dyn RowIdFilter>, +) -> io::Result<(Vec, Vec)> { + search_batch_reader_filter_with_observer(reader, queries, nq, k, nprobe, filter, |_| {}) +} + +fn search_batch_reader_filter_with_observer( + reader: &mut IVFPQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + filter: Option<&dyn RowIdFilter>, + mut observe_ephemeral_precomputed_lists: impl FnMut(usize), ) -> io::Result<(Vec, Vec)> { reader.ensure_loaded()?; let d = reader.d; @@ -1615,6 +1675,11 @@ pub fn search_batch_reader_filter( let use_precomputed = metric == MetricType::L2 && by_residual && !reader.precomputed_table.is_empty(); + let allow_ephemeral_precomputed = metric == MetricType::L2 + && by_residual + && !use_precomputed + && nq >= MIN_EPHEMERAL_PRECOMPUTE_QUERIES; + let mut pq_norms = None; let all_ip_tables: Vec> = if use_precomputed { (0..nq) @@ -1701,6 +1766,60 @@ pub fn search_batch_reader_filter( .iter() .map(|list| matching_rows(&list.ids, filter)) .collect::>(); + let use_ephemeral_precomputed = allow_ephemeral_precomputed && { + let mut matching_list_count = 0usize; + let mut probe_count = 0usize; + let mut active_query_count = 0usize; + + for rows in &matching_rows_by_list { + matching_list_count += usize::from(has_matching_rows(rows.as_ref())); + } + for probe_indices in &all_probe_indices { + let matching_probe_count = probe_indices + .iter() + .filter(|&&list_id| { + let position = list_positions[list_id]; + position != usize::MAX + && has_matching_rows(matching_rows_by_list[position].as_ref()) + }) + .count(); + probe_count += matching_probe_count; + active_query_count += usize::from(matching_probe_count > 0); + } + + should_use_ephemeral_precomputation( + matching_list_count, + active_query_count, + probe_count, + ) + }; + let ephemeral_precomputed_tables = if use_ephemeral_precomputed { + let pq_norms = pq_norms.get_or_insert_with(|| reader.pq.compute_centroid_norms()); + loaded_lists + .par_iter() + .zip(&matching_rows_by_list) + .map(|(list, rows)| { + let mut table = Vec::new(); + if has_matching_rows(rows.as_ref()) { + fill_list_precomputed_table( + &reader.quantizer_centroids[list.list_id * d..(list.list_id + 1) * d], + &reader.pq, + pq_norms, + &mut table, + ); + } + table + }) + .collect::>() + } else { + Vec::new() + }; + observe_ephemeral_precomputed_lists( + ephemeral_precomputed_tables + .iter() + .filter(|table| !table.is_empty()) + .count(), + ); let rows = (0..nq) .into_par_iter() @@ -1726,24 +1845,65 @@ pub fn search_batch_reader_filter( }; let mut heap = TopKHeap::new(k); let mut scratch = ReaderScanScratch::default(); + let query_uses_ephemeral_precomputed = use_ephemeral_precomputed + && all_probe_indices[qi].iter().any(|&list_id| { + let position = list_positions[list_id]; + position != usize::MAX && !ephemeral_precomputed_tables[position].is_empty() + }); + if query_uses_ephemeral_precomputed { + scratch.ip_table.resize(m * ksub, 0.0); + reader + .pq + .compute_inner_product_table(query, &mut scratch.ip_table); + } for (probe_rank, &list_id) in all_probe_indices[qi].iter().enumerate() { let position = list_positions[list_id]; if position == usize::MAX { continue; } - let dis0 = if use_precomputed { + let use_ephemeral_list = query_uses_ephemeral_precomputed + && !ephemeral_precomputed_tables[position].is_empty(); + let dis0 = if use_ephemeral_list { + fvec_l2sqr( + query, + &reader.quantizer_centroids[list_id * d..(list_id + 1) * d], + ) + } else if use_precomputed { all_coarse_dists[qi][probe_rank] } else { 0.0 }; - scan_reader_list( - &loaded_lists[position], - dis0, - &ctx, - matching_rows_by_list[position].as_ref(), - &mut scratch, - &mut heap, - ); + if use_ephemeral_list { + scratch.sim_table.resize(m * ksub, 0.0); + fvec_madd( + &ephemeral_precomputed_tables[position], + &scratch.ip_table, + -2.0, + &mut scratch.sim_table, + ); + scan_reader_codes( + &scratch.sim_table, + loaded_lists[position].codes(), + &loaded_lists[position].ids, + m, + ksub, + reader.pq.nbits, + reader.transposed_codes, + dis0, + matching_rows_by_list[position].as_ref(), + &mut scratch.distances, + &mut heap, + ); + } else { + scan_reader_list( + &loaded_lists[position], + dis0, + &ctx, + matching_rows_by_list[position].as_ref(), + &mut scratch, + &mut heap, + ); + } } heap.into_sorted() }) @@ -1967,6 +2127,56 @@ mod tests { data } + fn observed_ephemeral_precomputed_lists( + nq: usize, + nprobe: usize, + filter_step: Option, + apply_filter: bool, + seed: u64, + ) -> usize { + use crate::io::{write_index, IVFPQIndexReader, PosWriter}; + + let d = 16; + let nlist = 4; + let m = 4; + let n = 600; + let k = 5; + let data = generate_clustered_data(n, d, nlist, seed); + let ids = (0..n as i64).collect::>(); + let mut index = IVFPQIndex::new(d, nlist, m, MetricType::L2, false); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut bytes = Vec::new(); + write_index(&index, &mut PosWriter::new(&mut bytes)).unwrap(); + let filter = match filter_step { + Some(step) => ids.iter().copied().step_by(step).collect::>(), + None => HashSet::new(), + }; + let filter = if apply_filter { + Some(&filter as &dyn RowIdFilter) + } else { + None + }; + let precomputed_lists = AtomicUsize::new(0); + let mut reader = IVFPQIndexReader::open(Cursor::new(bytes)).unwrap(); + + search_batch_reader_filter_with_observer( + &mut reader, + &data[..nq * d], + nq, + k, + nprobe, + filter, + |count| { + precomputed_lists.fetch_add(count, Ordering::Relaxed); + }, + ) + .unwrap(); + + precomputed_lists.load(Ordering::Relaxed) + } + fn assert_invalid_merge(base: &IVFPQIndex, other: &IVFPQIndex, expected_message: &str) { let mut target = IVFPQIndex::from_trained(base); let before_ids = target.ids.clone(); @@ -2591,6 +2801,39 @@ mod tests { assert_eq!(actual, expected); } + #[test] + fn reader_list_precomputed_table_matches_index_table() { + let d = 16; + let nlist = 4; + let m = 4; + let n = 600; + let data = generate_clustered_data(n, d, nlist, 44); + let mut index = IVFPQIndex::new(d, nlist, m, MetricType::L2, false); + index.train(&data, n); + index.build_precomputed_table(); + let pq_norms = index.pq.compute_centroid_norms(); + + for list_id in 0..nlist { + let mut actual = Vec::new(); + fill_list_precomputed_table( + &index.quantizer_centroids[list_id * d..(list_id + 1) * d], + &index.pq, + &pq_norms, + &mut actual, + ); + let table_size = m * index.pq.ksub; + assert_eq!( + actual, + index.precomputed_table[list_id * table_size..(list_id + 1) * table_size] + ); + } + } + + #[test] + fn ephemeral_precomputation_requires_matching_probe_work() { + assert!(!should_use_ephemeral_precomputation(0, 0, 0)); + } + #[test] fn test_precomputed_table_matches_normal_search() { let d = 16; @@ -3168,6 +3411,187 @@ mod tests { ); } + #[test] + fn filtered_batch_reader_uses_ephemeral_list_precomputation() { + use crate::io::{write_index, IVFPQIndexReader, PosWriter}; + use std::sync::atomic::AtomicUsize; + + let d = 16; + let nlist = 4; + let m = 4; + let n = 600; + let nq = 64; + let k = 5; + let nprobe = nlist; + let data = generate_clustered_data(n, d, nlist, 45); + let ids = (0..n as i64).map(|id| 50_000 + id * 3).collect::>(); + let mut index = IVFPQIndex::new(d, nlist, m, MetricType::L2, false); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut bytes = Vec::new(); + write_index(&index, &mut PosWriter::new(&mut bytes)).unwrap(); + let filter = ids.iter().copied().step_by(5).collect::>(); + let queries = &data[..nq * d]; + let precomputed_lists = AtomicUsize::new(0); + let mut reader = IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + + let (batch_ids, batch_dists) = search_batch_reader_filter_with_observer( + &mut reader, + queries, + nq, + k, + nprobe, + Some(&filter), + |count| { + precomputed_lists.fetch_add(count, Ordering::Relaxed); + }, + ) + .unwrap(); + + assert_eq!(precomputed_lists.load(Ordering::Relaxed), nlist); + assert!( + reader.precomputed_table.is_empty(), + "batch-local precomputation must not remain resident on the reader" + ); + for query_index in 0..nq { + let mut single_reader = IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + let query = &queries[query_index * d..(query_index + 1) * d]; + let (single_ids, single_dists) = + search_with_reader_filter(&mut single_reader, query, k, nprobe, Some(&filter)) + .unwrap(); + assert_eq!( + &batch_ids[query_index * k..(query_index + 1) * k], + single_ids.as_slice() + ); + for (batch, single) in batch_dists[query_index * k..(query_index + 1) * k] + .iter() + .zip(&single_dists) + { + // The algebraically equivalent precomputed formula changes + // floating-point accumulation order. Allow a small absolute + // floor near zero plus a few ULPs for large distances. + let tolerance = 1e-4 + 4.0 * f32::EPSILON * single.abs(); + assert!( + (batch - single).abs() <= tolerance, + "ephemeral precomputation distance {batch} should match direct residual distance {single} within {tolerance}" + ); + } + } + } + + #[test] + fn unfiltered_batch_reader_uses_ephemeral_list_precomputation() { + use crate::io::{write_index, IVFPQIndexReader, PosWriter}; + use std::sync::atomic::AtomicUsize; + + let d = 16; + let nlist = 4; + let m = 4; + let n = 600; + let nq = 64; + let k = 5; + let nprobe = nlist; + let data = generate_clustered_data(n, d, nlist, 49); + let ids = (0..n as i64).map(|id| 70_000 + id * 3).collect::>(); + let mut index = IVFPQIndex::new(d, nlist, m, MetricType::L2, false); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut bytes = Vec::new(); + write_index(&index, &mut PosWriter::new(&mut bytes)).unwrap(); + let queries = &data[..nq * d]; + let precomputed_lists = AtomicUsize::new(0); + let mut reader = IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + + let (batch_ids, batch_dists) = search_batch_reader_filter_with_observer( + &mut reader, + queries, + nq, + k, + nprobe, + None, + |count| { + precomputed_lists.fetch_add(count, Ordering::Relaxed); + }, + ) + .unwrap(); + + assert_eq!(precomputed_lists.load(Ordering::Relaxed), nlist); + assert!( + reader.precomputed_table.is_empty(), + "batch-local precomputation must not remain resident on the reader" + ); + for query_index in 0..nq { + let mut single_reader = IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + let query = &queries[query_index * d..(query_index + 1) * d]; + let (single_ids, single_dists) = + search_with_reader_filter(&mut single_reader, query, k, nprobe, None).unwrap(); + assert_eq!( + &batch_ids[query_index * k..(query_index + 1) * k], + single_ids.as_slice() + ); + for (batch, single) in batch_dists[query_index * k..(query_index + 1) * k] + .iter() + .zip(&single_dists) + { + let tolerance = 1e-4 + 4.0 * f32::EPSILON * single.abs(); + assert!( + (batch - single).abs() <= tolerance, + "ephemeral precomputation distance {batch} should match direct residual distance {single} within {tolerance}" + ); + } + } + } + + #[test] + fn small_filtered_batch_reader_skips_ephemeral_list_precomputation() { + assert_eq!( + observed_ephemeral_precomputed_lists(4, 4, Some(5), true, 46), + 0, + "small batches should keep the direct residual-table path" + ); + } + + #[test] + fn single_probe_filtered_batch_reader_skips_ephemeral_list_precomputation() { + assert_eq!( + observed_ephemeral_precomputed_lists( + MIN_EPHEMERAL_PRECOMPUTE_QUERIES, + 1, + Some(5), + true, + 47, + ), + 0, + "single-probe batches cannot amortize list precomputation" + ); + } + + #[test] + fn empty_filtered_batch_reader_skips_ephemeral_list_precomputation() { + assert_eq!( + observed_ephemeral_precomputed_lists( + MIN_EPHEMERAL_PRECOMPUTE_QUERIES, + 4, + None, + true, + 48, + ), + 0, + "lists without matching rows should not be precomputed" + ); + } + + #[test] + fn small_unfiltered_batch_reader_skips_ephemeral_list_precomputation() { + assert_eq!( + observed_ephemeral_precomputed_lists(4, 4, None, false, 50), + 0, + "small unfiltered batches should keep the direct residual-table path" + ); + } + #[test] fn test_batch_reader_empty_roaring_filter_returns_empty_results() { use crate::io::{write_index, IVFPQIndexReader, PosWriter}; From 0ebb557808d825b90e1e77c84781cb941ff2b48e Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Tue, 28 Jul 2026 07:15:09 -0700 Subject: [PATCH 5/7] Add IVF-PQ batch table reuse modes --- core/src/index.rs | 41 +++- core/src/ivfpq.rs | 226 +++++++++++++++--- ffi/src/lib.rs | 11 +- .../vector/IvfPqBatchTableReuseMode.java | 68 ++++++ .../index/vector/VectorSearchParams.java | 48 +++- .../index/vector/VectorIndexJavaApiTest.java | 28 +++ jni/src/lib.rs | 40 +++- 7 files changed, 413 insertions(+), 49 deletions(-) create mode 100644 java/src/main/java/org/apache/paimon/index/vector/IvfPqBatchTableReuseMode.java diff --git a/core/src/index.rs b/core/src/index.rs index c2b4d10..39a6efc 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -33,9 +33,10 @@ use crate::ivfflat_io::{ search_batch_ivfflat_reader, search_batch_ivfflat_reader_roaring_filter, write_ivfflat_index, IVFFlatIndexReader, IVFFLAT_MAGIC, }; +pub use crate::ivfpq::IvfPqBatchTableReuseMode; use crate::ivfpq::{ - search_batch_reader, search_batch_reader_roaring_filter, search_with_reader, - search_with_reader_roaring_filter, IVFPQIndex, + search_batch_reader_roaring_filter_with_reuse_mode, search_batch_reader_with_reuse_mode, + search_with_reader, search_with_reader_roaring_filter, IVFPQIndex, }; use crate::ivfrq::IVFRQIndex; use crate::ivfrq_io::{ @@ -988,6 +989,7 @@ pub struct VectorSearchParams { pub top_k: usize, pub search_width: SearchWidth, pub width: usize, + pub ivfpq_batch_table_reuse: IvfPqBatchTableReuseMode, } impl VectorSearchParams { @@ -996,6 +998,7 @@ impl VectorSearchParams { top_k, search_width: SearchWidth::IvfNProbe, width: nprobe, + ivfpq_batch_table_reuse: IvfPqBatchTableReuseMode::Auto, } } @@ -1004,6 +1007,7 @@ impl VectorSearchParams { top_k, search_width: SearchWidth::DiskAnnLSearch, width: l_search, + ivfpq_batch_table_reuse: IvfPqBatchTableReuseMode::Auto, } } @@ -1012,9 +1016,15 @@ impl VectorSearchParams { top_k, search_width: SearchWidth::Auto, width: 0, + ivfpq_batch_table_reuse: IvfPqBatchTableReuseMode::Auto, } } + pub fn with_ivfpq_batch_table_reuse(mut self, mode: IvfPqBatchTableReuseMode) -> Self { + self.ivfpq_batch_table_reuse = mode; + self + } + pub fn configured_ivf_nprobe(self) -> Option { (self.search_width == SearchWidth::IvfNProbe).then_some(self.width) } @@ -1755,7 +1765,14 @@ impl VectorIndexReader { params.top_k, total_vectors, |nprobe| { - search_batch_reader(reader, queries, query_count, params.top_k, nprobe) + search_batch_reader_with_reuse_mode( + reader, + queries, + query_count, + params.top_k, + nprobe, + params.ivfpq_batch_table_reuse, + ) }, ) } @@ -1865,13 +1882,14 @@ impl VectorIndexReader { params.top_k, matching_count.unwrap_or(total_vectors), |nprobe| { - search_batch_reader_roaring_filter( + search_batch_reader_roaring_filter_with_reuse_mode( reader, queries, query_count, params.top_k, nprobe, roaring_filter_bytes, + params.ivfpq_batch_table_reuse, ) }, ) @@ -2870,6 +2888,21 @@ mod tests { .contains("cannot be used with a DiskANN")); } + #[test] + fn ivfpq_batch_table_reuse_is_auto_by_default_and_can_be_disabled() { + let params = VectorSearchParams::new(10, 4); + assert_eq!( + params.ivfpq_batch_table_reuse, + IvfPqBatchTableReuseMode::Auto + ); + assert_eq!( + params + .with_ivfpq_batch_table_reuse(IvfPqBatchTableReuseMode::Off) + .ivfpq_batch_table_reuse, + IvfPqBatchTableReuseMode::Off + ); + } + #[test] fn automatic_filtered_search_expands_until_results_are_filled() { let mut observed = Vec::new(); diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index c15f7cd..742b89b 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -925,6 +925,14 @@ fn has_matching_rows(matching_rows: Option<&MatchingRows>) -> bool { // per-query/list distance-table work. const MIN_EPHEMERAL_PRECOMPUTE_QUERIES: usize = 64; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u32)] +pub enum IvfPqBatchTableReuseMode { + Off = 0, + On = 1, + Auto = 2, +} + fn should_use_ephemeral_precomputation( matching_list_count: usize, active_query_count: usize, @@ -1570,7 +1578,25 @@ pub fn search_batch_reader( k: usize, nprobe: usize, ) -> io::Result<(Vec, Vec)> { - search_batch_reader_filter(reader, queries, nq, k, nprobe, None) + search_batch_reader_with_reuse_mode( + reader, + queries, + nq, + k, + nprobe, + IvfPqBatchTableReuseMode::Auto, + ) +} + +pub fn search_batch_reader_with_reuse_mode( + reader: &mut IVFPQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + reuse_mode: IvfPqBatchTableReuseMode, +) -> io::Result<(Vec, Vec)> { + search_batch_reader_filter_with_reuse_mode(reader, queries, nq, k, nprobe, None, reuse_mode) } /// Big batch search with an optional row-id filter. @@ -1582,9 +1608,39 @@ pub fn search_batch_reader_filter( nprobe: usize, filter: Option<&dyn RowIdFilter>, ) -> io::Result<(Vec, Vec)> { - search_batch_reader_filter_with_observer(reader, queries, nq, k, nprobe, filter, |_| {}) + search_batch_reader_filter_with_reuse_mode( + reader, + queries, + nq, + k, + nprobe, + filter, + IvfPqBatchTableReuseMode::Auto, + ) +} + +pub fn search_batch_reader_filter_with_reuse_mode( + reader: &mut IVFPQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + filter: Option<&dyn RowIdFilter>, + reuse_mode: IvfPqBatchTableReuseMode, +) -> io::Result<(Vec, Vec)> { + search_batch_reader_filter_with_reuse_mode_and_observer( + reader, + queries, + nq, + k, + nprobe, + filter, + reuse_mode, + |_| {}, + ) } +#[cfg(test)] fn search_batch_reader_filter_with_observer( reader: &mut IVFPQIndexReader, queries: &[f32], @@ -1593,6 +1649,28 @@ fn search_batch_reader_filter_with_observer( nprobe: usize, filter: Option<&dyn RowIdFilter>, mut observe_ephemeral_precomputed_lists: impl FnMut(usize), +) -> io::Result<(Vec, Vec)> { + search_batch_reader_filter_with_reuse_mode_and_observer( + reader, + queries, + nq, + k, + nprobe, + filter, + IvfPqBatchTableReuseMode::Auto, + &mut observe_ephemeral_precomputed_lists, + ) +} + +fn search_batch_reader_filter_with_reuse_mode_and_observer( + reader: &mut IVFPQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + filter: Option<&dyn RowIdFilter>, + reuse_mode: IvfPqBatchTableReuseMode, + mut observe_ephemeral_precomputed_lists: impl FnMut(usize), ) -> io::Result<(Vec, Vec)> { reader.ensure_loaded()?; let d = reader.d; @@ -1678,7 +1756,11 @@ fn search_batch_reader_filter_with_observer( let allow_ephemeral_precomputed = metric == MetricType::L2 && by_residual && !use_precomputed - && nq >= MIN_EPHEMERAL_PRECOMPUTE_QUERIES; + && match reuse_mode { + IvfPqBatchTableReuseMode::Off => false, + IvfPqBatchTableReuseMode::On => true, + IvfPqBatchTableReuseMode::Auto => nq >= MIN_EPHEMERAL_PRECOMPUTE_QUERIES, + }; let mut pq_norms = None; let all_ip_tables: Vec> = if use_precomputed { @@ -1766,33 +1848,34 @@ fn search_batch_reader_filter_with_observer( .iter() .map(|list| matching_rows(&list.ids, filter)) .collect::>(); - let use_ephemeral_precomputed = allow_ephemeral_precomputed && { - let mut matching_list_count = 0usize; - let mut probe_count = 0usize; - let mut active_query_count = 0usize; - - for rows in &matching_rows_by_list { - matching_list_count += usize::from(has_matching_rows(rows.as_ref())); - } - for probe_indices in &all_probe_indices { - let matching_probe_count = probe_indices - .iter() - .filter(|&&list_id| { - let position = list_positions[list_id]; - position != usize::MAX - && has_matching_rows(matching_rows_by_list[position].as_ref()) - }) - .count(); - probe_count += matching_probe_count; - active_query_count += usize::from(matching_probe_count > 0); - } + let use_ephemeral_precomputed = allow_ephemeral_precomputed + && (reuse_mode == IvfPqBatchTableReuseMode::On || { + let mut matching_list_count = 0usize; + let mut probe_count = 0usize; + let mut active_query_count = 0usize; + + for rows in &matching_rows_by_list { + matching_list_count += usize::from(has_matching_rows(rows.as_ref())); + } + for probe_indices in &all_probe_indices { + let matching_probe_count = probe_indices + .iter() + .filter(|&&list_id| { + let position = list_positions[list_id]; + position != usize::MAX + && has_matching_rows(matching_rows_by_list[position].as_ref()) + }) + .count(); + probe_count += matching_probe_count; + active_query_count += usize::from(matching_probe_count > 0); + } - should_use_ephemeral_precomputation( - matching_list_count, - active_query_count, - probe_count, - ) - }; + should_use_ephemeral_precomputation( + matching_list_count, + active_query_count, + probe_count, + ) + }); let ephemeral_precomputed_tables = if use_ephemeral_precomputed { let pq_norms = pq_norms.get_or_insert_with(|| reader.pq.compute_centroid_norms()); loaded_lists @@ -1938,9 +2021,37 @@ pub fn search_batch_reader_roaring_filter( k: usize, nprobe: usize, roaring_filter_bytes: &[u8], +) -> io::Result<(Vec, Vec)> { + search_batch_reader_roaring_filter_with_reuse_mode( + reader, + queries, + nq, + k, + nprobe, + roaring_filter_bytes, + IvfPqBatchTableReuseMode::Auto, + ) +} + +pub fn search_batch_reader_roaring_filter_with_reuse_mode( + reader: &mut IVFPQIndexReader, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + roaring_filter_bytes: &[u8], + reuse_mode: IvfPqBatchTableReuseMode, ) -> io::Result<(Vec, Vec)> { let filter = decode_roaring_filter(roaring_filter_bytes)?; - search_batch_reader_filter(reader, queries, nq, k, nprobe, Some(&filter)) + search_batch_reader_filter_with_reuse_mode( + reader, + queries, + nq, + k, + nprobe, + Some(&filter), + reuse_mode, + ) } // --- Top-K Heap --- @@ -2133,6 +2244,7 @@ mod tests { filter_step: Option, apply_filter: bool, seed: u64, + reuse_mode: IvfPqBatchTableReuseMode, ) -> usize { use crate::io::{write_index, IVFPQIndexReader, PosWriter}; @@ -2161,13 +2273,14 @@ mod tests { let precomputed_lists = AtomicUsize::new(0); let mut reader = IVFPQIndexReader::open(Cursor::new(bytes)).unwrap(); - search_batch_reader_filter_with_observer( + search_batch_reader_filter_with_reuse_mode_and_observer( &mut reader, &data[..nq * d], nq, k, nprobe, filter, + reuse_mode, |count| { precomputed_lists.fetch_add(count, Ordering::Relaxed); }, @@ -3547,7 +3660,14 @@ mod tests { #[test] fn small_filtered_batch_reader_skips_ephemeral_list_precomputation() { assert_eq!( - observed_ephemeral_precomputed_lists(4, 4, Some(5), true, 46), + observed_ephemeral_precomputed_lists( + 4, + 4, + Some(5), + true, + 46, + IvfPqBatchTableReuseMode::Auto, + ), 0, "small batches should keep the direct residual-table path" ); @@ -3562,6 +3682,7 @@ mod tests { Some(5), true, 47, + IvfPqBatchTableReuseMode::Auto, ), 0, "single-probe batches cannot amortize list precomputation" @@ -3577,6 +3698,7 @@ mod tests { None, true, 48, + IvfPqBatchTableReuseMode::Auto, ), 0, "lists without matching rows should not be precomputed" @@ -3586,12 +3708,50 @@ mod tests { #[test] fn small_unfiltered_batch_reader_skips_ephemeral_list_precomputation() { assert_eq!( - observed_ephemeral_precomputed_lists(4, 4, None, false, 50), + observed_ephemeral_precomputed_lists( + 4, + 4, + None, + false, + 50, + IvfPqBatchTableReuseMode::Auto, + ), 0, "small unfiltered batches should keep the direct residual-table path" ); } + #[test] + fn batch_table_reuse_off_never_precomputes_list_tables() { + assert_eq!( + observed_ephemeral_precomputed_lists( + MIN_EPHEMERAL_PRECOMPUTE_QUERIES, + 4, + Some(5), + true, + 51, + IvfPqBatchTableReuseMode::Off, + ), + 0, + "off mode must keep the direct residual-table path" + ); + } + + #[test] + fn batch_table_reuse_on_precomputes_for_small_batches() { + assert!( + observed_ephemeral_precomputed_lists( + 4, + 4, + Some(5), + true, + 52, + IvfPqBatchTableReuseMode::On, + ) > 0, + "on mode must bypass the automatic batch-size heuristic" + ); + } + #[test] fn test_batch_reader_empty_roaring_filter_returns_empty_results() { use crate::io::{write_index, IVFPQIndexReader, PosWriter}; diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index 7fa2c2c..7a5368b 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -19,9 +19,9 @@ use paimon_vindex_core::distance::MetricType; use paimon_vindex_core::index::{ - SearchWidth, VectorIndexConfig, VectorIndexMetadata, VectorIndexReadPlan, VectorIndexReader, - VectorIndexReaderOptions, VectorIndexTrainer, VectorIndexTraining, VectorIndexWriter, - VectorSearchParams, + IvfPqBatchTableReuseMode, SearchWidth, VectorIndexConfig, VectorIndexMetadata, + VectorIndexReadPlan, VectorIndexReader, VectorIndexReaderOptions, VectorIndexTrainer, + VectorIndexTraining, VectorIndexWriter, VectorSearchParams, }; use paimon_vindex_core::io::{ReadRequest, SeekRead, SeekReadCapabilities, SeekWrite}; use std::cell::RefCell; @@ -525,6 +525,7 @@ fn search_params_from_ffi(params: PaimonVindexSearchParams) -> Result Result Result Result { + match code { + 0 => Ok(IvfPqBatchTableReuseMode::Off), + 1 => Ok(IvfPqBatchTableReuseMode::On), + 2 => Ok(IvfPqBatchTableReuseMode::Auto), + value => Err(format!("invalid IVF-PQ batch table reuse mode: {value}")), + } +} + fn call_int_method(env: &mut JNIEnv, object: &JObject, name: &str) -> Result { env.call_method(object, name, "()I", &[]) .and_then(|value| value.i()) @@ -1023,3 +1035,25 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_fre } }) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ivfpq_batch_table_reuse_codes_map_to_core_modes() { + assert_eq!( + ivfpq_batch_table_reuse_mode(0).unwrap(), + IvfPqBatchTableReuseMode::Off + ); + assert_eq!( + ivfpq_batch_table_reuse_mode(1).unwrap(), + IvfPqBatchTableReuseMode::On + ); + assert_eq!( + ivfpq_batch_table_reuse_mode(2).unwrap(), + IvfPqBatchTableReuseMode::Auto + ); + assert!(ivfpq_batch_table_reuse_mode(3).is_err()); + } +} From 108893c5626f24607fa46fbf0dcb32b4b7152adf Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Wed, 29 Jul 2026 02:28:49 -0700 Subject: [PATCH 6/7] Fix 8-bit IVF-PQ batch table reuse stability --- core/src/ivfpq.rs | 363 ++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 318 insertions(+), 45 deletions(-) diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index 742b89b..bbeda05 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -16,10 +16,10 @@ // under the License. use crate::distance::{ - fvec_inner_product, fvec_l2sqr, fvec_madd, fvec_normalize, pq_distance_four_codes, - pq_distance_from_table, MetricType, + fvec_inner_product, fvec_madd, fvec_normalize, pq_distance_four_codes, pq_distance_from_table, + MetricType, }; -use crate::index_io_util::ivf_payload_is_oversized; +use crate::index_io_util::{ivf_payload_is_oversized, MAX_IVF_BATCH_READ_BYTES}; use crate::io::{IVFPQIndexReader, InvertedListPayload, SeekRead}; use crate::kmeans::{self, KMeansConfig}; use crate::opq::OPQMatrix; @@ -943,6 +943,25 @@ fn should_use_ephemeral_precomputation( setup_tables > 0 && probe_count >= setup_tables.saturating_mul(2) } +fn ephemeral_precomputed_table_fits_budget( + matching_list_count: usize, + query_scratch_count: usize, + m: usize, + ksub: usize, +) -> bool { + if matching_list_count == 0 { + return false; + } + matching_list_count + .checked_add(1) + .and_then(|tables| tables.checked_add(query_scratch_count)) + .and_then(|tables| tables.checked_mul(m)) + .and_then(|values| values.checked_mul(ksub)) + .and_then(|values| values.checked_mul(std::mem::size_of::())) + .is_some_and(|bytes| bytes <= MAX_IVF_BATCH_READ_BYTES) +} + +#[cfg(test)] fn fill_list_precomputed_table( coarse_centroid: &[f32], pq: &ProductQuantizer, @@ -969,6 +988,92 @@ fn fill_list_precomputed_table( } } +fn compute_stable_ephemeral_pq_norms(pq: &ProductQuantizer) -> Vec { + let mut norms = vec![0.0f64; pq.m * pq.ksub]; + for sub in 0..pq.m { + let range = pq.chunk_range(sub); + let chunk_dim = range.len(); + let pq_base = range.start * pq.ksub; + for code in 0..pq.ksub { + let pq_offset = pq_base + code * chunk_dim; + norms[sub * pq.ksub + code] = (0..chunk_dim) + .map(|dimension| { + let value = f64::from(pq.centroids[pq_offset + dimension]); + value * value + }) + .sum(); + } + } + norms +} + +fn fill_stable_ephemeral_list_table( + coarse_centroid: &[f32], + pq: &ProductQuantizer, + pq_norms: &[f64], + table: &mut Vec, +) { + table.resize(pq.m * pq.ksub, 0.0); + for sub in 0..pq.m { + let range = pq.chunk_range(sub); + let chunk_dim = range.len(); + let pq_base = range.start * pq.ksub; + for code in 0..pq.ksub { + let pq_offset = pq_base + code * chunk_dim; + let mut inner_product = 0.0f64; + for dimension in 0..chunk_dim { + let pq_value = f64::from(pq.centroids[pq_offset + dimension]); + inner_product += f64::from(coarse_centroid[range.start + dimension]) * pq_value; + } + let offset = sub * pq.ksub + code; + table[offset] = pq_norms[offset] + 2.0 * inner_product; + } + } +} + +fn fill_stable_ephemeral_query_table(query: &[f32], pq: &ProductQuantizer, table: &mut Vec) { + table.resize(pq.m * pq.ksub, 0.0); + for sub in 0..pq.m { + let range = pq.chunk_range(sub); + let chunk_dim = range.len(); + let pq_base = range.start * pq.ksub; + for code in 0..pq.ksub { + let pq_offset = pq_base + code * chunk_dim; + let mut inner_product = 0.0f64; + for dimension in 0..chunk_dim { + inner_product += f64::from(query[range.start + dimension]) + * f64::from(pq.centroids[pq_offset + dimension]); + } + table[sub * pq.ksub + code] = inner_product; + } + } +} + +fn combine_stable_ephemeral_tables( + list_table: &[f64], + query_table: &[f64], + query: &[f32], + coarse_centroid: &[f32], + pq: &ProductQuantizer, + sim_table: &mut Vec, +) { + sim_table.resize(pq.m * pq.ksub, 0.0); + for sub in 0..pq.m { + let range = pq.chunk_range(sub); + let mut residual_norm = 0.0f64; + for dimension in range { + let residual = f64::from(query[dimension]) - f64::from(coarse_centroid[dimension]); + residual_norm += residual * residual; + } + let table_base = sub * pq.ksub; + for code in 0..pq.ksub { + let offset = table_base + code; + sim_table[offset] = + (residual_norm + list_table[offset] - 2.0 * query_table[offset]).max(0.0) as f32; + } + } +} + /// Scan 4-bit packed codes using u8-domain accumulation. fn scan_codes_4bit( sim_table: &[f32], @@ -1221,7 +1326,7 @@ struct ReaderSearchContext<'a> { #[derive(Default)] struct ReaderScanScratch { sim_table: Vec, - ip_table: Vec, + ip_table: Vec, distances: Vec, } @@ -1751,8 +1856,10 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( } unique_lists.sort_unstable_by_key(|&list_id| reader.list_offsets[list_id]); - let use_precomputed = - metric == MetricType::L2 && by_residual && !reader.precomputed_table.is_empty(); + let use_precomputed = reuse_mode != IvfPqBatchTableReuseMode::Off + && metric == MetricType::L2 + && by_residual + && !reader.precomputed_table.is_empty(); let allow_ephemeral_precomputed = metric == MetricType::L2 && by_residual && !use_precomputed @@ -1761,8 +1868,6 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( IvfPqBatchTableReuseMode::On => true, IvfPqBatchTableReuseMode::Auto => nq >= MIN_EPHEMERAL_PRECOMPUTE_QUERIES, }; - let mut pq_norms = None; - let all_ip_tables: Vec> = if use_precomputed { (0..nq) .into_par_iter() @@ -1777,6 +1882,7 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( } else { Vec::new() }; + let mut stable_pq_norms = None; let mut heaps = (0..nq).map(|_| TopKHeap::new(k)).collect::>(); let mut batch_start = 0usize; @@ -1848,43 +1954,53 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( .iter() .map(|list| matching_rows(&list.ids, filter)) .collect::>(); + let matching_list_count = matching_rows_by_list + .iter() + .filter(|rows| has_matching_rows(rows.as_ref())) + .count(); + let (active_query_count, probe_count) = if allow_ephemeral_precomputed { + let mut probe_count = 0usize; + let mut active_query_count = 0usize; + for probe_indices in &all_probe_indices { + let matching_probe_count = probe_indices + .iter() + .filter(|&&list_id| { + let position = list_positions[list_id]; + position != usize::MAX + && has_matching_rows(matching_rows_by_list[position].as_ref()) + }) + .count(); + probe_count += matching_probe_count; + active_query_count += usize::from(matching_probe_count > 0); + } + (active_query_count, probe_count) + } else { + (0, 0) + }; + let query_scratch_count = active_query_count.min(rayon::current_num_threads()); let use_ephemeral_precomputed = allow_ephemeral_precomputed - && (reuse_mode == IvfPqBatchTableReuseMode::On || { - let mut matching_list_count = 0usize; - let mut probe_count = 0usize; - let mut active_query_count = 0usize; - - for rows in &matching_rows_by_list { - matching_list_count += usize::from(has_matching_rows(rows.as_ref())); - } - for probe_indices in &all_probe_indices { - let matching_probe_count = probe_indices - .iter() - .filter(|&&list_id| { - let position = list_positions[list_id]; - position != usize::MAX - && has_matching_rows(matching_rows_by_list[position].as_ref()) - }) - .count(); - probe_count += matching_probe_count; - active_query_count += usize::from(matching_probe_count > 0); - } - - should_use_ephemeral_precomputation( + && ephemeral_precomputed_table_fits_budget( + matching_list_count, + query_scratch_count, + m, + ksub, + ) + && (reuse_mode == IvfPqBatchTableReuseMode::On + || should_use_ephemeral_precomputation( matching_list_count, active_query_count, probe_count, - ) - }); + )); let ephemeral_precomputed_tables = if use_ephemeral_precomputed { - let pq_norms = pq_norms.get_or_insert_with(|| reader.pq.compute_centroid_norms()); + let pq_norms = stable_pq_norms + .get_or_insert_with(|| compute_stable_ephemeral_pq_norms(&reader.pq)); loaded_lists .par_iter() .zip(&matching_rows_by_list) .map(|(list, rows)| { let mut table = Vec::new(); if has_matching_rows(rows.as_ref()) { - fill_list_precomputed_table( + fill_stable_ephemeral_list_table( &reader.quantizer_centroids[list.list_id * d..(list.list_id + 1) * d], &reader.pq, pq_norms, @@ -1934,10 +2050,7 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( position != usize::MAX && !ephemeral_precomputed_tables[position].is_empty() }); if query_uses_ephemeral_precomputed { - scratch.ip_table.resize(m * ksub, 0.0); - reader - .pq - .compute_inner_product_table(query, &mut scratch.ip_table); + fill_stable_ephemeral_query_table(query, &reader.pq, &mut scratch.ip_table); } for (probe_rank, &list_id) in all_probe_indices[qi].iter().enumerate() { let position = list_positions[list_id]; @@ -1947,21 +2060,19 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( let use_ephemeral_list = query_uses_ephemeral_precomputed && !ephemeral_precomputed_tables[position].is_empty(); let dis0 = if use_ephemeral_list { - fvec_l2sqr( - query, - &reader.quantizer_centroids[list_id * d..(list_id + 1) * d], - ) + 0.0 } else if use_precomputed { all_coarse_dists[qi][probe_rank] } else { 0.0 }; if use_ephemeral_list { - scratch.sim_table.resize(m * ksub, 0.0); - fvec_madd( + combine_stable_ephemeral_tables( &ephemeral_precomputed_tables[position], &scratch.ip_table, - -2.0, + query, + &reader.quantizer_centroids[list_id * d..(list_id + 1) * d], + &reader.pq, &mut scratch.sim_table, ); scan_reader_codes( @@ -2947,6 +3058,32 @@ mod tests { assert!(!should_use_ephemeral_precomputation(0, 0, 0)); } + #[test] + fn ephemeral_precomputation_respects_batch_memory_budget() { + let max_values = + crate::index_io_util::MAX_IVF_BATCH_READ_BYTES / std::mem::size_of::(); + let max_list_values = max_values / 3; + assert!(ephemeral_precomputed_table_fits_budget( + 1, + 1, + 1, + max_list_values + )); + assert!(!ephemeral_precomputed_table_fits_budget( + 1, + 1, + 1, + max_list_values + 1 + )); + assert!(!ephemeral_precomputed_table_fits_budget(0, 1, 1, 1)); + assert!(!ephemeral_precomputed_table_fits_budget( + usize::MAX, + usize::MAX, + usize::MAX, + usize::MAX + )); + } + #[test] fn test_precomputed_table_matches_normal_search() { let d = 16; @@ -3657,6 +3794,95 @@ mod tests { } } + #[test] + fn forced_ephemeral_reuse_is_stable_for_8bit_large_offsets() { + use crate::io::{write_index, IVFPQIndexReader, PosWriter}; + + let d = 4; + let nlist = 4; + let m = 1; + let n = 512; + let nq = MIN_EPHEMERAL_PRECOMPUTE_QUERIES; + let k = 8; + let common = 1_000_000.0f32; + let spread = 500_000.0f32; + let data = (0..n) + .flat_map(|row| { + let cluster = row % nlist; + let point = row / nlist; + (0..d).map(move |dimension| { + common + + cluster as f32 * spread + + (((point * 17 + dimension * 13) % 31) as f32 - 15.0) * spread / 16.0 + }) + }) + .collect::>(); + let ids = (0..n as i64).collect::>(); + let mut index = IVFPQIndex::new(d, nlist, m, MetricType::L2, false); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut bytes = Vec::new(); + write_index(&index, &mut PosWriter::new(&mut bytes)).unwrap(); + let queries = &data[..nq * d]; + let mut reader = IVFPQIndexReader::open(Cursor::new(bytes)).unwrap(); + let (result_ids, result_distances) = search_batch_reader_with_reuse_mode( + &mut reader, + queries, + nq, + k, + nlist, + IvfPqBatchTableReuseMode::On, + ) + .unwrap(); + + let code_size = index.pq.code_size(); + let mut decoded = vec![0.0f32; d]; + for query_index in 0..nq { + let query = &queries[query_index * d..(query_index + 1) * d]; + for rank in 0..k { + let offset = query_index * k + rank; + let id = result_ids[offset]; + let reported = result_distances[offset]; + assert!( + reported >= 0.0, + "query {query_index} rank {rank} produced negative squared L2 distance {reported}" + ); + + let (list_id, position) = index + .ids + .iter() + .enumerate() + .find_map(|(list_id, list_ids)| { + list_ids + .iter() + .position(|candidate| *candidate == id) + .map(|position| (list_id, position)) + }) + .unwrap(); + let code_offset = position * code_size; + index.pq.decode( + &index.codes[list_id][code_offset..code_offset + code_size], + &mut decoded, + ); + let centroid = &index.quantizer_centroids[list_id * d..(list_id + 1) * d]; + let exact = (0..d) + .map(|dimension| { + let delta = f64::from(query[dimension]) + - f64::from(centroid[dimension]) + - f64::from(decoded[dimension]); + delta * delta + }) + .sum::(); + let tolerance = 1e-3 + 8.0 * f64::from(f32::EPSILON) * exact.abs().max(1.0); + assert!( + (f64::from(reported) - exact).abs() <= tolerance, + "query {query_index} rank {rank} reported {reported}, decoded oracle {exact}, tolerance {tolerance}" + ); + } + } + } + #[test] fn small_filtered_batch_reader_skips_ephemeral_list_precomputation() { assert_eq!( @@ -3737,6 +3963,53 @@ mod tests { ); } + #[test] + fn batch_table_reuse_off_ignores_resident_precomputed_tables() { + use crate::io::{write_index, IVFPQIndexReader, PosWriter}; + + let d = 16; + let nlist = 4; + let m = 4; + let n = 600; + let nq = 4; + let k = 5; + let data = generate_clustered_data(n, d, nlist, 53); + let ids = (0..n as i64).collect::>(); + let mut index = IVFPQIndex::new(d, nlist, m, MetricType::L2, false); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut bytes = Vec::new(); + write_index(&index, &mut PosWriter::new(&mut bytes)).unwrap(); + let queries = &data[..nq * d]; + let mut direct_reader = IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + let expected = search_batch_reader_with_reuse_mode( + &mut direct_reader, + queries, + nq, + k, + nlist, + IvfPqBatchTableReuseMode::Off, + ) + .unwrap(); + + let mut optimized_reader = IVFPQIndexReader::open(Cursor::new(bytes)).unwrap(); + optimized_reader.optimize_for_search().unwrap(); + assert!(!optimized_reader.precomputed_table.is_empty()); + optimized_reader.precomputed_table.fill(1_000_000_000.0); + let actual = search_batch_reader_with_reuse_mode( + &mut optimized_reader, + queries, + nq, + k, + nlist, + IvfPqBatchTableReuseMode::Off, + ) + .unwrap(); + + assert_eq!(actual, expected, "Off must ignore resident reuse tables"); + } + #[test] fn batch_table_reuse_on_precomputes_for_small_batches() { assert!( From 81e617025a396a1892a93912ced9d62fef28474c Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Wed, 29 Jul 2026 20:12:33 -0700 Subject: [PATCH 7/7] Restrict IVF-PQ batch table reuse to 8-bit --- core/src/ivfpq.rs | 93 ++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 91 insertions(+), 2 deletions(-) diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index bbeda05..f0c35ad 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -1860,7 +1860,8 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( && metric == MetricType::L2 && by_residual && !reader.precomputed_table.is_empty(); - let allow_ephemeral_precomputed = metric == MetricType::L2 + let allow_ephemeral_precomputed = reader.pq.nbits == 8 + && metric == MetricType::L2 && by_residual && !use_precomputed && match reuse_mode { @@ -2356,6 +2357,26 @@ mod tests { apply_filter: bool, seed: u64, reuse_mode: IvfPqBatchTableReuseMode, + ) -> usize { + observed_ephemeral_precomputed_lists_with_nbits( + 8, + nq, + nprobe, + filter_step, + apply_filter, + seed, + reuse_mode, + ) + } + + fn observed_ephemeral_precomputed_lists_with_nbits( + nbits: usize, + nq: usize, + nprobe: usize, + filter_step: Option, + apply_filter: bool, + seed: u64, + reuse_mode: IvfPqBatchTableReuseMode, ) -> usize { use crate::io::{write_index, IVFPQIndexReader, PosWriter}; @@ -2366,7 +2387,7 @@ mod tests { let k = 5; let data = generate_clustered_data(n, d, nlist, seed); let ids = (0..n as i64).collect::>(); - let mut index = IVFPQIndex::new(d, nlist, m, MetricType::L2, false); + let mut index = IVFPQIndex::with_nbits(d, nlist, m, nbits, MetricType::L2, false); index.train(&data, n); index.add(&data, &ids, n); @@ -4025,6 +4046,74 @@ mod tests { ); } + #[test] + fn four_bit_batch_table_reuse_modes_skip_ephemeral_precomputation() { + for reuse_mode in [IvfPqBatchTableReuseMode::Auto, IvfPqBatchTableReuseMode::On] { + assert_eq!( + observed_ephemeral_precomputed_lists_with_nbits( + 4, + MIN_EPHEMERAL_PRECOMPUTE_QUERIES, + 4, + None, + false, + 54, + reuse_mode, + ), + 0, + "4-bit {reuse_mode:?} must keep the existing scan path" + ); + } + } + + #[test] + fn four_bit_auto_batch_table_reuse_matches_off() { + use crate::io::{write_index, IVFPQIndexReader, PosWriter}; + + let d = 16; + let nlist = 4; + let m = 4; + let n = 600; + let nq = MIN_EPHEMERAL_PRECOMPUTE_QUERIES; + let k = 10; + let nprobe = nlist; + let data = generate_clustered_data(n, d, nlist, 55); + let ids = (0..n as i64).collect::>(); + let mut index = IVFPQIndex::with_nbits(d, nlist, m, 4, MetricType::L2, false); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut bytes = Vec::new(); + write_index(&index, &mut PosWriter::new(&mut bytes)).unwrap(); + let queries = &data[..nq * d]; + + let mut off_reader = IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + let expected = search_batch_reader_with_reuse_mode( + &mut off_reader, + queries, + nq, + k, + nprobe, + IvfPqBatchTableReuseMode::Off, + ) + .unwrap(); + + let mut auto_reader = IVFPQIndexReader::open(Cursor::new(bytes)).unwrap(); + let actual = search_batch_reader_with_reuse_mode( + &mut auto_reader, + queries, + nq, + k, + nprobe, + IvfPqBatchTableReuseMode::Auto, + ) + .unwrap(); + + assert_eq!( + actual, expected, + "4-bit Auto must preserve the Off path results" + ); + } + #[test] fn test_batch_reader_empty_roaring_filter_returns_empty_results() { use crate::io::{write_index, IVFPQIndexReader, PosWriter};