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,48 +117,62 @@ impl Index {
positions.push(union_positions.iter().collect()); positions.push(union_positions.iter().collect());
} }
// let positions = BestProximity::new(positions).next().unwrap_or_default(); let mut documents = Vec::new();
let _positions: Vec<Vec<u32>> = positions;
let positions = vec![0u32];
eprintln!("best proximity {:?}", positions);
let mut intersect_docids: Option<RoaringBitmap> = None; for (_proximity, positions) in BestProximity::new(positions) {
for ((word, is_prefix, dfa), pos) in words_positions.into_iter().zip(positions) { let mut same_proximity_union = RoaringBitmap::default();
let mut count = 0;
let mut union_docids = RoaringBitmap::default();
if false && word.len() <= 4 && is_prefix { for positions in positions {
let mut key = word.as_bytes()[..word.len().min(5)].to_vec(); let mut intersect_docids: Option<RoaringBitmap> = None;
key.extend_from_slice(&pos.to_be_bytes()); for ((word, is_prefix, dfa), pos) in words_positions.iter().zip(positions) {
if let Some(ids) = self.prefix_postings_ids.get(rtxn, &key)? { let mut count = 0;
let right = RoaringBitmap::deserialize_from(ids)?; let mut union_docids = RoaringBitmap::default();
union_docids.union_with(&right);
count = 1; // TODO re-enable the prefixes system
} if false && word.len() <= 4 && *is_prefix {
} else { let mut key = word.as_bytes()[..word.len().min(5)].to_vec();
let mut stream = fst.search(dfa).into_stream(); key.extend_from_slice(&pos.to_be_bytes());
while let Some(word) = stream.next() { if let Some(ids) = self.prefix_postings_ids.get(rtxn, &key)? {
let word = std::str::from_utf8(word)?; let right = RoaringBitmap::deserialize_from(ids)?;
let mut key = word.as_bytes().to_vec(); union_docids.union_with(&right);
key.extend_from_slice(&pos.to_be_bytes()); count = 1;
if let Some(attrs) = self.postings_ids.get(rtxn, &key)? { }
let right = RoaringBitmap::deserialize_from(attrs)?; } else {
union_docids.union_with(&right); let mut stream = fst.search(dfa).into_stream();
count += 1; while let Some(word) = stream.next() {
let word = std::str::from_utf8(word)?;
let mut key = word.as_bytes().to_vec();
key.extend_from_slice(&pos.to_be_bytes());
if let Some(attrs) = self.postings_ids.get(rtxn, &key)? {
let right = RoaringBitmap::deserialize_from(attrs)?;
union_docids.union_with(&right);
count += 1;
}
}
} }
let _ = count;
match &mut intersect_docids {
Some(left) => left.intersect_with(&union_docids),
None => intersect_docids = Some(union_docids),
}
}
if let Some(intersect_docids) = intersect_docids {
same_proximity_union.union_with(&intersect_docids);
} }
} }
let _ = count; documents.push(same_proximity_union);
match &mut intersect_docids { // We found enough documents we can stop here
Some(left) => left.intersect_with(&union_docids), if documents.iter().map(RoaringBitmap::len).sum::<u64>() >= 20 {
None => intersect_docids = Some(union_docids), break
} }
} }
eprintln!("{} candidates", intersect_docids.as_ref().map_or(0, |r| r.len())); eprintln!("{} candidates", documents.iter().map(RoaringBitmap::len).sum::<u64>());
Ok(documents.iter().flatten().take(20).collect())
Ok(intersect_docids.unwrap_or_default().iter().take(20).collect())
} }
} }