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 7237d93..656bb40 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}; @@ -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,74 @@ 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 has_matching_rows(matching_positions: Option<&[usize]>) -> bool { + match matching_positions { + Some(positions) => !positions.is_empty(), + None => true, + } +} + +fn should_scan_matching_positions(count: usize, matching_positions: &[usize]) -> bool { + matching_positions.len() <= count / 2 +} + +// Below this size, table construction and Rayon scheduling dominate the saved +// 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, + 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], @@ -756,19 +851,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 +894,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 +916,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 +986,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 +1008,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 +1067,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 +1087,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 +1149,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 +1158,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 +1167,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, @@ -982,6 +1181,7 @@ struct ReaderSearchContext<'a> { #[derive(Default)] struct ReaderScanScratch { sim_table: Vec, + ip_table: Vec, distances: Vec, } @@ -1088,6 +1288,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 +1298,7 @@ pub fn search_with_reader_filter( pq_nbits, transposed_codes, dis0, - filter, + positions.as_deref(), &mut scratch.distances, &mut heap, ); @@ -1132,7 +1333,6 @@ pub fn search_with_reader_filter( q: &q, ip_table: &ip_table, use_precomputed, - filter, d, m, ksub, @@ -1147,7 +1347,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 +1391,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 +1408,7 @@ fn scan_reader_list( ctx.pq.nbits, ctx.transposed_codes, dis0, - ctx.filter, + matching_positions, &mut scratch.distances, heap, ); @@ -1241,7 +1453,6 @@ fn reader_sim_table( q: query, ip_table, use_precomputed, - filter: None, d: reader.d, m: reader.m, ksub: reader.ksub, @@ -1267,22 +1478,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, + ); } } @@ -1295,7 +1547,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. @@ -1306,6 +1576,70 @@ pub fn search_batch_reader_filter( k: usize, nprobe: usize, filter: Option<&dyn RowIdFilter>, +) -> io::Result<(Vec, Vec)> { + 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], + nq: usize, + k: usize, + 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; @@ -1388,6 +1722,15 @@ 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 + && 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 { (0..nq) @@ -1442,6 +1785,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 +1796,7 @@ pub fn search_batch_reader_filter( pq_nbits, transposed_codes, *dis0, - filter, + positions.as_deref(), &mut distances, &mut heaps[*query_index], ); @@ -1469,6 +1813,67 @@ 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 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 positions in &matching_positions_by_list { + matching_list_count += usize::from(has_matching_rows(positions.as_deref())); + } + 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_positions_by_list[position].as_deref(), + ) + }) + .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_positions_by_list) + .map(|(list, positions)| { + let mut table = Vec::new(); + if has_matching_rows(positions.as_deref()) { + 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() @@ -1482,7 +1887,6 @@ pub fn search_batch_reader_filter( &[] }, use_precomputed, - filter, d, m, ksub, @@ -1495,17 +1899,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, &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_positions_by_list[position].as_deref(), + &mut scratch.distances, + &mut heap, + ); + } else { + scan_reader_list( + &loaded_lists[position], + dis0, + &ctx, + matching_positions_by_list[position].as_deref(), + &mut scratch, + &mut heap, + ); + } } heap.into_sorted() }) @@ -1540,9 +1992,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 --- @@ -1646,8 +2126,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, @@ -1709,6 +2209,58 @@ mod tests { data } + fn observed_ephemeral_precomputed_lists( + nq: usize, + nprobe: usize, + filter_step: Option, + apply_filter: bool, + seed: u64, + reuse_mode: IvfPqBatchTableReuseMode, + ) -> 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_reuse_mode_and_observer( + &mut reader, + &data[..nq * d], + nq, + k, + nprobe, + filter, + reuse_mode, + |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(); @@ -1804,6 +2356,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 +2745,59 @@ 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] + 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] @@ -2663,6 +3334,277 @@ 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 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, + IvfPqBatchTableReuseMode::Auto, + ), + 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, + IvfPqBatchTableReuseMode::Auto, + ), + 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, + IvfPqBatchTableReuseMode::Auto, + ), + 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, + 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()); + } +}