From 2006259a23c5081e40aa863df135a2d202af362f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <renault.cle@gmail.com>
Date: Sat, 14 Sep 2019 12:26:47 +0200
Subject: [PATCH] feat: Improve the number of documents counting

---
 meilidb-data/src/database/index/documents_index.rs  |  2 +-
 meilidb-data/src/database/index/mod.rs              | 13 ++++++++++---
 .../src/database/update/documents_addition.rs       |  2 +-
 .../src/database/update/documents_deletion.rs       |  2 +-
 .../src/database/update/synonyms_addition.rs        |  3 ++-
 .../src/database/update/synonyms_deletion.rs        |  3 ++-
 6 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/meilidb-data/src/database/index/documents_index.rs b/meilidb-data/src/database/index/documents_index.rs
index 83e95530e..b209f1ed0 100644
--- a/meilidb-data/src/database/index/documents_index.rs
+++ b/meilidb-data/src/database/index/documents_index.rs
@@ -54,7 +54,7 @@ impl DocumentsIndex {
         Ok(DocumentFieldsIter(iter))
     }
 
-    pub fn len(&self) -> RocksDbResult<usize> {
+    pub fn len(&self) -> RocksDbResult<u64> {
         let mut last_document_id = None;
         let mut count = 0;
 
diff --git a/meilidb-data/src/database/index/mod.rs b/meilidb-data/src/database/index/mod.rs
index e9992cb55..e3f205624 100644
--- a/meilidb-data/src/database/index/mod.rs
+++ b/meilidb-data/src/database/index/mod.rs
@@ -164,7 +164,7 @@ fn last_update_id(
 #[derive(Copy, Clone)]
 pub struct IndexStats {
     pub number_of_words: usize,
-    pub number_of_documents: usize,
+    pub number_of_documents: u64,
     pub number_attrs_in_ranked_map: usize,
 }
 
@@ -192,6 +192,7 @@ pub(crate) struct Cache {
     pub synonyms: Arc<fst::Set>,
     pub schema: Schema,
     pub ranked_map: RankedMap,
+    pub number_of_documents: u64,
 }
 
 impl Index {
@@ -241,7 +242,9 @@ impl Index {
             None => RankedMap::default(),
         };
 
-        let cache = Cache { words, synonyms, schema, ranked_map };
+        let number_of_documents = documents_index.len()?;
+
+        let cache = Cache { words, synonyms, schema, ranked_map, number_of_documents };
         let cache = Arc::new(ArcSwap::from_pointee(cache));
 
         let last_update_id = last_update_id(&updates_index, &updates_results_index)?;
@@ -280,7 +283,7 @@ impl Index {
         let cache = self.cache.load();
         Ok(IndexStats {
             number_of_words: cache.words.len(),
-            number_of_documents: self.documents_index.len()?,
+            number_of_documents: cache.number_of_documents,
             number_attrs_in_ranked_map: cache.ranked_map.len(),
         })
     }
@@ -319,6 +322,10 @@ impl Index {
         self.custom_settings_index.clone()
     }
 
+    pub fn number_of_documents(&self) -> u64 {
+        self.cache.load().number_of_documents
+    }
+
     pub fn documents_addition<D>(&self) -> DocumentsAddition<D> {
         DocumentsAddition::new(self)
     }
diff --git a/meilidb-data/src/database/update/documents_addition.rs b/meilidb-data/src/database/update/documents_addition.rs
index 3daeb278c..0ff04221c 100644
--- a/meilidb-data/src/database/update/documents_addition.rs
+++ b/meilidb-data/src/database/update/documents_addition.rs
@@ -132,7 +132,7 @@ pub fn apply_documents_addition(
     let synonyms = cache.synonyms.clone();
     let schema = cache.schema.clone();
 
-    let cache = Cache { words, synonyms, schema, ranked_map };
+    let cache = Cache { words, synonyms, schema, ranked_map, number_of_documents };
     index.cache.store(Arc::new(cache));
 
     Ok(())
diff --git a/meilidb-data/src/database/update/documents_deletion.rs b/meilidb-data/src/database/update/documents_deletion.rs
index 7c7526a64..99dc35334 100644
--- a/meilidb-data/src/database/update/documents_deletion.rs
+++ b/meilidb-data/src/database/update/documents_deletion.rs
@@ -143,7 +143,7 @@ pub fn apply_documents_deletion(
     let synonyms = cache.synonyms.clone();
     let schema = cache.schema.clone();
 
-    let cache = Cache { words, synonyms, schema, ranked_map };
+    let cache = Cache { words, synonyms, schema, ranked_map, number_of_documents };
     index.cache.store(Arc::new(cache));
 
     Ok(())
diff --git a/meilidb-data/src/database/update/synonyms_addition.rs b/meilidb-data/src/database/update/synonyms_addition.rs
index 1995adc5a..95a650fb7 100644
--- a/meilidb-data/src/database/update/synonyms_addition.rs
+++ b/meilidb-data/src/database/update/synonyms_addition.rs
@@ -85,8 +85,9 @@ pub fn apply_synonyms_addition(
     let ranked_map = cache.ranked_map.clone();
     let synonyms = Arc::new(synonyms);
     let schema = cache.schema.clone();
+    let number_of_documents = cache.number_of_documents;
 
-    let cache = Cache { words, synonyms, schema, ranked_map };
+    let cache = Cache { words, synonyms, schema, ranked_map, number_of_documents };
     index.cache.store(Arc::new(cache));
 
     Ok(())
diff --git a/meilidb-data/src/database/update/synonyms_deletion.rs b/meilidb-data/src/database/update/synonyms_deletion.rs
index d472cc73c..bc036cb06 100644
--- a/meilidb-data/src/database/update/synonyms_deletion.rs
+++ b/meilidb-data/src/database/update/synonyms_deletion.rs
@@ -128,8 +128,9 @@ pub fn apply_synonyms_deletion(
     let ranked_map = cache.ranked_map.clone();
     let synonyms = Arc::new(synonyms);
     let schema = cache.schema.clone();
+    let number_of_documents = cache.number_of_documents;
 
-    let cache = Cache { words, synonyms, schema, ranked_map };
+    let cache = Cache { words, synonyms, schema, ranked_map, number_of_documents };
     index.cache.store(Arc::new(cache));
 
     Ok(())