Improve the Mtbl heed codec to only encode MTBL databases

src/bin/indexer.rs

@@ -579,13 +579,16 @@ fn main() -> anyhow::Result<()> {
     builder.extend(docs_stores);
     builder.build().write_into(&mut writer)?;
     let file = writer.into_inner()?;
+
+    // Read back the documents MTBL database from the file.
     let documents_mmap = unsafe { memmap::Mmap::map(&file)? };
+    let documents = Reader::new(documents_mmap)?;
 
     debug!("We are writing the postings lists and documents into LMDB on disk...");
     // We merge the postings lists into LMDB.
     let mut wtxn = env.write_txn()?;
     merge_into_lmdb(stores, |k, v| lmdb_writer(&mut wtxn, &index, k, v))?;
-    index.put_documents(&mut wtxn, &documents_mmap)?;
+    index.put_documents(&mut wtxn, &documents)?;
     let count = index.number_of_documents(&wtxn)?;
     wtxn.commit()?;
 

src/heed_codec/mtbl_codec.rs

@@ -1,9 +1,10 @@
 use std::borrow::Cow;
+use std::marker::PhantomData;
 use oxidized_mtbl::Reader;
 
-pub struct MtblCodec;
+pub struct MtblCodec<A>(PhantomData<A>);
 
-impl<'a> heed::BytesDecode<'a> for MtblCodec {
+impl<'a> heed::BytesDecode<'a> for MtblCodec<&'a [u8]> {
     type DItem = Reader<&'a [u8]>;
 
     fn bytes_decode(bytes: &'a [u8]) -> Option<Self::DItem> {
@@ -11,10 +12,10 @@ impl<'a> heed::BytesDecode<'a> for MtblCodec {
     }
 }
 
-impl heed::BytesEncode<'_> for MtblCodec {
-    type EItem = [u8];
+impl<'a, A: AsRef<[u8]> + 'a> heed::BytesEncode<'a> for MtblCodec<A> {
+    type EItem = Reader<A>;
 
     fn bytes_encode(item: &Self::EItem) -> Option<Cow<[u8]>> {
-        Some(Cow::Borrowed(item))
+        Some(Cow::Borrowed(item.as_bytes()))
     }
 }

src/lib.rs

@@ -12,6 +12,7 @@ use anyhow::{bail, Context};
 use fxhash::{FxHasher32, FxHasher64};
 use heed::types::*;
 use heed::{PolyDatabase, Database};
+use oxidized_mtbl as omtbl;
 
 pub use self::search::{Search, SearchResult};
 pub use self::criterion::{Criterion, default_criteria};
@@ -90,7 +91,7 @@ impl Index {
         iter: impl IntoIterator<Item=DocumentId>,
     ) -> anyhow::Result<Vec<(DocumentId, Vec<u8>)>>
     {
-        match self.main.get::<_, Str, MtblCodec>(rtxn, DOCUMENTS_KEY)? {
+        match self.main.get::<_, Str, MtblCodec<&[u8]>>(rtxn, DOCUMENTS_KEY)? {
             Some(documents) => {
                 iter.into_iter().map(|id| {
                     let key = id.to_be_bytes();
@@ -103,13 +104,13 @@ impl Index {
         }
     }
 
-    pub fn put_documents(&self, wtxn: &mut heed::RwTxn, documents: &[u8]) -> anyhow::Result<()> {
-        Ok(self.main.put::<_, Str, MtblCodec>(wtxn, DOCUMENTS_KEY, documents)?)
+    pub fn put_documents<A: AsRef<[u8]>>(&self, wtxn: &mut heed::RwTxn, documents: &omtbl::Reader<A>) -> anyhow::Result<()> {
+        Ok(self.main.put::<_, Str, MtblCodec<A>>(wtxn, DOCUMENTS_KEY, documents)?)
     }
 
     /// Returns the number of documents indexed in the database.
     pub fn number_of_documents<'t>(&self, rtxn: &'t heed::RoTxn) -> anyhow::Result<usize> {
-        match self.main.get::<_, Str, MtblCodec>(rtxn, DOCUMENTS_KEY)? {
+        match self.main.get::<_, Str, MtblCodec<&[u8]>>(rtxn, DOCUMENTS_KEY)? {
             Some(documents) => Ok(documents.metadata().count_entries as usize),
             None => return Ok(0),
         }