MeiliSearch/milli/src/search/distinct/mod.rs

156 lines
5.5 KiB
Rust
Raw Normal View History

mod facet_distinct;
mod noop_distinct;
2021-06-16 18:33:33 +02:00
pub use facet_distinct::FacetDistinct;
pub use noop_distinct::NoopDistinct;
use roaring::RoaringBitmap;
use crate::{DocumentId, Result};
2021-04-14 12:00:45 +02:00
/// A trait implemented by document interators that are returned by calls to `Distinct::distinct`.
/// It provides a way to get back the ownership to the excluded set.
pub trait DocIter: Iterator<Item = Result<DocumentId>> {
2021-04-14 12:00:45 +02:00
/// Returns ownership on the internal exluded set.
fn into_excluded(self) -> RoaringBitmap;
}
2021-04-14 12:00:45 +02:00
/// A trait that is implemented by structs that perform a distinct on `candidates`. Calling distinct
/// must return an iterator containing only distinct documents, and add the discarded documents to
/// the excluded set. The excluded set can later be retrieved by calling `DocIter::excluded` on the
/// returned iterator.
pub trait Distinct {
type Iter: DocIter;
fn distinct(&mut self, candidates: RoaringBitmap, excluded: RoaringBitmap) -> Self::Iter;
}
2021-04-15 15:29:37 +02:00
#[cfg(test)]
mod test {
2021-05-03 15:58:47 +02:00
use std::collections::HashSet;
use std::io::Cursor;
2021-04-15 15:29:37 +02:00
use once_cell::sync::Lazy;
2021-06-16 18:33:33 +02:00
use rand::seq::SliceRandom;
use rand::Rng;
2021-04-15 15:29:37 +02:00
use roaring::RoaringBitmap;
use serde_json::{json, Value};
use crate::documents::{DocumentsBatchBuilder, DocumentsBatchReader};
2021-06-16 18:33:33 +02:00
use crate::index::tests::TempIndex;
use crate::index::Index;
use crate::update::{
IndexDocuments, IndexDocumentsConfig, IndexDocumentsMethod, IndexerConfig, Settings,
};
2021-06-16 18:33:33 +02:00
use crate::{DocumentId, FieldId, BEU32};
2021-04-15 15:29:37 +02:00
static JSON: Lazy<Vec<u8>> = Lazy::new(|| {
2021-04-15 15:29:37 +02:00
let mut rng = rand::thread_rng();
let num_docs = rng.gen_range(10..30);
let mut builder = DocumentsBatchBuilder::new(Vec::new());
let txts = ["Toto", "Titi", "Tata"];
2021-04-15 15:29:37 +02:00
let cats = (1..10).map(|i| i.to_string()).collect::<Vec<_>>();
let cat_ints = (1..10).collect::<Vec<_>>();
for i in 0..num_docs {
let txt = txts.choose(&mut rng).unwrap();
let mut sample_txts = cats.clone();
sample_txts.shuffle(&mut rng);
let mut sample_ints = cat_ints.clone();
sample_ints.shuffle(&mut rng);
let json = json!({
2021-04-15 15:29:37 +02:00
"id": i,
"txt": txt,
"cat-int": rng.gen_range(0..3),
"txts": sample_txts[..(rng.gen_range(0..3))],
"cat-ints": sample_ints[..(rng.gen_range(0..3))],
});
2021-10-24 14:41:36 +02:00
let object = match json {
Value::Object(object) => object,
_ => panic!(),
};
builder.append_json_object(&object).unwrap();
2021-04-15 15:29:37 +02:00
}
builder.into_inner().unwrap()
});
2021-04-15 15:29:37 +02:00
/// Returns a temporary index populated with random test documents, the FieldId for the
/// distinct attribute, and the RoaringBitmap with the document ids.
pub(crate) fn generate_index(distinct: &str) -> (TempIndex, FieldId, RoaringBitmap) {
2021-04-15 15:29:37 +02:00
let index = TempIndex::new();
let mut txn = index.write_txn().unwrap();
// set distinct and faceted attributes for the index.
let config = IndexerConfig::default();
let mut update = Settings::new(&mut txn, &index, &config);
update.set_distinct_field(distinct.to_string());
update.execute(|_| (), || false).unwrap();
2021-04-15 15:29:37 +02:00
// add documents to the index
let config = IndexerConfig::default();
let indexing_config = IndexDocumentsConfig {
update_method: IndexDocumentsMethod::ReplaceDocuments,
..Default::default()
};
2022-06-15 14:35:19 +02:00
let addition =
IndexDocuments::new(&mut txn, &index, &config, indexing_config, |_| (), || false)
.unwrap();
2021-04-15 15:29:37 +02:00
let reader =
crate::documents::DocumentsBatchReader::from_reader(Cursor::new(JSON.as_slice()))
.unwrap();
2022-06-15 14:35:19 +02:00
let (addition, user_error) = addition.add_documents(reader).unwrap();
user_error.unwrap();
addition.execute().unwrap();
2021-04-15 15:29:37 +02:00
let fields_map = index.fields_ids_map(&txn).unwrap();
let fid = fields_map.id(&distinct).unwrap();
let documents = DocumentsBatchReader::from_reader(Cursor::new(JSON.as_slice())).unwrap();
let map = (0..documents.documents_count() as u32).collect();
2021-04-15 15:29:37 +02:00
txn.commit().unwrap();
(index, fid, map)
}
/// Checks that all the candidates are distinct, and returns the candidates number.
pub(crate) fn validate_distinct_candidates(
candidates: impl Iterator<Item = crate::Result<DocumentId>>,
2021-04-15 15:29:37 +02:00
distinct: FieldId,
index: &Index,
2021-06-16 18:33:33 +02:00
) -> usize {
2021-04-15 15:29:37 +02:00
fn test(seen: &mut HashSet<String>, value: &Value) {
match value {
Value::Null | Value::Object(_) | Value::Bool(_) => (),
Value::Number(_) | Value::String(_) => {
let s = value.to_string();
assert!(seen.insert(s));
}
2021-06-16 18:33:33 +02:00
Value::Array(values) => values.into_iter().for_each(|value| test(seen, value)),
2021-04-15 15:29:37 +02:00
}
}
let mut seen = HashSet::<String>::new();
let txn = index.read_txn().unwrap();
let mut count = 0;
for candidate in candidates {
count += 1;
let candidate = candidate.unwrap();
let id = BEU32::new(candidate);
let document = index.documents.get(&txn, &id).unwrap().unwrap();
let value = document.get(distinct).unwrap();
let value = serde_json::from_slice(value).unwrap();
test(&mut seen, &value);
}
count
}
}