use std::fs::File; use std::io::BufReader; use heed::types::{Bytes, DecodeIgnore}; use heed::{BytesDecode, Error, RoTxn, RwTxn}; use obkv::KvReader; use roaring::RoaringBitmap; use crate::facet::FacetType; use crate::heed_codec::facet::{ FacetGroupKey, FacetGroupKeyCodec, FacetGroupValue, FacetGroupValueCodec, }; use crate::heed_codec::BytesRefCodec; use crate::search::facet::get_highest_level; use crate::update::del_add::DelAdd; use crate::update::index_documents::valid_lmdb_key; use crate::{CboRoaringBitmapCodec, Index, Result}; enum InsertionResult { InPlace, Expand, Insert, } enum DeletionResult { InPlace, Reduce { next: Option> }, Remove { next: Option> }, } /// Algorithm to incrementally insert and delete elememts into the /// `facet_id_(string/f64)_docids` databases. pub struct FacetsUpdateIncremental { inner: FacetsUpdateIncrementalInner, delta_data: grenad::Reader>, } impl FacetsUpdateIncremental { pub fn new( index: &Index, facet_type: FacetType, delta_data: grenad::Reader>, group_size: u8, min_level_size: u8, max_group_size: u8, ) -> Self { FacetsUpdateIncremental { inner: FacetsUpdateIncrementalInner { db: match facet_type { FacetType::String => index .facet_id_string_docids .remap_key_type::>(), FacetType::Number => index .facet_id_f64_docids .remap_key_type::>(), }, group_size, max_group_size, min_level_size, }, delta_data, } } pub fn execute(self, wtxn: &mut RwTxn) -> crate::Result<()> { let mut cursor = self.delta_data.into_cursor()?; while let Some((key, value)) = cursor.move_on_next()? { if !valid_lmdb_key(key) { continue; } let key = FacetGroupKeyCodec::::bytes_decode(key) .map_err(heed::Error::Encoding)?; let value = KvReader::new(value); let docids_to_delete = value .get(DelAdd::Deletion) .map(CboRoaringBitmapCodec::bytes_decode) .map(|o| o.map_err(heed::Error::Encoding)); let docids_to_add = value .get(DelAdd::Addition) .map(CboRoaringBitmapCodec::bytes_decode) .map(|o| o.map_err(heed::Error::Encoding)); if let Some(docids_to_delete) = docids_to_delete { let docids_to_delete = docids_to_delete?; self.inner.delete(wtxn, key.field_id, key.left_bound, &docids_to_delete)?; } if let Some(docids_to_add) = docids_to_add { let docids_to_add = docids_to_add?; self.inner.insert(wtxn, key.field_id, key.left_bound, &docids_to_add)?; } } Ok(()) } } /// Implementation of `FacetsUpdateIncremental` that is independent of milli's `Index` type pub struct FacetsUpdateIncrementalInner { pub db: heed::Database, FacetGroupValueCodec>, pub group_size: u8, pub min_level_size: u8, pub max_group_size: u8, } impl FacetsUpdateIncrementalInner { /// Find the `FacetGroupKey`/`FacetGroupValue` in the database that /// should be used to insert the new `facet_value` for the given `field_id` and `level` /// where `level` must be strictly greater than 0. /// /// For example, when inserting the facet value `4`, there are two possibilities: /// /// 1. We find a key whose lower bound is 3 followed by a key whose lower bound is 6. Therefore, /// we know that the implicit range of the first key is 3..6, which contains 4. /// So the new facet value belongs in that first key/value pair. /// /// 2. The first key of the level has a lower bound of `5`. We return this key/value pair /// but will need to change the lowerbound of this key to `4` in order to insert this facet value. fn find_insertion_key_value( &self, field_id: u16, level: u8, facet_value: &[u8], txn: &RoTxn, ) -> Result<(FacetGroupKey>, FacetGroupValue)> { assert!(level > 0); match self.db.get_lower_than_or_equal_to( txn, &FacetGroupKey { field_id, level, left_bound: facet_value }, )? { Some((key, value)) => { if key.level != level { let mut prefix = vec![]; prefix.extend_from_slice(&field_id.to_be_bytes()); prefix.push(level); let mut iter = self .db .remap_types::() .prefix_iter(txn, prefix.as_slice())?; let (key_bytes, value) = iter.next().unwrap()?; Ok(( FacetGroupKeyCodec::::bytes_decode(key_bytes) .map_err(Error::Encoding)? .into_owned(), value, )) } else { Ok((key.into_owned(), value)) } } None => { // We checked that the level is > 0 // Since all keys of level 1 are greater than those of level 0, // we are guaranteed that db.get_lower_than_or_equal_to(key) exists panic!() } } } /// Insert the given facet value and corresponding document ids in the level 0 of the database /// /// ## Return /// See documentation of `insert_in_level` fn insert_in_level_0( &self, txn: &mut RwTxn, field_id: u16, facet_value: &[u8], docids: &RoaringBitmap, ) -> Result { let key = FacetGroupKey { field_id, level: 0, left_bound: facet_value }; let value = FacetGroupValue { bitmap: docids.clone(), size: 1 }; let mut level0_prefix = vec![]; level0_prefix.extend_from_slice(&field_id.to_be_bytes()); level0_prefix.push(0); let mut iter = self.db.remap_types::().prefix_iter(txn, &level0_prefix)?; if iter.next().is_none() { drop(iter); self.db.put(txn, &key, &value)?; Ok(InsertionResult::Insert) } else { drop(iter); let old_value = self.db.get(txn, &key)?; match old_value { Some(mut updated_value) => { // now merge the two updated_value.bitmap |= value.bitmap; self.db.put(txn, &key, &updated_value)?; Ok(InsertionResult::InPlace) } None => { self.db.put(txn, &key, &value)?; Ok(InsertionResult::Insert) } } } } /// Insert the given facet value and corresponding document ids in all the levels of the database up to the given `level`. /// This function works recursively. /// /// ## Return /// Returns the effect of adding the facet value to the database on the given `level`. /// /// - `InsertionResult::InPlace` means that inserting the `facet_value` into the `level` did not have /// an effect on the number of keys in that level. Therefore, it did not increase the number of children /// of the parent node. /// /// - `InsertionResult::Insert` means that inserting the `facet_value` into the `level` resulted /// in the addition of a new key in that level, and that therefore the number of children /// of the parent node should be incremented. fn insert_in_level( &self, txn: &mut RwTxn, field_id: u16, level: u8, facet_value: &[u8], docids: &RoaringBitmap, ) -> Result { if level == 0 { return self.insert_in_level_0(txn, field_id, facet_value, docids); } let max_group_size = self.max_group_size; let result = self.insert_in_level(txn, field_id, level - 1, facet_value, docids)?; // level below inserted an element let (insertion_key, insertion_value) = self.find_insertion_key_value(field_id, level, facet_value, txn)?; match result { // because we know that we inserted in place, the facet_value is not a new one // thus it doesn't extend a group, and thus the insertion key computed above is // still correct InsertionResult::InPlace => { let mut updated_value = insertion_value; updated_value.bitmap |= docids; self.db.put(txn, &insertion_key.as_ref(), &updated_value)?; return Ok(InsertionResult::InPlace); } InsertionResult::Expand => {} InsertionResult::Insert => {} } // Here we know that inserting the facet value in the level below resulted in the creation // of a new key. Therefore, it may be the case that we need to modify the left bound of the // insertion key (see documentation of `find_insertion_key_value` for an example of when that // could happen). let (insertion_key, insertion_key_was_modified) = { let mut new_insertion_key = insertion_key.clone(); let mut key_should_be_modified = false; if facet_value < insertion_key.left_bound.as_slice() { new_insertion_key.left_bound = facet_value.to_vec(); key_should_be_modified = true; } if key_should_be_modified { let is_deleted = self.db.delete(txn, &insertion_key.as_ref())?; assert!(is_deleted); self.db.put(txn, &new_insertion_key.as_ref(), &insertion_value)?; } (new_insertion_key, key_should_be_modified) }; // Now we know that the insertion key contains the `facet_value`. // We still need to update the insertion value by: // 1. Incrementing the number of children (since the recursive call returned `InsertionResult::Insert`) // 2. Merge the previous docids with the new one let mut updated_value = insertion_value; if matches!(result, InsertionResult::Insert) { updated_value.size += 1; } if updated_value.size < max_group_size { updated_value.bitmap |= docids; self.db.put(txn, &insertion_key.as_ref(), &updated_value)?; if insertion_key_was_modified { return Ok(InsertionResult::Expand); } else { return Ok(InsertionResult::InPlace); } } // We've increased the group size of the value and realised it has become greater than or equal to `max_group_size` // Therefore it must be split into two nodes. let size_left = updated_value.size / 2; let size_right = updated_value.size - size_left; let level_below = level - 1; let start_key = FacetGroupKey { field_id, level: level_below, left_bound: insertion_key.left_bound.as_slice(), }; let mut iter = self.db.range(txn, &(start_key..))?.take((size_left as usize) + (size_right as usize)); let group_left = { let mut values_left = RoaringBitmap::new(); let mut i = 0; for next in iter.by_ref() { let (_key, value) = next?; i += 1; values_left |= &value.bitmap; if i == size_left { break; } } let key = FacetGroupKey { field_id, level, left_bound: insertion_key.left_bound.clone() }; let value = FacetGroupValue { size: size_left, bitmap: values_left }; (key, value) }; let group_right = { let ( FacetGroupKey { left_bound: right_left_bound, .. }, FacetGroupValue { bitmap: mut values_right, .. }, ) = iter.next().unwrap()?; for next in iter.by_ref() { let (_, value) = next?; values_right |= &value.bitmap; } let key = FacetGroupKey { field_id, level, left_bound: right_left_bound.to_vec() }; let value = FacetGroupValue { size: size_right, bitmap: values_right }; (key, value) }; drop(iter); let _ = self.db.delete(txn, &insertion_key.as_ref())?; self.db.put(txn, &group_left.0.as_ref(), &group_left.1)?; self.db.put(txn, &group_right.0.as_ref(), &group_right.1)?; Ok(InsertionResult::Insert) } /// Insert the given facet value and corresponding document ids in the database. pub fn insert( &self, txn: &mut RwTxn, field_id: u16, facet_value: &[u8], docids: &RoaringBitmap, ) -> Result<()> { if docids.is_empty() { return Ok(()); } let group_size = self.group_size; let highest_level = get_highest_level(txn, self.db, field_id)?; let result = self.insert_in_level(txn, field_id, highest_level, facet_value, docids)?; match result { InsertionResult::InPlace => return Ok(()), InsertionResult::Expand => return Ok(()), InsertionResult::Insert => {} } // Here we check whether the highest level has exceeded `min_level_size` * `self.group_size`. // If it has, we must build an addition level above it. let mut highest_level_prefix = vec![]; highest_level_prefix.extend_from_slice(&field_id.to_be_bytes()); highest_level_prefix.push(highest_level); let size_highest_level = self.db.remap_types::().prefix_iter(txn, &highest_level_prefix)?.count(); if size_highest_level < self.group_size as usize * self.min_level_size as usize { return Ok(()); } let mut groups_iter = self .db .remap_types::() .prefix_iter(txn, &highest_level_prefix)?; let nbr_new_groups = size_highest_level / self.group_size as usize; let nbr_leftover_elements = size_highest_level % self.group_size as usize; let mut to_add = vec![]; for _ in 0..nbr_new_groups { let mut first_key = None; let mut values = RoaringBitmap::new(); for _ in 0..group_size { let (key_bytes, value_i) = groups_iter.next().unwrap()?; let key_i = FacetGroupKeyCodec::::bytes_decode(key_bytes) .map_err(Error::Encoding)?; if first_key.is_none() { first_key = Some(key_i); } values |= value_i.bitmap; } let key = FacetGroupKey { field_id, level: highest_level + 1, left_bound: first_key.unwrap().left_bound, }; let value = FacetGroupValue { size: group_size, bitmap: values }; to_add.push((key.into_owned(), value)); } // now we add the rest of the level, in case its size is > group_size * min_level_size // this can indeed happen if the min_level_size parameter changes between two calls to `insert` if nbr_leftover_elements > 0 { let mut first_key = None; let mut values = RoaringBitmap::new(); for _ in 0..nbr_leftover_elements { let (key_bytes, value_i) = groups_iter.next().unwrap()?; let key_i = FacetGroupKeyCodec::::bytes_decode(key_bytes) .map_err(Error::Encoding)?; if first_key.is_none() { first_key = Some(key_i); } values |= value_i.bitmap; } let key = FacetGroupKey { field_id, level: highest_level + 1, left_bound: first_key.unwrap().left_bound, }; // Note: nbr_leftover_elements can be casted to a u8 since it is bounded by `max_group_size` // when it is created above. let value = FacetGroupValue { size: nbr_leftover_elements as u8, bitmap: values }; to_add.push((key.into_owned(), value)); } drop(groups_iter); for (key, value) in to_add { self.db.put(txn, &key.as_ref(), &value)?; } Ok(()) } /// Delete the given document id from the given facet value in the database, from level 0 to the /// the given level. /// /// ## Return /// Returns the effect of removing the document id from the database on the given `level`. /// /// - `DeletionResult::InPlace` means that deleting the document id did not have /// an effect on the keys in that level. /// /// - `DeletionResult::Reduce` means that deleting the document id resulted in a change in the /// number of keys in the level. For example, removing a document id from the facet value `3` could /// cause it to have no corresponding document in level 0 anymore, and therefore the key was deleted /// entirely. In that case, `DeletionResult::Remove` is returned. The parent of the deleted key must /// then adjust its group size. If its group size falls to 0, then it will need to be deleted as well. /// /// - `DeletionResult::Reduce` means that deleting the document id resulted in a change in the /// bounds of the keys of the level. For example, removing a document id from the facet value /// `3` might have caused the facet value `3` to have no corresponding document in level 0. Therefore, /// in level 1, the key with the left bound `3` had to be changed to the next facet value (e.g. 4). /// In that case `DeletionResult::Reduce` is returned. The parent of the reduced key may need to adjust /// its left bound as well. fn delete_in_level( &self, txn: &mut RwTxn, field_id: u16, level: u8, facet_value: &[u8], docids: &RoaringBitmap, ) -> Result { if level == 0 { return self.delete_in_level_0(txn, field_id, facet_value, docids); } let (deletion_key, mut bitmap) = self.find_insertion_key_value(field_id, level, facet_value, txn)?; let result = self.delete_in_level(txn, field_id, level - 1, facet_value, docids)?; let mut decrease_size = false; let next_key = match result { DeletionResult::InPlace => { bitmap.bitmap -= docids; self.db.put(txn, &deletion_key.as_ref(), &bitmap)?; return Ok(DeletionResult::InPlace); } DeletionResult::Reduce { next } => next, DeletionResult::Remove { next } => { decrease_size = true; next } }; // If either DeletionResult::Reduce or DeletionResult::Remove was returned, // then we may need to adjust the left_bound of the deletion key. // If DeletionResult::Remove was returned, then we need to decrease the group // size of the deletion key. let mut updated_value = bitmap; if decrease_size { updated_value.size -= 1; } if updated_value.size == 0 { self.db.delete(txn, &deletion_key.as_ref())?; Ok(DeletionResult::Remove { next: next_key }) } else { let mut updated_deletion_key = deletion_key.clone(); let reduced_range = facet_value == deletion_key.left_bound; if reduced_range { updated_deletion_key.left_bound = next_key.clone().unwrap(); } updated_value.bitmap -= docids; let _ = self.db.delete(txn, &deletion_key.as_ref())?; self.db.put(txn, &updated_deletion_key.as_ref(), &updated_value)?; if reduced_range { Ok(DeletionResult::Reduce { next: next_key }) } else { Ok(DeletionResult::InPlace) } } } fn delete_in_level_0( &self, txn: &mut RwTxn, field_id: u16, facet_value: &[u8], docids: &RoaringBitmap, ) -> Result { let key = FacetGroupKey { field_id, level: 0, left_bound: facet_value }; let mut bitmap = self.db.get(txn, &key)?.unwrap().bitmap; bitmap -= docids; if bitmap.is_empty() { let mut next_key = None; if let Some((next, _)) = self.db.remap_data_type::().get_greater_than(txn, &key)? { if next.field_id == field_id && next.level == 0 { next_key = Some(next.left_bound.to_vec()); } } self.db.delete(txn, &key)?; Ok(DeletionResult::Remove { next: next_key }) } else { self.db.put(txn, &key, &FacetGroupValue { size: 1, bitmap })?; Ok(DeletionResult::InPlace) } } pub fn delete( &self, txn: &mut RwTxn, field_id: u16, facet_value: &[u8], docids: &RoaringBitmap, ) -> Result<()> { if self .db .remap_data_type::() .get(txn, &FacetGroupKey { field_id, level: 0, left_bound: facet_value })? .is_none() { return Ok(()); } let highest_level = get_highest_level(txn, self.db, field_id)?; let result = self.delete_in_level(txn, field_id, highest_level, facet_value, docids)?; match result { DeletionResult::InPlace => return Ok(()), DeletionResult::Reduce { .. } => return Ok(()), DeletionResult::Remove { .. } => {} } // if we either removed a key from the highest level, its size may have fallen // below `min_level_size`, in which case we need to remove the entire level let mut highest_level_prefix = vec![]; highest_level_prefix.extend_from_slice(&field_id.to_be_bytes()); highest_level_prefix.push(highest_level); if highest_level == 0 || self .db .remap_types::() .prefix_iter(txn, &highest_level_prefix)? .count() >= self.min_level_size as usize { return Ok(()); } let mut to_delete = vec![]; let mut iter = self.db.remap_types::().prefix_iter(txn, &highest_level_prefix)?; for el in iter.by_ref() { let (k, _) = el?; to_delete.push( FacetGroupKeyCodec::::bytes_decode(k) .map_err(Error::Encoding)? .into_owned(), ); } drop(iter); for k in to_delete { self.db.delete(txn, &k.as_ref())?; } Ok(()) } } impl<'a> FacetGroupKey<&'a [u8]> { pub fn into_owned(self) -> FacetGroupKey> { FacetGroupKey { field_id: self.field_id, level: self.level, left_bound: self.left_bound.to_vec(), } } } impl FacetGroupKey> { pub fn as_ref(&self) -> FacetGroupKey<&[u8]> { FacetGroupKey { field_id: self.field_id, level: self.level, left_bound: self.left_bound.as_slice(), } } } #[cfg(test)] mod tests { use rand::seq::SliceRandom; use rand::{Rng, SeedableRng}; use roaring::RoaringBitmap; use crate::heed_codec::facet::OrderedF64Codec; use crate::heed_codec::StrRefCodec; use crate::milli_snap; use crate::update::facet::test_helpers::FacetIndex; #[test] fn append() { let index = FacetIndex::::new(4, 8, 5); for i in 0..256u16 { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i as u32); let mut txn = index.env.write_txn().unwrap(); index.insert(&mut txn, 0, &(i as f64), &bitmap); txn.commit().unwrap(); } let txn = index.env.read_txn().unwrap(); index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}")); } #[test] fn many_field_ids_append() { let index = FacetIndex::::new(4, 8, 5); for i in 0..256u16 { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i as u32); let mut txn = index.env.write_txn().unwrap(); index.insert(&mut txn, 0, &(i as f64), &bitmap); txn.commit().unwrap(); } for i in 0..256u16 { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i as u32); let mut txn = index.env.write_txn().unwrap(); index.insert(&mut txn, 2, &(i as f64), &bitmap); txn.commit().unwrap(); } for i in 0..256u16 { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i as u32); let mut txn = index.env.write_txn().unwrap(); index.insert(&mut txn, 1, &(i as f64), &bitmap); txn.commit().unwrap(); } let txn = index.env.read_txn().unwrap(); index.verify_structure_validity(&txn, 0); index.verify_structure_validity(&txn, 1); index.verify_structure_validity(&txn, 2); txn.commit().unwrap(); milli_snap!(format!("{index}")); } #[test] fn many_field_ids_prepend() { let index = FacetIndex::::new(4, 8, 5); for i in (0..256).rev() { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i as u32); let mut txn = index.env.write_txn().unwrap(); index.insert(&mut txn, 0, &(i as f64), &bitmap); txn.commit().unwrap(); } for i in (0..256).rev() { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i as u32); let mut txn = index.env.write_txn().unwrap(); index.insert(&mut txn, 2, &(i as f64), &bitmap); txn.commit().unwrap(); } for i in (0..256).rev() { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i as u32); let mut txn = index.env.write_txn().unwrap(); index.insert(&mut txn, 1, &(i as f64), &bitmap); txn.commit().unwrap(); } let txn = index.env.read_txn().unwrap(); index.verify_structure_validity(&txn, 0); index.verify_structure_validity(&txn, 1); index.verify_structure_validity(&txn, 2); txn.commit().unwrap(); milli_snap!(format!("{index}")); } #[test] fn prepend() { let index = FacetIndex::::new(4, 8, 5); let mut txn = index.env.write_txn().unwrap(); for i in (0..256).rev() { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i); index.insert(&mut txn, 0, &(i as f64), &bitmap); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}")); } #[test] fn shuffled() { let index = FacetIndex::::new(4, 8, 5); let mut txn = index.env.write_txn().unwrap(); let mut keys = (0..256).collect::>(); let mut rng = rand::rngs::SmallRng::from_seed([0; 32]); keys.shuffle(&mut rng); for (_i, key) in keys.into_iter().enumerate() { let mut bitmap = RoaringBitmap::new(); bitmap.insert(key); index.insert(&mut txn, 0, &(key as f64), &bitmap); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}")); } #[test] fn merge_values() { let index = FacetIndex::::new(4, 8, 5); let mut txn = index.env.write_txn().unwrap(); let mut keys = (0..256).collect::>(); let mut rng = rand::rngs::SmallRng::from_seed([0; 32]); keys.shuffle(&mut rng); for (_i, key) in keys.into_iter().enumerate() { let mut bitmap = RoaringBitmap::new(); bitmap.insert(key); bitmap.insert(rng.gen_range(256..512)); index.verify_structure_validity(&txn, 0); index.insert(&mut txn, 0, &(key as f64), &bitmap); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}")); } #[test] fn delete_from_end() { let index = FacetIndex::::new(4, 8, 5); let mut txn = index.env.write_txn().unwrap(); for i in 0..256 { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i); index.verify_structure_validity(&txn, 0); index.insert(&mut txn, 0, &(i as f64), &bitmap); } for i in (200..256).rev() { index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(i as f64), i as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 200); let mut txn = index.env.write_txn().unwrap(); for i in (150..200).rev() { index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(i as f64), i as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 150); let mut txn = index.env.write_txn().unwrap(); for i in (100..150).rev() { index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(i as f64), i as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 100); let mut txn = index.env.write_txn().unwrap(); for i in (17..100).rev() { index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(i as f64), i as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 17); let mut txn = index.env.write_txn().unwrap(); for i in (15..17).rev() { index.delete_single_docid(&mut txn, 0, &(i as f64), i as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 15); let mut txn = index.env.write_txn().unwrap(); for i in (0..15).rev() { index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(i as f64), i as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 0); } #[test] fn delete_from_start() { let index = FacetIndex::::new(4, 8, 5); let mut txn = index.env.write_txn().unwrap(); for i in 0..256 { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i); index.verify_structure_validity(&txn, 0); index.insert(&mut txn, 0, &(i as f64), &bitmap); } for i in 0..128 { index.delete_single_docid(&mut txn, 0, &(i as f64), i as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 127); let mut txn = index.env.write_txn().unwrap(); for i in 128..216 { index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(i as f64), i as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 215); let mut txn = index.env.write_txn().unwrap(); for i in 216..256 { index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(i as f64), i as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 255); } #[test] #[allow(clippy::needless_range_loop)] fn delete_shuffled() { let index = FacetIndex::::new(4, 8, 5); let mut txn = index.env.write_txn().unwrap(); for i in 0..256 { let mut bitmap = RoaringBitmap::new(); bitmap.insert(i); index.verify_structure_validity(&txn, 0); index.insert(&mut txn, 0, &(i as f64), &bitmap); } let mut keys = (0..256).collect::>(); let mut rng = rand::rngs::SmallRng::from_seed([0; 32]); keys.shuffle(&mut rng); for i in 0..128 { let key = keys[i]; index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(key as f64), key as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 127); let mut txn = index.env.write_txn().unwrap(); for i in 128..216 { let key = keys[i]; index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(key as f64), key as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); let mut txn = index.env.write_txn().unwrap(); milli_snap!(format!("{index}"), 215); for i in 216..256 { let key = keys[i]; index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(key as f64), key as u32); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), 255); } #[test] fn in_place_level0_insert() { let index = FacetIndex::::new(4, 8, 5); let mut txn = index.env.write_txn().unwrap(); let mut keys = (0..16).collect::>(); let mut rng = rand::rngs::SmallRng::from_seed([0; 32]); keys.shuffle(&mut rng); for i in 0..4 { for &key in keys.iter() { let mut bitmap = RoaringBitmap::new(); bitmap.insert(rng.gen_range(i * 256..(i + 1) * 256)); index.verify_structure_validity(&txn, 0); index.insert(&mut txn, 0, &(key as f64), &bitmap); } } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}")); } #[test] fn in_place_level0_delete() { let index = FacetIndex::::new(4, 8, 5); let mut txn = index.env.write_txn().unwrap(); let mut keys = (0..64).collect::>(); let mut rng = rand::rngs::SmallRng::from_seed([0; 32]); keys.shuffle(&mut rng); for &key in keys.iter() { let mut bitmap = RoaringBitmap::new(); bitmap.insert(key); bitmap.insert(key + 100); index.verify_structure_validity(&txn, 0); index.insert(&mut txn, 0, &(key as f64), &bitmap); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), "before_delete"); let mut txn = index.env.write_txn().unwrap(); for &key in keys.iter() { index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &(key as f64), key + 100); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), "after_delete"); } #[test] fn shuffle_merge_string_and_delete() { let index = FacetIndex::::new(4, 8, 5); let mut txn = index.env.write_txn().unwrap(); let mut keys = (1000..1064).collect::>(); let mut rng = rand::rngs::SmallRng::from_seed([0; 32]); keys.shuffle(&mut rng); for &key in keys.iter() { let mut bitmap = RoaringBitmap::new(); bitmap.insert(key); bitmap.insert(key + 100); index.verify_structure_validity(&txn, 0); index.insert(&mut txn, 0, &format!("{key:x}").as_str(), &bitmap); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), "before_delete"); let mut txn = index.env.write_txn().unwrap(); for &key in keys.iter() { index.verify_structure_validity(&txn, 0); index.delete_single_docid(&mut txn, 0, &format!("{key:x}").as_str(), key + 100); } index.verify_structure_validity(&txn, 0); txn.commit().unwrap(); milli_snap!(format!("{index}"), "after_delete"); } } // fuzz tests #[cfg(all(test, fuzzing))] /** Fuzz test for the incremental indxer. The fuzz test uses fuzzcheck, a coverage-guided fuzzer. See https://github.com/loiclec/fuzzcheck-rs and https://fuzzcheck.neocities.org for more information. It is only run when using the `cargo fuzzcheck` command line tool, which can be installed with: ```sh cargo install cargo-fuzzcheck ``` To start the fuzz test, run (from the base folder or from milli/): ```sh cargo fuzzcheck update::facet::incremental::fuzz::fuzz ``` and wait a couple minutes to make sure the code was thoroughly tested, then hit `Ctrl-C` to stop the fuzzer. The corpus generated by the fuzzer is located in milli/fuzz. To work on this module with rust-analyzer working properly, add the following to your .cargo/config.toml file: ```toml [build] rustflags = ["--cfg", "fuzzing"] ``` The fuzz test generates sequences of additions and deletions to the facet database and ensures that: 1. its structure is still internally valid 2. its content is the same as a trivially correct implementation of the same database */ mod fuzz { use std::collections::{BTreeMap, HashMap}; use std::iter::FromIterator; use std::rc::Rc; use fuzzcheck::mutators::integer::U8Mutator; use fuzzcheck::mutators::integer_within_range::{U16WithinRangeMutator, U8WithinRangeMutator}; use fuzzcheck::mutators::vector::VecMutator; use fuzzcheck::DefaultMutator; use roaring::RoaringBitmap; use tempfile::TempDir; use super::*; use crate::update::facet::test_helpers::FacetIndex; #[derive(Default)] pub struct TrivialDatabase { pub elements: BTreeMap>, } impl TrivialDatabase where T: Ord + Clone + Eq + std::fmt::Debug, { #[no_coverage] pub fn insert(&mut self, field_id: u16, new_key: &T, new_values: &RoaringBitmap) { if new_values.is_empty() { return; } let values_field_id = self.elements.entry(field_id).or_default(); let values = values_field_id.entry(new_key.clone()).or_default(); *values |= new_values; } #[no_coverage] pub fn delete(&mut self, field_id: u16, key: &T, values_to_remove: &RoaringBitmap) { if let Some(values_field_id) = self.elements.get_mut(&field_id) { if let Some(values) = values_field_id.get_mut(&key) { *values -= values_to_remove; if values.is_empty() { values_field_id.remove(&key); } } if values_field_id.is_empty() { self.elements.remove(&field_id); } } } } #[derive(Clone, DefaultMutator, serde::Serialize, serde::Deserialize)] struct Operation { #[field_mutator(VecMutator = { VecMutator::new(u8::default_mutator(), 0 ..= 5) })] key: Vec, #[field_mutator(U8WithinRangeMutator = { U8WithinRangeMutator::new(..32) })] group_size: u8, #[field_mutator(U8WithinRangeMutator = { U8WithinRangeMutator::new(..32) })] max_group_size: u8, #[field_mutator(U8WithinRangeMutator = { U8WithinRangeMutator::new(..32) })] min_level_size: u8, #[field_mutator(U16WithinRangeMutator = { U16WithinRangeMutator::new(..=3) })] field_id: u16, kind: OperationKind, } #[derive(Clone, DefaultMutator, serde::Serialize, serde::Deserialize)] enum OperationKind { Insert( #[field_mutator(VecMutator = { VecMutator::new(U8Mutator::default(), 0 ..= 10) })] Vec, ), Delete( #[field_mutator(VecMutator = { VecMutator::new(U8Mutator::default(), 0 ..= 10) })] Vec, ), } #[no_coverage] fn compare_with_trivial_database(tempdir: Rc, operations: &[Operation]) { let index = FacetIndex::::open_from_tempdir(tempdir, 4, 8, 5); // dummy params, they'll be overwritten let mut txn = index.env.write_txn().unwrap(); let mut trivial_db = TrivialDatabase::>::default(); let mut value_to_keys = HashMap::>>::new(); for Operation { key, group_size, max_group_size, min_level_size, field_id, kind } in operations { index.set_group_size(*group_size); index.set_max_group_size(*max_group_size); index.set_min_level_size(*min_level_size); match kind { OperationKind::Insert(values) => { let mut bitmap = RoaringBitmap::new(); for value in values { bitmap.insert(*value as u32); value_to_keys.entry(*value).or_default().push(key.clone()); } index.insert(&mut txn, *field_id, &key.as_slice(), &bitmap); trivial_db.insert(*field_id, &key, &bitmap); } OperationKind::Delete(values) => { let values = RoaringBitmap::from_iter(values.iter().copied().map(|x| x as u32)); let mut values_per_key = HashMap::new(); for value in values { if let Some(keys) = value_to_keys.get(&(value as u8)) { for key in keys { let values: &mut RoaringBitmap = values_per_key.entry(key).or_default(); values.insert(value); } } } for (key, values) in values_per_key { index.delete(&mut txn, *field_id, &key.as_slice(), &values); trivial_db.delete(*field_id, &key, &values); } } } } for (field_id, values_field_id) in trivial_db.elements.iter() { let level0iter = index .content .as_polymorph() .prefix_iter::<_, Bytes, FacetGroupValueCodec>(&mut txn, &field_id.to_be_bytes()) .unwrap(); for ((key, values), group) in values_field_id.iter().zip(level0iter) { let (group_key, group_values) = group.unwrap(); let group_key = FacetGroupKeyCodec::::bytes_decode(group_key).unwrap(); assert_eq!(key, &group_key.left_bound); assert_eq!(values, &group_values.bitmap); } } for (field_id, values_field_id) in trivial_db.elements.iter() { let level0iter = index .content .as_polymorph() .prefix_iter::<_, Bytes, FacetGroupValueCodec>(&txn, &field_id.to_be_bytes()) .unwrap(); for ((key, values), group) in values_field_id.iter().zip(level0iter) { let (group_key, group_values) = group.unwrap(); let group_key = FacetGroupKeyCodec::::bytes_decode(group_key).unwrap(); assert_eq!(key, &group_key.left_bound); assert_eq!(values, &group_values.bitmap); } index.verify_structure_validity(&txn, *field_id); } txn.abort().unwrap(); } #[test] #[no_coverage] fn fuzz() { let tempdir = Rc::new(TempDir::new().unwrap()); let tempdir_cloned = tempdir.clone(); let result = fuzzcheck::fuzz_test(move |operations: &[Operation]| { compare_with_trivial_database(tempdir_cloned.clone(), operations) }) .default_mutator() .serde_serializer() .default_sensor_and_pool_with_custom_filter(|file, function| { file == std::path::Path::new("milli/src/update/facet/incremental.rs") && !function.contains("serde") && !function.contains("tests::") && !function.contains("fuzz::") && !function.contains("display_bitmap") }) .arguments_from_cargo_fuzzcheck() .launch(); assert!(!result.found_test_failure); } }