Change the project to become a workspace with milli as a default-member

milli/src/update/available_documents_ids.rs

@@ -0,0 +1,67 @@
+use std::iter::{Chain, FromIterator};
+use std::ops::RangeInclusive;
+use roaring::bitmap::{RoaringBitmap, IntoIter};
+
+pub struct AvailableDocumentsIds {
+    iter: Chain<IntoIter, RangeInclusive<u32>>,
+}
+
+impl AvailableDocumentsIds {
+    pub fn from_documents_ids(docids: &RoaringBitmap) -> AvailableDocumentsIds {
+        match docids.max() {
+            Some(last_id) => {
+                let mut available = RoaringBitmap::from_iter(0..last_id);
+                available.difference_with(&docids);
+
+                let iter = match last_id.checked_add(1) {
+                    Some(id) => id..=u32::max_value(),
+                    None => 1..=0, // empty range iterator
+                };
+
+                AvailableDocumentsIds {
+                    iter: available.into_iter().chain(iter),
+                }
+            },
+            None => {
+                let empty = RoaringBitmap::new().into_iter();
+                AvailableDocumentsIds {
+                    iter: empty.chain(0..=u32::max_value()),
+                }
+            },
+        }
+    }
+}
+
+impl Iterator for AvailableDocumentsIds {
+    type Item = u32;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter.next()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn empty() {
+        let base = RoaringBitmap::new();
+        let left = AvailableDocumentsIds::from_documents_ids(&base);
+        let right = 0..=u32::max_value();
+        left.zip(right).take(500).for_each(|(l, r)| assert_eq!(l, r));
+    }
+
+    #[test]
+    fn scattered() {
+        let mut base = RoaringBitmap::new();
+        base.insert(0);
+        base.insert(10);
+        base.insert(100);
+        base.insert(405);
+
+        let left = AvailableDocumentsIds::from_documents_ids(&base);
+        let right = (0..=u32::max_value()).filter(|&n| n != 0 && n != 10 && n != 100 && n != 405);
+        left.zip(right).take(500).for_each(|(l, r)| assert_eq!(l, r));
+    }
+}

milli/src/update/clear_documents.rs

@@ -0,0 +1,56 @@
+use roaring::RoaringBitmap;
+use crate::{ExternalDocumentsIds, Index};
+
+pub struct ClearDocuments<'t, 'u, 'i> {
+    wtxn: &'t mut heed::RwTxn<'i, 'u>,
+    index: &'i Index,
+    _update_id: u64,
+}
+
+impl<'t, 'u, 'i> ClearDocuments<'t, 'u, 'i> {
+    pub fn new(
+        wtxn: &'t mut heed::RwTxn<'i, 'u>,
+        index: &'i Index,
+        update_id: u64
+    ) -> ClearDocuments<'t, 'u, 'i> {
+
+        ClearDocuments { wtxn, index, _update_id: update_id }
+    }
+
+    pub fn execute(self) -> anyhow::Result<usize> {
+        let Index {
+            env: _env,
+            main: _main,
+            word_docids,
+            docid_word_positions,
+            word_pair_proximity_docids,
+            facet_field_id_value_docids,
+            field_id_docid_facet_values,
+            documents,
+        } = self.index;
+
+        // We retrieve the number of documents ids that we are deleting.
+        let number_of_documents = self.index.number_of_documents(self.wtxn)?;
+        let faceted_fields = self.index.faceted_fields_ids(self.wtxn)?;
+
+        // We clean some of the main engine datastructures.
+        self.index.put_words_fst(self.wtxn, &fst::Set::default())?;
+        self.index.put_external_documents_ids(self.wtxn, &ExternalDocumentsIds::default())?;
+        self.index.put_documents_ids(self.wtxn, &RoaringBitmap::default())?;
+
+        // We clean all the faceted documents ids.
+        for (field_id, _) in faceted_fields {
+            self.index.put_faceted_documents_ids(self.wtxn, field_id, &RoaringBitmap::default())?;
+        }
+
+        // Clear the other databases.
+        word_docids.clear(self.wtxn)?;
+        docid_word_positions.clear(self.wtxn)?;
+        word_pair_proximity_docids.clear(self.wtxn)?;
+        facet_field_id_value_docids.clear(self.wtxn)?;
+        field_id_docid_facet_values.clear(self.wtxn)?;
+        documents.clear(self.wtxn)?;
+
+        Ok(number_of_documents)
+    }
+}

milli/src/update/delete_documents.rs

@@ -0,0 +1,250 @@
+use fst::IntoStreamer;
+use heed::types::ByteSlice;
+use roaring::RoaringBitmap;
+
+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<DeleteDocuments<'t, 'u, 'i>>
+    {
+        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<u32> {
+        let docid = self.external_documents_ids.get(external_id)?;
+        self.delete_document(docid);
+        Some(docid)
+    }
+
+    pub fn execute(self) -> anyhow::Result<usize> {
+        // 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,
+            docid_word_positions,
+            word_pair_proximity_docids,
+            facet_field_id_value_docids,
+            field_id_docid_facet_values,
+            documents,
+        } = self.index;
+
+        // 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()? {
+                if let Some(content) = obkv.get(id_field) {
+                    let external_id: SmallString32 = serde_json::from_slice(content).unwrap();
+                    external_ids.push(external_id);
+                }
+                iter.del_current()?;
+            }
+            drop(iter);
+
+            // We iterate througt 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()?;
+            }
+        }
+
+        // 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 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::<ByteSlice>().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::<FieldDocIdFacetStringCodec>();
+                    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::<FieldDocIdFacetF64Codec>();
+                    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::<FieldDocIdFacetI64Codec>();
+                    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() as usize)
+    }
+}

milli/src/update/facets.rs

@@ -0,0 +1,262 @@
+use std::cmp;
+use std::fs::File;
+use std::num::NonZeroUsize;
+
+use grenad::{CompressionType, Reader, Writer, FileFuse};
+use heed::types::{ByteSlice, DecodeIgnore};
+use heed::{BytesEncode, Error};
+use log::debug;
+use num_traits::{Bounded, Zero};
+use roaring::RoaringBitmap;
+
+use crate::facet::FacetType;
+use crate::heed_codec::CboRoaringBitmapCodec;
+use crate::heed_codec::facet::{FacetLevelValueI64Codec, FacetLevelValueF64Codec};
+use crate::Index;
+use crate::update::index_documents::WriteMethod;
+use crate::update::index_documents::{create_writer, writer_into_reader, write_into_lmdb_database};
+
+pub struct Facets<'t, 'u, 'i> {
+    wtxn: &'t mut heed::RwTxn<'i, 'u>,
+    index: &'i Index,
+    pub(crate) chunk_compression_type: CompressionType,
+    pub(crate) chunk_compression_level: Option<u32>,
+    pub(crate) chunk_fusing_shrink_size: Option<u64>,
+    level_group_size: NonZeroUsize,
+    min_level_size: NonZeroUsize,
+    _update_id: u64,
+}
+
+impl<'t, 'u, 'i> Facets<'t, 'u, 'i> {
+    pub fn new(
+        wtxn: &'t mut heed::RwTxn<'i, 'u>,
+        index: &'i Index,
+        update_id: u64,
+    ) -> Facets<'t, 'u, 'i> {
+        Facets {
+            wtxn,
+            index,
+            chunk_compression_type: CompressionType::None,
+            chunk_compression_level: None,
+            chunk_fusing_shrink_size: None,
+            level_group_size: NonZeroUsize::new(4).unwrap(),
+            min_level_size: NonZeroUsize::new(5).unwrap(),
+            _update_id: update_id,
+        }
+    }
+
+    pub fn level_group_size(&mut self, value: NonZeroUsize) -> &mut Self {
+        self.level_group_size = NonZeroUsize::new(cmp::max(value.get(), 2)).unwrap();
+        self
+    }
+
+    pub fn min_level_size(&mut self, value: NonZeroUsize) -> &mut Self {
+        self.min_level_size = value;
+        self
+    }
+
+    pub fn execute(self) -> anyhow::Result<()> {
+        // We get the faceted fields to be able to create the facet levels.
+        let faceted_fields = self.index.faceted_fields_ids(self.wtxn)?;
+
+        debug!("Computing and writing the facet values levels docids into LMDB on disk...");
+        for (field_id, facet_type) in faceted_fields {
+            let (content, documents_ids) = match facet_type {
+                FacetType::Integer => {
+                    clear_field_levels::<i64, FacetLevelValueI64Codec>(
+                        self.wtxn,
+                        self.index.facet_field_id_value_docids,
+                        field_id,
+                    )?;
+
+                    let documents_ids = compute_faceted_documents_ids(
+                        self.wtxn,
+                        self.index.facet_field_id_value_docids,
+                        field_id,
+                    )?;
+
+                    let content = compute_facet_levels::<i64, FacetLevelValueI64Codec>(
+                        self.wtxn,
+                        self.index.facet_field_id_value_docids,
+                        self.chunk_compression_type,
+                        self.chunk_compression_level,
+                        self.chunk_fusing_shrink_size,
+                        self.level_group_size,
+                        self.min_level_size,
+                        field_id,
+                    )?;
+
+                    (Some(content), documents_ids)
+                },
+                FacetType::Float => {
+                    clear_field_levels::<f64, FacetLevelValueF64Codec>(
+                        self.wtxn,
+                        self.index.facet_field_id_value_docids,
+                        field_id,
+                    )?;
+
+                    let documents_ids = compute_faceted_documents_ids(
+                        self.wtxn,
+                        self.index.facet_field_id_value_docids,
+                        field_id,
+                    )?;
+
+                    let content = compute_facet_levels::<f64, FacetLevelValueF64Codec>(
+                        self.wtxn,
+                        self.index.facet_field_id_value_docids,
+                        self.chunk_compression_type,
+                        self.chunk_compression_level,
+                        self.chunk_fusing_shrink_size,
+                        self.level_group_size,
+                        self.min_level_size,
+                        field_id,
+                    )?;
+
+                    (Some(content), documents_ids)
+                },
+                FacetType::String => {
+                    let documents_ids = compute_faceted_documents_ids(
+                        self.wtxn,
+                        self.index.facet_field_id_value_docids,
+                        field_id,
+                    )?;
+
+                    (None, documents_ids)
+                },
+            };
+
+            if let Some(content) = content {
+                write_into_lmdb_database(
+                    self.wtxn,
+                    *self.index.facet_field_id_value_docids.as_polymorph(),
+                    content,
+                    |_, _| anyhow::bail!("invalid facet level merging"),
+                    WriteMethod::GetMergePut,
+                )?;
+            }
+
+            self.index.put_faceted_documents_ids(self.wtxn, field_id, &documents_ids)?;
+        }
+
+        Ok(())
+    }
+}
+
+fn clear_field_levels<'t, T: 't, KC>(
+    wtxn: &'t mut heed::RwTxn,
+    db: heed::Database<ByteSlice, CboRoaringBitmapCodec>,
+    field_id: u8,
+) -> heed::Result<()>
+where
+    T: Copy + Bounded,
+    KC: heed::BytesDecode<'t, DItem = (u8, u8, T, T)>,
+    KC: for<'x> heed::BytesEncode<'x, EItem = (u8, u8, T, T)>,
+{
+    let left = (field_id, 1, T::min_value(), T::min_value());
+    let right = (field_id, u8::MAX, T::max_value(), T::max_value());
+    let range = left..=right;
+    db.remap_key_type::<KC>().delete_range(wtxn, &range).map(drop)
+}
+
+fn compute_facet_levels<'t, T: 't, KC>(
+    rtxn: &'t heed::RoTxn,
+    db: heed::Database<ByteSlice, CboRoaringBitmapCodec>,
+    compression_type: CompressionType,
+    compression_level: Option<u32>,
+    shrink_size: Option<u64>,
+    level_group_size: NonZeroUsize,
+    min_level_size: NonZeroUsize,
+    field_id: u8,
+) -> anyhow::Result<Reader<FileFuse>>
+where
+    T: Copy + PartialEq + PartialOrd + Bounded + Zero,
+    KC: heed::BytesDecode<'t, DItem = (u8, u8, T, T)>,
+    KC: for<'x> heed::BytesEncode<'x, EItem = (u8, u8, T, T)>,
+{
+    let first_level_size = db.prefix_iter(rtxn, &[field_id])?
+        .remap_types::<DecodeIgnore, DecodeIgnore>()
+        .fold(Ok(0usize), |count, result| result.and(count).map(|c| c + 1))?;
+
+    // It is forbidden to keep a cursor and write in a database at the same time with LMDB
+    // therefore we write the facet levels entries into a grenad file before transfering them.
+    let mut writer = tempfile::tempfile().and_then(|file| {
+        create_writer(compression_type, compression_level, file)
+    })?;
+
+    let level_0_range = {
+        let left = (field_id, 0, T::min_value(), T::min_value());
+        let right = (field_id, 0, T::max_value(), T::max_value());
+        left..=right
+    };
+
+    // Groups sizes are always a power of the original level_group_size and therefore a group
+    // always maps groups of the previous level and never splits previous levels groups in half.
+    let group_size_iter = (1u8..)
+        .map(|l| (l, level_group_size.get().pow(l as u32)))
+        .take_while(|(_, s)| first_level_size / *s >= min_level_size.get());
+
+    for (level, group_size) in group_size_iter {
+        let mut left = T::zero();
+        let mut right = T::zero();
+        let mut group_docids = RoaringBitmap::new();
+
+        let db = db.remap_key_type::<KC>();
+        for (i, result) in db.range(rtxn, &level_0_range)?.enumerate() {
+            let ((_field_id, _level, value, _right), docids) = result?;
+
+            if i == 0 {
+                left = value;
+            } else if i % group_size == 0 {
+                // we found the first bound of the next group, we must store the left
+                // and right bounds associated with the docids.
+                write_entry::<T, KC>(&mut writer, field_id, level, left, right, &group_docids)?;
+
+                // We save the left bound for the new group and also reset the docids.
+                group_docids = RoaringBitmap::new();
+                left = value;
+            }
+
+            // The right bound is always the bound we run through.
+            group_docids.union_with(&docids);
+            right = value;
+        }
+
+        if !group_docids.is_empty() {
+            write_entry::<T, KC>(&mut writer, field_id, level, left, right, &group_docids)?;
+        }
+    }
+
+    writer_into_reader(writer, shrink_size)
+}
+
+fn compute_faceted_documents_ids(
+    rtxn: &heed::RoTxn,
+    db: heed::Database<ByteSlice, CboRoaringBitmapCodec>,
+    field_id: u8,
+) -> anyhow::Result<RoaringBitmap>
+{
+    let mut documents_ids = RoaringBitmap::new();
+    for result in db.prefix_iter(rtxn, &[field_id])? {
+        let (_key, docids) = result?;
+        documents_ids.union_with(&docids);
+    }
+    Ok(documents_ids)
+}
+
+fn write_entry<T, KC>(
+    writer: &mut Writer<File>,
+    field_id: u8,
+    level: u8,
+    left: T,
+    right: T,
+    ids: &RoaringBitmap,
+) -> anyhow::Result<()>
+where
+    KC: for<'x> heed::BytesEncode<'x, EItem = (u8, u8, T, T)>,
+{
+    let key = (field_id, level, left, right);
+    let key = KC::bytes_encode(&key).ok_or(Error::Encoding)?;
+    let data = CboRoaringBitmapCodec::bytes_encode(&ids).ok_or(Error::Encoding)?;
+    writer.insert(&key, &data)?;
+    Ok(())
+}

milli/src/update/index_documents/merge_function.rs

@@ -0,0 +1,105 @@
+use std::borrow::Cow;
+
+use anyhow::{bail, ensure, Context};
+use bstr::ByteSlice as _;
+use fst::IntoStreamer;
+use roaring::RoaringBitmap;
+
+use crate::heed_codec::CboRoaringBitmapCodec;
+
+const WORDS_FST_KEY: &[u8] = crate::index::WORDS_FST_KEY.as_bytes();
+const FIELDS_IDS_MAP_KEY: &[u8] = crate::index::FIELDS_IDS_MAP_KEY.as_bytes();
+const DOCUMENTS_IDS_KEY: &[u8] = crate::index::DOCUMENTS_IDS_KEY.as_bytes();
+
+pub fn main_merge(key: &[u8], values: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    match key {
+        WORDS_FST_KEY => {
+            let fsts: Vec<_> = values.iter().map(|v| fst::Set::new(v).unwrap()).collect();
+
+            // Union of the FSTs
+            let mut op = fst::set::OpBuilder::new();
+            fsts.iter().for_each(|fst| op.push(fst.into_stream()));
+            let op = op.r#union();
+
+            let mut build = fst::SetBuilder::memory();
+            build.extend_stream(op.into_stream()).unwrap();
+            Ok(build.into_inner().unwrap())
+        },
+        FIELDS_IDS_MAP_KEY => {
+            ensure!(values.windows(2).all(|vs| vs[0] == vs[1]), "fields ids map doesn't match");
+            Ok(values[0].to_vec())
+        },
+        DOCUMENTS_IDS_KEY => roaring_bitmap_merge(values),
+        otherwise => bail!("wut {:?}", otherwise),
+    }
+}
+
+pub fn word_docids_merge(_key: &[u8], values: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    roaring_bitmap_merge(values)
+}
+
+pub fn docid_word_positions_merge(key: &[u8], _values: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    bail!("merging docid word positions is an error ({:?})", key.as_bstr())
+}
+
+pub fn field_id_docid_facet_values_merge(_key: &[u8], values: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    let first = values.first().context("no value to merge")?;
+    ensure!(values.iter().all(|v| v == first), "invalid field id docid facet value merging");
+    Ok(first.to_vec())
+}
+
+pub fn words_pairs_proximities_docids_merge(_key: &[u8], values: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    cbo_roaring_bitmap_merge(values)
+}
+
+pub fn facet_field_value_docids_merge(_key: &[u8], values: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    cbo_roaring_bitmap_merge(values)
+}
+
+pub fn documents_merge(key: &[u8], _values: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    bail!("merging documents is an error ({:?})", key.as_bstr())
+}
+
+pub fn merge_two_obkvs(base: obkv::KvReader, update: obkv::KvReader, buffer: &mut Vec<u8>) {
+    use itertools::merge_join_by;
+    use itertools::EitherOrBoth::{Both, Left, Right};
+
+    buffer.clear();
+
+    let mut writer = obkv::KvWriter::new(buffer);
+    for eob in merge_join_by(base.iter(), update.iter(), |(b, _), (u, _)| b.cmp(u)) {
+        match eob {
+            Both(_, (k, v)) | Left((k, v)) | Right((k, v)) => writer.insert(k, v).unwrap(),
+        }
+    }
+
+    writer.finish().unwrap();
+}
+
+fn roaring_bitmap_merge(values: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    let (head, tail) = values.split_first().unwrap();
+    let mut head = RoaringBitmap::deserialize_from(&head[..])?;
+
+    for value in tail {
+        let bitmap = RoaringBitmap::deserialize_from(&value[..])?;
+        head.union_with(&bitmap);
+    }
+
+    let mut vec = Vec::with_capacity(head.serialized_size());
+    head.serialize_into(&mut vec)?;
+    Ok(vec)
+}
+
+fn cbo_roaring_bitmap_merge(values: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    let (head, tail) = values.split_first().unwrap();
+    let mut head = CboRoaringBitmapCodec::deserialize_from(&head[..])?;
+
+    for value in tail {
+        let bitmap = CboRoaringBitmapCodec::deserialize_from(&value[..])?;
+        head.union_with(&bitmap);
+    }
+
+    let mut vec = Vec::new();
+    CboRoaringBitmapCodec::serialize_into(&head, &mut vec)?;
+    Ok(vec)
+}

milli/src/update/index_documents/store.rs

@@ -0,0 +1,690 @@
+use std::borrow::Cow;
+use std::collections::{BTreeMap, HashMap, HashSet};
+use std::convert::{TryFrom, TryInto};
+use std::fs::File;
+use std::iter::FromIterator;
+use std::time::Instant;
+use std::{cmp, iter};
+
+use anyhow::{bail, Context};
+use bstr::ByteSlice as _;
+use fst::Set;
+use grenad::{Reader, FileFuse, Writer, Sorter, CompressionType};
+use heed::BytesEncode;
+use linked_hash_map::LinkedHashMap;
+use log::{debug, info};
+use meilisearch_tokenizer::{Analyzer, AnalyzerConfig};
+use ordered_float::OrderedFloat;
+use roaring::RoaringBitmap;
+use serde_json::Value;
+use tempfile::tempfile;
+
+use crate::facet::FacetType;
+use crate::heed_codec::facet::{FacetValueStringCodec, FacetLevelValueF64Codec, FacetLevelValueI64Codec};
+use crate::heed_codec::facet::{FieldDocIdFacetStringCodec, FieldDocIdFacetF64Codec, FieldDocIdFacetI64Codec};
+use crate::heed_codec::{BoRoaringBitmapCodec, CboRoaringBitmapCodec};
+use crate::update::UpdateIndexingStep;
+use crate::{json_to_string, SmallVec8, SmallVec32, SmallString32, Position, DocumentId, FieldId};
+
+use super::{MergeFn, create_writer, create_sorter, writer_into_reader};
+use super::merge_function::{
+    main_merge, word_docids_merge, words_pairs_proximities_docids_merge,
+    facet_field_value_docids_merge, field_id_docid_facet_values_merge,
+};
+
+const LMDB_MAX_KEY_LENGTH: usize = 511;
+const ONE_KILOBYTE: usize = 1024 * 1024;
+
+const MAX_POSITION: usize = 1000;
+const WORDS_FST_KEY: &[u8] = crate::index::WORDS_FST_KEY.as_bytes();
+
+pub struct Readers {
+    pub main: Reader<FileFuse>,
+    pub word_docids: Reader<FileFuse>,
+    pub docid_word_positions: Reader<FileFuse>,
+    pub words_pairs_proximities_docids: Reader<FileFuse>,
+    pub facet_field_value_docids: Reader<FileFuse>,
+    pub field_id_docid_facet_values: Reader<FileFuse>,
+    pub documents: Reader<FileFuse>,
+}
+
+pub struct Store<'s, A> {
+    // Indexing parameters
+    searchable_fields: HashSet<FieldId>,
+    faceted_fields: HashMap<FieldId, FacetType>,
+    // Caches
+    word_docids: LinkedHashMap<SmallVec32<u8>, RoaringBitmap>,
+    word_docids_limit: usize,
+    words_pairs_proximities_docids: LinkedHashMap<(SmallVec32<u8>, SmallVec32<u8>, u8), RoaringBitmap>,
+    words_pairs_proximities_docids_limit: usize,
+    facet_field_value_docids: LinkedHashMap<(u8, FacetValue), RoaringBitmap>,
+    facet_field_value_docids_limit: usize,
+    // MTBL parameters
+    chunk_compression_type: CompressionType,
+    chunk_compression_level: Option<u32>,
+    chunk_fusing_shrink_size: Option<u64>,
+    // MTBL sorters
+    main_sorter: Sorter<MergeFn>,
+    word_docids_sorter: Sorter<MergeFn>,
+    words_pairs_proximities_docids_sorter: Sorter<MergeFn>,
+    facet_field_value_docids_sorter: Sorter<MergeFn>,
+    field_id_docid_facet_values_sorter: Sorter<MergeFn>,
+    // MTBL writers
+    docid_word_positions_writer: Writer<File>,
+    documents_writer: Writer<File>,
+    // tokenizer
+    analyzer: Analyzer<'s, A>,
+}
+
+impl<'s, A: AsRef<[u8]>> Store<'s, A> {
+    pub fn new(
+        searchable_fields: HashSet<FieldId>,
+        faceted_fields: HashMap<FieldId, FacetType>,
+        linked_hash_map_size: Option<usize>,
+        max_nb_chunks: Option<usize>,
+        max_memory: Option<usize>,
+        chunk_compression_type: CompressionType,
+        chunk_compression_level: Option<u32>,
+        chunk_fusing_shrink_size: Option<u64>,
+        stop_words: &'s Set<A>,
+    ) -> anyhow::Result<Self>
+    {
+        // We divide the max memory by the number of sorter the Store have.
+        let max_memory = max_memory.map(|mm| cmp::max(ONE_KILOBYTE, mm / 4));
+        let linked_hash_map_size = linked_hash_map_size.unwrap_or(500);
+
+        let main_sorter = create_sorter(
+            main_merge,
+            chunk_compression_type,
+            chunk_compression_level,
+            chunk_fusing_shrink_size,
+            max_nb_chunks,
+            max_memory,
+        );
+        let word_docids_sorter = create_sorter(
+            word_docids_merge,
+            chunk_compression_type,
+            chunk_compression_level,
+            chunk_fusing_shrink_size,
+            max_nb_chunks,
+            max_memory,
+        );
+        let words_pairs_proximities_docids_sorter = create_sorter(
+            words_pairs_proximities_docids_merge,
+            chunk_compression_type,
+            chunk_compression_level,
+            chunk_fusing_shrink_size,
+            max_nb_chunks,
+            max_memory,
+        );
+        let facet_field_value_docids_sorter = create_sorter(
+            facet_field_value_docids_merge,
+            chunk_compression_type,
+            chunk_compression_level,
+            chunk_fusing_shrink_size,
+            max_nb_chunks,
+            max_memory,
+        );
+        let field_id_docid_facet_values_sorter = create_sorter(
+            field_id_docid_facet_values_merge,
+            chunk_compression_type,
+            chunk_compression_level,
+            chunk_fusing_shrink_size,
+            max_nb_chunks,
+            Some(1024 * 1024 * 1024), // 1MB
+        );
+
+        let documents_writer = tempfile().and_then(|f| {
+            create_writer(chunk_compression_type, chunk_compression_level, f)
+        })?;
+        let docid_word_positions_writer = tempfile().and_then(|f| {
+            create_writer(chunk_compression_type, chunk_compression_level, f)
+        })?;
+
+        let analyzer = Analyzer::new(AnalyzerConfig::default_with_stopwords(stop_words));
+
+        Ok(Store {
+            // Indexing parameters.
+            searchable_fields,
+            faceted_fields,
+            // Caches
+            word_docids: LinkedHashMap::with_capacity(linked_hash_map_size),
+            word_docids_limit: linked_hash_map_size,
+            words_pairs_proximities_docids: LinkedHashMap::with_capacity(linked_hash_map_size),
+            words_pairs_proximities_docids_limit: linked_hash_map_size,
+            facet_field_value_docids: LinkedHashMap::with_capacity(linked_hash_map_size),
+            facet_field_value_docids_limit: linked_hash_map_size,
+            // MTBL parameters
+            chunk_compression_type,
+            chunk_compression_level,
+            chunk_fusing_shrink_size,
+            // MTBL sorters
+            main_sorter,
+            word_docids_sorter,
+            words_pairs_proximities_docids_sorter,
+            facet_field_value_docids_sorter,
+            field_id_docid_facet_values_sorter,
+            // MTBL writers
+            docid_word_positions_writer,
+            documents_writer,
+            // tokenizer
+            analyzer,
+        })
+    }
+
+    // Save the documents ids under the position and word we have seen it.
+    fn insert_word_docid(&mut self, word: &str, id: DocumentId) -> anyhow::Result<()> {
+        // if get_refresh finds the element it is assured to be at the end of the linked hash map.
+        match self.word_docids.get_refresh(word.as_bytes()) {
+            Some(old) => { old.insert(id); },
+            None => {
+                let word_vec = SmallVec32::from(word.as_bytes());
+                // A newly inserted element is append at the end of the linked hash map.
+                self.word_docids.insert(word_vec, RoaringBitmap::from_iter(Some(id)));
+                // If the word docids just reached it's capacity we must make sure to remove
+                // one element, this way next time we insert we doesn't grow the capacity.
+                if self.word_docids.len() == self.word_docids_limit {
+                    // Removing the front element is equivalent to removing the LRU element.
+                    let lru = self.word_docids.pop_front();
+                    Self::write_word_docids(&mut self.word_docids_sorter, lru)?;
+                }
+            }
+        }
+        Ok(())
+    }
+
+    // Save the documents ids under the facet field id and value we have seen it.
+    fn insert_facet_values_docid(
+        &mut self,
+        field_id: FieldId,
+        field_value: FacetValue,
+        id: DocumentId,
+    ) -> anyhow::Result<()>
+    {
+        Self::write_field_id_docid_facet_value(&mut self.field_id_docid_facet_values_sorter, field_id, id, &field_value)?;
+
+        let key = (field_id, field_value);
+        // if get_refresh finds the element it is assured to be at the end of the linked hash map.
+        match self.facet_field_value_docids.get_refresh(&key) {
+            Some(old) => { old.insert(id); },
+            None => {
+                // A newly inserted element is append at the end of the linked hash map.
+                self.facet_field_value_docids.insert(key, RoaringBitmap::from_iter(Some(id)));
+                // If the word docids just reached it's capacity we must make sure to remove
+                // one element, this way next time we insert we doesn't grow the capacity.
+                if self.facet_field_value_docids.len() == self.facet_field_value_docids_limit {
+                    // Removing the front element is equivalent to removing the LRU element.
+                    Self::write_facet_field_value_docids(
+                        &mut self.facet_field_value_docids_sorter,
+                        self.facet_field_value_docids.pop_front(),
+                    )?;
+                }
+            }
+        }
+        Ok(())
+    }
+
+    // Save the documents ids under the words pairs proximities that it contains.
+    fn insert_words_pairs_proximities_docids<'a>(
+        &mut self,
+        words_pairs_proximities: impl IntoIterator<Item=((&'a str, &'a str), u8)>,
+        id: DocumentId,
+    ) -> anyhow::Result<()>
+    {
+        for ((w1, w2), prox) in words_pairs_proximities {
+            let w1 = SmallVec32::from(w1.as_bytes());
+            let w2 = SmallVec32::from(w2.as_bytes());
+            let key = (w1, w2, prox);
+            // if get_refresh finds the element it is assured
+            // to be at the end of the linked hash map.
+            match self.words_pairs_proximities_docids.get_refresh(&key) {
+                Some(old) => { old.insert(id); },
+                None => {
+                    // A newly inserted element is append at the end of the linked hash map.
+                    let ids = RoaringBitmap::from_iter(Some(id));
+                    self.words_pairs_proximities_docids.insert(key, ids);
+                }
+            }
+        }
+
+        // If the linked hashmap is over capacity we must remove the overflowing elements.
+        let len = self.words_pairs_proximities_docids.len();
+        let overflow = len.checked_sub(self.words_pairs_proximities_docids_limit);
+        if let Some(overflow) = overflow {
+            let mut lrus = Vec::with_capacity(overflow);
+            // Removing front elements is equivalent to removing the LRUs.
+            let iter = iter::from_fn(|| self.words_pairs_proximities_docids.pop_front());
+            iter.take(overflow).for_each(|x| lrus.push(x));
+            Self::write_words_pairs_proximities(&mut self.words_pairs_proximities_docids_sorter, lrus)?;
+        }
+
+        Ok(())
+    }
+
+    fn write_document(
+        &mut self,
+        document_id: DocumentId,
+        words_positions: &mut HashMap<String, SmallVec32<Position>>,
+        facet_values: &mut HashMap<FieldId, SmallVec8<FacetValue>>,
+        record: &[u8],
+    ) -> anyhow::Result<()>
+    {
+        // We compute the list of words pairs proximities (self-join) and write it directly to disk.
+        let words_pair_proximities = compute_words_pair_proximities(&words_positions);
+        self.insert_words_pairs_proximities_docids(words_pair_proximities, document_id)?;
+
+        // We store document_id associated with all the words the record contains.
+        for (word, _) in words_positions.drain() {
+            self.insert_word_docid(&word, document_id)?;
+        }
+
+        self.documents_writer.insert(document_id.to_be_bytes(), record)?;
+        Self::write_docid_word_positions(&mut self.docid_word_positions_writer, document_id, words_positions)?;
+
+        // We store document_id associated with all the field id and values.
+        for (field, values) in facet_values.drain() {
+            for value in values {
+                self.insert_facet_values_docid(field, value, document_id)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn write_words_pairs_proximities(
+        sorter: &mut Sorter<MergeFn>,
+        iter: impl IntoIterator<Item=((SmallVec32<u8>, SmallVec32<u8>, u8), RoaringBitmap)>,
+    ) -> anyhow::Result<()>
+    {
+        let mut key = Vec::new();
+        let mut buffer = Vec::new();
+
+        for ((w1, w2, min_prox), docids) in iter {
+            key.clear();
+            key.extend_from_slice(w1.as_bytes());
+            key.push(0);
+            key.extend_from_slice(w2.as_bytes());
+            // Storing the minimun proximity found between those words
+            key.push(min_prox);
+            // We serialize the document ids into a buffer
+            buffer.clear();
+            buffer.reserve(CboRoaringBitmapCodec::serialized_size(&docids));
+            CboRoaringBitmapCodec::serialize_into(&docids, &mut buffer)?;
+            // that we write under the generated key into MTBL
+            if lmdb_key_valid_size(&key) {
+                sorter.insert(&key, &buffer)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn write_docid_word_positions(
+        writer: &mut Writer<File>,
+        id: DocumentId,
+        words_positions: &HashMap<String, SmallVec32<Position>>,
+    ) -> anyhow::Result<()>
+    {
+        // We prefix the words by the document id.
+        let mut key = id.to_be_bytes().to_vec();
+        let base_size = key.len();
+
+        // We order the words lexicographically, this way we avoid passing by a sorter.
+        let words_positions = BTreeMap::from_iter(words_positions);
+
+        for (word, positions) in words_positions {
+            key.truncate(base_size);
+            key.extend_from_slice(word.as_bytes());
+            // We serialize the positions into a buffer.
+            let positions = RoaringBitmap::from_iter(positions.iter().cloned());
+            let bytes = BoRoaringBitmapCodec::bytes_encode(&positions)
+                .with_context(|| "could not serialize positions")?;
+            // that we write under the generated key into MTBL
+            if lmdb_key_valid_size(&key) {
+                writer.insert(&key, &bytes)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn write_facet_field_value_docids<I>(
+        sorter: &mut Sorter<MergeFn>,
+        iter: I,
+    ) -> anyhow::Result<()>
+    where I: IntoIterator<Item=((FieldId, FacetValue), RoaringBitmap)>
+    {
+        use FacetValue::*;
+
+        for ((field_id, value), docids) in iter {
+            let result = match value {
+                String(s) => FacetValueStringCodec::bytes_encode(&(field_id, &s)).map(Cow::into_owned),
+                Float(f) => FacetLevelValueF64Codec::bytes_encode(&(field_id, 0, *f, *f)).map(Cow::into_owned),
+                Integer(i) => FacetLevelValueI64Codec::bytes_encode(&(field_id, 0, i, i)).map(Cow::into_owned),
+            };
+            let key = result.context("could not serialize facet key")?;
+            let bytes = CboRoaringBitmapCodec::bytes_encode(&docids)
+                .context("could not serialize docids")?;
+            if lmdb_key_valid_size(&key) {
+                sorter.insert(&key, &bytes)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn write_field_id_docid_facet_value(
+        sorter: &mut Sorter<MergeFn>,
+        field_id: FieldId,
+        document_id: DocumentId,
+        value: &FacetValue,
+    ) -> anyhow::Result<()>
+    {
+        use FacetValue::*;
+
+        let result = match value {
+            String(s) => FieldDocIdFacetStringCodec::bytes_encode(&(field_id, document_id, s)).map(Cow::into_owned),
+            Float(f) => FieldDocIdFacetF64Codec::bytes_encode(&(field_id, document_id, **f)).map(Cow::into_owned),
+            Integer(i) => FieldDocIdFacetI64Codec::bytes_encode(&(field_id, document_id, *i)).map(Cow::into_owned),
+        };
+
+        let key = result.context("could not serialize facet key")?;
+        if lmdb_key_valid_size(&key) {
+            sorter.insert(&key, &[])?;
+        }
+
+        Ok(())
+    }
+
+    fn write_word_docids<I>(sorter: &mut Sorter<MergeFn>, iter: I) -> anyhow::Result<()>
+    where I: IntoIterator<Item=(SmallVec32<u8>, RoaringBitmap)>
+    {
+        let mut key = Vec::new();
+        let mut buffer = Vec::new();
+
+        for (word, ids) in iter {
+            key.clear();
+            key.extend_from_slice(&word);
+            // We serialize the document ids into a buffer
+            buffer.clear();
+            let ids = RoaringBitmap::from_iter(ids);
+            buffer.reserve(ids.serialized_size());
+            ids.serialize_into(&mut buffer)?;
+            // that we write under the generated key into MTBL
+            if lmdb_key_valid_size(&key) {
+                sorter.insert(&key, &buffer)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    pub fn index<F>(
+        mut self,
+        mut documents: grenad::Reader<&[u8]>,
+        documents_count: usize,
+        thread_index: usize,
+        num_threads: usize,
+        log_every_n: Option<usize>,
+        mut progress_callback: F,
+    ) -> anyhow::Result<Readers>
+    where F: FnMut(UpdateIndexingStep),
+    {
+        debug!("{:?}: Indexing in a Store...", thread_index);
+
+        let mut before = Instant::now();
+        let mut words_positions = HashMap::new();
+        let mut facet_values = HashMap::new();
+
+        let mut count: usize = 0;
+        while let Some((key, value)) = documents.next()? {
+            let document_id = key.try_into().map(u32::from_be_bytes).unwrap();
+            let document = obkv::KvReader::new(value);
+
+            // We skip documents that must not be indexed by this thread.
+            if count % num_threads == thread_index {
+                // This is a log routine that we do every `log_every_n` documents.
+                if thread_index == 0 && log_every_n.map_or(false, |len| count % len == 0) {
+                    info!("We have seen {} documents so far ({:.02?}).", format_count(count), before.elapsed());
+                    progress_callback(UpdateIndexingStep::IndexDocuments {
+                        documents_seen: count,
+                        total_documents: documents_count,
+                    });
+                    before = Instant::now();
+                }
+
+                for (attr, content) in document.iter() {
+                    if self.faceted_fields.contains_key(&attr) || self.searchable_fields.contains(&attr) {
+                        let value = serde_json::from_slice(content)?;
+
+                        if let Some(ftype) = self.faceted_fields.get(&attr) {
+                            let mut values = parse_facet_value(*ftype, &value).with_context(|| {
+                                format!("extracting facets from the value {}", value)
+                            })?;
+                            facet_values.entry(attr).or_insert_with(SmallVec8::new).extend(values.drain(..));
+                        }
+
+                        if self.searchable_fields.contains(&attr) {
+                            let content = match json_to_string(&value) {
+                                Some(content) => content,
+                                None => continue,
+                            };
+
+                            let analyzed = self.analyzer.analyze(&content);
+                            let tokens = analyzed
+                                .tokens()
+                                .filter(|t| t.is_word())
+                                .map(|t| t.text().to_string());
+
+                            for (pos, word) in tokens.enumerate().take(MAX_POSITION) {
+                                let position = (attr as usize * MAX_POSITION + pos) as u32;
+                                words_positions.entry(word).or_insert_with(SmallVec32::new).push(position);
+                            }
+                        }
+                    }
+                }
+
+                // We write the document in the documents store.
+                self.write_document(document_id, &mut words_positions, &mut facet_values, value)?;
+            }
+
+            // Compute the document id of the next document.
+            count += 1;
+        }
+
+        progress_callback(UpdateIndexingStep::IndexDocuments {
+            documents_seen: count,
+            total_documents: documents_count,
+        });
+
+        let readers = self.finish()?;
+        debug!("{:?}: Store created!", thread_index);
+        Ok(readers)
+    }
+
+    fn finish(mut self) -> anyhow::Result<Readers> {
+        let comp_type = self.chunk_compression_type;
+        let comp_level = self.chunk_compression_level;
+        let shrink_size = self.chunk_fusing_shrink_size;
+
+        Self::write_word_docids(&mut self.word_docids_sorter, self.word_docids)?;
+        Self::write_words_pairs_proximities(
+            &mut self.words_pairs_proximities_docids_sorter,
+            self.words_pairs_proximities_docids,
+        )?;
+        Self::write_facet_field_value_docids(
+            &mut self.facet_field_value_docids_sorter,
+            self.facet_field_value_docids,
+        )?;
+
+        let mut word_docids_wtr = tempfile().and_then(|f| create_writer(comp_type, comp_level, f))?;
+        let mut builder = fst::SetBuilder::memory();
+
+        let mut iter = self.word_docids_sorter.into_iter()?;
+        while let Some((word, val)) = iter.next()? {
+            // This is a lexicographically ordered word position
+            // we use the key to construct the words fst.
+            builder.insert(word)?;
+            word_docids_wtr.insert(word, val)?;
+        }
+
+        let fst = builder.into_set();
+        self.main_sorter.insert(WORDS_FST_KEY, fst.as_fst().as_bytes())?;
+
+        let mut main_wtr = tempfile().and_then(|f| create_writer(comp_type, comp_level, f))?;
+        self.main_sorter.write_into(&mut main_wtr)?;
+
+        let mut words_pairs_proximities_docids_wtr = tempfile().and_then(|f| create_writer(comp_type, comp_level, f))?;
+        self.words_pairs_proximities_docids_sorter.write_into(&mut words_pairs_proximities_docids_wtr)?;
+
+        let mut facet_field_value_docids_wtr = tempfile().and_then(|f| create_writer(comp_type, comp_level, f))?;
+        self.facet_field_value_docids_sorter.write_into(&mut facet_field_value_docids_wtr)?;
+
+        let mut field_id_docid_facet_values_wtr = tempfile().and_then(|f| create_writer(comp_type, comp_level, f))?;
+        self.field_id_docid_facet_values_sorter.write_into(&mut field_id_docid_facet_values_wtr)?;
+
+        let main = writer_into_reader(main_wtr, shrink_size)?;
+        let word_docids = writer_into_reader(word_docids_wtr, shrink_size)?;
+        let words_pairs_proximities_docids = writer_into_reader(words_pairs_proximities_docids_wtr, shrink_size)?;
+        let facet_field_value_docids = writer_into_reader(facet_field_value_docids_wtr, shrink_size)?;
+        let field_id_docid_facet_values = writer_into_reader(field_id_docid_facet_values_wtr, shrink_size)?;
+        let docid_word_positions = writer_into_reader(self.docid_word_positions_writer, shrink_size)?;
+        let documents = writer_into_reader(self.documents_writer, shrink_size)?;
+
+        Ok(Readers {
+            main,
+            word_docids,
+            docid_word_positions,
+            words_pairs_proximities_docids,
+            facet_field_value_docids,
+            field_id_docid_facet_values,
+            documents,
+        })
+    }
+}
+
+/// Outputs a list of all pairs of words with the shortest proximity between 1 and 7 inclusive.
+///
+/// This list is used by the engine to calculate the documents containing words that are
+/// close to each other.
+fn compute_words_pair_proximities(
+    word_positions: &HashMap<String, SmallVec32<Position>>,
+) -> HashMap<(&str, &str), u8>
+{
+    use itertools::Itertools;
+
+    let mut words_pair_proximities = HashMap::new();
+    for ((w1, ps1), (w2, ps2)) in word_positions.iter().cartesian_product(word_positions) {
+        let mut min_prox = None;
+        for (ps1, ps2) in ps1.iter().cartesian_product(ps2) {
+            let prox = crate::proximity::positions_proximity(*ps1, *ps2);
+            let prox = u8::try_from(prox).unwrap();
+            // We don't care about a word that appear at the
+            // same position or too far from the other.
+            if prox >= 1 && prox <= 7 && min_prox.map_or(true, |mp| prox < mp) {
+                min_prox = Some(prox)
+            }
+        }
+
+        if let Some(min_prox) = min_prox {
+            words_pair_proximities.insert((w1.as_str(), w2.as_str()), min_prox);
+        }
+    }
+
+    words_pair_proximities
+}
+
+fn format_count(n: usize) -> String {
+    human_format::Formatter::new().with_decimals(1).with_separator("").format(n as f64)
+}
+
+fn lmdb_key_valid_size(key: &[u8]) -> bool {
+    !key.is_empty() && key.len() <= LMDB_MAX_KEY_LENGTH
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+enum FacetValue {
+    String(SmallString32),
+    Float(OrderedFloat<f64>),
+    Integer(i64),
+}
+
+fn parse_facet_value(ftype: FacetType, value: &Value) -> anyhow::Result<SmallVec8<FacetValue>> {
+    use FacetValue::*;
+
+    fn inner_parse_facet_value(
+        ftype: FacetType,
+        value: &Value,
+        can_recurse: bool,
+        output: &mut SmallVec8<FacetValue>,
+    ) -> anyhow::Result<()>
+    {
+        match value {
+            Value::Null => Ok(()),
+            Value::Bool(b) => {
+                output.push(Integer(*b as i64));
+                Ok(())
+            },
+            Value::Number(number) => match ftype {
+                FacetType::String => {
+                    let string = SmallString32::from(number.to_string());
+                    output.push(String(string));
+                    Ok(())
+                },
+                FacetType::Float => match number.as_f64() {
+                    Some(float) => {
+                        output.push(Float(OrderedFloat(float)));
+                        Ok(())
+                    },
+                    None => bail!("invalid facet type, expecting {} found integer", ftype),
+                },
+                FacetType::Integer => match number.as_i64() {
+                    Some(integer) => {
+                        output.push(Integer(integer));
+                        Ok(())
+                    },
+                    None => if number.is_f64() {
+                        bail!("invalid facet type, expecting {} found float", ftype)
+                    } else {
+                        bail!("invalid facet type, expecting {} found out-of-bound integer (64bit)", ftype)
+                    },
+                },
+            },
+            Value::String(string) => {
+                let string = string.trim().to_lowercase();
+                if string.is_empty() { return Ok(()) }
+                match ftype {
+                    FacetType::String => {
+                        let string = SmallString32::from(string);
+                        output.push(String(string));
+                        Ok(())
+                    },
+                    FacetType::Float => match string.parse() {
+                        Ok(float) => {
+                            output.push(Float(OrderedFloat(float)));
+                            Ok(())
+                        },
+                        Err(_err) => bail!("invalid facet type, expecting {} found string", ftype),
+                    },
+                    FacetType::Integer => match string.parse() {
+                        Ok(integer) => {
+                            output.push(Integer(integer));
+                            Ok(())
+                        },
+                        Err(_err) => bail!("invalid facet type, expecting {} found string", ftype),
+                    },
+                }
+            },
+            Value::Array(values) => if can_recurse {
+                values.iter().map(|v| inner_parse_facet_value(ftype, v, false, output)).collect()
+            } else {
+                bail!("invalid facet type, expecting {} found sub-array ()", ftype)
+            },
+            Value::Object(_) => bail!("invalid facet type, expecting {} found object", ftype),
+        }
+    }
+
+    let mut facet_values = SmallVec8::new();
+    inner_parse_facet_value(ftype, value, true, &mut facet_values)?;
+    Ok(facet_values)
+}

milli/src/update/index_documents/transform.rs

@@ -0,0 +1,641 @@
+use std::borrow::Cow;
+use std::fs::File;
+use std::io::{Read, Seek, SeekFrom};
+use std::iter::Peekable;
+use std::time::Instant;
+
+use anyhow::{anyhow, Context};
+use grenad::CompressionType;
+use log::info;
+use roaring::RoaringBitmap;
+use serde_json::{Map, Value};
+
+use crate::{Index, BEU32, MergeFn, FieldsIdsMap, ExternalDocumentsIds, FieldId};
+use crate::update::{AvailableDocumentsIds, UpdateIndexingStep};
+use super::merge_function::merge_two_obkvs;
+use super::{create_writer, create_sorter, IndexDocumentsMethod};
+
+const DEFAULT_PRIMARY_KEY_NAME: &str = "id";
+
+pub struct TransformOutput {
+    pub primary_key: String,
+    pub fields_ids_map: FieldsIdsMap,
+    pub external_documents_ids: ExternalDocumentsIds<'static>,
+    pub new_documents_ids: RoaringBitmap,
+    pub replaced_documents_ids: RoaringBitmap,
+    pub documents_count: usize,
+    pub documents_file: File,
+}
+
+/// Extract the external ids, deduplicate and compute the new internal documents ids
+/// and fields ids, writing all the documents under their internal ids into a final file.
+///
+/// Outputs the new `FieldsIdsMap`, the new `UsersIdsDocumentsIds` map, the new documents ids,
+/// the replaced documents ids, the number of documents in this update and the file
+/// containing all those documents.
+pub struct Transform<'t, 'i> {
+    pub rtxn: &'t heed::RoTxn<'i>,
+    pub index: &'i Index,
+    pub log_every_n: Option<usize>,
+    pub chunk_compression_type: CompressionType,
+    pub chunk_compression_level: Option<u32>,
+    pub chunk_fusing_shrink_size: Option<u64>,
+    pub max_nb_chunks: Option<usize>,
+    pub max_memory: Option<usize>,
+    pub index_documents_method: IndexDocumentsMethod,
+    pub autogenerate_docids: bool,
+}
+
+impl Transform<'_, '_> {
+    pub fn output_from_json<R, F>(self, reader: R, progress_callback: F) -> anyhow::Result<TransformOutput>
+    where
+        R: Read,
+        F: Fn(UpdateIndexingStep) + Sync,
+    {
+        self.output_from_generic_json(reader, false, progress_callback)
+    }
+
+    pub fn output_from_json_stream<R, F>(self, reader: R, progress_callback: F) -> anyhow::Result<TransformOutput>
+    where
+        R: Read,
+        F: Fn(UpdateIndexingStep) + Sync,
+    {
+        self.output_from_generic_json(reader, true, progress_callback)
+    }
+
+    fn output_from_generic_json<R, F>(
+        self,
+        reader: R,
+        is_stream: bool,
+        progress_callback: F,
+    ) -> anyhow::Result<TransformOutput>
+    where
+        R: Read,
+        F: Fn(UpdateIndexingStep) + Sync,
+    {
+        let mut fields_ids_map = self.index.fields_ids_map(self.rtxn)?;
+        let external_documents_ids = self.index.external_documents_ids(self.rtxn).unwrap();
+
+        // Deserialize the whole batch of documents in memory.
+        let mut documents: Peekable<Box<dyn Iterator<Item=serde_json::Result<Map<String, Value>>>>> = if is_stream {
+            let iter = serde_json::Deserializer::from_reader(reader).into_iter();
+            let iter = Box::new(iter) as Box<dyn Iterator<Item=_>>;
+            iter.peekable()
+        } else {
+            let vec: Vec<_> = serde_json::from_reader(reader)?;
+            let iter = vec.into_iter().map(Ok);
+            let iter = Box::new(iter) as Box<dyn Iterator<Item=_>>;
+            iter.peekable()
+        };
+
+        // We extract the primary key from the first document in
+        // the batch if it hasn't already been defined in the index
+        let first = documents.peek().and_then(|r| r.as_ref().ok());
+        let alternative_name = first.and_then(|doc| doc.keys().find(|k| k.contains(DEFAULT_PRIMARY_KEY_NAME)).cloned());
+        let (primary_key_id, primary_key) = compute_primary_key_pair(
+            self.index.primary_key(self.rtxn)?,
+            &mut fields_ids_map,
+            alternative_name,
+            self.autogenerate_docids
+        )?;
+
+        if documents.peek().is_none() {
+            return Ok(TransformOutput {
+                primary_key,
+                fields_ids_map,
+                external_documents_ids: ExternalDocumentsIds::default(),
+                new_documents_ids: RoaringBitmap::new(),
+                replaced_documents_ids: RoaringBitmap::new(),
+                documents_count: 0,
+                documents_file: tempfile::tempfile()?,
+            });
+        }
+
+        // We must choose the appropriate merge function for when two or more documents
+        // with the same user id must be merged or fully replaced in the same batch.
+        let merge_function = match self.index_documents_method {
+            IndexDocumentsMethod::ReplaceDocuments => keep_latest_obkv,
+            IndexDocumentsMethod::UpdateDocuments => merge_obkvs,
+        };
+
+        // We initialize the sorter with the user indexing settings.
+        let mut sorter = create_sorter(
+            merge_function,
+            self.chunk_compression_type,
+            self.chunk_compression_level,
+            self.chunk_fusing_shrink_size,
+            self.max_nb_chunks,
+            self.max_memory,
+        );
+
+        let mut json_buffer = Vec::new();
+        let mut obkv_buffer = Vec::new();
+        let mut uuid_buffer = [0; uuid::adapter::Hyphenated::LENGTH];
+        let mut documents_count = 0;
+
+        for result in documents {
+            let document = result?;
+
+            if self.log_every_n.map_or(false, |len| documents_count % len == 0) {
+                progress_callback(UpdateIndexingStep::TransformFromUserIntoGenericFormat {
+                    documents_seen: documents_count,
+                });
+            }
+
+            obkv_buffer.clear();
+            let mut writer = obkv::KvWriter::new(&mut obkv_buffer);
+
+            // We prepare the fields ids map with the documents keys.
+            for (key, _value) in &document {
+                fields_ids_map.insert(&key).context("field id limit reached")?;
+            }
+
+            // We retrieve the user id from the document based on the primary key name,
+            // if the document id isn't present we generate a uuid.
+            let external_id = match document.get(&primary_key) {
+                Some(value) => match value {
+                    Value::String(string) => Cow::Borrowed(string.as_str()),
+                    Value::Number(number) => Cow::Owned(number.to_string()),
+                    _ => return Err(anyhow!("documents ids must be either strings or numbers")),
+                },
+                None => {
+                    if !self.autogenerate_docids {
+                        return Err(anyhow!("missing primary key"));
+                    }
+                    let uuid = uuid::Uuid::new_v4().to_hyphenated().encode_lower(&mut uuid_buffer);
+                    Cow::Borrowed(uuid)
+                },
+            };
+
+            // We iterate in the fields ids ordered.
+            for (field_id, name) in fields_ids_map.iter() {
+                json_buffer.clear();
+
+                // We try to extract the value from the document and if we don't find anything
+                // and this should be the document id we return the one we generated.
+                if let Some(value) = document.get(name) {
+                    // We serialize the attribute values.
+                    serde_json::to_writer(&mut json_buffer, value)?;
+                    writer.insert(field_id, &json_buffer)?;
+                }
+                else if field_id == primary_key_id {
+                    // We validate the document id [a-zA-Z0-9\-_].
+                    let external_id = match validate_document_id(&external_id) {
+                        Some(valid) => valid,
+                        None => return Err(anyhow!("invalid document id: {:?}", external_id)),
+                    };
+
+                    // We serialize the document id.
+                    serde_json::to_writer(&mut json_buffer, &external_id)?;
+                    writer.insert(field_id, &json_buffer)?;
+                }
+            }
+
+            // We use the extracted/generated user id as the key for this document.
+            sorter.insert(external_id.as_bytes(), &obkv_buffer)?;
+            documents_count += 1;
+        }
+
+        progress_callback(UpdateIndexingStep::TransformFromUserIntoGenericFormat {
+            documents_seen: documents_count,
+        });
+
+        // Now that we have a valid sorter that contains the user id and the obkv we
+        // give it to the last transforming function which returns the TransformOutput.
+        self.output_from_sorter(
+            sorter,
+            primary_key,
+            fields_ids_map,
+            documents_count,
+            external_documents_ids,
+            progress_callback,
+        )
+    }
+
+    pub fn output_from_csv<R, F>(self, reader: R, progress_callback: F) -> anyhow::Result<TransformOutput>
+    where
+        R: Read,
+        F: Fn(UpdateIndexingStep) + Sync,
+    {
+        let mut fields_ids_map = self.index.fields_ids_map(self.rtxn)?;
+        let external_documents_ids = self.index.external_documents_ids(self.rtxn).unwrap();
+
+        let mut csv = csv::Reader::from_reader(reader);
+        let headers = csv.headers()?;
+
+        let mut fields_ids = Vec::new();
+        // Generate the new fields ids based on the current fields ids and this CSV headers.
+        for (i, header) in headers.iter().enumerate() {
+            let id = fields_ids_map.insert(header).context("field id limit reached)")?;
+            fields_ids.push((id, i));
+        }
+
+        // Extract the position of the primary key in the current headers, None if not found.
+        let primary_key_pos = match self.index.primary_key(self.rtxn)? {
+            Some(primary_key) => {
+               // The primary key is known so we must find the position in the CSV headers.
+               headers.iter().position(|h| h == primary_key)
+            },
+            None => headers.iter().position(|h| h.contains("id")),
+        };
+
+        // Returns the field id in the fields ids map, create an "id" field
+        // in case it is not in the current headers.
+        let alternative_name = primary_key_pos.map(|pos| headers[pos].to_string());
+        let (primary_key_id, _) = compute_primary_key_pair(
+            self.index.primary_key(self.rtxn)?,
+            &mut fields_ids_map,
+            alternative_name,
+            self.autogenerate_docids
+        )?;
+
+        // The primary key field is not present in the header, so we need to create it.
+        if primary_key_pos.is_none() {
+            fields_ids.push((primary_key_id, usize::max_value()));
+        }
+
+        // We sort the fields ids by the fields ids map id, this way we are sure to iterate over
+        // the records fields in the fields ids map order and correctly generate the obkv.
+        fields_ids.sort_unstable_by_key(|(field_id, _)| *field_id);
+
+        // We initialize the sorter with the user indexing settings.
+        let mut sorter = create_sorter(
+            keep_latest_obkv,
+            self.chunk_compression_type,
+            self.chunk_compression_level,
+            self.chunk_fusing_shrink_size,
+            self.max_nb_chunks,
+            self.max_memory,
+        );
+
+        // We write into the sorter to merge and deduplicate the documents
+        // based on the external ids.
+        let mut json_buffer = Vec::new();
+        let mut obkv_buffer = Vec::new();
+        let mut uuid_buffer = [0; uuid::adapter::Hyphenated::LENGTH];
+        let mut documents_count = 0;
+
+        let mut record = csv::StringRecord::new();
+        while csv.read_record(&mut record)? {
+            obkv_buffer.clear();
+            let mut writer = obkv::KvWriter::new(&mut obkv_buffer);
+
+            if self.log_every_n.map_or(false, |len| documents_count % len == 0) {
+                progress_callback(UpdateIndexingStep::TransformFromUserIntoGenericFormat {
+                    documents_seen: documents_count,
+                });
+            }
+
+            // We extract the user id if we know where it is or generate an UUID V4 otherwise.
+            let external_id = match primary_key_pos {
+                Some(pos) => {
+                    let external_id = &record[pos];
+                    // We validate the document id [a-zA-Z0-9\-_].
+                    match validate_document_id(&external_id) {
+                        Some(valid) => valid,
+                        None => return Err(anyhow!("invalid document id: {:?}", external_id)),
+                    }
+                },
+                None => uuid::Uuid::new_v4().to_hyphenated().encode_lower(&mut uuid_buffer),
+            };
+
+            // When the primary_key_field_id is found in the fields ids list
+            // we return the generated document id instead of the record field.
+            let iter = fields_ids.iter()
+                .map(|(fi, i)| {
+                    let field = if *fi == primary_key_id { external_id } else { &record[*i] };
+                    (fi, field)
+                });
+
+            // We retrieve the field id based on the fields ids map fields ids order.
+            for (field_id, field) in iter {
+                // We serialize the attribute values as JSON strings.
+                json_buffer.clear();
+                serde_json::to_writer(&mut json_buffer, &field)?;
+                writer.insert(*field_id, &json_buffer)?;
+            }
+
+            // We use the extracted/generated user id as the key for this document.
+            sorter.insert(external_id, &obkv_buffer)?;
+            documents_count += 1;
+        }
+
+        progress_callback(UpdateIndexingStep::TransformFromUserIntoGenericFormat {
+            documents_seen: documents_count,
+        });
+
+        // Now that we have a valid sorter that contains the user id and the obkv we
+        // give it to the last transforming function which returns the TransformOutput.
+        let primary_key_name = fields_ids_map
+            .name(primary_key_id)
+            .map(String::from)
+            .expect("Primary key must be present in fields id map");
+        self.output_from_sorter(
+            sorter,
+            primary_key_name,
+            fields_ids_map,
+            documents_count,
+            external_documents_ids,
+            progress_callback,
+        )
+    }
+
+    /// Generate the `TransformOutput` based on the given sorter that can be generated from any
+    /// format like CSV, JSON or JSON stream. This sorter must contain a key that is the document
+    /// id for the user side and the value must be an obkv where keys are valid fields ids.
+    fn output_from_sorter<F>(
+        self,
+        sorter: grenad::Sorter<MergeFn>,
+        primary_key: String,
+        fields_ids_map: FieldsIdsMap,
+        approximate_number_of_documents: usize,
+        mut external_documents_ids: ExternalDocumentsIds<'_>,
+        progress_callback: F,
+    ) -> anyhow::Result<TransformOutput>
+    where
+        F: Fn(UpdateIndexingStep) + Sync,
+    {
+        let documents_ids = self.index.documents_ids(self.rtxn)?;
+        let mut available_documents_ids = AvailableDocumentsIds::from_documents_ids(&documents_ids);
+
+        // Once we have sort and deduplicated the documents we write them into a final file.
+        let mut final_sorter = create_sorter(
+            |_docid, _obkvs| Err(anyhow!("cannot merge two documents")),
+            self.chunk_compression_type,
+            self.chunk_compression_level,
+            self.chunk_fusing_shrink_size,
+            self.max_nb_chunks,
+            self.max_memory,
+        );
+        let mut new_external_documents_ids_builder = fst::MapBuilder::memory();
+        let mut replaced_documents_ids = RoaringBitmap::new();
+        let mut new_documents_ids = RoaringBitmap::new();
+        let mut obkv_buffer = Vec::new();
+
+        // While we write into final file we get or generate the internal documents ids.
+        let mut documents_count = 0;
+        let mut iter = sorter.into_iter()?;
+        while let Some((external_id, update_obkv)) = iter.next()? {
+
+            if self.log_every_n.map_or(false, |len| documents_count % len == 0) {
+                progress_callback(UpdateIndexingStep::ComputeIdsAndMergeDocuments {
+                    documents_seen: documents_count,
+                    total_documents: approximate_number_of_documents,
+                });
+            }
+
+            let (docid, obkv) = match external_documents_ids.get(external_id) {
+                Some(docid) => {
+                    // If we find the user id in the current external documents ids map
+                    // we use it and insert it in the list of replaced documents.
+                    replaced_documents_ids.insert(docid);
+
+                    // Depending on the update indexing method we will merge
+                    // the document update with the current document or not.
+                    match self.index_documents_method {
+                        IndexDocumentsMethod::ReplaceDocuments => (docid, update_obkv),
+                        IndexDocumentsMethod::UpdateDocuments => {
+                            let key = BEU32::new(docid);
+                            let base_obkv = self.index.documents.get(&self.rtxn, &key)?
+                                .context("document not found")?;
+                            let update_obkv = obkv::KvReader::new(update_obkv);
+                            merge_two_obkvs(base_obkv, update_obkv, &mut obkv_buffer);
+                            (docid, obkv_buffer.as_slice())
+                        }
+                    }
+                },
+                None => {
+                    // If this user id is new we add it to the external documents ids map
+                    // for new ids and into the list of new documents.
+                    let new_docid = available_documents_ids.next()
+                        .context("no more available documents ids")?;
+                    new_external_documents_ids_builder.insert(external_id, new_docid as u64)?;
+                    new_documents_ids.insert(new_docid);
+                    (new_docid, update_obkv)
+                },
+            };
+
+            // We insert the document under the documents ids map into the final file.
+            final_sorter.insert(docid.to_be_bytes(), obkv)?;
+            documents_count += 1;
+        }
+
+        progress_callback(UpdateIndexingStep::ComputeIdsAndMergeDocuments {
+            documents_seen: documents_count,
+            total_documents: documents_count,
+        });
+
+        // We create a final writer to write the new documents in order from the sorter.
+        let file = tempfile::tempfile()?;
+        let mut writer = create_writer(self.chunk_compression_type, self.chunk_compression_level, file)?;
+
+        // Once we have written all the documents into the final sorter, we write the documents
+        // into this writer, extract the file and reset the seek to be able to read it again.
+        final_sorter.write_into(&mut writer)?;
+        let mut documents_file = writer.into_inner()?;
+        documents_file.seek(SeekFrom::Start(0))?;
+
+        let before_docids_merging = Instant::now();
+        // We merge the new external ids with existing external documents ids.
+        let new_external_documents_ids = new_external_documents_ids_builder.into_map();
+        external_documents_ids.insert_ids(&new_external_documents_ids)?;
+
+        info!("Documents external merging took {:.02?}", before_docids_merging.elapsed());
+
+        Ok(TransformOutput {
+            primary_key,
+            fields_ids_map,
+            external_documents_ids: external_documents_ids.into_static(),
+            new_documents_ids,
+            replaced_documents_ids,
+            documents_count,
+            documents_file,
+        })
+    }
+
+    /// Returns a `TransformOutput` with a file that contains the documents of the index
+    /// with the attributes reordered accordingly to the `FieldsIdsMap` given as argument.
+    // TODO this can be done in parallel by using the rayon `ThreadPool`.
+    pub fn remap_index_documents(
+        self,
+        primary_key: String,
+        old_fields_ids_map: FieldsIdsMap,
+        new_fields_ids_map: FieldsIdsMap,
+    ) -> anyhow::Result<TransformOutput>
+    {
+        let external_documents_ids = self.index.external_documents_ids(self.rtxn)?;
+        let documents_ids = self.index.documents_ids(self.rtxn)?;
+        let documents_count = documents_ids.len() as usize;
+
+        // We create a final writer to write the new documents in order from the sorter.
+        let file = tempfile::tempfile()?;
+        let mut writer = create_writer(self.chunk_compression_type, self.chunk_compression_level, file)?;
+
+        let mut obkv_buffer = Vec::new();
+        for result in self.index.documents.iter(self.rtxn)? {
+            let (docid, obkv) = result?;
+            let docid = docid.get();
+
+            obkv_buffer.clear();
+            let mut obkv_writer = obkv::KvWriter::new(&mut obkv_buffer);
+
+            // We iterate over the new `FieldsIdsMap` ids in order and construct the new obkv.
+            for (id, name) in new_fields_ids_map.iter() {
+                if let Some(val) = old_fields_ids_map.id(name).and_then(|id| obkv.get(id)) {
+                    obkv_writer.insert(id, val)?;
+                }
+            }
+
+            let buffer = obkv_writer.into_inner()?;
+            writer.insert(docid.to_be_bytes(), buffer)?;
+        }
+
+        // Once we have written all the documents, we extract
+        // the file and reset the seek to be able to read it again.
+        let mut documents_file = writer.into_inner()?;
+        documents_file.seek(SeekFrom::Start(0))?;
+
+        Ok(TransformOutput {
+            primary_key,
+            fields_ids_map: new_fields_ids_map,
+            external_documents_ids: external_documents_ids.into_static(),
+            new_documents_ids: documents_ids,
+            replaced_documents_ids: RoaringBitmap::default(),
+            documents_count,
+            documents_file,
+        })
+    }
+}
+
+/// Given an optional primary key and an optional alternative name, returns the (field_id, attr_name)
+/// for the primary key according to the following rules:
+/// - if primary_key is `Some`, returns the id and the name, else
+/// - if alternative_name is Some, adds alternative to the fields_ids_map, and returns the pair, else
+/// - if autogenerate_docids is true, insert the default id value in the field ids map ("id") and
+/// returns the pair, else
+/// - returns an error.
+fn compute_primary_key_pair(
+    primary_key: Option<&str>,
+    fields_ids_map: &mut FieldsIdsMap,
+    alternative_name: Option<String>,
+    autogenerate_docids: bool,
+) -> anyhow::Result<(FieldId, String)> {
+    match primary_key {
+        Some(primary_key) => {
+            let id = fields_ids_map.id(primary_key).expect("primary key must be present in the fields id map");
+            Ok((id, primary_key.to_string()))
+        }
+        None => {
+            let name = match alternative_name {
+                Some(key) => key,
+                None => {
+                    if !autogenerate_docids {
+                        // If there is no primary key in the current document batch, we must
+                        // return an error and not automatically generate any document id.
+                        anyhow::bail!("missing primary key")
+                    }
+                    DEFAULT_PRIMARY_KEY_NAME.to_string()
+                },
+            };
+            let id = fields_ids_map.insert(&name).context("field id limit reached")?;
+            Ok((id, name))
+        },
+    }
+}
+
+/// Only the last value associated with an id is kept.
+fn keep_latest_obkv(_key: &[u8], obkvs: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    obkvs.last().context("no last value").map(|last| last.clone().into_owned())
+}
+
+/// Merge all the obks in the order we see them.
+fn merge_obkvs(_key: &[u8], obkvs: &[Cow<[u8]>]) -> anyhow::Result<Vec<u8>> {
+    let mut iter = obkvs.iter();
+    let first = iter.next().map(|b| b.clone().into_owned()).context("no first value")?;
+    Ok(iter.fold(first, |acc, current| {
+        let first = obkv::KvReader::new(&acc);
+        let second = obkv::KvReader::new(current);
+        let mut buffer = Vec::new();
+        merge_two_obkvs(first, second, &mut buffer);
+        buffer
+    }))
+}
+
+fn validate_document_id(document_id: &str) -> Option<&str> {
+    let document_id = document_id.trim();
+    Some(document_id).filter(|id| {
+        !id.is_empty() && id.chars().all(|c| {
+            matches!(c, 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_')
+        })
+    })
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    mod compute_primary_key {
+        use super::compute_primary_key_pair;
+        use super::FieldsIdsMap;
+
+        #[test]
+        #[should_panic]
+        fn should_panic_primary_key_not_in_map() {
+            let mut fields_map = FieldsIdsMap::new();
+            let _result = compute_primary_key_pair(
+                Some("toto"),
+                &mut fields_map,
+                None,
+                false);
+        }
+
+        #[test]
+        fn should_return_primary_key_if_is_some() {
+            let mut fields_map = FieldsIdsMap::new();
+            fields_map.insert("toto").unwrap();
+            let result = compute_primary_key_pair(
+                Some("toto"),
+                &mut fields_map,
+                Some("tata".to_string()),
+                false);
+            assert_eq!(result.unwrap(), (0u8, "toto".to_string()));
+            assert_eq!(fields_map.len(), 1);
+        }
+
+        #[test]
+        fn should_return_alternative_if_primary_is_none() {
+            let mut fields_map = FieldsIdsMap::new();
+            let result = compute_primary_key_pair(
+                None,
+                &mut fields_map,
+                Some("tata".to_string()),
+                false);
+            assert_eq!(result.unwrap(), (0u8, "tata".to_string()));
+            assert_eq!(fields_map.len(), 1);
+        }
+
+        #[test]
+        fn should_return_default_if_both_are_none() {
+            let mut fields_map = FieldsIdsMap::new();
+            let result = compute_primary_key_pair(
+                None,
+                &mut fields_map,
+                None,
+                true);
+            assert_eq!(result.unwrap(), (0u8, "id".to_string()));
+            assert_eq!(fields_map.len(), 1);
+        }
+
+        #[test]
+        fn should_return_err_if_both_are_none_and_recompute_is_false(){
+            let mut fields_map = FieldsIdsMap::new();
+            let result = compute_primary_key_pair(
+                None,
+                &mut fields_map,
+                None,
+                false);
+            assert!(result.is_err());
+            assert_eq!(fields_map.len(), 0);
+        }
+    }
+}

milli/src/update/mod.rs

@@ -0,0 +1,17 @@
+mod available_documents_ids;
+mod clear_documents;
+mod delete_documents;
+mod facets;
+mod index_documents;
+mod settings;
+mod update_builder;
+mod update_step;
+
+pub use self::available_documents_ids::AvailableDocumentsIds;
+pub use self::clear_documents::ClearDocuments;
+pub use self::delete_documents::DeleteDocuments;
+pub use self::index_documents::{IndexDocuments, IndexDocumentsMethod, UpdateFormat, DocumentAdditionResult};
+pub use self::facets::Facets;
+pub use self::settings::Settings;
+pub use self::update_builder::UpdateBuilder;
+pub use self::update_step::UpdateIndexingStep;

milli/src/update/settings.rs

@@ -0,0 +1,522 @@
+use std::collections::HashMap;
+use std::str::FromStr;
+
+use anyhow::Context;
+use grenad::CompressionType;
+use itertools::Itertools;
+use rayon::ThreadPool;
+
+use crate::criterion::Criterion;
+use crate::facet::FacetType;
+use crate::update::index_documents::{Transform, IndexDocumentsMethod};
+use crate::update::{ClearDocuments, IndexDocuments, UpdateIndexingStep};
+use crate::{Index, FieldsIdsMap};
+
+pub struct Settings<'a, 't, 'u, 'i> {
+    wtxn: &'t mut heed::RwTxn<'i, 'u>,
+    index: &'i Index,
+    pub(crate) log_every_n: Option<usize>,
+    pub(crate) max_nb_chunks: Option<usize>,
+    pub(crate) max_memory: Option<usize>,
+    pub(crate) linked_hash_map_size: Option<usize>,
+    pub(crate) chunk_compression_type: CompressionType,
+    pub(crate) chunk_compression_level: Option<u32>,
+    pub(crate) chunk_fusing_shrink_size: Option<u64>,
+    pub(crate) thread_pool: Option<&'a ThreadPool>,
+    update_id: u64,
+
+    // If a struct field is set to `None` it means that it hasn't been set by the user,
+    // however if it is `Some(None)` it means that the user forced a reset of the setting.
+    searchable_fields: Option<Option<Vec<String>>>,
+    displayed_fields: Option<Option<Vec<String>>>,
+    faceted_fields: Option<Option<HashMap<String, String>>>,
+    criteria: Option<Option<Vec<String>>>,
+}
+
+impl<'a, 't, 'u, 'i> Settings<'a, 't, 'u, 'i> {
+    pub fn new(
+        wtxn: &'t mut heed::RwTxn<'i, 'u>,
+        index: &'i Index,
+        update_id: u64,
+    ) -> Settings<'a, 't, 'u, 'i> {
+        Settings {
+            wtxn,
+            index,
+            log_every_n: None,
+            max_nb_chunks: None,
+            max_memory: None,
+            linked_hash_map_size: None,
+            chunk_compression_type: CompressionType::None,
+            chunk_compression_level: None,
+            chunk_fusing_shrink_size: None,
+            thread_pool: None,
+            searchable_fields: None,
+            displayed_fields: None,
+            faceted_fields: None,
+            criteria: None,
+            update_id,
+        }
+    }
+
+    pub fn reset_searchable_fields(&mut self) {
+        self.searchable_fields = Some(None);
+    }
+
+    pub fn set_searchable_fields(&mut self, names: Vec<String>) {
+        self.searchable_fields = Some(Some(names));
+    }
+
+    pub fn reset_displayed_fields(&mut self) {
+        self.displayed_fields = Some(None);
+    }
+
+    pub fn set_displayed_fields(&mut self, names: Vec<String>) {
+        self.displayed_fields = Some(Some(names));
+    }
+
+    pub fn set_faceted_fields(&mut self, names_facet_types: HashMap<String, String>) {
+        self.faceted_fields = Some(Some(names_facet_types));
+    }
+
+    pub fn reset_faceted_fields(&mut self) {
+        self.faceted_fields = Some(None);
+    }
+
+    pub fn reset_criteria(&mut self) {
+        self.criteria = Some(None);
+    }
+
+    pub fn set_criteria(&mut self, criteria: Vec<String>) {
+        self.criteria = Some(Some(criteria));
+    }
+
+    fn reindex<F>(&mut self, cb: &F, old_fields_ids_map: FieldsIdsMap) -> anyhow::Result<()>
+    where
+        F: Fn(UpdateIndexingStep, u64) + Sync
+    {
+        let fields_ids_map = self.index.fields_ids_map(self.wtxn)?;
+        let update_id = self.update_id;
+        let cb = |step| cb(step, update_id);
+        // if the settings are set before any document update, we don't need to do anything, and
+        // will set the primary key during the first document addition.
+        if self.index.number_of_documents(&self.wtxn)? == 0 {
+            return Ok(())
+        }
+
+        let transform = Transform {
+            rtxn: &self.wtxn,
+            index: self.index,
+            log_every_n: self.log_every_n,
+            chunk_compression_type: self.chunk_compression_type,
+            chunk_compression_level: self.chunk_compression_level,
+            chunk_fusing_shrink_size: self.chunk_fusing_shrink_size,
+            max_nb_chunks: self.max_nb_chunks,
+            max_memory: self.max_memory,
+            index_documents_method: IndexDocumentsMethod::ReplaceDocuments,
+            autogenerate_docids: false,
+        };
+
+        // There already has been a document addition, the primary key should be set by now.
+        let primary_key = self.index.primary_key(&self.wtxn)?.context("Index must have a primary key")?;
+
+        // We remap the documents fields based on the new `FieldsIdsMap`.
+        let output = transform.remap_index_documents(
+            primary_key.to_string(),
+            old_fields_ids_map,
+            fields_ids_map.clone())?;
+
+        // We clear the full database (words-fst, documents ids and documents content).
+        ClearDocuments::new(self.wtxn, self.index, self.update_id).execute()?;
+
+        // We index the generated `TransformOutput` which must contain
+        // all the documents with fields in the newly defined searchable order.
+        let mut indexing_builder = IndexDocuments::new(self.wtxn, self.index, self.update_id);
+        indexing_builder.log_every_n = self.log_every_n;
+        indexing_builder.max_nb_chunks = self.max_nb_chunks;
+        indexing_builder.max_memory = self.max_memory;
+        indexing_builder.linked_hash_map_size = self.linked_hash_map_size;
+        indexing_builder.chunk_compression_type = self.chunk_compression_type;
+        indexing_builder.chunk_compression_level = self.chunk_compression_level;
+        indexing_builder.chunk_fusing_shrink_size = self.chunk_fusing_shrink_size;
+        indexing_builder.thread_pool = self.thread_pool;
+        indexing_builder.execute_raw(output, &cb)?;
+        Ok(())
+    }
+
+    fn update_displayed(&mut self) -> anyhow::Result<bool> {
+        match self.displayed_fields {
+            Some(Some(ref fields)) => {
+                let mut fields_ids_map = self.index.fields_ids_map(self.wtxn)?;
+                // fields are deduplicated, only the first occurrence is taken into account
+                let names: Vec<_> = fields
+                    .iter()
+                    .unique()
+                    .map(String::as_str)
+                    .collect();
+
+                for name in names.iter() {
+                    fields_ids_map
+                        .insert(name)
+                        .context("field id limit exceeded")?;
+                }
+                self.index.put_displayed_fields(self.wtxn, &names)?;
+                self.index.put_fields_ids_map(self.wtxn, &fields_ids_map)?;
+            }
+            Some(None) => { self.index.delete_displayed_fields(self.wtxn)?; },
+            None => return Ok(false),
+        }
+        Ok(true)
+    }
+
+    /// Udpates the index's searchable attributes. This causes the field map to be recomputed to
+    /// reflect the order of the searchable attributes.
+    fn update_searchable(&mut self) -> anyhow::Result<bool> {
+        match self.searchable_fields {
+            Some(Some(ref fields)) => {
+                // every time the searchable attributes are updated, we need to update the
+                // ids for any settings that uses the facets. (displayed_fields,
+                // faceted_fields)
+                let old_fields_ids_map = self.index.fields_ids_map(self.wtxn)?;
+
+                let mut new_fields_ids_map = FieldsIdsMap::new();
+                // fields are deduplicated, only the first occurrence is taken into account
+                let names = fields
+                    .iter()
+                    .unique()
+                    .map(String::as_str)
+                    .collect::<Vec<_>>();
+
+                // Add all the searchable attributes to the field map, and then add the
+                // remaining fields from the old field map to the new one
+                for name in names.iter() {
+                    new_fields_ids_map
+                        .insert(&name)
+                        .context("field id limit exceeded")?;
+                }
+
+                for (_, name) in old_fields_ids_map.iter() {
+                    new_fields_ids_map
+                        .insert(&name)
+                        .context("field id limit exceeded")?;
+                }
+
+                self.index.put_searchable_fields(self.wtxn, &names)?;
+                self.index.put_fields_ids_map(self.wtxn, &new_fields_ids_map)?;
+            }
+            Some(None) => { self.index.delete_searchable_fields(self.wtxn)?; },
+            None => return Ok(false),
+        }
+        Ok(true)
+    }
+
+    fn update_facets(&mut self) -> anyhow::Result<bool> {
+        match self.faceted_fields {
+            Some(Some(ref fields)) => {
+                let mut fields_ids_map = self.index.fields_ids_map(self.wtxn)?;
+                let mut new_facets = HashMap::new();
+                for (name, ty) in fields {
+                    fields_ids_map.insert(name).context("field id limit exceeded")?;
+                    let ty = FacetType::from_str(&ty)?;
+                    new_facets.insert(name.clone(), ty);
+                }
+                self.index.put_faceted_fields(self.wtxn, &new_facets)?;
+                self.index.put_fields_ids_map(self.wtxn, &fields_ids_map)?;
+            }
+            Some(None) => { self.index.delete_faceted_fields(self.wtxn)?; },
+            None => return Ok(false)
+        }
+        Ok(true)
+    }
+
+    fn update_criteria(&mut self) -> anyhow::Result<()> {
+        match self.criteria {
+            Some(Some(ref fields)) => {
+                let faceted_fields = self.index.faceted_fields(&self.wtxn)?;
+                let mut new_criteria = Vec::new();
+                for name in fields {
+                    let criterion = Criterion::from_str(&faceted_fields, &name)?;
+                    new_criteria.push(criterion);
+                }
+                self.index.put_criteria(self.wtxn, &new_criteria)?;
+            }
+            Some(None) => { self.index.delete_criteria(self.wtxn)?; }
+            None => (),
+        }
+        Ok(())
+    }
+
+    pub fn execute<F>(mut self, progress_callback: F) -> anyhow::Result<()>
+    where
+        F: Fn(UpdateIndexingStep, u64) + Sync
+        {
+            let old_fields_ids_map = self.index.fields_ids_map(&self.wtxn)?;
+            self.update_displayed()?;
+            let facets_updated = self.update_facets()?;
+            // update_criteria MUST be called after update_facets, since criterion fields must be set
+            // as facets.
+            self.update_criteria()?;
+            let searchable_updated = self.update_searchable()?;
+
+            if facets_updated || searchable_updated {
+                self.reindex(&progress_callback, old_fields_ids_map)?;
+            }
+            Ok(())
+        }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use heed::EnvOpenOptions;
+    use maplit::hashmap;
+
+    use crate::facet::FacetType;
+    use crate::update::{IndexDocuments, UpdateFormat};
+
+    #[test]
+    fn set_and_reset_searchable_fields() {
+        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 ids from 1 to 3.
+        let mut wtxn = index.write_txn().unwrap();
+        let content = &b"name,age\nkevin,23\nkevina,21\nbenoit,34\n"[..];
+        let mut builder = IndexDocuments::new(&mut wtxn, &index, 0);
+        builder.update_format(UpdateFormat::Csv);
+        builder.execute(content, |_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // We change the searchable fields to be the "name" field only.
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = Settings::new(&mut wtxn, &index, 1);
+        builder.set_searchable_fields(vec!["name".into()]);
+        builder.execute(|_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the searchable field is correctly set to "name" only.
+        let rtxn = index.read_txn().unwrap();
+        // When we search for something that is not in
+        // the searchable fields it must not return any document.
+        let result = index.search(&rtxn).query("23").execute().unwrap();
+        assert!(result.documents_ids.is_empty());
+
+        // When we search for something that is in the searchable fields
+        // we must find the appropriate document.
+        let result = index.search(&rtxn).query(r#""kevin""#).execute().unwrap();
+        let documents = index.documents(&rtxn, result.documents_ids).unwrap();
+        assert_eq!(documents.len(), 1);
+        assert_eq!(documents[0].1.get(0), Some(&br#""kevin""#[..]));
+        drop(rtxn);
+
+        // We change the searchable fields to be the "name" field only.
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = Settings::new(&mut wtxn, &index, 2);
+        builder.reset_searchable_fields();
+        builder.execute(|_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the searchable field have been reset and documents are found now.
+        let rtxn = index.read_txn().unwrap();
+        let searchable_fields = index.searchable_fields(&rtxn).unwrap();
+        assert_eq!(searchable_fields, None);
+        let result = index.search(&rtxn).query("23").execute().unwrap();
+        assert_eq!(result.documents_ids.len(), 1);
+        let documents = index.documents(&rtxn, result.documents_ids).unwrap();
+        assert_eq!(documents[0].1.get(0), Some(&br#""kevin""#[..]));
+        drop(rtxn);
+    }
+
+    #[test]
+    fn mixup_searchable_with_displayed_fields() {
+        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 ids from 1 to 3.
+        let mut wtxn = index.write_txn().unwrap();
+        let content = &b"name,age\nkevin,23\nkevina,21\nbenoit,34\n"[..];
+        let mut builder = IndexDocuments::new(&mut wtxn, &index, 0);
+        builder.update_format(UpdateFormat::Csv);
+        builder.execute(content, |_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // In the same transaction we change the displayed fields to be only the "age".
+        // We also change the searchable fields to be the "name" field only.
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = Settings::new(&mut wtxn, &index, 1);
+        builder.set_displayed_fields(vec!["age".into()]);
+        builder.set_searchable_fields(vec!["name".into()]);
+        builder.execute(|_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the displayed fields are correctly set to `None` (default value).
+        let rtxn = index.read_txn().unwrap();
+        let fields_ids = index.displayed_fields(&rtxn).unwrap();
+        assert_eq!(fields_ids.unwrap(), (&["age"][..]));
+        drop(rtxn);
+
+        // We change the searchable fields to be the "name" field only.
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = Settings::new(&mut wtxn, &index, 2);
+        builder.reset_searchable_fields();
+        builder.execute(|_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the displayed fields always contains only the "age" field.
+        let rtxn = index.read_txn().unwrap();
+        let fields_ids = index.displayed_fields(&rtxn).unwrap();
+        assert_eq!(fields_ids.unwrap(), &["age"][..]);
+        drop(rtxn);
+    }
+
+    #[test]
+    fn default_displayed_fields() {
+        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 ids from 1 to 3.
+        let mut wtxn = index.write_txn().unwrap();
+        let content = &b"name,age\nkevin,23\nkevina,21\nbenoit,34\n"[..];
+        let mut builder = IndexDocuments::new(&mut wtxn, &index, 0);
+        builder.update_format(UpdateFormat::Csv);
+        builder.execute(content, |_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the displayed fields are correctly set to `None` (default value).
+        let rtxn = index.read_txn().unwrap();
+        let fields_ids = index.displayed_fields(&rtxn).unwrap();
+        assert_eq!(fields_ids, None);
+        drop(rtxn);
+    }
+
+    #[test]
+    fn set_and_reset_displayed_field() {
+        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 ids from 1 to 3.
+        let mut wtxn = index.write_txn().unwrap();
+        let content = &b"name,age\nkevin,23\nkevina,21\nbenoit,34\n"[..];
+        let mut builder = IndexDocuments::new(&mut wtxn, &index, 0);
+        builder.update_format(UpdateFormat::Csv);
+        builder.execute(content, |_, _| ()).unwrap();
+
+        // In the same transaction we change the displayed fields to be only the age.
+        let mut builder = Settings::new(&mut wtxn, &index, 0);
+        builder.set_displayed_fields(vec!["age".into()]);
+        builder.execute(|_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the displayed fields are correctly set to only the "age" field.
+        let rtxn = index.read_txn().unwrap();
+        let fields_ids = index.displayed_fields(&rtxn).unwrap();
+        assert_eq!(fields_ids.unwrap(), &["age"][..]);
+        drop(rtxn);
+
+        // We reset the fields ids to become `None`, the default value.
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = Settings::new(&mut wtxn, &index, 0);
+        builder.reset_displayed_fields();
+        builder.execute(|_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the displayed fields are correctly set to `None` (default value).
+        let rtxn = index.read_txn().unwrap();
+        let fields_ids = index.displayed_fields(&rtxn).unwrap();
+        assert_eq!(fields_ids, None);
+        drop(rtxn);
+    }
+
+    #[test]
+    fn set_faceted_fields() {
+        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();
+
+        // Set the faceted fields to be the age.
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = Settings::new(&mut wtxn, &index, 0);
+        builder.set_faceted_fields(hashmap!{ "age".into() => "integer".into() });
+        builder.execute(|_, _| ()).unwrap();
+
+        // Then index some documents.
+        let content = &b"name,age\nkevin,23\nkevina,21\nbenoit,34\n"[..];
+        let mut builder = IndexDocuments::new(&mut wtxn, &index, 1);
+        builder.update_format(UpdateFormat::Csv);
+        builder.execute(content, |_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the displayed fields are correctly set.
+        let rtxn = index.read_txn().unwrap();
+        let fields_ids = index.faceted_fields(&rtxn).unwrap();
+        assert_eq!(fields_ids, hashmap!{ "age".to_string() => FacetType::Integer });
+        // Only count the field_id 0 and level 0 facet values.
+        let count = index.facet_field_id_value_docids.prefix_iter(&rtxn, &[0, 0]).unwrap().count();
+        assert_eq!(count, 3);
+        drop(rtxn);
+
+        // Index a little more documents with new and current facets values.
+        let mut wtxn = index.write_txn().unwrap();
+        let content = &b"name,age\nkevin2,23\nkevina2,21\nbenoit2,35\n"[..];
+        let mut builder = IndexDocuments::new(&mut wtxn, &index, 2);
+        builder.update_format(UpdateFormat::Csv);
+        builder.execute(content, |_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        let rtxn = index.read_txn().unwrap();
+        // Only count the field_id 0 and level 0 facet values.
+        let count = index.facet_field_id_value_docids.prefix_iter(&rtxn, &[0, 0]).unwrap().count();
+        assert_eq!(count, 4);
+        drop(rtxn);
+    }
+
+    #[test]
+    fn setting_searchable_recomputes_other_settings() {
+        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();
+
+        // Set all the settings except searchable
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = Settings::new(&mut wtxn, &index, 0);
+        builder.set_displayed_fields(vec!["hello".to_string()]);
+        builder.set_faceted_fields(hashmap!{
+            "age".into() => "integer".into(),
+            "toto".into() => "integer".into(),
+        });
+        builder.set_criteria(vec!["asc(toto)".to_string()]);
+        builder.execute(|_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // check the output
+        let rtxn = index.read_txn().unwrap();
+        assert_eq!(&["hello"][..], index.displayed_fields(&rtxn).unwrap().unwrap());
+        // since no documents have been pushed the primary key is still unset
+        assert!(index.primary_key(&rtxn).unwrap().is_none());
+        assert_eq!(vec![Criterion::Asc("toto".to_string())], index.criteria(&rtxn).unwrap());
+        drop(rtxn);
+
+        // We set toto and age as searchable to force reordering of the fields
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = Settings::new(&mut wtxn, &index, 1);
+        builder.set_searchable_fields(vec!["toto".to_string(), "age".to_string()]);
+        builder.execute(|_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        let rtxn = index.read_txn().unwrap();
+        assert_eq!(&["hello"][..], index.displayed_fields(&rtxn).unwrap().unwrap());
+        assert!(index.primary_key(&rtxn).unwrap().is_none());
+        assert_eq!(vec![Criterion::Asc("toto".to_string())], index.criteria(&rtxn).unwrap());
+        drop(rtxn);
+    }
+}

milli/src/update/update_builder.rs

@@ -0,0 +1,138 @@
+use grenad::CompressionType;
+use rayon::ThreadPool;
+
+use crate::Index;
+use super::{ClearDocuments, DeleteDocuments, IndexDocuments, Settings, Facets};
+
+pub struct UpdateBuilder<'a> {
+    pub(crate) log_every_n: Option<usize>,
+    pub(crate) max_nb_chunks: Option<usize>,
+    pub(crate) max_memory: Option<usize>,
+    pub(crate) linked_hash_map_size: Option<usize>,
+    pub(crate) chunk_compression_type: CompressionType,
+    pub(crate) chunk_compression_level: Option<u32>,
+    pub(crate) chunk_fusing_shrink_size: Option<u64>,
+    pub(crate) thread_pool: Option<&'a ThreadPool>,
+    pub(crate) update_id: u64,
+}
+
+impl<'a> UpdateBuilder<'a> {
+    pub fn new(update_id: u64) -> UpdateBuilder<'a> {
+        UpdateBuilder {
+            log_every_n: None,
+            max_nb_chunks: None,
+            max_memory: None,
+            linked_hash_map_size: None,
+            chunk_compression_type: CompressionType::None,
+            chunk_compression_level: None,
+            chunk_fusing_shrink_size: None,
+            thread_pool: None,
+            update_id,
+        }
+    }
+
+    pub fn log_every_n(&mut self, log_every_n: usize) {
+        self.log_every_n = Some(log_every_n);
+    }
+
+    pub fn max_nb_chunks(&mut self, max_nb_chunks: usize) {
+        self.max_nb_chunks = Some(max_nb_chunks);
+    }
+
+    pub fn max_memory(&mut self, max_memory: usize) {
+        self.max_memory = Some(max_memory);
+    }
+
+    pub fn linked_hash_map_size(&mut self, linked_hash_map_size: usize) {
+        self.linked_hash_map_size = Some(linked_hash_map_size);
+    }
+
+    pub fn chunk_compression_type(&mut self, chunk_compression_type: CompressionType) {
+        self.chunk_compression_type = chunk_compression_type;
+    }
+
+    pub fn chunk_compression_level(&mut self, chunk_compression_level: u32) {
+        self.chunk_compression_level = Some(chunk_compression_level);
+    }
+
+    pub fn chunk_fusing_shrink_size(&mut self, chunk_fusing_shrink_size: u64) {
+        self.chunk_fusing_shrink_size = Some(chunk_fusing_shrink_size);
+    }
+
+    pub fn thread_pool(&mut self, thread_pool: &'a ThreadPool) {
+        self.thread_pool = Some(thread_pool);
+    }
+
+    pub fn clear_documents<'t, 'u, 'i>(
+        self,
+        wtxn: &'t mut heed::RwTxn<'i, 'u>,
+        index: &'i Index,
+    ) -> ClearDocuments<'t, 'u, 'i>
+    {
+        ClearDocuments::new(wtxn, index, self.update_id)
+    }
+
+    pub fn delete_documents<'t, 'u, 'i>(
+        self,
+        wtxn: &'t mut heed::RwTxn<'i, 'u>,
+        index: &'i Index,
+    ) -> anyhow::Result<DeleteDocuments<'t, 'u, 'i>>
+    {
+        DeleteDocuments::new(wtxn, index, self.update_id)
+    }
+
+    pub fn index_documents<'t, 'u, 'i>(
+        self,
+        wtxn: &'t mut heed::RwTxn<'i, 'u>,
+        index: &'i Index,
+    ) -> IndexDocuments<'t, 'u, 'i, 'a>
+    {
+        let mut builder = IndexDocuments::new(wtxn, index, self.update_id);
+
+        builder.log_every_n = self.log_every_n;
+        builder.max_nb_chunks = self.max_nb_chunks;
+        builder.max_memory = self.max_memory;
+        builder.linked_hash_map_size = self.linked_hash_map_size;
+        builder.chunk_compression_type = self.chunk_compression_type;
+        builder.chunk_compression_level = self.chunk_compression_level;
+        builder.chunk_fusing_shrink_size = self.chunk_fusing_shrink_size;
+        builder.thread_pool = self.thread_pool;
+
+        builder
+    }
+
+    pub fn settings<'t, 'u, 'i>(
+        self,
+        wtxn: &'t mut heed::RwTxn<'i, 'u>,
+        index: &'i Index,
+    ) -> Settings<'a, 't, 'u, 'i>
+    {
+        let mut builder = Settings::new(wtxn, index, self.update_id);
+
+        builder.log_every_n = self.log_every_n;
+        builder.max_nb_chunks = self.max_nb_chunks;
+        builder.max_memory = self.max_memory;
+        builder.linked_hash_map_size = self.linked_hash_map_size;
+        builder.chunk_compression_type = self.chunk_compression_type;
+        builder.chunk_compression_level = self.chunk_compression_level;
+        builder.chunk_fusing_shrink_size = self.chunk_fusing_shrink_size;
+        builder.thread_pool = self.thread_pool;
+
+        builder
+    }
+
+    pub fn facets<'t, 'u, 'i>(
+        self,
+        wtxn: &'t mut heed::RwTxn<'i, 'u>,
+        index: &'i Index,
+    ) -> Facets<'t, 'u, 'i>
+    {
+        let mut builder = Facets::new(wtxn, index, self.update_id);
+
+        builder.chunk_compression_type = self.chunk_compression_type;
+        builder.chunk_compression_level = self.chunk_compression_level;
+        builder.chunk_fusing_shrink_size = self.chunk_fusing_shrink_size;
+
+        builder
+    }
+}

milli/src/update/update_step.rs

@@ -0,0 +1,36 @@
+use UpdateIndexingStep::*;
+
+#[derive(Debug, Clone, Copy)]
+pub enum UpdateIndexingStep {
+    /// Transform from the original user given format (CSV, JSON, JSON lines)
+    /// into a generic format based on the obkv and grenad crates. This step also
+    /// deduplicate potential documents in this batch update by merging or replacing them.
+    TransformFromUserIntoGenericFormat { documents_seen: usize },
+
+    /// This step check the external document id, computes the internal ids and merge
+    /// the documents that are already present in the database.
+    ComputeIdsAndMergeDocuments { documents_seen: usize, total_documents: usize },
+
+    /// Extract the documents words using the tokenizer and compute the documents
+    /// facets. Stores those words, facets and documents ids on disk.
+    IndexDocuments { documents_seen: usize, total_documents: usize },
+
+    /// Merge the previously extracted data (words and facets) into the final LMDB database.
+    /// These extracted data are split into multiple databases.
+    MergeDataIntoFinalDatabase { databases_seen: usize, total_databases: usize },
+}
+
+impl UpdateIndexingStep {
+    pub const fn step(&self) -> usize {
+        match self {
+            TransformFromUserIntoGenericFormat { .. } => 0,
+            ComputeIdsAndMergeDocuments { .. } => 1,
+            IndexDocuments { .. } => 2,
+            MergeDataIntoFinalDatabase { .. } => 3,
+        }
+    }
+
+    pub const fn number_of_steps(&self) -> usize {
+        4
+    }
+}