use std::collections::HashMap; use std::collections::hash_map::Entry; use anyhow::anyhow; use chrono::Utc; use fst::IntoStreamer; use heed::types::ByteSlice; use roaring::RoaringBitmap; use serde_json::Value; use crate::facet::FacetType; use crate::{Index, BEU32, SmallString32, ExternalDocumentsIds}; use crate::heed_codec::facet::{FieldDocIdFacetStringCodec, FieldDocIdFacetF64Codec, FieldDocIdFacetI64Codec}; use super::ClearDocuments; pub struct DeleteDocuments<'t, 'u, 'i> { wtxn: &'t mut heed::RwTxn<'i, 'u>, index: &'i Index, external_documents_ids: ExternalDocumentsIds<'static>, documents_ids: RoaringBitmap, update_id: u64, } impl<'t, 'u, 'i> DeleteDocuments<'t, 'u, 'i> { pub fn new( wtxn: &'t mut heed::RwTxn<'i, 'u>, index: &'i Index, update_id: u64, ) -> anyhow::Result> { let external_documents_ids = index .external_documents_ids(wtxn)? .into_static(); Ok(DeleteDocuments { wtxn, index, external_documents_ids, documents_ids: RoaringBitmap::new(), update_id, }) } pub fn delete_document(&mut self, docid: u32) { self.documents_ids.insert(docid); } pub fn delete_documents(&mut self, docids: &RoaringBitmap) { self.documents_ids.union_with(docids); } pub fn delete_external_id(&mut self, external_id: &str) -> Option { let docid = self.external_documents_ids.get(external_id)?; self.delete_document(docid); Some(docid) } pub fn execute(self) -> anyhow::Result { self.index.set_updated_at(self.wtxn, &Utc::now())?; // We retrieve the current documents ids that are in the database. let mut documents_ids = self.index.documents_ids(self.wtxn)?; // We can and must stop removing documents in a database that is empty. if documents_ids.is_empty() { return Ok(0); } // We remove the documents ids that we want to delete // from the documents in the database and write them back. let current_documents_ids_len = documents_ids.len(); documents_ids.difference_with(&self.documents_ids); self.index.put_documents_ids(self.wtxn, &documents_ids)?; // We can execute a ClearDocuments operation when the number of documents // to delete is exactly the number of documents in the database. if current_documents_ids_len == self.documents_ids.len() { return ClearDocuments::new(self.wtxn, self.index, self.update_id).execute(); } let fields_ids_map = self.index.fields_ids_map(self.wtxn)?; let id_field = fields_ids_map.id("id").expect(r#"the field "id" to be present"#); let Index { env: _env, main: _main, word_docids, word_prefix_docids, docid_word_positions, word_pair_proximity_docids, word_prefix_pair_proximity_docids, facet_field_id_value_docids, field_id_docid_facet_values, documents, } = self.index; // Number of fields for each document that has been deleted. let mut fields_ids_distribution_diff = HashMap::new(); // Retrieve the words and the external documents ids contained in the documents. let mut words = Vec::new(); let mut external_ids = Vec::new(); for docid in &self.documents_ids { // We create an iterator to be able to get the content and delete the document // content itself. It's faster to acquire a cursor to get and delete, // as we avoid traversing the LMDB B-Tree two times but only once. let key = BEU32::new(docid); let mut iter = documents.range_mut(self.wtxn, &(key..=key))?; if let Some((_key, obkv)) = iter.next().transpose()? { for (field_id, _) in obkv.iter() { *fields_ids_distribution_diff.entry(field_id).or_default() += 1; } if let Some(content) = obkv.get(id_field) { let external_id = match serde_json::from_slice(content).unwrap() { Value::String(string) => SmallString32::from(string.as_str()), Value::Number(number) => SmallString32::from(number.to_string()), _ => return Err(anyhow!("documents ids must be either strings or numbers")), }; external_ids.push(external_id); } iter.del_current()?; } drop(iter); // We iterate through the words positions of the document id, // retrieve the word and delete the positions. let mut iter = docid_word_positions.prefix_iter_mut(self.wtxn, &(docid, ""))?; while let Some(result) = iter.next() { let ((_docid, word), _positions) = result?; // This boolean will indicate if we must remove this word from the words FST. words.push((SmallString32::from(word), false)); iter.del_current()?; } } let mut fields_distribution = self.index.fields_distribution(self.wtxn)?; // We use pre-calculated number of fields occurrences that needs to be deleted // to reflect deleted documents. // If all field occurrences are removed, delete the entry from distribution. // Otherwise, insert new number of occurrences (current_count - count_diff). for (field_id, count_diff) in fields_ids_distribution_diff { let field_name = fields_ids_map.name(field_id).unwrap(); if let Entry::Occupied(mut entry) = fields_distribution.entry(field_name.to_string()) { match entry.get().checked_sub(count_diff) { Some(0) | None => entry.remove(), Some(count) => entry.insert(count) }; } } self.index.put_fields_distribution(self.wtxn, &fields_distribution)?; // We create the FST map of the external ids that we must delete. external_ids.sort_unstable(); let external_ids_to_delete = fst::Set::from_iter(external_ids.iter().map(AsRef::as_ref))?; // We acquire the current external documents ids map... let mut new_external_documents_ids = self.index.external_documents_ids(self.wtxn)?; // ...and remove the to-delete external ids. new_external_documents_ids.delete_ids(external_ids_to_delete)?; // We write the new external ids into the main database. let new_external_documents_ids = new_external_documents_ids.into_static(); self.index.put_external_documents_ids(self.wtxn, &new_external_documents_ids)?; // Maybe we can improve the get performance of the words // if we sort the words first, keeping the LMDB pages in cache. words.sort_unstable(); // We iterate over the words and delete the documents ids // from the word docids database. for (word, must_remove) in &mut words { // We create an iterator to be able to get the content and delete the word docids. // It's faster to acquire a cursor to get and delete or put, as we avoid traversing // the LMDB B-Tree two times but only once. let mut iter = word_docids.prefix_iter_mut(self.wtxn, &word)?; if let Some((key, mut docids)) = iter.next().transpose()? { if key == word.as_ref() { let previous_len = docids.len(); docids.difference_with(&self.documents_ids); if docids.is_empty() { iter.del_current()?; *must_remove = true; } else if docids.len() != previous_len { iter.put_current(key, &docids)?; } } } } // We construct an FST set that contains the words to delete from the words FST. let words_to_delete = words.iter().filter_map(|(word, must_remove)| { if *must_remove { Some(word.as_ref()) } else { None } }); let words_to_delete = fst::Set::from_iter(words_to_delete)?; let new_words_fst = { // We retrieve the current words FST from the database. let words_fst = self.index.words_fst(self.wtxn)?; let difference = words_fst.op().add(&words_to_delete).difference(); // We stream the new external ids that does no more contains the to-delete external ids. let mut new_words_fst_builder = fst::SetBuilder::memory(); new_words_fst_builder.extend_stream(difference.into_stream())?; // We create an words FST set from the above builder. new_words_fst_builder.into_set() }; // We write the new words FST into the main database. self.index.put_words_fst(self.wtxn, &new_words_fst)?; // We iterate over the word prefix docids database and remove the deleted documents ids // from every docids lists. We register the empty prefixes in an fst Set for futur deletion. let mut prefixes_to_delete = fst::SetBuilder::memory(); let mut iter = word_prefix_docids.iter_mut(self.wtxn)?; while let Some(result) = iter.next() { let (prefix, mut docids) = result?; let previous_len = docids.len(); docids.difference_with(&self.documents_ids); if docids.is_empty() { iter.del_current()?; prefixes_to_delete.insert(prefix)?; } else if docids.len() != previous_len { iter.put_current(prefix, &docids)?; } } drop(iter); // We compute the new prefix FST and write it only if there is a change. let prefixes_to_delete = prefixes_to_delete.into_set(); if !prefixes_to_delete.is_empty() { let new_words_prefixes_fst = { // We retrieve the current words prefixes FST from the database. let words_prefixes_fst = self.index.words_prefixes_fst(self.wtxn)?; let difference = words_prefixes_fst.op().add(&prefixes_to_delete).difference(); // We stream the new external ids that does no more contains the to-delete external ids. let mut new_words_prefixes_fst_builder = fst::SetBuilder::memory(); new_words_prefixes_fst_builder.extend_stream(difference.into_stream())?; // We create an words FST set from the above builder. new_words_prefixes_fst_builder.into_set() }; // We write the new words prefixes FST into the main database. self.index.put_words_prefixes_fst(self.wtxn, &new_words_prefixes_fst)?; } // We delete the documents ids from the word prefix pair proximity database docids // and remove the empty pairs too. let db = word_prefix_pair_proximity_docids.remap_key_type::(); let mut iter = db.iter_mut(self.wtxn)?; while let Some(result) = iter.next() { let (key, mut docids) = result?; let previous_len = docids.len(); docids.difference_with(&self.documents_ids); if docids.is_empty() { iter.del_current()?; } else if docids.len() != previous_len { iter.put_current(key, &docids)?; } } drop(iter); // We delete the documents ids that are under the pairs of words, // it is faster and use no memory to iterate over all the words pairs than // to compute the cartesian product of every words of the deleted documents. let mut iter = word_pair_proximity_docids.remap_key_type::().iter_mut(self.wtxn)?; while let Some(result) = iter.next() { let (bytes, mut docids) = result?; let previous_len = docids.len(); docids.difference_with(&self.documents_ids); if docids.is_empty() { iter.del_current()?; } else if docids.len() != previous_len { iter.put_current(bytes, &docids)?; } } drop(iter); // Remove the documents ids from the faceted documents ids. let faceted_fields = self.index.faceted_fields_ids(self.wtxn)?; for (field_id, facet_type) in faceted_fields { let mut docids = self.index.faceted_documents_ids(self.wtxn, field_id)?; docids.difference_with(&self.documents_ids); self.index.put_faceted_documents_ids(self.wtxn, field_id, &docids)?; // We delete the entries that are part of the documents ids. let iter = field_id_docid_facet_values.prefix_iter_mut(self.wtxn, &[field_id])?; match facet_type { FacetType::String => { let mut iter = iter.remap_key_type::(); while let Some(result) = iter.next() { let ((_fid, docid, _value), ()) = result?; if self.documents_ids.contains(docid) { iter.del_current()?; } } }, FacetType::Float => { let mut iter = iter.remap_key_type::(); while let Some(result) = iter.next() { let ((_fid, docid, _value), ()) = result?; if self.documents_ids.contains(docid) { iter.del_current()?; } } }, FacetType::Integer => { let mut iter = iter.remap_key_type::(); while let Some(result) = iter.next() { let ((_fid, docid, _value), ()) = result?; if self.documents_ids.contains(docid) { iter.del_current()?; } } }, } } // We delete the documents ids that are under the facet field id values. let mut iter = facet_field_id_value_docids.iter_mut(self.wtxn)?; while let Some(result) = iter.next() { let (bytes, mut docids) = result?; let previous_len = docids.len(); docids.difference_with(&self.documents_ids); if docids.is_empty() { iter.del_current()?; } else if docids.len() != previous_len { iter.put_current(bytes, &docids)?; } } drop(iter); Ok(self.documents_ids.len()) } } #[cfg(test)] mod tests { use heed::EnvOpenOptions; use crate::update::{IndexDocuments, UpdateFormat}; use super::*; #[test] fn delete_documents_with_numbers_as_primary_key() { let path = tempfile::tempdir().unwrap(); let mut options = EnvOpenOptions::new(); options.map_size(10 * 1024 * 1024); // 10 MB let index = Index::new(options, &path).unwrap(); // First we send 3 documents with an id for only one of them. let mut wtxn = index.write_txn().unwrap(); let content = &br#"[ { "id": 0, "name": "kevin", "object": { "key1": "value1", "key2": "value2" } }, { "id": 1, "name": "kevina", "array": ["I", "am", "fine"] }, { "id": 2, "name": "benoit", "array_of_object": [{ "wow": "amazing" }] } ]"#[..]; let mut builder = IndexDocuments::new(&mut wtxn, &index, 0); builder.update_format(UpdateFormat::Json); builder.execute(content, |_, _| ()).unwrap(); // delete those documents, ids are synchronous therefore 0, 1, and 2. let mut builder = DeleteDocuments::new(&mut wtxn, &index, 1).unwrap(); builder.delete_document(0); builder.delete_document(1); builder.delete_document(2); builder.execute().unwrap(); wtxn.commit().unwrap(); let rtxn = index.read_txn().unwrap(); assert!(index.fields_distribution(&rtxn).unwrap().is_empty()); } }