From 5fe7c4545a57d7356d25668e7ecbdcf1616082b2 Mon Sep 17 00:00:00 2001 From: Vivek Kumar Date: Thu, 19 Oct 2023 16:48:45 +0530 Subject: [PATCH] compute all candidates correctly when skipping --- milli/src/search/new/bucket_sort.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/milli/src/search/new/bucket_sort.rs b/milli/src/search/new/bucket_sort.rs index cf2f08cce..df9c14c7d 100644 --- a/milli/src/search/new/bucket_sort.rs +++ b/milli/src/search/new/bucket_sort.rs @@ -46,9 +46,8 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>( if let Some(distinct_fid) = distinct_fid { let mut excluded = RoaringBitmap::new(); let mut results = vec![]; - let mut skip = 0; for docid in universe.iter() { - if results.len() >= length { + if results.len() >= from + length { break; } if excluded.contains(docid) { @@ -56,16 +55,16 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>( } distinct_single_docid(ctx.index, ctx.txn, distinct_fid, docid, &mut excluded)?; - skip += 1; - if skip <= from { - continue; - } - results.push(docid); } let mut all_candidates = universe - excluded; all_candidates.extend(results.iter().copied()); + if results.len() >= from { + results.drain(..from); + } else { + results.clear(); + } return Ok(BucketSortOutput { scores: vec![Default::default(); results.len()],