This is an automated email from the ASF dual-hosted git repository.
jerry-024 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 cf0506d ivf: incrementally scan expanded batch probe ranges (#73)
cf0506d is described below
commit cf0506dac7fc1700d5a9898c00a41cc42a7c4868
Author: shyjsarah <[email protected]>
AuthorDate: Wed Aug 12 15:44:15 2026 +0800
ivf: incrementally scan expanded batch probe ranges (#73)
Co-authored-by: shaoyijie <[email protected]>
---
core/src/index.rs | 340 ++++++++++++++++++++++++++++++++++++-------------
core/src/ivfflat_io.rs | 102 ++++++++++++++-
core/src/ivfpq.rs | 304 ++++++++++++++++++++++++++++++++++++++++---
core/src/ivfrq_io.rs | 107 +++++++++++++++-
core/src/ivfsq_io.rs | 103 ++++++++++++++-
5 files changed, 835 insertions(+), 121 deletions(-)
diff --git a/core/src/index.rs b/core/src/index.rs
index 6454bd1..60c7762 100644
--- a/core/src/index.rs
+++ b/core/src/index.rs
@@ -30,25 +30,25 @@ use crate::distance::MetricType;
use crate::io::{write_index, IVFPQIndexReader, ReadRequest, SeekRead,
SeekWrite, MAGIC};
use crate::ivfflat::IVFFlatIndex;
use crate::ivfflat_io::{
- search_batch_ivfflat_reader, search_batch_ivfflat_reader_roaring_filter,
write_ivfflat_index,
- IVFFlatIndexReader, IVFFLAT_MAGIC,
+ search_batch_ivfflat_reader_filter_range,
search_batch_ivfflat_reader_roaring_filter_range,
+ write_ivfflat_index, IVFFlatIndexReader, IVFFLAT_MAGIC,
};
use crate::ivfpq::{
- search_batch_reader_roaring_filter_with_reuse_mode_and_budget,
- search_batch_reader_with_reuse_mode_and_budget, search_with_reader,
+ search_batch_reader_roaring_filter_with_reuse_mode_and_budget_range,
+ search_batch_reader_with_reuse_mode_and_budget_range, search_with_reader,
search_with_reader_roaring_filter, IVFPQIndex,
};
pub use crate::ivfpq::{IvfPqBatchTableReuseMode,
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES};
use crate::ivfrq::IVFRQIndex;
pub use crate::ivfrq_io::IVFRQSearchStats;
use crate::ivfrq_io::{
- search_batch_ivfrq_reader, search_batch_ivfrq_reader_roaring_filter,
write_ivfrq_index,
- IVFRQIndexReader, IVF_RQ_MAGIC,
+ search_batch_ivfrq_reader_filter_range,
search_batch_ivfrq_reader_roaring_filter_range,
+ write_ivfrq_index, IVFRQIndexReader, IVF_RQ_MAGIC,
};
use crate::ivfsq::IVFSQIndex;
use crate::ivfsq_io::{
- search_batch_ivfsq_reader, search_batch_ivfsq_reader_roaring_filter,
write_ivfsq_index,
- IVFSQIndexReader, IVF_SQ_MAGIC,
+ search_batch_ivfsq_reader_filter_range,
search_batch_ivfsq_reader_roaring_filter_range,
+ write_ivfsq_index, IVFSQIndexReader, IVF_SQ_MAGIC,
};
pub use crate::read_options::{DeploymentProfile, VectorIndexReadPlan,
VectorIndexReaderOptions};
use crate::rq::{is_supported_rq_bits, padded_dimension, DEFAULT_RQ_BITS};
@@ -1784,13 +1784,22 @@ impl<R: SeekRead> VectorIndexReader<R> {
query_count,
params.top_k,
total_vectors,
- |active_queries, active_query_count, nprobe| {
- search_batch_ivfflat_reader(
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+ search_batch_ivfflat_reader_filter_range(
reader,
active_queries,
active_query_count,
params.top_k,
- nprobe,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
+ None,
)
},
)
@@ -1807,13 +1816,22 @@ impl<R: SeekRead> VectorIndexReader<R> {
query_count,
params.top_k,
total_vectors,
- |active_queries, active_query_count, nprobe| {
- search_batch_ivfsq_reader(
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+ search_batch_ivfsq_reader_filter_range(
reader,
active_queries,
active_query_count,
params.top_k,
- nprobe,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
+ None,
)
},
)
@@ -1830,13 +1848,21 @@ impl<R: SeekRead> VectorIndexReader<R> {
query_count,
params.top_k,
total_vectors,
- |active_queries, active_query_count, nprobe| {
- search_batch_reader_with_reuse_mode_and_budget(
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+ search_batch_reader_with_reuse_mode_and_budget_range(
reader,
active_queries,
active_query_count,
params.top_k,
- nprobe,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
params.ivfpq_batch_table_reuse,
params.ivfpq_batch_table_reuse_max_bytes,
)
@@ -1856,13 +1882,22 @@ impl<R: SeekRead> VectorIndexReader<R> {
query_count,
params.top_k,
total_vectors,
- |active_queries, active_query_count, nprobe| {
- let result = search_batch_ivfrq_reader(
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+ let result = search_batch_ivfrq_reader_filter_range(
reader,
active_queries,
active_query_count,
params.top_k,
- nprobe,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
+ None,
);
if result.is_ok() {
aggregate_stats.merge(reader.last_search_stats());
@@ -1912,13 +1947,21 @@ impl<R: SeekRead> VectorIndexReader<R> {
query_count,
params.top_k,
matching_count.unwrap_or(total_vectors),
- |active_queries, active_query_count, nprobe| {
- search_batch_ivfflat_reader_roaring_filter(
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+ search_batch_ivfflat_reader_roaring_filter_range(
reader,
active_queries,
active_query_count,
params.top_k,
- nprobe,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
roaring_filter_bytes,
)
},
@@ -1937,13 +1980,21 @@ impl<R: SeekRead> VectorIndexReader<R> {
query_count,
params.top_k,
matching_count.unwrap_or(total_vectors),
- |active_queries, active_query_count, nprobe| {
- search_batch_ivfsq_reader_roaring_filter(
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+ search_batch_ivfsq_reader_roaring_filter_range(
reader,
active_queries,
active_query_count,
params.top_k,
- nprobe,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
roaring_filter_bytes,
)
},
@@ -1962,13 +2013,21 @@ impl<R: SeekRead> VectorIndexReader<R> {
query_count,
params.top_k,
matching_count.unwrap_or(total_vectors),
- |active_queries, active_query_count, nprobe| {
-
search_batch_reader_roaring_filter_with_reuse_mode_and_budget(
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+
search_batch_reader_roaring_filter_with_reuse_mode_and_budget_range(
reader,
active_queries,
active_query_count,
params.top_k,
- nprobe,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
roaring_filter_bytes,
params.ivfpq_batch_table_reuse,
params.ivfpq_batch_table_reuse_max_bytes,
@@ -1990,13 +2049,21 @@ impl<R: SeekRead> VectorIndexReader<R> {
query_count,
params.top_k,
matching_count.unwrap_or(total_vectors),
- |active_queries, active_query_count, nprobe| {
- let result = search_batch_ivfrq_reader_roaring_filter(
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+ let result =
search_batch_ivfrq_reader_roaring_filter_range(
reader,
active_queries,
active_query_count,
params.top_k,
- nprobe,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
roaring_filter_bytes,
);
if result.is_ok() {
@@ -2060,8 +2127,8 @@ fn progressive_ivf_search(
/// Runs automatic IVF batch expansion independently for each query.
///
/// Queries which already produced the required number of results are removed
from later rounds.
-/// Retried queries are still searched from the first probed list; reusing
work across rounds is
-/// intentionally left to the incremental-search implementation.
+/// Each callback scans only the half-open probe range passed to it. In
expansion rounds it receives
+/// the active queries' accumulated Top-K as a seed and returns the updated
Top-K.
fn progressive_ivf_batch_search(
params: VectorSearchParams,
nlist: usize,
@@ -2070,7 +2137,7 @@ fn progressive_ivf_batch_search(
query_count: usize,
top_k: usize,
available_matches: usize,
- search: impl FnMut(&[f32], usize, usize) -> io::Result<(Vec<i64>,
Vec<f32>)>,
+ search: impl FnMut(&[f32], usize, usize, usize, &[i64], &[f32]) ->
io::Result<(Vec<i64>, Vec<f32>)>,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
progressive_ivf_batch_search_with_retry_buffer_limit(
params,
@@ -2095,16 +2162,23 @@ fn progressive_ivf_batch_search_with_retry_buffer_limit(
top_k: usize,
available_matches: usize,
retry_buffer_limit_bytes: usize,
- mut search: impl FnMut(&[f32], usize, usize) -> io::Result<(Vec<i64>,
Vec<f32>)>,
+ mut search: impl FnMut(
+ &[f32],
+ usize,
+ usize,
+ usize,
+ &[i64],
+ &[f32],
+ ) -> io::Result<(Vec<i64>, Vec<f32>)>,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
if params.search_width != SearchWidth::Auto {
- return search(queries, query_count, initial_nprobe);
+ return search(queries, query_count, 0, initial_nprobe, &[], &[]);
}
let dimension = queries.len() / query_count;
let required_per_query = top_k.min(available_matches);
let mut nprobe = initial_nprobe;
- let (mut result_ids, mut result_distances) = search(queries, query_count,
nprobe)?;
+ let (mut result_ids, mut result_distances) = search(queries, query_count,
0, nprobe, &[], &[])?;
let mut active_queries = result_distances
.chunks_exact(top_k)
.take(query_count)
@@ -2119,23 +2193,8 @@ fn progressive_ivf_batch_search_with_retry_buffer_limit(
return Ok((result_ids, result_distances));
}
+ let previous_nprobe = nprobe;
nprobe = nprobe.saturating_mul(2).min(nlist);
- if active_queries.len() == query_count {
- drop(result_ids);
- drop(result_distances);
- (result_ids, result_distances) = search(queries, query_count,
nprobe)?;
- active_queries = result_distances
- .chunks_exact(top_k)
- .take(query_count)
- .enumerate()
- .filter_map(|(query_index, distances)| {
- (!ivf_search_result_is_complete(distances,
required_per_query))
- .then_some(query_index)
- })
- .collect();
- continue;
- }
-
let bytes_per_retry_query = dimension
.saturating_mul(std::mem::size_of::<f32>())
.saturating_add(
@@ -2148,12 +2207,25 @@ fn progressive_ivf_batch_search_with_retry_buffer_limit(
let mut next_active_queries = Vec::new();
for active_chunk in active_queries.chunks(retry_chunk_size) {
let mut packed_queries = Vec::with_capacity(active_chunk.len() *
dimension);
+ let mut seed_ids = Vec::with_capacity(active_chunk.len() * top_k);
+ let mut seed_distances = Vec::with_capacity(active_chunk.len() *
top_k);
for &query_index in active_chunk {
- let start = query_index * dimension;
- packed_queries.extend_from_slice(&queries[start..start +
dimension]);
+ let query_start = query_index * dimension;
+
packed_queries.extend_from_slice(&queries[query_start..query_start +
dimension]);
+ let result_start = query_index * top_k;
+
seed_ids.extend_from_slice(&result_ids[result_start..result_start + top_k]);
+ seed_distances
+
.extend_from_slice(&result_distances[result_start..result_start + top_k]);
}
- let (round_ids, round_distances) = search(&packed_queries,
active_chunk.len(), nprobe)?;
+ let (round_ids, round_distances) = search(
+ &packed_queries,
+ active_chunk.len(),
+ previous_nprobe,
+ nprobe,
+ &seed_ids,
+ &seed_distances,
+ )?;
for (round_index, &query_index) in active_chunk.iter().enumerate()
{
let round_start = round_index * top_k;
let result_start = query_index * top_k;
@@ -3253,27 +3325,39 @@ mod tests {
}
#[test]
- fn automatic_batch_search_retries_only_incomplete_queries() {
+ fn automatic_batch_search_scans_only_new_probe_ranges_and_merges_results()
{
let queries = vec![10.0, 20.0, 30.0];
let mut observed = Vec::new();
let result = progressive_ivf_batch_search(
- VectorSearchParams::automatic(2),
+ VectorSearchParams::automatic(3),
8,
2,
&queries,
3,
- 2,
+ 3,
10,
- |active_queries, active_query_count, nprobe| {
- observed.push((nprobe, active_query_count,
active_queries.to_vec()));
- match nprobe {
- 2 => Ok((
- vec![100, 101, 200, -1, 300, 301],
- vec![1.0, 2.0, 1.0, f32::MAX, 1.0, 2.0],
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+ observed.push((
+ probe_start,
+ probe_end,
+ active_query_count,
+ active_queries.to_vec(),
+ seed_ids.to_vec(),
+ seed_distances.to_vec(),
+ ));
+ match (probe_start, probe_end) {
+ (0, 2) => Ok((
+ vec![100, 101, 102, 200, -1, -1, 300, 301, 302],
+ vec![1.0, 2.0, 3.0, 5.0, f32::MAX, f32::MAX, 1.0, 2.0,
3.0],
)),
- 4 => Ok((vec![200, -1], vec![1.0, f32::MAX])),
- 8 => Ok((vec![200, 201], vec![1.0, 2.0])),
- _ => unreachable!("unexpected nprobe {nprobe}"),
+ (2, 4) => Ok((vec![201, 200, -1], vec![4.0, 5.0,
f32::MAX])),
+ (4, 8) => Ok((vec![202, 201, 200], vec![3.0, 4.0, 5.0])),
+ _ => unreachable!("unexpected probe range
{probe_start}..{probe_end}"),
}
},
)
@@ -3282,17 +3366,31 @@ mod tests {
assert_eq!(
observed,
vec![
- (2, 3, vec![10.0, 20.0, 30.0]),
- (4, 1, vec![20.0]),
- (8, 1, vec![20.0]),
+ (0, 2, 3, vec![10.0, 20.0, 30.0], vec![], vec![]),
+ (
+ 2,
+ 4,
+ 1,
+ vec![20.0],
+ vec![200, -1, -1],
+ vec![5.0, f32::MAX, f32::MAX]
+ ),
+ (
+ 4,
+ 8,
+ 1,
+ vec![20.0],
+ vec![201, 200, -1],
+ vec![4.0, 5.0, f32::MAX]
+ ),
]
);
- assert_eq!(result.0, vec![100, 101, 200, 201, 300, 301]);
- assert_eq!(result.1, vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0]);
+ assert_eq!(result.0, vec![100, 101, 102, 202, 201, 200, 300, 301,
302]);
+ assert_eq!(result.1, vec![1.0, 2.0, 3.0, 3.0, 4.0, 5.0, 1.0, 2.0,
3.0]);
}
#[test]
- fn automatic_batch_search_chunks_partial_retries_to_bound_memory() {
+ fn automatic_batch_search_chunks_queries_and_seeds_to_bound_retry_memory()
{
let queries = vec![10.0, 11.0, 20.0, 21.0, 30.0, 31.0, 40.0, 41.0];
let mut observed = Vec::new();
let result = progressive_ivf_batch_search_with_retry_buffer_limit(
@@ -3304,9 +3402,21 @@ mod tests {
2,
10,
32,
- |active_queries, active_query_count, nprobe| {
- observed.push((nprobe, active_query_count,
active_queries.to_vec()));
- if nprobe == 2 {
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+ observed.push((
+ probe_start,
+ probe_end,
+ active_query_count,
+ active_queries.to_vec(),
+ seed_ids.to_vec(),
+ seed_distances.to_vec(),
+ ));
+ if probe_start == 0 {
return Ok((
vec![100, 101, 200, -1, 300, -1, 400, -1],
vec![1.0, 2.0, 1.0, f32::MAX, 1.0, f32::MAX, 1.0,
f32::MAX],
@@ -3322,10 +3432,31 @@ mod tests {
assert_eq!(
observed,
vec![
- (2, 4, queries),
- (4, 1, vec![20.0, 21.0]),
- (4, 1, vec![30.0, 31.0]),
- (4, 1, vec![40.0, 41.0]),
+ (0, 2, 4, queries, vec![], vec![]),
+ (
+ 2,
+ 4,
+ 1,
+ vec![20.0, 21.0],
+ vec![200, -1],
+ vec![1.0, f32::MAX],
+ ),
+ (
+ 2,
+ 4,
+ 1,
+ vec![30.0, 31.0],
+ vec![300, -1],
+ vec![1.0, f32::MAX],
+ ),
+ (
+ 2,
+ 4,
+ 1,
+ vec![40.0, 41.0],
+ vec![400, -1],
+ vec![1.0, f32::MAX],
+ ),
]
);
assert_eq!(result.0, vec![100, 101, 20, 21, 30, 31, 40, 41]);
@@ -3343,8 +3474,20 @@ mod tests {
3,
2,
10,
- |active_queries, active_query_count, nprobe| {
- observed.push((nprobe, active_query_count,
active_queries.to_vec()));
+ |active_queries,
+ active_query_count,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances| {
+ observed.push((
+ probe_start,
+ probe_end,
+ active_query_count,
+ active_queries.to_vec(),
+ seed_ids.to_vec(),
+ seed_distances.to_vec(),
+ ));
Ok((
vec![100, 101, 200, 201, 300, 301],
vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0],
@@ -3353,11 +3496,38 @@ mod tests {
)
.unwrap();
- assert_eq!(observed, vec![(4, 3, queries)]);
+ assert_eq!(observed, vec![(0, 4, 3, queries, vec![], vec![])]);
assert_eq!(result.0, vec![100, 101, 200, 201, 300, 301]);
assert_eq!(result.1, vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0]);
}
+ #[test]
+ fn progressive_batch_seed_preserves_negative_one_row_ids() {
+ let result = progressive_ivf_batch_search(
+ VectorSearchParams::automatic(2),
+ 4,
+ 2,
+ &[1.0],
+ 1,
+ 2,
+ 10,
+ |_, _, probe_start, probe_end, seed_ids, seed_distances| {
+ if probe_start == 0 {
+ Ok((vec![-1, -1], vec![0.5, f32::MAX]))
+ } else {
+ assert_eq!((probe_start, probe_end), (2, 4));
+ assert_eq!(seed_ids, &[-1, -1]);
+ assert_eq!(seed_distances, &[0.5, f32::MAX]);
+ Ok((vec![-1, 8], vec![0.5, 2.0]))
+ }
+ },
+ )
+ .unwrap();
+
+ assert_eq!(result.0, vec![-1, 8]);
+ assert_eq!(result.1, vec![0.5, 2.0]);
+ }
+
#[test]
fn capped_automatic_filtered_search_can_expand_past_the_initial_cap() {
let params =
VectorSearchParams::automatic(2).with_max_initial_filter_expansion_factor(4);
diff --git a/core/src/ivfflat_io.rs b/core/src/ivfflat_io.rs
index 1409caa..10ddb5c 100644
--- a/core/src/ivfflat_io.rs
+++ b/core/src/ivfflat_io.rs
@@ -767,6 +767,20 @@ pub fn search_batch_ivfflat_reader_filter<R: SeekRead>(
k: usize,
nprobe: usize,
filter: Option<&dyn RowIdFilter>,
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_ivfflat_reader_filter_range(reader, queries, nq, k, 0,
nprobe, &[], &[], filter)
+}
+
+pub(crate) fn search_batch_ivfflat_reader_filter_range<R: SeekRead>(
+ reader: &mut IVFFlatIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ probe_start: usize,
+ probe_end: usize,
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ filter: Option<&dyn RowIdFilter>,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
reader.ensure_loaded()?;
let d = reader.d;
@@ -799,12 +813,13 @@ pub fn search_batch_ivfflat_reader_filter<R: SeekRead>(
"k must be greater than 0",
));
}
- if nprobe == 0 {
+ if probe_start >= probe_end {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
- "nprobe must be greater than 0",
+ "probe range must be non-empty",
));
}
+ validate_batch_seed(seed_ids, seed_distances, nq, k)?;
let mut processed = queries[..expected_query_len].to_vec();
if reader.metric == MetricType::Cosine {
@@ -819,13 +834,13 @@ pub fn search_batch_ivfflat_reader_filter<R: SeekRead>(
&reader.quantizer_centroids,
reader.nlist,
d,
- nprobe,
+ probe_end,
);
let mut list_to_queries = vec![Vec::new(); reader.nlist];
let mut unique_lists = Vec::new();
for (qi, probe_indices) in all_probe_indices.iter().enumerate() {
- for &list_id in probe_indices {
+ for &list_id in probe_indices.iter().skip(probe_start) {
if list_to_queries[list_id].is_empty() {
unique_lists.push(list_id);
}
@@ -834,6 +849,7 @@ pub fn search_batch_ivfflat_reader_filter<R: SeekRead>(
}
let mut heaps: Vec<ReaderTopKHeap> = (0..nq).map(|_|
ReaderTopKHeap::new(k)).collect();
+ seed_flat_heaps(&mut heaps, seed_ids, seed_distances, k);
let mut batch_start = 0usize;
while batch_start < unique_lists.len() {
let first_list = unique_lists[batch_start];
@@ -914,9 +930,85 @@ pub fn search_batch_ivfflat_reader_roaring_filter<R:
SeekRead>(
k: usize,
nprobe: usize,
roaring_filter_bytes: &[u8],
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_ivfflat_reader_roaring_filter_range(
+ reader,
+ queries,
+ nq,
+ k,
+ 0,
+ nprobe,
+ &[],
+ &[],
+ roaring_filter_bytes,
+ )
+}
+
+pub(crate) fn search_batch_ivfflat_reader_roaring_filter_range<R: SeekRead>(
+ reader: &mut IVFFlatIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ probe_start: usize,
+ probe_end: usize,
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ roaring_filter_bytes: &[u8],
) -> io::Result<(Vec<i64>, Vec<f32>)> {
let filter = decode_roaring_filter(roaring_filter_bytes)?;
- search_batch_ivfflat_reader_filter(reader, queries, nq, k, nprobe,
Some(&filter))
+ search_batch_ivfflat_reader_filter_range(
+ reader,
+ queries,
+ nq,
+ k,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
+ Some(&filter),
+ )
+}
+
+fn validate_batch_seed(
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ nq: usize,
+ k: usize,
+) -> io::Result<()> {
+ let expected = nq
+ .checked_mul(k)
+ .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "nq * k
overflows usize"))?;
+ if (seed_ids.is_empty() && seed_distances.is_empty())
+ || (seed_ids.len() == expected && seed_distances.len() == expected)
+ {
+ Ok(())
+ } else {
+ Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "seed result lengths must both equal nq * k",
+ ))
+ }
+}
+
+fn seed_flat_heaps(
+ heaps: &mut [ReaderTopKHeap],
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ k: usize,
+) {
+ for (query_index, heap) in heaps.iter_mut().enumerate() {
+ let start = query_index * k;
+ for (&id, &distance) in seed_ids
+ .get(start..start + k)
+ .unwrap_or_default()
+ .iter()
+ .zip(seed_distances.get(start..start + k).unwrap_or_default())
+ {
+ if distance != f32::MAX {
+ heap.push(distance, id);
+ }
+ }
+ }
}
fn scan_flat_list(
diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs
index 75afcf1..7cfd70f 100644
--- a/core/src/ivfpq.rs
+++ b/core/src/ivfpq.rs
@@ -1769,18 +1769,78 @@ pub fn
search_batch_reader_with_reuse_mode_and_budget<R: SeekRead>(
reuse_mode: IvfPqBatchTableReuseMode,
reuse_max_bytes: usize,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
- search_batch_reader_filter_with_reuse_mode_and_budget(
+ search_batch_reader_with_reuse_mode_and_budget_range(
reader,
queries,
nq,
k,
+ 0,
nprobe,
+ &[],
+ &[],
+ reuse_mode,
+ reuse_max_bytes,
+ )
+}
+
+pub(crate) fn search_batch_reader_with_reuse_mode_and_budget_range<R:
SeekRead>(
+ reader: &mut IVFPQIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ probe_start: usize,
+ probe_end: usize,
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ reuse_mode: IvfPqBatchTableReuseMode,
+ reuse_max_bytes: usize,
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_reader_filter_with_reuse_mode_and_budget_range(
+ reader,
+ queries,
+ nq,
+ k,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
None,
reuse_mode,
reuse_max_bytes,
)
}
+fn search_batch_reader_filter_with_reuse_mode_and_budget_range<R: SeekRead>(
+ reader: &mut IVFPQIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ probe_start: usize,
+ probe_end: usize,
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ filter: Option<&dyn RowIdFilter>,
+ reuse_mode: IvfPqBatchTableReuseMode,
+ reuse_max_bytes: usize,
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_reader_filter_with_reuse_mode_and_observer(
+ reader,
+ queries,
+ nq,
+ k,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
+ filter,
+ reuse_mode,
+ reuse_max_bytes,
+ |_| {},
+ #[cfg(test)]
+ None,
+ )
+}
+
/// Big batch search with an optional row-id filter.
pub fn search_batch_reader_filter<R: SeekRead>(
reader: &mut IVFPQIndexReader<R>,
@@ -1832,18 +1892,18 @@ pub fn
search_batch_reader_filter_with_reuse_mode_and_budget<R: SeekRead>(
reuse_mode: IvfPqBatchTableReuseMode,
reuse_max_bytes: usize,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
- search_batch_reader_filter_with_reuse_mode_and_observer(
+ search_batch_reader_filter_with_reuse_mode_and_budget_range(
reader,
queries,
nq,
k,
+ 0,
nprobe,
+ &[],
+ &[],
filter,
reuse_mode,
reuse_max_bytes,
- |_| {},
- #[cfg(test)]
- None,
)
}
@@ -1862,7 +1922,10 @@ fn search_batch_reader_filter_with_observer<R: SeekRead>(
queries,
nq,
k,
+ 0,
nprobe,
+ &[],
+ &[],
filter,
IvfPqBatchTableReuseMode::Auto,
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
@@ -1876,7 +1939,10 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
queries: &[f32],
nq: usize,
k: usize,
- nprobe: usize,
+ probe_start: usize,
+ probe_end: usize,
+ seed_ids: &[i64],
+ seed_distances: &[f32],
filter: Option<&dyn RowIdFilter>,
reuse_mode: IvfPqBatchTableReuseMode,
reuse_max_bytes: usize,
@@ -1913,12 +1979,13 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
"k must be greater than 0",
));
}
- if nprobe == 0 {
+ if probe_start >= probe_end {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
- "nprobe must be greater than 0",
+ "probe range must be non-empty",
));
}
+ validate_batch_seed(seed_ids, seed_distances, nq, k)?;
let m = reader.m;
let ksub = reader.ksub;
@@ -1945,7 +2012,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
&reader.quantizer_centroids,
reader.nlist,
d,
- nprobe,
+ probe_end,
);
// Step 3: Read every probed list once. Queries share the decoded list
@@ -1953,7 +2020,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
let mut seen = vec![false; reader.nlist];
let mut unique_lists = Vec::new();
for probe_indices in &all_probe_indices {
- for &list_id in probe_indices {
+ for &list_id in probe_indices.iter().skip(probe_start) {
if !seen[list_id] && reader.list_counts[list_id] > 0 {
seen[list_id] = true;
unique_lists.push(list_id);
@@ -2000,7 +2067,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
// probed list for that query can share one table.
let reuse_non_residual_tables = reader.pq.nbits == 8
&& !by_residual
- && nprobe > 1
+ && probe_end - probe_start > 1
&& match reuse_mode {
IvfPqBatchTableReuseMode::Off => false,
IvfPqBatchTableReuseMode::On => true,
@@ -2029,6 +2096,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
let mut stable_pq_norms = None;
let mut heaps = (0..nq).map(|_| TopKHeap::new(k)).collect::<Vec<_>>();
+ seed_heaps(&mut heaps, seed_ids, seed_distances, k);
let mut batch_start = 0usize;
while batch_start < unique_lists.len() {
let first_list = unique_lists[batch_start];
@@ -2037,7 +2105,11 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
.filter_map(|query_index| {
all_probe_indices[query_index]
.iter()
- .position(|&list_id| list_id == first_list)
+ .enumerate()
+ .skip(probe_start)
+ .find_map(|(probe_rank, &list_id)| {
+ (list_id == first_list).then_some(probe_rank)
+ })
.map(|probe_rank| {
let query = &processed[query_index *
d..(query_index + 1) * d];
let sim_table =
(!reuse_non_residual_tables).then(|| {
@@ -2128,6 +2200,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
for probe_indices in &all_probe_indices {
let matching_probe_count = probe_indices
.iter()
+ .skip(probe_start)
.filter(|&&list_id| {
let position = list_positions[list_id];
position != usize::MAX
@@ -2213,14 +2286,20 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
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()
- });
+ && all_probe_indices[qi]
+ .iter()
+ .skip(probe_start)
+ .any(|&list_id| {
+ let position = list_positions[list_id];
+ position != usize::MAX
+ &&
!ephemeral_precomputed_tables[position].is_empty()
+ });
if query_uses_ephemeral_precomputed {
fill_stable_ephemeral_query_table(query, &reader.pq, &mut
scratch.ip_table);
}
- for (probe_rank, &list_id) in
all_probe_indices[qi].iter().enumerate() {
+ for (probe_rank, &list_id) in
+ all_probe_indices[qi].iter().enumerate().skip(probe_start)
+ {
let position = list_positions[list_id];
if position == usize::MAX {
continue;
@@ -2300,7 +2379,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
std::io::stderr().lock(),
"[paimon-vindex] ivfpq_batch_table_reuse
strategy=non_residual_query_table \
mode={reuse_mode:?} enabled={reuse_non_residual_tables} used={}
metric={} \
- pq_bits={} nq={nq} nprobe={nprobe} unique_lists={} filtered={}
required_bytes={:?} \
+ pq_bits={} nq={nq} nprobe={probe_end} unique_lists={} filtered={}
required_bytes={:?} \
budget_bytes={reuse_max_bytes} tables_built={tables_built}",
tables_built > 0,
metric.as_str(),
@@ -2365,13 +2444,44 @@ pub fn
search_batch_reader_roaring_filter_with_reuse_mode_and_budget<R: SeekRead
reuse_mode: IvfPqBatchTableReuseMode,
reuse_max_bytes: usize,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
- let filter = decode_roaring_filter(roaring_filter_bytes)?;
- search_batch_reader_filter_with_reuse_mode_and_budget(
+ search_batch_reader_roaring_filter_with_reuse_mode_and_budget_range(
reader,
queries,
nq,
k,
+ 0,
nprobe,
+ &[],
+ &[],
+ roaring_filter_bytes,
+ reuse_mode,
+ reuse_max_bytes,
+ )
+}
+
+pub(crate) fn
search_batch_reader_roaring_filter_with_reuse_mode_and_budget_range<R:
SeekRead>(
+ reader: &mut IVFPQIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ probe_start: usize,
+ probe_end: usize,
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ roaring_filter_bytes: &[u8],
+ reuse_mode: IvfPqBatchTableReuseMode,
+ reuse_max_bytes: usize,
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ let filter = decode_roaring_filter(roaring_filter_bytes)?;
+ search_batch_reader_filter_with_reuse_mode_and_budget_range(
+ reader,
+ queries,
+ nq,
+ k,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
Some(&filter),
reuse_mode,
reuse_max_bytes,
@@ -2418,6 +2528,43 @@ impl TopKHeap {
}
}
+fn validate_batch_seed(
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ nq: usize,
+ k: usize,
+) -> io::Result<()> {
+ let expected = nq
+ .checked_mul(k)
+ .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "nq * k
overflows usize"))?;
+ if (seed_ids.is_empty() && seed_distances.is_empty())
+ || (seed_ids.len() == expected && seed_distances.len() == expected)
+ {
+ Ok(())
+ } else {
+ Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "seed result lengths must both equal nq * k",
+ ))
+ }
+}
+
+fn seed_heaps(heaps: &mut [TopKHeap], seed_ids: &[i64], seed_distances:
&[f32], k: usize) {
+ for (query_index, heap) in heaps.iter_mut().enumerate() {
+ let start = query_index * k;
+ for (&id, &distance) in seed_ids
+ .get(start..start + k)
+ .unwrap_or_default()
+ .iter()
+ .zip(seed_distances.get(start..start + k).unwrap_or_default())
+ {
+ if distance != f32::MAX {
+ heap.push(distance, id);
+ }
+ }
+ }
+}
+
// --- Utilities ---
fn compute_residuals(
@@ -2622,7 +2769,10 @@ mod tests {
&data[..nq * d],
nq,
k,
+ 0,
nprobe,
+ &[],
+ &[],
filter,
reuse_mode,
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
@@ -3793,6 +3943,102 @@ mod tests {
}
}
+ #[test]
+ fn incremental_batch_probe_ranges_partition_the_one_shot_lists() {
+ use crate::io::{write_index, PosWriter};
+
+ let d = 16;
+ let nlist = 4;
+ let m = 4;
+ let n = 128;
+ let nq = 2;
+ let k = n;
+ let data = generate_clustered_data(n, d, nlist, 2_026);
+ let ids = (0..n as i64).collect::<Vec<_>>();
+ let queries = &data[..nq * d];
+
+ 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 mut full_reader =
IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap();
+ let (full_ids, _) = search_batch_reader(&mut full_reader, queries, nq,
k, nlist).unwrap();
+
+ let mut incremental_reader =
IVFPQIndexReader::open(Cursor::new(bytes)).unwrap();
+ let (first_ids, first_distances) =
search_batch_reader_with_reuse_mode_and_budget_range(
+ &mut incremental_reader,
+ queries,
+ nq,
+ k,
+ 0,
+ 2,
+ &[],
+ &[],
+ IvfPqBatchTableReuseMode::Auto,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ )
+ .unwrap();
+ let (second_delta_ids, _) =
search_batch_reader_with_reuse_mode_and_budget_range(
+ &mut incremental_reader,
+ queries,
+ nq,
+ k,
+ 2,
+ nlist,
+ &[],
+ &[],
+ IvfPqBatchTableReuseMode::Auto,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ )
+ .unwrap();
+ let (second_ids, _) =
search_batch_reader_with_reuse_mode_and_budget_range(
+ &mut incremental_reader,
+ queries,
+ nq,
+ k,
+ 2,
+ nlist,
+ &first_ids,
+ &first_distances,
+ IvfPqBatchTableReuseMode::Auto,
+ DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
+ )
+ .unwrap();
+
+ for query_index in 0..nq {
+ let result = query_index * k..(query_index + 1) * k;
+ let full = full_ids[result.clone()]
+ .iter()
+ .copied()
+ .filter(|&id| id != -1)
+ .collect::<HashSet<_>>();
+ let first = first_ids[result.clone()]
+ .iter()
+ .copied()
+ .filter(|&id| id != -1)
+ .collect::<HashSet<_>>();
+ let second_delta = second_delta_ids[result.clone()]
+ .iter()
+ .copied()
+ .filter(|&id| id != -1)
+ .collect::<HashSet<_>>();
+ let second = second_ids[result]
+ .iter()
+ .copied()
+ .filter(|&id| id != -1)
+ .collect::<HashSet<_>>();
+
+ assert!(first.is_disjoint(&second_delta));
+ assert_eq!(
+ first.union(&second_delta).copied().collect::<HashSet<_>>(),
+ full
+ );
+ assert_eq!(second, full);
+ }
+ }
+
#[test]
fn
inner_product_batch_table_reuse_modes_preserve_results_and_control_table_builds()
{
use crate::io::{write_index, IVFPQIndexReader, PosWriter};
@@ -3820,7 +4066,10 @@ mod tests {
&data[..nq * d],
nq,
k,
+ 0,
nprobe,
+ &[],
+ &[],
None,
IvfPqBatchTableReuseMode::On,
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
@@ -3838,7 +4087,10 @@ mod tests {
&data[..nq * d],
nq,
k,
+ 0,
nprobe,
+ &[],
+ &[],
None,
IvfPqBatchTableReuseMode::On,
nq * m * 256 * std::mem::size_of::<f32>() - 1,
@@ -3883,7 +4135,10 @@ mod tests {
&data[..large_nq * d],
large_nq,
k,
+ 0,
nprobe,
+ &[],
+ &[],
None,
IvfPqBatchTableReuseMode::Auto,
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
@@ -3904,7 +4159,10 @@ mod tests {
&data[..nq * d],
nq,
k,
+ 0,
nprobe,
+ &[],
+ &[],
None,
IvfPqBatchTableReuseMode::Auto,
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
@@ -3926,7 +4184,10 @@ mod tests {
&data[..nq * d],
nq,
k,
+ 0,
nprobe,
+ &[],
+ &[],
Some(&empty_filter),
IvfPqBatchTableReuseMode::On,
DEFAULT_IVFPQ_BATCH_TABLE_REUSE_MAX_BYTES,
@@ -4406,7 +4667,10 @@ mod tests {
queries,
nq,
k,
+ 0,
nlist,
+ &[],
+ &[],
None,
IvfPqBatchTableReuseMode::On,
1,
diff --git a/core/src/ivfrq_io.rs b/core/src/ivfrq_io.rs
index 328a332..fc8d6d5 100644
--- a/core/src/ivfrq_io.rs
+++ b/core/src/ivfrq_io.rs
@@ -636,9 +636,30 @@ pub fn search_batch_ivfrq_reader_filter<R: SeekRead>(
k: usize,
nprobe: usize,
filter: Option<&dyn RowIdFilter>,
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_ivfrq_reader_filter_range(reader, queries, nq, k, 0, nprobe,
&[], &[], filter)
+}
+
+pub(crate) fn search_batch_ivfrq_reader_filter_range<R: SeekRead>(
+ reader: &mut IVFRQIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ probe_start: usize,
+ probe_end: usize,
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ filter: Option<&dyn RowIdFilter>,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
reader.ensure_loaded()?;
- validate_search_inputs(queries, nq, reader.d, k, nprobe)?;
+ validate_search_inputs(queries, nq, reader.d, k, probe_end)?;
+ if probe_start >= probe_end {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "probe range must be non-empty",
+ ));
+ }
+ validate_batch_seed(seed_ids, seed_distances, nq, k)?;
let processed = preprocess_vectors(queries, nq, reader.d, reader.metric);
let (all_probe_indices, all_probe_distances) =
kmeans::find_topk_batch_with_centroid_norms(
&processed,
@@ -647,7 +668,7 @@ pub fn search_batch_ivfrq_reader_filter<R: SeekRead>(
&reader.quantizer_centroid_norms,
reader.nlist,
reader.d,
- nprobe,
+ probe_end,
);
let query_norms = processed
.chunks_exact(reader.d)
@@ -684,7 +705,7 @@ pub fn search_batch_ivfrq_reader_filter<R: SeekRead>(
let mut seen_lists = vec![false; reader.nlist];
let mut unique_lists = Vec::new();
for probes in &all_probe_indices {
- for &list_id in probes {
+ for &list_id in probes.iter().skip(probe_start) {
if !seen_lists[list_id] {
seen_lists[list_id] = true;
unique_lists.push(list_id);
@@ -693,6 +714,7 @@ pub fn search_batch_ivfrq_reader_filter<R: SeekRead>(
}
let mut heaps: Vec<TopKHeap> = (0..nq).map(|_| TopKHeap::new(k)).collect();
+ seed_heaps(&mut heaps, seed_ids, seed_distances, k);
let mut query_stats = vec![IVFRQSearchStats::default(); nq];
let mut aggregate_stats = IVFRQSearchStats {
query_count: nq,
@@ -778,8 +800,10 @@ pub fn search_batch_ivfrq_reader_filter<R: SeekRead>(
.zip(query_stats.par_iter_mut())
.enumerate()
.for_each(|(query_index, (heap, stats))| {
- for (probe_position, &list_id) in
- all_probe_indices[query_index].iter().enumerate()
+ for (probe_position, &list_id) in
all_probe_indices[query_index]
+ .iter()
+ .enumerate()
+ .skip(probe_start)
{
let position = list_positions[list_id];
if position == usize::MAX {
@@ -824,9 +848,80 @@ pub fn search_batch_ivfrq_reader_roaring_filter<R:
SeekRead>(
k: usize,
nprobe: usize,
roaring_filter_bytes: &[u8],
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_ivfrq_reader_roaring_filter_range(
+ reader,
+ queries,
+ nq,
+ k,
+ 0,
+ nprobe,
+ &[],
+ &[],
+ roaring_filter_bytes,
+ )
+}
+
+pub(crate) fn search_batch_ivfrq_reader_roaring_filter_range<R: SeekRead>(
+ reader: &mut IVFRQIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ probe_start: usize,
+ probe_end: usize,
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ roaring_filter_bytes: &[u8],
) -> io::Result<(Vec<i64>, Vec<f32>)> {
let filter = decode_roaring_filter(roaring_filter_bytes)?;
- search_batch_ivfrq_reader_filter(reader, queries, nq, k, nprobe,
Some(&filter))
+ search_batch_ivfrq_reader_filter_range(
+ reader,
+ queries,
+ nq,
+ k,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
+ Some(&filter),
+ )
+}
+
+fn validate_batch_seed(
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ nq: usize,
+ k: usize,
+) -> io::Result<()> {
+ let expected = nq
+ .checked_mul(k)
+ .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "nq * k
overflows usize"))?;
+ if (seed_ids.is_empty() && seed_distances.is_empty())
+ || (seed_ids.len() == expected && seed_distances.len() == expected)
+ {
+ Ok(())
+ } else {
+ Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "seed result lengths must both equal nq * k",
+ ))
+ }
+}
+
+fn seed_heaps(heaps: &mut [TopKHeap], seed_ids: &[i64], seed_distances:
&[f32], k: usize) {
+ for (query_index, heap) in heaps.iter_mut().enumerate() {
+ let start = query_index * k;
+ for (&id, &distance) in seed_ids
+ .get(start..start + k)
+ .unwrap_or_default()
+ .iter()
+ .zip(seed_distances.get(start..start + k).unwrap_or_default())
+ {
+ if distance != f32::MAX {
+ heap.push(distance, id);
+ }
+ }
+ }
}
fn scan_blocked_list(
diff --git a/core/src/ivfsq_io.rs b/core/src/ivfsq_io.rs
index 541abee..0bc5819 100644
--- a/core/src/ivfsq_io.rs
+++ b/core/src/ivfsq_io.rs
@@ -642,9 +642,30 @@ pub fn search_batch_ivfsq_reader_filter<R: SeekRead>(
k: usize,
nprobe: usize,
filter: Option<&dyn RowIdFilter>,
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_ivfsq_reader_filter_range(reader, queries, nq, k, 0, nprobe,
&[], &[], filter)
+}
+
+pub(crate) fn search_batch_ivfsq_reader_filter_range<R: SeekRead>(
+ reader: &mut IVFSQIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ probe_start: usize,
+ probe_end: usize,
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ filter: Option<&dyn RowIdFilter>,
) -> io::Result<(Vec<i64>, Vec<f32>)> {
reader.ensure_loaded()?;
- validate_search_inputs(queries, nq, reader.d, k, nprobe)?;
+ validate_search_inputs(queries, nq, reader.d, k, probe_end)?;
+ if probe_start >= probe_end {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "probe range must be non-empty",
+ ));
+ }
+ validate_batch_seed(seed_ids, seed_distances, nq, k)?;
let processed = preprocess_vectors(queries, nq, reader.d, reader.metric);
let (all_probe_indices, _) = kmeans::find_topk_batch(
&processed,
@@ -652,12 +673,12 @@ pub fn search_batch_ivfsq_reader_filter<R: SeekRead>(
&reader.quantizer_centroids,
reader.nlist,
reader.d,
- nprobe,
+ probe_end,
);
let mut seen = vec![false; reader.nlist];
let mut unique_lists = Vec::new();
for list_ids in &all_probe_indices {
- for &list_id in list_ids {
+ for &list_id in list_ids.iter().skip(probe_start) {
if !seen[list_id] {
seen[list_id] = true;
unique_lists.push(list_id);
@@ -666,13 +687,14 @@ pub fn search_batch_ivfsq_reader_filter<R: SeekRead>(
}
let mut list_to_queries = vec![Vec::new(); reader.nlist];
for (query_index, list_ids) in all_probe_indices.iter().enumerate() {
- for &list_id in list_ids {
+ for &list_id in list_ids.iter().skip(probe_start) {
list_to_queries[list_id].push(query_index);
}
}
let d = reader.d;
let metric = reader.metric;
let mut heaps = (0..nq).map(|_| TopKHeap::new(k)).collect::<Vec<_>>();
+ seed_heaps(&mut heaps, seed_ids, seed_distances, k);
// Oversized-list chunks are scanned query-by-query, so one reusable
// distance buffer is sufficient regardless of the batch width.
let mut stream_scratch = SqScanScratch::default();
@@ -764,9 +786,80 @@ pub fn search_batch_ivfsq_reader_roaring_filter<R:
SeekRead>(
k: usize,
nprobe: usize,
roaring_filter_bytes: &[u8],
+) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ search_batch_ivfsq_reader_roaring_filter_range(
+ reader,
+ queries,
+ nq,
+ k,
+ 0,
+ nprobe,
+ &[],
+ &[],
+ roaring_filter_bytes,
+ )
+}
+
+pub(crate) fn search_batch_ivfsq_reader_roaring_filter_range<R: SeekRead>(
+ reader: &mut IVFSQIndexReader<R>,
+ queries: &[f32],
+ nq: usize,
+ k: usize,
+ probe_start: usize,
+ probe_end: usize,
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ roaring_filter_bytes: &[u8],
) -> io::Result<(Vec<i64>, Vec<f32>)> {
let filter = decode_roaring_filter(roaring_filter_bytes)?;
- search_batch_ivfsq_reader_filter(reader, queries, nq, k, nprobe,
Some(&filter))
+ search_batch_ivfsq_reader_filter_range(
+ reader,
+ queries,
+ nq,
+ k,
+ probe_start,
+ probe_end,
+ seed_ids,
+ seed_distances,
+ Some(&filter),
+ )
+}
+
+fn validate_batch_seed(
+ seed_ids: &[i64],
+ seed_distances: &[f32],
+ nq: usize,
+ k: usize,
+) -> io::Result<()> {
+ let expected = nq
+ .checked_mul(k)
+ .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "nq * k
overflows usize"))?;
+ if (seed_ids.is_empty() && seed_distances.is_empty())
+ || (seed_ids.len() == expected && seed_distances.len() == expected)
+ {
+ Ok(())
+ } else {
+ Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "seed result lengths must both equal nq * k",
+ ))
+ }
+}
+
+fn seed_heaps(heaps: &mut [TopKHeap], seed_ids: &[i64], seed_distances:
&[f32], k: usize) {
+ for (query_index, heap) in heaps.iter_mut().enumerate() {
+ let start = query_index * k;
+ for (&id, &distance) in seed_ids
+ .get(start..start + k)
+ .unwrap_or_default()
+ .iter()
+ .zip(seed_distances.get(start..start + k).unwrap_or_default())
+ {
+ if distance != f32::MAX {
+ heap.push(distance, id);
+ }
+ }
+ }
}
pub struct SqListData {
