Introduce a proximity based documents retriever

This commit is contained in:
Kerollmops 2020-06-10 16:54:22 +02:00
parent 78f27c0465
commit 66a4b26811
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4

View File

@ -117,17 +117,19 @@ impl Index {
positions.push(union_positions.iter().collect());
}
// let positions = BestProximity::new(positions).next().unwrap_or_default();
let _positions: Vec<Vec<u32>> = positions;
let positions = vec![0u32];
eprintln!("best proximity {:?}", positions);
let mut documents = Vec::new();
for (_proximity, positions) in BestProximity::new(positions) {
let mut same_proximity_union = RoaringBitmap::default();
for positions in positions {
let mut intersect_docids: Option<RoaringBitmap> = None;
for ((word, is_prefix, dfa), pos) in words_positions.into_iter().zip(positions) {
for ((word, is_prefix, dfa), pos) in words_positions.iter().zip(positions) {
let mut count = 0;
let mut union_docids = RoaringBitmap::default();
if false && word.len() <= 4 && is_prefix {
// TODO re-enable the prefixes system
if false && word.len() <= 4 && *is_prefix {
let mut key = word.as_bytes()[..word.len().min(5)].to_vec();
key.extend_from_slice(&pos.to_be_bytes());
if let Some(ids) = self.prefix_postings_ids.get(rtxn, &key)? {
@ -157,8 +159,20 @@ impl Index {
}
}
eprintln!("{} candidates", intersect_docids.as_ref().map_or(0, |r| r.len()));
if let Some(intersect_docids) = intersect_docids {
same_proximity_union.union_with(&intersect_docids);
}
}
Ok(intersect_docids.unwrap_or_default().iter().take(20).collect())
documents.push(same_proximity_union);
// We found enough documents we can stop here
if documents.iter().map(RoaringBitmap::len).sum::<u64>() >= 20 {
break
}
}
eprintln!("{} candidates", documents.iter().map(RoaringBitmap::len).sum::<u64>());
Ok(documents.iter().flatten().take(20).collect())
}
}