use std::collections::{HashMap, HashSet}; use grenad::CompressionType; use heed::types::{Bytes, Str}; use heed::Database; use crate::update::del_add::{deladd_serialize_add_side, DelAdd, KvWriterDelAdd}; use crate::update::index_documents::{ create_sorter, merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap, valid_lmdb_key, write_sorter_into_database, CursorClonableMmap, MergeDeladdCboRoaringBitmaps, }; use crate::{CboRoaringBitmapCodec, Result}; pub struct WordPrefixDocids<'t, 'i> { wtxn: &'t mut heed::RwTxn<'i>, word_docids: Database, word_prefix_docids: Database, pub(crate) chunk_compression_type: CompressionType, pub(crate) chunk_compression_level: Option, pub(crate) max_nb_chunks: Option, pub(crate) max_memory: Option, } impl<'t, 'i> WordPrefixDocids<'t, 'i> { pub fn new( wtxn: &'t mut heed::RwTxn<'i>, word_docids: Database, word_prefix_docids: Database, ) -> WordPrefixDocids<'t, 'i> { WordPrefixDocids { wtxn, word_docids, word_prefix_docids, chunk_compression_type: CompressionType::None, chunk_compression_level: None, max_nb_chunks: None, max_memory: None, } } #[tracing::instrument( level = "trace", skip_all, target = "indexing::prefix", name = "word_prefix_docids" )] pub fn execute( self, new_word_docids: grenad::Merger, new_prefix_fst_words: &[String], common_prefix_fst_words: &[&[String]], del_prefix_fst_words: &HashSet>, ) -> Result<()> { // It is forbidden to keep a mutable reference into the database // and write into it at the same time, therefore we write into another file. let mut prefix_docids_sorter = create_sorter( grenad::SortAlgorithm::Unstable, MergeDeladdCboRoaringBitmaps, self.chunk_compression_type, self.chunk_compression_level, self.max_nb_chunks, self.max_memory, ); if !common_prefix_fst_words.is_empty() { let mut current_prefixes: Option<&&[String]> = None; let mut prefixes_cache = HashMap::new(); let mut new_word_docids_iter = new_word_docids.into_stream_merger_iter()?; while let Some((word, data)) = new_word_docids_iter.next()? { current_prefixes = match current_prefixes.take() { Some(prefixes) if word.starts_with(prefixes[0].as_bytes()) => Some(prefixes), _otherwise => { write_prefixes_in_sorter(&mut prefixes_cache, &mut prefix_docids_sorter)?; common_prefix_fst_words .iter() .find(|prefixes| word.starts_with(prefixes[0].as_bytes())) } }; if let Some(prefixes) = current_prefixes { for prefix in prefixes.iter() { if word.starts_with(prefix.as_bytes()) { match prefixes_cache.get_mut(prefix.as_bytes()) { Some(value) => value.push(data.to_owned()), None => { prefixes_cache .insert(prefix.clone().into(), vec![data.to_owned()]); } } } } } } write_prefixes_in_sorter(&mut prefixes_cache, &mut prefix_docids_sorter)?; } // We fetch the docids associated to the newly added word prefix fst only. let db = self.word_docids.remap_data_type::(); let mut buffer = Vec::new(); for prefix in new_prefix_fst_words { let prefix = std::str::from_utf8(prefix.as_bytes())?; for result in db.prefix_iter(self.wtxn, prefix)? { let (_word, data) = result?; buffer.clear(); let mut writer = KvWriterDelAdd::new(&mut buffer); writer.insert(DelAdd::Addition, data)?; prefix_docids_sorter.insert(prefix, writer.into_inner()?)?; } } // We remove all the entries that are no more required in this word prefix docids database. let mut iter = self.word_prefix_docids.iter_mut(self.wtxn)?.lazily_decode_data(); while let Some((prefix, _)) = iter.next().transpose()? { if del_prefix_fst_words.contains(prefix.as_bytes()) { unsafe { iter.del_current()? }; } } drop(iter); let database_is_empty = self.word_prefix_docids.is_empty(self.wtxn)?; // We finally write the word prefix docids into the LMDB database. write_sorter_into_database( prefix_docids_sorter, &self.word_prefix_docids, self.wtxn, database_is_empty, deladd_serialize_add_side, merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap, )?; Ok(()) } } fn write_prefixes_in_sorter( prefixes: &mut HashMap, Vec>>, sorter: &mut grenad::Sorter, ) -> Result<()> { for (key, data_slices) in prefixes.drain() { for data in data_slices { if valid_lmdb_key(&key) { sorter.insert(&key, data)?; } } } Ok(()) }