Extract and write the documents and words fst in the database

This commit is contained in:
Clément Renault 2024-09-04 09:59:19 +02:00
parent 52d32b4ee9
commit 27b4cab857
No known key found for this signature in database
GPG Key ID: F250A4C4E3AE5F5F
5 changed files with 225 additions and 38 deletions

View File

@ -5,6 +5,7 @@ use grenad::Merger;
use heed::types::Bytes;
use super::StdResult;
use crate::index::main_key::{DOCUMENTS_IDS_KEY, WORDS_FST_KEY};
use crate::update::new::KvReaderFieldId;
use crate::update::MergeDeladdCboRoaringBitmaps;
use crate::{DocumentId, Index};
@ -22,12 +23,14 @@ pub fn extractors_merger_channels(cap: usize) -> ExtractorsMergerChannels {
ExtractorsMergerChannels {
merger_receiver: MergerReceiver(receiver),
deladd_cbo_roaring_bitmap_sender: DeladdCboRoaringBitmapSender(sender.clone()),
extracted_documents_sender: ExtractedDocumentsSender(sender.clone()),
}
}
pub struct ExtractorsMergerChannels {
pub merger_receiver: MergerReceiver,
pub deladd_cbo_roaring_bitmap_sender: DeladdCboRoaringBitmapSender,
pub extracted_documents_sender: ExtractedDocumentsSender,
}
pub struct KeyValueEntry {
@ -95,18 +98,37 @@ impl DocumentEntry {
}
}
pub enum WriterOperation {
WordDocids(EntryOperation),
Document(DocumentEntry),
pub struct DocumentDeletionEntry(DocumentId);
impl DocumentDeletionEntry {
pub fn key(&self) -> [u8; 4] {
self.0.to_be_bytes()
}
}
pub struct WriterOperation {
database: Database,
entry: EntryOperation,
}
pub enum Database {
WordDocids,
Documents,
Main,
}
impl WriterOperation {
pub fn database(&self, index: &Index) -> heed::Database<Bytes, Bytes> {
match self {
WriterOperation::WordDocids(_) => index.word_docids.remap_types(),
WriterOperation::Document(_) => index.documents.remap_types(),
match self.database {
Database::Main => index.main.remap_types(),
Database::Documents => index.documents.remap_types(),
Database::WordDocids => index.word_docids.remap_types(),
}
}
pub fn entry(self) -> EntryOperation {
self.entry
}
}
pub struct WriterReceiver(Receiver<WriterOperation>);
@ -123,37 +145,93 @@ impl IntoIterator for WriterReceiver {
pub struct MergerSender(Sender<WriterOperation>);
impl MergerSender {
pub fn main(&self) -> MainSender<'_> {
MainSender(&self.0)
}
pub fn word_docids(&self) -> WordDocidsSender<'_> {
WordDocidsSender(&self.0)
}
pub fn documents(&self) -> DocumentsSender<'_> {
DocumentsSender(&self.0)
}
pub fn send_documents_ids(&self, bitmap: &[u8]) -> StdResult<(), SendError<()>> {
let entry = EntryOperation::Write(KeyValueEntry::from_key_value(
DOCUMENTS_IDS_KEY.as_bytes(),
bitmap,
));
match self.0.send(WriterOperation { database: Database::Main, entry }) {
Ok(()) => Ok(()),
Err(SendError(_)) => Err(SendError(())),
}
}
}
pub struct MainSender<'a>(&'a Sender<WriterOperation>);
impl MainSender<'_> {
pub fn write_words_fst(&self, value: &[u8]) -> StdResult<(), SendError<()>> {
let entry =
EntryOperation::Write(KeyValueEntry::from_key_value(WORDS_FST_KEY.as_bytes(), value));
match self.0.send(WriterOperation { database: Database::Main, entry }) {
Ok(()) => Ok(()),
Err(SendError(_)) => Err(SendError(())),
}
}
pub fn delete(&self, key: &[u8]) -> StdResult<(), SendError<()>> {
let entry = EntryOperation::Delete(KeyEntry::from_key(key));
match self.0.send(WriterOperation { database: Database::Main, entry }) {
Ok(()) => Ok(()),
Err(SendError(_)) => Err(SendError(())),
}
}
}
pub struct WordDocidsSender<'a>(&'a Sender<WriterOperation>);
impl WordDocidsSender<'_> {
pub fn write(&self, key: &[u8], value: &[u8]) -> StdResult<(), SendError<()>> {
let operation = EntryOperation::Write(KeyValueEntry::from_key_value(key, value));
match self.0.send(WriterOperation::WordDocids(operation)) {
let entry = EntryOperation::Write(KeyValueEntry::from_key_value(key, value));
match self.0.send(WriterOperation { database: Database::WordDocids, entry }) {
Ok(()) => Ok(()),
Err(SendError(_)) => Err(SendError(())),
}
}
pub fn delete(&self, key: &[u8]) -> StdResult<(), SendError<()>> {
let operation = EntryOperation::Delete(KeyEntry::from_key(key));
match self.0.send(WriterOperation::WordDocids(operation)) {
let entry = EntryOperation::Delete(KeyEntry::from_key(key));
match self.0.send(WriterOperation { database: Database::WordDocids, entry }) {
Ok(()) => Ok(()),
Err(SendError(_)) => Err(SendError(())),
}
}
}
#[derive(Clone)]
pub struct DocumentSender(Sender<WriterOperation>);
pub struct DocumentsSender<'a>(&'a Sender<WriterOperation>);
impl DocumentSender {
pub fn send(&self, document: DocumentEntry) -> StdResult<(), SendError<()>> {
match self.0.send(WriterOperation::Document(document)) {
impl DocumentsSender<'_> {
/// TODO do that efficiently
pub fn uncompressed(
&self,
docid: DocumentId,
document: &KvReaderFieldId,
) -> StdResult<(), SendError<()>> {
let entry = EntryOperation::Write(KeyValueEntry::from_key_value(
&docid.to_be_bytes(),
document.as_bytes(),
));
match self.0.send(WriterOperation { database: Database::Documents, entry }) {
Ok(()) => Ok(()),
Err(SendError(_)) => Err(SendError(())),
}
}
pub fn delete(&self, docid: DocumentId) -> StdResult<(), SendError<()>> {
let entry = EntryOperation::Delete(KeyEntry::from_key(&docid.to_be_bytes()));
match self.0.send(WriterOperation { database: Database::Documents, entry }) {
Ok(()) => Ok(()),
Err(SendError(_)) => Err(SendError(())),
}
@ -162,6 +240,8 @@ impl DocumentSender {
pub enum MergerOperation {
WordDocidsMerger(Merger<File, MergeDeladdCboRoaringBitmaps>),
InsertDocument { docid: DocumentId, document: Box<KvReaderFieldId> },
DeleteDocument { docid: DocumentId },
}
pub struct MergerReceiver(Receiver<MergerOperation>);
@ -190,3 +270,26 @@ impl DeladdCboRoaringBitmapSender {
}
}
}
#[derive(Clone)]
pub struct ExtractedDocumentsSender(Sender<MergerOperation>);
impl ExtractedDocumentsSender {
pub fn insert(
&self,
docid: DocumentId,
document: Box<KvReaderFieldId>,
) -> StdResult<(), SendError<()>> {
match self.0.send(MergerOperation::InsertDocument { docid, document }) {
Ok(()) => Ok(()),
Err(SendError(_)) => Err(SendError(())),
}
}
pub fn delete(&self, docid: DocumentId) -> StdResult<(), SendError<()>> {
match self.0.send(MergerOperation::DeleteDocument { docid }) {
Ok(()) => Ok(()),
Err(SendError(_)) => Err(SendError(())),
}
}
}

View File

@ -1,15 +1,15 @@
use std::fs::File;
use charabia::TokenizerBuilder;
use grenad::{Merger, ReaderCursor};
use grenad::Merger;
use heed::RoTxn;
use rayon::iter::{IntoParallelIterator, ParallelBridge, ParallelIterator};
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use super::cache::CachedSorter;
use super::tokenize_document::DocumentTokenizer;
use crate::update::new::{DocumentChange, ItemsPool};
use crate::update::{create_sorter, GrenadParameters, MergeDeladdCboRoaringBitmaps};
use crate::{FieldsIdsMap, GlobalFieldsIdsMap, Index, Result, MAX_POSITION_PER_ATTRIBUTE};
use crate::{GlobalFieldsIdsMap, Index, Result, MAX_POSITION_PER_ATTRIBUTE};
pub trait SearchableExtractor {
fn run_extraction(

View File

@ -1,13 +1,11 @@
use std::collections::HashMap;
use charabia::{SeparatorKind, Token, TokenKind, Tokenizer, TokenizerBuilder};
use heed::RoTxn;
use charabia::{SeparatorKind, Token, TokenKind, Tokenizer};
use serde_json::Value;
use crate::update::new::KvReaderFieldId;
use crate::{
FieldId, FieldsIdsMap, GlobalFieldsIdsMap, Index, InternalError, LocalizedAttributesRule,
Result, MAX_POSITION_PER_ATTRIBUTE, MAX_WORD_LENGTH,
FieldId, GlobalFieldsIdsMap, InternalError, LocalizedAttributesRule, Result, MAX_WORD_LENGTH,
};
pub struct DocumentTokenizer<'a> {
@ -239,6 +237,8 @@ mod test {
use serde_json::json;
use super::*;
use crate::FieldsIdsMap;
#[test]
fn test_tokenize_document() {
let mut fields_ids_map = FieldsIdsMap::new();

View File

@ -13,7 +13,6 @@ pub use update_by_function::UpdateByFunction;
use super::channel::{
extractors_merger_channels, merger_writer_channel, EntryOperation, ExtractorsMergerChannels,
WriterOperation,
};
use super::document_change::DocumentChange;
use super::extract::{SearchableExtractor, WordDocidsExtractor};
@ -57,8 +56,11 @@ where
PI::Iter: Clone,
{
let (merger_sender, writer_receiver) = merger_writer_channel(100);
let ExtractorsMergerChannels { merger_receiver, deladd_cbo_roaring_bitmap_sender } =
extractors_merger_channels(100);
let ExtractorsMergerChannels {
merger_receiver,
deladd_cbo_roaring_bitmap_sender,
extracted_documents_sender,
} = extractors_merger_channels(100);
let fields_ids_map_lock = RwLock::new(fields_ids_map);
let global_fields_ids_map = GlobalFieldsIdsMap::new(&fields_ids_map_lock);
@ -68,6 +70,28 @@ where
let handle = Builder::new().name(S("indexer-extractors")).spawn_scoped(s, move || {
pool.in_place_scope(|_s| {
let document_changes = document_changes.into_par_iter();
// document but we need to create a function that collects and compresses documents.
document_changes.clone().into_par_iter().try_for_each(|result| {
match result? {
DocumentChange::Deletion(deletion) => {
let docid = deletion.docid();
extracted_documents_sender.delete(docid).unwrap();
}
DocumentChange::Update(update) => {
let docid = update.docid();
let content = update.new();
extracted_documents_sender.insert(docid, content.boxed()).unwrap();
}
DocumentChange::Insertion(insertion) => {
let docid = insertion.docid();
let content = insertion.new();
extracted_documents_sender.insert(docid, content.boxed()).unwrap();
}
}
Ok(()) as Result<_>
})?;
// word docids
let merger = WordDocidsExtractor::run_extraction(
index,
@ -90,15 +114,15 @@ where
merge_grenad_entries(merger_receiver, merger_sender, &rtxn, index)
})?;
// TODO Split this code into another function
for operation in writer_receiver {
let database = operation.database(index);
match operation {
WriterOperation::WordDocids(operation) => match operation {
EntryOperation::Delete(e) => database.delete(wtxn, e.entry()).map(drop)?,
match operation.entry() {
EntryOperation::Delete(e) => {
if !database.delete(wtxn, e.entry())? {
unreachable!("We tried to delete an unknown key")
}
}
EntryOperation::Write(e) => database.put(wtxn, e.key(), e.value())?,
},
WriterOperation::Document(e) => database.put(wtxn, &e.key(), e.content())?,
}
}

View File

@ -1,9 +1,14 @@
use fst::set::OpBuilder;
use fst::{Set, SetBuilder};
use heed::types::Bytes;
use heed::RoTxn;
use memmap2::Mmap;
use roaring::RoaringBitmap;
use tempfile::tempfile;
use super::channel::{MergerReceiver, MergerSender};
use super::KvReaderDelAdd;
use crate::index::main_key::WORDS_FST_KEY;
use crate::update::del_add::DelAdd;
use crate::update::new::channel::MergerOperation;
use crate::{CboRoaringBitmapCodec, Index, Result};
@ -16,12 +21,15 @@ pub fn merge_grenad_entries(
index: &Index,
) -> Result<()> {
let mut buffer = Vec::new();
let mut documents_ids = index.documents_ids(rtxn)?;
for merger_operation in receiver {
match merger_operation {
MergerOperation::WordDocidsMerger(merger) => {
let sender = sender.word_docids();
let word_docids_sender = sender.word_docids();
let database = index.word_docids.remap_types::<Bytes, Bytes>();
let mut add_words_fst = SetBuilder::new(tempfile()?)?;
let mut del_words_fst = SetBuilder::new(tempfile()?)?;
/// TODO manage the error correctly
let mut merger_iter = merger.into_stream_merger_iter().unwrap();
@ -35,17 +43,62 @@ pub fn merge_grenad_entries(
match merge_cbo_bitmaps(current, del, add)? {
Operation::Write(bitmap) => {
let value = cbo_serialize_into_vec(&bitmap, &mut buffer);
sender.write(key, value).unwrap();
let value = cbo_bitmap_serialize_into_vec(&bitmap, &mut buffer);
word_docids_sender.write(key, value).unwrap();
add_words_fst.insert(key)?;
}
Operation::Delete => {
word_docids_sender.delete(key).unwrap();
del_words_fst.insert(key)?;
}
Operation::Delete => sender.delete(key).unwrap(),
Operation::Ignore => (),
}
}
// Move that into a dedicated function
let words_fst = index.words_fst(rtxn)?;
let add_words_fst_file = add_words_fst.into_inner()?;
let add_words_fst_mmap = unsafe { Mmap::map(&add_words_fst_file)? };
let add_words_fst = Set::new(&add_words_fst_mmap)?;
let del_words_fst_file = del_words_fst.into_inner()?;
let del_words_fst_mmap = unsafe { Mmap::map(&del_words_fst_file)? };
let del_words_fst = Set::new(&del_words_fst_mmap)?;
// TO BE IMPROVED @many
let diff = words_fst.op().add(&del_words_fst).difference();
let stream = add_words_fst.op().add(diff).union();
let mut words_fst = SetBuilder::new(tempfile()?)?;
words_fst.extend_stream(stream)?;
let words_fst_file = words_fst.into_inner()?;
let words_fst_mmap = unsafe { Mmap::map(&words_fst_file)? };
// PLEASE SEND THIS AS AN MMAP
let main_sender = sender.main();
main_sender.write_words_fst(&words_fst_mmap).unwrap();
}
MergerOperation::InsertDocument { docid, document } => {
documents_ids.insert(docid);
sender.documents().uncompressed(docid, &document).unwrap();
}
MergerOperation::DeleteDocument { docid } => {
if !documents_ids.remove(docid) {
unreachable!("Tried deleting a document that we do not know about");
}
sender.documents().delete(docid).unwrap();
}
}
}
// Send the documents ids unionized with the current one
/// TODO return the slice of bytes directly
serialize_bitmap_into_vec(&documents_ids, &mut buffer);
sender.send_documents_ids(&buffer).unwrap();
// ...
Ok(())
}
@ -86,9 +139,16 @@ fn merge_cbo_bitmaps(
}
}
/// Return the slice directly from the serialize_into method
fn cbo_serialize_into_vec<'b>(bitmap: &RoaringBitmap, buffer: &'b mut Vec<u8>) -> &'b [u8] {
/// TODO Return the slice directly from the serialize_into method
fn cbo_bitmap_serialize_into_vec<'b>(bitmap: &RoaringBitmap, buffer: &'b mut Vec<u8>) -> &'b [u8] {
buffer.clear();
CboRoaringBitmapCodec::serialize_into(bitmap, buffer);
buffer.as_slice()
}
/// TODO Return the slice directly from the serialize_into method
fn serialize_bitmap_into_vec<'b>(bitmap: &RoaringBitmap, buffer: &'b mut Vec<u8>) {
buffer.clear();
bitmap.serialize_into(buffer).unwrap();
// buffer.as_slice()
}