jerry-024 commented on code in PR #72:
URL:
https://github.com/apache/paimon-vector-index/pull/72#discussion_r3746984631
##########
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:
**[major] Avoid retaining the full-batch and retry buffers at the same time**
At this call, `result_ids` and `result_distances` for the original full
batch are still live. A partial retry also owns `packed_queries`, and `search`
allocates another IDs/distances pair. Peak memory therefore grows by
`O(active_queries * (dimension + top_k))` on top of the full result buffers.
For 100k queries at dimension 768 and `top_k = 100`, that is roughly 427 MB of
additional live data before search-internal heaps and probe metadata, which can
turn automatic expansion into an OOM.
Please avoid keeping all three representations live together—for example,
replace the previous result directly when all queries remain active, and
process partial retries in bounded chunks or pass active-query indices.
--
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]
