From 353a1f17e5b85c87ba2417ee9026ba2c84829a7e Mon Sep 17 00:00:00 2001 From: shaoyijie Date: Mon, 27 Jul 2026 03:30:04 -0700 Subject: [PATCH 1/3] 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/3] 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/3] 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;