shyjsarah commented on code in PR #72:
URL:
https://github.com/apache/paimon-vector-index/pull/72#discussion_r3763250981
##########
core/src/index.rs:
##########
@@ -2020,21 +2028,97 @@ fn progressive_ivf_search(
.1
.chunks_exact(top_k)
.take(query_count)
- .all(|distances| {
- distances
- .iter()
- .filter(|&&distance| distance != f32::MAX)
- .take(required_per_query)
- .count()
- >= required_per_query
- })
+ .all(|distances| ivf_search_result_is_complete(distances,
required_per_query))
{
return Ok(result);
}
nprobe = nprobe.saturating_mul(2).min(nlist);
}
}
+/// 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.
+fn progressive_ivf_batch_search(
+ params: VectorSearchParams,
+ nlist: usize,
+ initial_nprobe: usize,
+ queries: &[f32],
+ query_count: usize,
+ top_k: usize,
+ available_matches: usize,
+ mut search: impl FnMut(&[f32], usize, usize) -> 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);
+ }
+
+ 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 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::<Vec<_>>();
+
+ loop {
+ if nprobe >= nlist || required_per_query == 0 ||
active_queries.is_empty() {
+ return Ok((result_ids, result_distances));
+ }
+
+ nprobe = nprobe.saturating_mul(2).min(nlist);
+ let packed_queries;
+ let round_queries = if active_queries.len() == query_count {
+ queries
+ } else {
+ packed_queries = {
+ let mut packed = Vec::with_capacity(active_queries.len() *
dimension);
+ for &query_index in &active_queries {
+ let start = query_index * dimension;
+ packed.extend_from_slice(&queries[start..start +
dimension]);
+ }
+ packed
+ };
+ &packed_queries
+ };
+
+ let (round_ids, round_distances) = search(round_queries,
active_queries.len(), nprobe)?;
Review Comment:
Fixed in ce1bad4. Full-batch retries now drop the previous result buffers
before searching again. Partial retries are processed in chunks bounded by a 64
MiB query/result buffer budget, with a regression test covering chunked
retry/scatter behavior.
##########
core/src/index.rs:
##########
@@ -1843,18 +1846,19 @@ impl<R: SeekRead> VectorIndexReader<R> {
let total_vectors = usize::try_from(reader.total_vectors)
.map_err(|_| invalid_input("negative IVF vector count"))?;
let nprobe = params.resolve_ivf_nprobe(reader.nlist,
total_vectors, None)?;
- progressive_ivf_search(
+ progressive_ivf_batch_search(
params,
reader.nlist,
nprobe,
+ queries,
query_count,
params.top_k,
total_vectors,
- |nprobe| {
+ |active_queries, active_query_count, nprobe| {
Review Comment:
Fixed in ce1bad4. IVF-RQ statistics are accumulated across all internal
progressive rounds, and the published query_count is restored to the original
logical batch size for both filtered and unfiltered paths. Added a filtered
progressive-search regression test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
