feat: add a method to get an iterator over all documents ids

This commit is contained in:
qdequele 2019-09-18 15:04:13 +02:00
parent 4f71219e17
commit a36c991897
No known key found for this signature in database
GPG Key ID: B3F0A000EBF11745
3 changed files with 53 additions and 0 deletions

View File

@ -55,6 +55,11 @@ impl DocumentsIndex {
Ok(DocumentFieldsIter(iter))
}
pub fn documents_ids(&self) -> RocksDbResult<DocumentsIdsIter> {
let iter = DocumentsKeysIter(self.0.iter()?);
Ok(DocumentsIdsIter { inner: iter, last: None })
}
pub fn documents_fields_repartition(&self, schema: Schema) -> RocksDbResult<HashMap<String, u64>> {
let iter = self.0.iter()?;
let mut repartition_attributes_id = HashMap::new();
@ -120,3 +125,22 @@ impl Iterator for DocumentsKeysIter<'_> {
}
}
}
pub struct DocumentsIdsIter<'a> {
inner: DocumentsKeysIter<'a>,
last: Option<DocumentId>,
}
impl Iterator for DocumentsIdsIter<'_> {
type Item = DocumentId;
fn next(&mut self) -> Option<Self::Item> {
for DocumentAttrKey { document_id, .. } in &mut self.inner {
if self.last != Some(document_id) {
self.last = Some(document_id);
return Some(document_id)
}
}
None
}
}

View File

@ -19,6 +19,7 @@ use crate::serde::{Deserializer, DeserializerError};
pub use self::custom_settings_index::{CustomSettingsIndex, RankingOrdering, StopWords, RankingOrder, DistinctField, RankingRules};
pub use self::common_index::CommonIndex;
pub use self::documents_index::DocumentsIdsIter;
use self::docs_words_index::DocsWordsIndex;
use self::documents_index::DocumentsIndex;
use self::main_index::MainIndex;
@ -374,6 +375,10 @@ impl Index {
Ok(self.update_status(update_id)?.unwrap())
}
pub fn documents_ids(&self) -> Result<DocumentsIdsIter, Error> {
Ok(self.documents_index.documents_ids()?)
}
pub fn document<T>(
&self,
fields: Option<&HashSet<&str>>,

View File

@ -189,3 +189,27 @@ fn custom_settings() {
assert_eq!(ret_distinct_field, distinct_field);
assert_eq!(ret_ranking_rules, ranking_rules);
}
#[test]
fn documents_ids() {
let tmp_dir = tempfile::tempdir().unwrap();
let database = Database::open(&tmp_dir).unwrap();
let schema = simple_schema();
let index = database.create_index("hello", schema).unwrap();
let doc1 = json!({ "objectId": 123, "title": "hello" });
let doc2 = json!({ "objectId": 456, "title": "world" });
let doc3 = json!({ "objectId": 789 });
let mut addition = index.documents_addition();
addition.update_document(&doc1);
addition.update_document(&doc2);
addition.update_document(&doc3);
let update_id = addition.finalize().unwrap();
let status = index.update_status_blocking(update_id).unwrap();
assert!(status.result.is_ok());
let documents_ids_count = index.documents_ids().unwrap().count();
assert_eq!(documents_ids_count, 3);
}