load index dump

meilisearch-http/src/index/dump.rs

@@ -0,0 +1,120 @@
+use std::{fs::{create_dir_all, File}, path::Path, sync::Arc};
+
+use anyhow::Context;
+use heed::RoTxn;
+use indexmap::IndexMap;
+use milli::update::{IndexDocumentsMethod, UpdateFormat::JsonStream};
+use serde::{Deserialize, Serialize};
+use anyhow::bail;
+
+use crate::option::IndexerOpts;
+
+use super::update_handler::UpdateHandler;
+use super::{Checked, Index, Settings};
+
+#[derive(Serialize, Deserialize)]
+struct DumpMeta {
+    settings: Settings<Checked>,
+    primary_key: Option<String>,
+}
+
+const META_FILE_NAME: &'static str = "meta.json";
+const DATA_FILE_NAME: &'static str = "documents.jsonl";
+
+impl Index {
+    pub fn dump(&self, path: impl AsRef<Path>) -> anyhow::Result<()> {
+        // acquire write txn make sure any ongoing write is finnished before we start.
+        let txn = self.env.write_txn()?;
+
+        self.dump_documents(&txn, &path)?;
+        self.dump_meta(&txn, &path)?;
+
+        Ok(())
+    }
+
+    fn dump_documents(&self, txn: &RoTxn, path: impl AsRef<Path>) -> anyhow::Result<()> {
+        println!("dumping documents");
+        let document_file_path = path.as_ref().join(DATA_FILE_NAME);
+        let mut document_file = File::create(&document_file_path)?;
+
+        let documents = self.all_documents(txn)?;
+        let fields_ids_map = self.fields_ids_map(txn)?;
+
+        // dump documents
+        let mut json_map = IndexMap::new();
+        for document in documents {
+            let (_, reader) = document?;
+
+            for (fid, bytes) in reader.iter() {
+                if let Some(name) = fields_ids_map.name(fid) {
+                    json_map.insert(name, serde_json::from_slice::<serde_json::Value>(bytes)?);
+                }
+            }
+
+            serde_json::to_writer(&mut document_file, &json_map)?;
+            std::io::Write::write(&mut document_file, b"\n")?;
+
+            json_map.clear();
+        }
+
+        Ok(())
+    }
+
+    fn dump_meta(&self, txn: &RoTxn, path: impl AsRef<Path>) -> anyhow::Result<()> {
+        println!("dumping settings");
+        let meta_file_path = path.as_ref().join(META_FILE_NAME);
+        let mut meta_file = File::create(&meta_file_path)?;
+
+        let settings = self.settings_txn(txn)?;
+        let primary_key = self.primary_key(txn)?.map(String::from);
+        let meta = DumpMeta { settings, primary_key };
+
+        serde_json::to_writer(&mut meta_file, &meta)?;
+
+        Ok(())
+    }
+
+    pub fn load_dump(
+        src: impl AsRef<Path>,
+        dst: impl AsRef<Path>,
+        size: u64,
+        indexing_options: &IndexerOpts,
+    ) -> anyhow::Result<()> {
+        let dir_name = src
+            .as_ref()
+            .file_name()
+            .with_context(|| format!("invalid dump index: {}", src.as_ref().display()))?;
+        let dst_dir_path = dst.as_ref().join(dir_name);
+        create_dir_all(&dst_dir_path)?;
+
+        let meta_path = src.as_ref().join(META_FILE_NAME);
+        let mut meta_file = File::open(meta_path)?;
+        let DumpMeta { settings, primary_key } = serde_json::from_reader(&mut meta_file)?;
+        let index = Self::open(&dst_dir_path, size as usize)?;
+        let mut txn = index.write_txn()?;
+
+        let handler = UpdateHandler::new(&indexing_options)?;
+
+        index.update_settings_txn(&mut txn, &settings, handler.update_builder(0))?;
+
+        let document_file_path = src.as_ref().join(DATA_FILE_NAME);
+        let document_file = File::open(&document_file_path)?;
+        index.update_documents_txn(
+            &mut txn,
+            JsonStream,
+            IndexDocumentsMethod::UpdateDocuments,
+            Some(document_file),
+            handler.update_builder(0),
+            primary_key.as_deref(),
+        )?;
+
+        txn.commit()?;
+
+        match Arc::try_unwrap(index.0) {
+            Ok(inner) => inner.prepare_for_closing().wait(),
+            Err(_) => bail!("Could not close index properly."),
+        }
+
+        Ok(())
+    }
+}

meilisearch-http/src/index/mod.rs

@@ -1,11 +1,9 @@
-use std::{collections::{BTreeSet, HashSet}, io::Write, marker::PhantomData, path::{Path, PathBuf}};
+use std::{collections::{BTreeSet, HashSet}, marker::PhantomData, path::Path};
 use std::ops::Deref;
 use std::sync::Arc;
-use std::fs::File;
 
 use anyhow::{bail, Context};
-use heed::RoTxn;
-use indexmap::IndexMap;
+use heed::{EnvOpenOptions, RoTxn};
 use milli::obkv_to_json;
 use serde_json::{Map, Value};
 
@@ -16,6 +14,8 @@ use serde::{de::Deserializer, Deserialize};
 
 mod search;
 mod updates;
+mod dump;
+pub mod update_handler;
 
 pub type Document = Map<String, Value>;
 
@@ -39,6 +39,14 @@ where
 }
 
 impl Index {
+    pub fn open(path: impl AsRef<Path>, size: usize) -> anyhow::Result<Self> {
+        std::fs::create_dir_all(&path)?;
+        let mut options = EnvOpenOptions::new();
+        options.map_size(size);
+        let index = milli::Index::new(options, &path)?;
+        Ok(Index(Arc::new(index)))
+    }
+
     pub fn settings(&self) -> anyhow::Result<Settings<Checked>> {
         let txn = self.read_txn()?;
         self.settings_txn(&txn)
@@ -167,57 +175,4 @@ impl Index {
         displayed_fields_ids.retain(|fid| attributes_to_retrieve_ids.contains(fid));
         Ok(displayed_fields_ids)
     }
-
-    pub fn dump(&self, path: PathBuf) -> anyhow::Result<()> {
-        // acquire write txn make sure any ongoing write is finnished before we start.
-        let txn = self.env.write_txn()?;
-
-        self.dump_documents(&txn, &path)?;
-        self.dump_meta(&txn, &path)?;
-
-        Ok(())
-    }
-
-    fn dump_documents(&self, txn: &RoTxn, path: impl AsRef<Path>) -> anyhow::Result<()> {
-        println!("dumping documents");
-        let document_file_path = path.as_ref().join("documents.jsonl");
-        let mut document_file = File::create(&document_file_path)?;
-
-        let documents = self.all_documents(txn)?;
-        let fields_ids_map = self.fields_ids_map(txn)?;
-
-        // dump documents
-        let mut json_map = IndexMap::new();
-        for document in documents {
-            let (_, reader) = document?;
-
-            for (fid, bytes) in reader.iter() {
-                if let Some(name) = fields_ids_map.name(fid) {
-                    json_map.insert(name, serde_json::from_slice::<serde_json::Value>(bytes)?);
-                }
-            }
-
-            serde_json::to_writer(&mut document_file, &json_map)?;
-            document_file.write(b"\n")?;
-
-            json_map.clear();
-        }
-
-        Ok(())
-    }
-
-    fn dump_meta(&self, txn: &RoTxn, path: impl AsRef<Path>) -> anyhow::Result<()> {
-        println!("dumping settings");
-        let meta_file_path = path.as_ref().join("meta.json");
-        let mut meta_file = File::create(&meta_file_path)?;
-
-        let settings = self.settings_txn(txn)?;
-        let json = serde_json::json!({
-            "settings": settings,
-        });
-
-        serde_json::to_writer(&mut meta_file, &json)?;
-
-        Ok(())
-    }
 }

meilisearch-http/src/index/update_handler.rs

@@ -0,0 +1,93 @@
+use std::fs::File;
+
+use crate::index::Index;
+use anyhow::Result;
+use grenad::CompressionType;
+use milli::update::UpdateBuilder;
+use rayon::ThreadPool;
+
+use crate::index_controller::UpdateMeta;
+use crate::index_controller::{Failed, Processed, Processing};
+use crate::option::IndexerOpts;
+
+pub struct UpdateHandler {
+    max_nb_chunks: Option<usize>,
+    chunk_compression_level: Option<u32>,
+    thread_pool: ThreadPool,
+    log_frequency: usize,
+    max_memory: usize,
+    linked_hash_map_size: usize,
+    chunk_compression_type: CompressionType,
+    chunk_fusing_shrink_size: u64,
+}
+
+impl UpdateHandler {
+    pub fn new(opt: &IndexerOpts) -> anyhow::Result<Self> {
+        let thread_pool = rayon::ThreadPoolBuilder::new()
+            .num_threads(opt.indexing_jobs.unwrap_or(0))
+            .build()?;
+        Ok(Self {
+            max_nb_chunks: opt.max_nb_chunks,
+            chunk_compression_level: opt.chunk_compression_level,
+            thread_pool,
+            log_frequency: opt.log_every_n,
+            max_memory: opt.max_memory.get_bytes() as usize,
+            linked_hash_map_size: opt.linked_hash_map_size,
+            chunk_compression_type: opt.chunk_compression_type,
+            chunk_fusing_shrink_size: opt.chunk_fusing_shrink_size.get_bytes(),
+        })
+    }
+
+    pub fn update_builder(&self, update_id: u64) -> UpdateBuilder {
+        // We prepare the update by using the update builder.
+        let mut update_builder = UpdateBuilder::new(update_id);
+        if let Some(max_nb_chunks) = self.max_nb_chunks {
+            update_builder.max_nb_chunks(max_nb_chunks);
+        }
+        if let Some(chunk_compression_level) = self.chunk_compression_level {
+            update_builder.chunk_compression_level(chunk_compression_level);
+        }
+        update_builder.thread_pool(&self.thread_pool);
+        update_builder.log_every_n(self.log_frequency);
+        update_builder.max_memory(self.max_memory);
+        update_builder.linked_hash_map_size(self.linked_hash_map_size);
+        update_builder.chunk_compression_type(self.chunk_compression_type);
+        update_builder.chunk_fusing_shrink_size(self.chunk_fusing_shrink_size);
+        update_builder
+    }
+
+    pub fn handle_update(
+        &self,
+        meta: Processing,
+        content: Option<File>,
+        index: Index,
+    ) -> Result<Processed, Failed> {
+        use UpdateMeta::*;
+
+        let update_id = meta.id();
+
+        let update_builder = self.update_builder(update_id);
+
+        let result = match meta.meta() {
+            DocumentsAddition {
+                method,
+                format,
+                primary_key,
+            } => index.update_documents(
+                *format,
+                *method,
+                content,
+                update_builder,
+                primary_key.as_deref(),
+            ),
+            ClearDocuments => index.clear_documents(update_builder),
+            DeleteDocuments => index.delete_documents(content, update_builder),
+            Settings(settings) => index.update_settings(settings, update_builder),
+        };
+
+        match result {
+            Ok(result) => Ok(meta.process(result)),
+            Err(e) => Err(meta.fail(e.to_string())),
+        }
+    }
+}