diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index 7237d93..0f25512 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_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 { + probed_lists[list_id] = true; + } + } + self.ids + .iter() + .zip(probed_lists) + .map(|(ids, probed)| { + if probed { + matching_rows(ids, Some(filter)).unwrap() + } else { + MatchingRows::Sparse(Vec::new()) + } + }) + .collect::>() + }); let results: Vec> = (0..nq) .into_par_iter() @@ -398,6 +417,10 @@ impl IVFPQIndex { if count == 0 { continue; } + 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. @@ -428,13 +451,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(rows) = matching_rows { + for position in rows.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 +469,7 @@ impl IVFPQIndex { m, ksub, dis0, - filter, + matching_rows, &mut heap, ); } else { @@ -457,7 +481,7 @@ impl IVFPQIndex { m, ksub, dis0, - filter, + matching_rows, &mut heap, ); } @@ -747,6 +771,149 @@ fn invalid_merge_input(message: impl Into) -> io::Error { io::Error::new(io::ErrorKind::InvalidInput, message.into()) } +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| { + 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), + } + }) +} + +// 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. fn scan_codes_4bit( sim_table: &[f32], @@ -756,19 +923,20 @@ fn scan_codes_4bit( m: usize, _ksub: usize, dis0: f32, - filter: Option<&dyn RowIdFilter>, + matching_rows: Option<&MatchingRows>, heap: &mut TopKHeap, ) { 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(rows) = matching_rows { + for position in rows.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 +949,7 @@ fn scan_codes_4bit_transposed( count: usize, m: usize, dis0: f32, - filter: Option<&dyn RowIdFilter>, + matching_rows: Option<&MatchingRows>, heap: &mut TopKHeap, ) { let cs = m / 2; @@ -835,13 +1003,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(rows) = matching_rows { + for position in rows.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 +1025,24 @@ fn scan_codes_transposed_with_scratch( m: usize, ksub: usize, dis0: f32, - filter: Option<&dyn RowIdFilter>, + matching_rows: Option<&MatchingRows>, heap: &mut TopKHeap, dists: &mut Vec, ) { debug_assert!(m > 0); + if let Some(rows) = + matching_rows.filter(|rows| should_scan_sparse(count, rows, TRANSPOSED_SPARSE_SCAN_DIVISOR)) + { + 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]; + } + heap.push(distance, ids[row]); + } + return; + } + dists.resize(count, 0.0); let table = &sim_table[..ksub]; let column = &codes[..count]; @@ -902,13 +1084,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(rows) = matching_rows { + for position in rows.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 +1104,20 @@ fn scan_codes_batched( m: usize, ksub: usize, dis0: f32, - filter: Option<&dyn RowIdFilter>, + matching_rows: Option<&MatchingRows>, heap: &mut TopKHeap, ) { + if let Some(rows) = + matching_rows.filter(|rows| should_scan_sparse(count, rows, ROW_MAJOR_SPARSE_SCAN_DIVISOR)) + { + 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]); + } + return; + } + let mut i = 0; while i + 4 <= count { @@ -937,28 +1131,19 @@ 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; - } + if matching_rows.is_none_or(|rows| rows.contains(idx)) { + heap.push(dis0 + dists[j], ids[idx]); } - heap.push(dis0 + dists[j], id); } 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]; - if let Some(f) = filter { - if !f.contains(id) { - i += 1; - continue; - } + 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]); } - heap.push(dist, id); i += 1; } } @@ -967,7 +1152,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 +1272,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_rows(ids, filter); scan_reader_codes( &sim_table, codes, @@ -1097,7 +1282,7 @@ pub fn search_with_reader_filter( pq_nbits, transposed_codes, dis0, - filter, + positions.as_ref(), &mut scratch.distances, &mut heap, ); @@ -1132,7 +1317,6 @@ pub fn search_with_reader_filter( q: &q, ip_table: &ip_table, use_precomputed, - filter, d, m, ksub, @@ -1147,7 +1331,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_rows(&entry.ids, filter); + scan_reader_list( + entry, + *dis0, + &ctx, + positions.as_ref(), + scratch, + &mut local_heap, + ); local_heap.into_sorted() }) .collect::>(); @@ -1183,9 +1375,13 @@ fn scan_reader_list( entry: &InvertedListPayload, dis0: f32, ctx: &ReaderSearchContext<'_>, + matching_rows: Option<&MatchingRows>, scratch: &mut ReaderScanScratch, heap: &mut TopKHeap, ) { + if matching_rows.is_some_and(MatchingRows::is_empty) { + return; + } fill_reader_sim_table(entry.list_id, ctx, &mut scratch.sim_table); scan_reader_codes( &scratch.sim_table, @@ -1196,7 +1392,7 @@ fn scan_reader_list( ctx.pq.nbits, ctx.transposed_codes, dis0, - ctx.filter, + matching_rows, &mut scratch.distances, heap, ); @@ -1241,7 +1437,6 @@ fn reader_sim_table( q: query, ip_table, use_precomputed, - filter: None, d: reader.d, m: reader.m, ksub: reader.ksub, @@ -1267,22 +1462,54 @@ fn scan_reader_codes( pq_nbits: usize, transposed_codes: bool, dis0: f32, - filter: Option<&dyn RowIdFilter>, + matching_rows: Option<&MatchingRows>, distances: &mut Vec, heap: &mut TopKHeap, ) { + 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, filter, heap); + scan_codes_4bit_transposed(sim_table, codes, ids, count, m, dis0, matching_rows, 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_rows, + 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_rows, + 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_rows, + heap, + ); } } @@ -1442,6 +1669,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_rows(ids, filter); for (query_index, dis0, sim_table) in &query_tables { scan_reader_codes( sim_table, @@ -1452,7 +1680,7 @@ pub fn search_batch_reader_filter( pq_nbits, transposed_codes, *dis0, - filter, + positions.as_ref(), &mut distances, &mut heaps[*query_index], ); @@ -1469,6 +1697,10 @@ pub fn search_batch_reader_filter( for (position, list) in loaded_lists.iter().enumerate() { list_positions[list.list_id] = position; } + let matching_rows_by_list = loaded_lists + .iter() + .map(|list| matching_rows(&list.ids, filter)) + .collect::>(); let rows = (0..nq) .into_par_iter() @@ -1482,7 +1714,6 @@ pub fn search_batch_reader_filter( &[] }, use_precomputed, - filter, d, m, ksub, @@ -1505,7 +1736,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_rows_by_list[position].as_ref(), + &mut scratch, + &mut heap, + ); } heap.into_sorted() }) @@ -1646,8 +1884,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 +2062,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 +2451,144 @@ 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 matching_rows = MatchingRows::Sparse(matching_positions); + let mut filtered_heap = TopKHeap::new(matching_rows.len()); + scan_codes_transposed_with_scratch( + &table, + &codes, + &ids, + count, + m, + ksub, + dis0, + Some(&matching_rows), + &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] + fn row_major_sparse_scan_matches_exact_distances() { + let count = 400; + let m = 8; + let ksub = 256; + let dis0 = 1.25; + let ids = (10_000..10_000 + count as i64).collect::>(); + 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 = matching_positions + .iter() + .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_batched( + &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] @@ -2663,6 +3125,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};