This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-vector-index.git
The following commit(s) were added to refs/heads/main by this push:
new a722d21 ivfpq: Push down row filters for 8-bit batch scans (#64)
a722d21 is described below
commit a722d213b95a923157292bd8595b324ed2cbad65
Author: shyjsarah <[email protected]>
AuthorDate: Wed Jul 29 22:12:53 2026 +0800
ivfpq: Push down row filters for 8-bit batch scans (#64)
---
core/src/ivfpq.rs | 623 ++++++++++++++++++++++++++++++++++++++++++++++++------
1 file changed, 564 insertions(+), 59 deletions(-)
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::<Vec<_>>()
+ });
let results: Vec<Vec<(f32, i64)>> = (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<String>) ->
io::Error {
io::Error::new(io::ErrorKind::InvalidInput, message.into())
}
+enum MatchingRows {
+ Sparse(Vec<usize>),
+ Bitmap { words: Vec<u64>, 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::<usize>(),
+ Self::Bitmap { words, .. } => words.capacity() *
std::mem::size_of::<u64>(),
+ }
+ }
+}
+
+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<Self::Item> {
+ 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<MatchingRows> {
+ filter.map(|filter| {
+ let bitmap_words = ids.len().div_ceil(64);
+ let sparse_limit =
+ bitmap_words.saturating_mul(std::mem::size_of::<u64>()) /
std::mem::size_of::<usize>();
+ let mut positions = Vec::new();
+ let mut bitmap = None::<Vec<u64>>;
+ 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<f32>,
) {
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<R: SeekRead>(
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<R: SeekRead>(
pq_nbits,
transposed_codes,
dis0,
- filter,
+ positions.as_ref(),
&mut scratch.distances,
&mut heap,
);
@@ -1132,7 +1317,6 @@ pub fn search_with_reader_filter<R: SeekRead>(
q: &q,
ip_table: &ip_table,
use_precomputed,
- filter,
d,
m,
ksub,
@@ -1147,7 +1331,15 @@ pub fn search_with_reader_filter<R: SeekRead>(
.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::<Vec<_>>();
@@ -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<R: SeekRead>(
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<f32>,
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<R: SeekRead>(
// 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<R: SeekRead>(
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<R: SeekRead>(
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::<Vec<_>>();
let rows = (0..nq)
.into_par_iter()
@@ -1482,7 +1714,6 @@ pub fn search_batch_reader_filter<R: SeekRead>(
&[]
},
use_precomputed,
- filter,
d,
m,
ksub,
@@ -1505,7 +1736,14 @@ pub fn search_batch_reader_filter<R: SeekRead>(
} 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::<Vec<_>>();
+ 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::<usize>();
+
+ 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::<Vec<_>>();
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::<Vec<_>>();
+ 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::<Vec<_>>();
+ 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::<Vec<_>>();
+ let codes = (0..count * m)
+ .map(|index| ((index * 37 + 11) % ksub) as u8)
+ .collect::<Vec<_>>();
+ let table = (0..m * ksub)
+ .map(|index| ((index * 29 + 7) % 113) as f32 * 0.03125)
+ .collect::<Vec<_>>();
+ let matching_positions = (0..count).step_by(3).collect::<Vec<_>>();
+ 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::<Vec<_>>();
+ 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::<Vec<_>>();
+ let sparse_filter = [3i64, 511,
900].into_iter().collect::<HashSet<_>>();
+ let sparse = matching_rows(&ids, Some(&sparse_filter)).unwrap();
+ assert!(matches!(sparse, MatchingRows::Sparse(_)));
+ assert_eq!(sparse.positions().collect::<Vec<_>>(), vec![3, 511, 900]);
+
+ let dense_filter = ids.iter().copied().collect::<HashSet<_>>();
+ 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::<Vec<_>>(),
+ (0..ids.len()).collect::<Vec<_>>()
+ );
+ assert!(
+ dense.storage_bytes() <= ids.len().div_ceil(64) * size_of::<u64>(),
+ "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::<Vec<_>>();
+ let codes = (0..count * m)
+ .map(|index| ((index * 73 + 19) % ksub) as u8)
+ .collect::<Vec<_>>();
+ let table = (0..m * ksub)
+ .map(|index| ((index * 31 + 5) % 127) as f32 * 0.015625)
+ .collect::<Vec<_>>();
+ let filter = ids
+ .iter()
+ .copied()
+ .filter(|id| id % 4 != 0)
+ .collect::<HashSet<_>>();
+ 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::<Vec<_>>();
+ 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<i64> = (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};
