Fix a bug in facet_range_search and add documentation

This commit is contained in:
Loïc Lecrenier 2022-09-21 14:39:11 +02:00 committed by Loïc Lecrenier
parent a2270b7432
commit d0109627b9
14 changed files with 173 additions and 28 deletions

View File

@ -9,6 +9,8 @@ use crate::heed_codec::facet::{
}; };
use crate::Result; use crate::Result;
/// Find all the document ids for which the given field contains a value contained within
/// the two bounds.
pub fn find_docids_of_facet_within_bounds<'t, BoundCodec>( pub fn find_docids_of_facet_within_bounds<'t, BoundCodec>(
rtxn: &'t heed::RoTxn<'t>, rtxn: &'t heed::RoTxn<'t>,
db: heed::Database<FacetGroupKeyCodec<BoundCodec>, FacetGroupValueCodec>, db: heed::Database<FacetGroupKeyCodec<BoundCodec>, FacetGroupValueCodec>,
@ -24,11 +26,11 @@ where
let inner; let inner;
let left = match left { let left = match left {
Bound::Included(left) => { Bound::Included(left) => {
inner = BoundCodec::bytes_encode(left).unwrap(); inner = BoundCodec::bytes_encode(left).ok_or(heed::Error::Encoding)?;
Bound::Included(inner.as_ref()) Bound::Included(inner.as_ref())
} }
Bound::Excluded(left) => { Bound::Excluded(left) => {
inner = BoundCodec::bytes_encode(left).unwrap(); inner = BoundCodec::bytes_encode(left).ok_or(heed::Error::Encoding)?;
Bound::Excluded(inner.as_ref()) Bound::Excluded(inner.as_ref())
} }
Bound::Unbounded => Bound::Unbounded, Bound::Unbounded => Bound::Unbounded,
@ -36,11 +38,11 @@ where
let inner; let inner;
let right = match right { let right = match right {
Bound::Included(right) => { Bound::Included(right) => {
inner = BoundCodec::bytes_encode(right).unwrap(); inner = BoundCodec::bytes_encode(right).ok_or(heed::Error::Encoding)?;
Bound::Included(inner.as_ref()) Bound::Included(inner.as_ref())
} }
Bound::Excluded(right) => { Bound::Excluded(right) => {
inner = BoundCodec::bytes_encode(right).unwrap(); inner = BoundCodec::bytes_encode(right).ok_or(heed::Error::Encoding)?;
Bound::Excluded(inner.as_ref()) Bound::Excluded(inner.as_ref())
} }
Bound::Unbounded => Bound::Unbounded, Bound::Unbounded => Bound::Unbounded,
@ -49,9 +51,11 @@ where
let mut f = FacetRangeSearch { rtxn, db, field_id, left, right, docids }; let mut f = FacetRangeSearch { rtxn, db, field_id, left, right, docids };
let highest_level = get_highest_level(rtxn, db, field_id)?; let highest_level = get_highest_level(rtxn, db, field_id)?;
if let Some(first_bound) = get_first_facet_value::<ByteSliceRef>(rtxn, db, field_id)? { if let Some(starting_left_bound) = get_first_facet_value::<ByteSliceRef>(rtxn, db, field_id)? {
let last_bound = get_last_facet_value::<ByteSliceRef>(rtxn, db, field_id)?.unwrap(); let rightmost_bound =
f.run(highest_level, first_bound, Bound::Included(last_bound), usize::MAX)?; Bound::Included(get_last_facet_value::<ByteSliceRef>(rtxn, db, field_id)?.unwrap()); // will not fail because get_first_facet_value succeeded
let group_size = usize::MAX;
f.run(highest_level, starting_left_bound, rightmost_bound, group_size)?;
Ok(()) Ok(())
} else { } else {
return Ok(()); return Ok(());
@ -107,7 +111,25 @@ impl<'t, 'b, 'bitmap> FacetRangeSearch<'t, 'b, 'bitmap> {
Ok(()) Ok(())
} }
/// Recursive part of the algorithm for level > 0 /// Recursive part of the algorithm for level > 0.
///
/// It works by visiting a slice of a level and checking whether the range asscociated
/// with each visited element is contained within the bounds.
///
/// 1. So long as the element's range is less than the left bound, we do nothing and keep iterating
/// 2. If the element's range is fully contained by the bounds, then all of its docids are added to
/// the roaring bitmap.
/// 3. If the element's range merely intersects the bounds, then we call the algorithm recursively
/// on the children of the element from the level below.
/// 4. If the element's range is greater than the right bound, we do nothing and stop iterating.
/// Note that the right bound is found through either the `left_bound` of the *next* element,
/// or from the `rightmost_bound` argument
///
/// ## Arguments
/// - `level`: the level being visited
/// - `starting_left_bound`: the left_bound of the first element to visit
/// - `rightmost_bound`: the right bound of the last element that should be visited
/// - `group_size`: the number of elements that should be visited
fn run( fn run(
&mut self, &mut self,
level: u8, level: u8,
@ -123,13 +145,14 @@ impl<'t, 'b, 'bitmap> FacetRangeSearch<'t, 'b, 'bitmap> {
FacetGroupKey { field_id: self.field_id, level, left_bound: starting_left_bound }; FacetGroupKey { field_id: self.field_id, level, left_bound: starting_left_bound };
let mut iter = self.db.range(&self.rtxn, &(left_key..))?.take(group_size); let mut iter = self.db.range(&self.rtxn, &(left_key..))?.take(group_size);
// We iterate over the range while keeping in memory the previous value
let (mut previous_key, mut previous_value) = iter.next().unwrap()?; let (mut previous_key, mut previous_value) = iter.next().unwrap()?;
for el in iter { for el in iter {
let (next_key, next_value) = el?; let (next_key, next_value) = el?;
// the right of the iter range is unbounded, so we need to make sure that we are not iterating // the right of the iter range is potentially unbounded (e.g. if `group_size` is usize::MAX),
// on the next field id // so we need to make sure that we are not iterating on the next field id
if next_key.field_id != self.field_id { if next_key.field_id != self.field_id {
return Ok(()); break;
} }
// now, do we skip, stop, or visit? // now, do we skip, stop, or visit?
let should_skip = { let should_skip = {
@ -176,6 +199,8 @@ impl<'t, 'b, 'bitmap> FacetRangeSearch<'t, 'b, 'bitmap> {
previous_value = next_value; previous_value = next_value;
continue; continue;
} }
// from here, we should visit the children of the previous element and
// call the function recursively
let level = level - 1; let level = level - 1;
let starting_left_bound = previous_key.left_bound; let starting_left_bound = previous_key.left_bound;
@ -187,7 +212,7 @@ impl<'t, 'b, 'bitmap> FacetRangeSearch<'t, 'b, 'bitmap> {
previous_key = next_key; previous_key = next_key;
previous_value = next_value; previous_value = next_value;
} }
// previous_key/previous_value are the last element // previous_key/previous_value are the last element's key/value
// now, do we skip, stop, or visit? // now, do we skip, stop, or visit?
let should_skip = { let should_skip = {
@ -224,18 +249,41 @@ impl<'t, 'b, 'bitmap> FacetRangeSearch<'t, 'b, 'bitmap> {
Bound::Unbounded => true, Bound::Unbounded => true,
}; };
let right_condition = match (self.right, rightmost_bound) { let right_condition = match (self.right, rightmost_bound) {
(Bound::Included(right), Bound::Included(rightmost)) => rightmost <= right, (Bound::Included(right), Bound::Included(rightmost)) => {
(Bound::Included(right), Bound::Excluded(rightmost)) => rightmost < right, // we need to stay within the bound ..=right
// e.g. x < 8 and rightmost is <= y // the element's range goes to ..=righmost
// condition met if rightmost < 8 // so the element fits entirely within the bound if rightmost <= right
(Bound::Excluded(right), Bound::Included(rightmost)) => rightmost < right, rightmost <= right
// e.g. x < 8 and rightmost is < y }
// condition met only if y <= 8? (Bound::Included(right), Bound::Excluded(rightmost)) => {
(Bound::Excluded(right), Bound::Excluded(rightmost)) => rightmost <= right, // we need to stay within the bound ..=right
// e.g. x < inf. , so yes we take the whole thing // the element's range goes to ..righmost
(Bound::Unbounded, _) => true, // so the element fits entirely within the bound if rightmost <= right
// e.g. x < 7 , righmost is inf rightmost <= right
(_, Bound::Unbounded) => false, // panic? }
(Bound::Excluded(right), Bound::Included(rightmost)) => {
// we need to stay within the bound ..right
// the element's range goes to ..=righmost
// so the element fits entirely within the bound if rightmost < right
rightmost < right
}
(Bound::Excluded(right), Bound::Excluded(rightmost)) => {
// we need to stay within the bound ..right
// the element's range goes to ..righmost
// so the element fits entirely within the bound if rightmost <= right
rightmost <= right
}
(Bound::Unbounded, _) => {
// we need to stay within the bound ..inf
// so the element always fits entirely within the bound
true
}
(_, Bound::Unbounded) => {
// we need to stay within a finite bound
// but the element's range goes to ..inf
// so the element never fits entirely within the bound
false
}
}; };
left_condition && right_condition left_condition && right_condition
}; };
@ -262,7 +310,10 @@ mod tests {
use super::find_docids_of_facet_within_bounds; use super::find_docids_of_facet_within_bounds;
use crate::heed_codec::facet::{FacetGroupKeyCodec, OrderedF64Codec}; use crate::heed_codec::facet::{FacetGroupKeyCodec, OrderedF64Codec};
use crate::milli_snap; use crate::milli_snap;
use crate::search::facet::tests::{get_random_looking_index, get_simple_index}; use crate::search::facet::tests::{
get_random_looking_index, get_random_looking_index_with_multiple_field_ids,
get_simple_index, get_simple_index_with_multiple_field_ids,
};
use crate::snapshot_tests::display_bitmap; use crate::snapshot_tests::display_bitmap;
#[test] #[test]
@ -272,7 +323,12 @@ mod tests {
} }
#[test] #[test]
fn filter_range_increasing() { fn filter_range_increasing() {
let indexes = [get_simple_index(), get_random_looking_index()]; let indexes = [
get_simple_index(),
get_random_looking_index(),
get_simple_index_with_multiple_field_ids(),
get_random_looking_index_with_multiple_field_ids(),
];
for (i, index) in indexes.iter().enumerate() { for (i, index) in indexes.iter().enumerate() {
let txn = index.env.read_txn().unwrap(); let txn = index.env.read_txn().unwrap();
let mut results = String::new(); let mut results = String::new();
@ -316,7 +372,12 @@ mod tests {
} }
#[test] #[test]
fn filter_range_decreasing() { fn filter_range_decreasing() {
let indexes = [get_simple_index(), get_random_looking_index()]; let indexes = [
get_simple_index(),
get_random_looking_index(),
get_simple_index_with_multiple_field_ids(),
get_random_looking_index_with_multiple_field_ids(),
];
for (i, index) in indexes.iter().enumerate() { for (i, index) in indexes.iter().enumerate() {
let txn = index.env.read_txn().unwrap(); let txn = index.env.read_txn().unwrap();
@ -367,7 +428,12 @@ mod tests {
} }
#[test] #[test]
fn filter_range_pinch() { fn filter_range_pinch() {
let indexes = [get_simple_index(), get_random_looking_index()]; let indexes = [
get_simple_index(),
get_random_looking_index(),
get_simple_index_with_multiple_field_ids(),
get_random_looking_index_with_multiple_field_ids(),
];
for (i, index) in indexes.iter().enumerate() { for (i, index) in indexes.iter().enumerate() {
let txn = index.env.read_txn().unwrap(); let txn = index.env.read_txn().unwrap();

View File

@ -119,4 +119,35 @@ pub(crate) mod tests {
txn.commit().unwrap(); txn.commit().unwrap();
index index
} }
pub fn get_simple_index_with_multiple_field_ids() -> FacetIndex<OrderedF64Codec> {
let index = FacetIndex::<OrderedF64Codec>::new(4, 8, 5);
let mut txn = index.env.write_txn().unwrap();
for fid in 0..2 {
for i in 0..256u16 {
let mut bitmap = RoaringBitmap::new();
bitmap.insert(i as u32);
index.insert(&mut txn, fid, &(i as f64), &bitmap);
}
}
txn.commit().unwrap();
index
}
pub fn get_random_looking_index_with_multiple_field_ids() -> FacetIndex<OrderedF64Codec> {
let index = FacetIndex::<OrderedF64Codec>::new(4, 8, 5);
let mut txn = index.env.write_txn().unwrap();
let mut rng = rand::rngs::SmallRng::from_seed([0; 32]);
let keys =
std::iter::from_fn(|| Some(rng.gen_range(0..256))).take(128).collect::<Vec<u32>>();
for fid in 0..2 {
for (_i, &key) in keys.iter().enumerate() {
let mut bitmap = RoaringBitmap::new();
bitmap.insert(key);
bitmap.insert(key + 100);
index.insert(&mut txn, fid, &(key as f64), &bitmap);
}
}
txn.commit().unwrap();
index
}
} }

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
fcedc563a82c1c61f50174a5f3f982b6

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
6cc26e77fc6bd9145deedf14cf422b03

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
57d35cfa419a19a1a1f8d7c8ef096e0f

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
3dbe0547b42759795e9b16989df72cee

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
c1c7a0bb91d53d33724583b6d4a99f16

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
12213d3f1047a0c3d08e4670a7d688e7

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
ca59f20e043a4d52c49e15b10adf96bb

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
cb69e0fe10fb299bafe77514204379cb

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
3456db9a1bb94c33c1e9f656184ee711

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
2127cd818b457e0611e0c8e1a871602a

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
b976551ceff412bfb2ec9bfbda320bbb

View File

@ -0,0 +1,4 @@
---
source: milli/src/search/facet/facet_range_search.rs
---
7620ca1a96882c7147d3fd996570f9b3