Move the documents MTBL database inside the Index

src/lib.rs

@@ -2,12 +2,17 @@ mod best_proximity;
 mod heed_codec;
 mod iter_shortest_paths;
 mod query_tokens;
+mod transitive_arc;
 
 use std::borrow::Cow;
 use std::collections::{HashSet, HashMap};
+use std::fs::{File, OpenOptions};
 use std::hash::BuildHasherDefault;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
 use std::time::Instant;
 
+use anyhow::Context;
 use cow_utils::CowUtils;
 use fst::{IntoStreamer, Streamer};
 use fxhash::{FxHasher32, FxHasher64};
@@ -15,12 +20,15 @@ use heed::types::*;
 use heed::{PolyDatabase, Database};
 use levenshtein_automata::LevenshteinAutomatonBuilder as LevBuilder;
 use log::debug;
+use memmap::Mmap;
 use once_cell::sync::Lazy;
+use oxidized_mtbl as omtbl;
 use roaring::RoaringBitmap;
 
 use self::best_proximity::BestProximity;
 use self::heed_codec::RoaringBitmapCodec;
 use self::query_tokens::{QueryTokens, QueryToken};
+use self::transitive_arc::TransitiveArc;
 
 // Building these factories is not free.
 static LEVDIST0: Lazy<LevBuilder> = Lazy::new(|| LevBuilder::new(0, true));
@@ -39,6 +47,8 @@ pub type Position = u32;
 
 #[derive(Clone)]
 pub struct Index {
+    // The database path, where the LMDB and MTBL files are.
+    path: PathBuf,
     /// Contains many different types (e.g. the documents CSV headers).
     pub main: PolyDatabase,
     /// A word and all the positions where it appears in the whole dataset.
@@ -49,20 +59,40 @@ pub struct Index {
     pub prefix_word_position_docids: Database<ByteSlice, RoaringBitmapCodec>,
     /// Maps a word and an attribute (u32) to all the documents ids that it appears in.
     pub word_attribute_docids: Database<ByteSlice, RoaringBitmapCodec>,
+    /// The MTBL store that contains the documents content.
+    documents: omtbl::Reader<TransitiveArc<Mmap>>,
 }
 
 impl Index {
-    pub fn new(env: &heed::Env) -> heed::Result<Index> {
+    pub fn new<P: AsRef<Path>>(env: &heed::Env, path: P) -> anyhow::Result<Index> {
+        let documents_path = path.as_ref().join("documents.mtbl");
+        let mut documents = OpenOptions::new().create(true).write(true).read(true).open(documents_path)?;
+        // If the file is empty we must initialize it like an empty MTBL database.
+        if documents.metadata()?.len() == 0 {
+            omtbl::Writer::new(&mut documents).finish()?;
+        }
+        let documents = unsafe { memmap::Mmap::map(&documents)? };
+
         Ok(Index {
+            path: path.as_ref().to_path_buf(),
             main: env.create_poly_database(None)?,
             word_positions: env.create_database(Some("word-positions"))?,
             prefix_word_positions: env.create_database(Some("prefix-word-positions"))?,
             word_position_docids: env.create_database(Some("word-position-docids"))?,
             prefix_word_position_docids: env.create_database(Some("prefix-word-position-docids"))?,
             word_attribute_docids: env.create_database(Some("word-attribute-docids"))?,
+            documents: omtbl::Reader::new(TransitiveArc(Arc::new(documents)))?,
         })
     }
 
+    pub fn refresh_documents(&mut self) -> anyhow::Result<()> {
+        let documents_path = self.path.join("documents.mtbl");
+        let documents = File::open(&documents_path)?;
+        let documents = unsafe { memmap::Mmap::map(&documents)? };
+        self.documents = omtbl::Reader::new(TransitiveArc(Arc::new(documents)))?;
+        Ok(())
+    }
+
     pub fn put_headers(&self, wtxn: &mut heed::RwTxn, headers: &[u8]) -> anyhow::Result<()> {
         Ok(self.main.put::<_, Str, ByteSlice>(wtxn, "headers", headers)?)
     }
@@ -93,6 +123,21 @@ impl Index {
         }
     }
 
+    /// Returns a [`Vec`] of the requested documents. Returns an error if a document is missing.
+    pub fn documents<I: IntoIterator<Item=DocumentId>>(&self, iter: I) -> anyhow::Result<Vec<(DocumentId, Vec<u8>)>> {
+        iter.into_iter().map(|id| {
+            let key = id.to_be_bytes();
+            let content = self.documents.clone().get(&key)?.with_context(|| format!("Could not find document {}.", id))?;
+            Ok((id, content.as_ref().to_vec()))
+        })
+        .collect()
+    }
+
+    /// Returns the number of documents indexed in the database.
+    pub fn number_of_documents(&self) -> usize {
+        self.documents.metadata().count_entries as usize
+    }
+
     pub fn search(&self, rtxn: &heed::RoTxn, query: &str) -> anyhow::Result<(HashSet<String>, Vec<DocumentId>)> {
         let fst = match self.fst(rtxn)? {
             Some(fst) => fst,