2022-08-30 14:17:40 +02:00
|
|
|
use crate::heed_codec::facet::new::{
|
|
|
|
FacetGroupValue, FacetGroupValueCodec, FacetKey, FacetKeyCodec, MyByteSlice,
|
|
|
|
};
|
2022-08-31 09:36:19 +02:00
|
|
|
use heed::Result;
|
2022-08-30 15:22:39 +02:00
|
|
|
use roaring::RoaringBitmap;
|
2022-08-30 14:17:40 +02:00
|
|
|
|
|
|
|
use super::{get_first_facet_value, get_highest_level};
|
|
|
|
|
|
|
|
pub fn ascending_facet_sort<'t>(
|
|
|
|
rtxn: &'t heed::RoTxn<'t>,
|
2022-08-31 09:36:19 +02:00
|
|
|
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
|
2022-08-30 14:17:40 +02:00
|
|
|
field_id: u16,
|
|
|
|
candidates: RoaringBitmap,
|
2022-08-31 09:36:19 +02:00
|
|
|
) -> Result<Box<dyn Iterator<Item = Result<RoaringBitmap>> + 't>> {
|
|
|
|
let highest_level = get_highest_level(rtxn, db, field_id)?;
|
|
|
|
if let Some(first_bound) = get_first_facet_value::<MyByteSlice>(rtxn, db, field_id)? {
|
2022-08-30 14:17:40 +02:00
|
|
|
let first_key = FacetKey { field_id, level: highest_level, left_bound: first_bound };
|
|
|
|
let iter = db.range(rtxn, &(first_key..)).unwrap().take(usize::MAX);
|
|
|
|
|
2022-08-30 15:22:39 +02:00
|
|
|
Ok(Box::new(AscendingFacetSort { rtxn, db, field_id, stack: vec![(candidates, iter)] }))
|
2022-08-30 14:17:40 +02:00
|
|
|
} else {
|
2022-08-30 15:22:39 +02:00
|
|
|
Ok(Box::new(std::iter::empty()))
|
2022-08-30 14:17:40 +02:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
struct AscendingFacetSort<'t, 'e> {
|
|
|
|
rtxn: &'t heed::RoTxn<'e>,
|
2022-08-31 09:36:19 +02:00
|
|
|
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
|
2022-08-30 14:17:40 +02:00
|
|
|
field_id: u16,
|
|
|
|
stack: Vec<(
|
|
|
|
RoaringBitmap,
|
|
|
|
std::iter::Take<heed::RoRange<'t, FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>>,
|
|
|
|
)>,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<'t, 'e> Iterator for AscendingFacetSort<'t, 'e> {
|
2022-08-31 09:36:19 +02:00
|
|
|
type Item = Result<RoaringBitmap>;
|
2022-08-30 14:17:40 +02:00
|
|
|
|
|
|
|
fn next(&mut self) -> Option<Self::Item> {
|
|
|
|
'outer: loop {
|
|
|
|
let (documents_ids, deepest_iter) = self.stack.last_mut()?;
|
|
|
|
for result in deepest_iter {
|
|
|
|
let (
|
|
|
|
FacetKey { level, left_bound, field_id },
|
|
|
|
FacetGroupValue { size: group_size, mut bitmap },
|
|
|
|
) = result.unwrap();
|
|
|
|
// The range is unbounded on the right and the group size for the highest level is MAX,
|
|
|
|
// so we need to check that we are not iterating over the next field id
|
|
|
|
if field_id != self.field_id {
|
|
|
|
return None;
|
|
|
|
}
|
|
|
|
|
|
|
|
// If the last iterator found an empty set of documents it means
|
|
|
|
// that we found all the documents in the sub level iterations already,
|
|
|
|
// we can pop this level iterator.
|
|
|
|
if documents_ids.is_empty() {
|
|
|
|
break;
|
|
|
|
}
|
|
|
|
|
|
|
|
bitmap &= &*documents_ids;
|
|
|
|
if !bitmap.is_empty() {
|
|
|
|
*documents_ids -= &bitmap;
|
|
|
|
|
|
|
|
if level == 0 {
|
2022-08-31 09:36:19 +02:00
|
|
|
return Some(Ok(bitmap));
|
2022-08-30 14:17:40 +02:00
|
|
|
}
|
|
|
|
let starting_key_below =
|
|
|
|
FacetKey { field_id: self.field_id, level: level - 1, left_bound };
|
2022-08-30 15:22:39 +02:00
|
|
|
let iter = match self.db.range(&self.rtxn, &(starting_key_below..)) {
|
|
|
|
Ok(iter) => iter,
|
|
|
|
Err(e) => return Some(Err(e.into())),
|
|
|
|
}
|
|
|
|
.take(group_size as usize);
|
2022-08-30 14:17:40 +02:00
|
|
|
|
|
|
|
self.stack.push((bitmap, iter));
|
|
|
|
continue 'outer;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
self.stack.pop();
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[cfg(test)]
|
|
|
|
mod tests {
|
2022-08-30 15:22:39 +02:00
|
|
|
use rand::Rng;
|
|
|
|
use rand::SeedableRng;
|
2022-08-30 14:17:40 +02:00
|
|
|
use roaring::RoaringBitmap;
|
|
|
|
|
2022-08-30 15:22:39 +02:00
|
|
|
use crate::{
|
|
|
|
heed_codec::facet::new::ordered_f64_codec::OrderedF64Codec,
|
|
|
|
search::facet::{facet_sort_ascending::ascending_facet_sort, test::FacetIndex},
|
|
|
|
snapshot_tests::display_bitmap,
|
|
|
|
};
|
|
|
|
|
|
|
|
fn get_simple_index() -> FacetIndex<OrderedF64Codec> {
|
|
|
|
let index = FacetIndex::<OrderedF64Codec>::new(4, 8);
|
2022-08-30 14:17:40 +02:00
|
|
|
let mut txn = index.env.write_txn().unwrap();
|
|
|
|
for i in 0..256u16 {
|
|
|
|
let mut bitmap = RoaringBitmap::new();
|
|
|
|
bitmap.insert(i as u32);
|
2022-08-31 14:19:52 +02:00
|
|
|
index.insert(&mut txn, 0, &(i as f64), &bitmap);
|
2022-08-30 14:17:40 +02:00
|
|
|
}
|
|
|
|
txn.commit().unwrap();
|
|
|
|
index
|
|
|
|
}
|
2022-08-30 15:22:39 +02:00
|
|
|
fn get_random_looking_index() -> FacetIndex<OrderedF64Codec> {
|
|
|
|
let index = FacetIndex::<OrderedF64Codec>::new(4, 8);
|
2022-08-30 14:17:40 +02:00
|
|
|
let mut txn = index.env.write_txn().unwrap();
|
|
|
|
|
2022-08-31 14:19:52 +02:00
|
|
|
let mut rng = rand::rngs::SmallRng::from_seed([0; 32]);
|
2022-08-30 15:22:39 +02:00
|
|
|
let keys =
|
|
|
|
std::iter::from_fn(|| Some(rng.gen_range(0..256))).take(128).collect::<Vec<u32>>();
|
2022-08-30 14:17:40 +02:00
|
|
|
|
|
|
|
for (_i, key) in keys.into_iter().enumerate() {
|
|
|
|
let mut bitmap = RoaringBitmap::new();
|
|
|
|
bitmap.insert(key);
|
|
|
|
bitmap.insert(key + 100);
|
2022-08-30 15:22:39 +02:00
|
|
|
index.insert(&mut txn, 0, &(key as f64), &bitmap);
|
2022-08-30 14:17:40 +02:00
|
|
|
}
|
|
|
|
txn.commit().unwrap();
|
|
|
|
index
|
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn random_looking_index_snap() {
|
|
|
|
let index = get_random_looking_index();
|
|
|
|
insta::assert_display_snapshot!(index)
|
|
|
|
}
|
|
|
|
#[test]
|
|
|
|
fn filter_sort() {
|
|
|
|
let indexes = [get_simple_index(), get_random_looking_index()];
|
2022-08-31 14:19:52 +02:00
|
|
|
for (i, index) in indexes.iter().enumerate() {
|
2022-08-30 14:17:40 +02:00
|
|
|
let txn = index.env.read_txn().unwrap();
|
|
|
|
let candidates = (200..=300).into_iter().collect::<RoaringBitmap>();
|
|
|
|
let mut results = String::new();
|
2022-08-31 14:19:52 +02:00
|
|
|
let iter = ascending_facet_sort(&txn, index.db.content, 0, candidates).unwrap();
|
|
|
|
for el in iter {
|
|
|
|
let docids = el.unwrap();
|
|
|
|
results.push_str(&display_bitmap(&docids));
|
2022-08-30 14:17:40 +02:00
|
|
|
}
|
|
|
|
insta::assert_snapshot!(format!("filter_sort_{i}_ascending"), results);
|
|
|
|
|
|
|
|
txn.commit().unwrap();
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|