Refine the facet distribution to use both databases

This commit is contained in:
Clément Renault 2021-05-03 15:17:24 +02:00 committed by Kerollmops
parent 79efded841
commit 02c655ff1a
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4
4 changed files with 106 additions and 114 deletions

View File

@ -50,7 +50,6 @@ impl<'t> AscDesc<'t> {
Self::new(index, rtxn, parent, field_name, false) Self::new(index, rtxn, parent, field_name, false)
} }
fn new( fn new(
index: &'t Index, index: &'t Index,
rtxn: &'t heed::RoTxn, rtxn: &'t heed::RoTxn,
@ -59,7 +58,6 @@ impl<'t> AscDesc<'t> {
ascending: bool, ascending: bool,
) -> anyhow::Result<Self> { ) -> anyhow::Result<Self> {
let fields_ids_map = index.fields_ids_map(rtxn)?; let fields_ids_map = index.fields_ids_map(rtxn)?;
let faceted_fields = index.faceted_fields(rtxn)?;
let field_id = fields_ids_map let field_id = fields_ids_map
.id(&field_name) .id(&field_name)
.with_context(|| format!("field {:?} isn't registered", field_name))?; .with_context(|| format!("field {:?} isn't registered", field_name))?;

View File

@ -55,27 +55,45 @@ pub enum FacetCondition {
And(Box<Self>, Box<Self>), And(Box<Self>, Box<Self>),
} }
fn field_id<'a>( fn field_id(
fields_ids_map: &FieldsIdsMap, fields_ids_map: &FieldsIdsMap,
faceted_fields: &HashSet<FieldId>, faceted_fields: &HashSet<FieldId>,
items: &mut Pairs<'a, Rule>, items: &mut Pairs<Rule>,
) -> Result<FieldId, PestError<Rule>> ) -> Result<FieldId, PestError<Rule>>
{ {
// lexing ensures that we at least have a key // lexing ensures that we at least have a key
let key = items.next().unwrap(); let key = items.next().unwrap();
match fields_ids_map.id(key.as_str()) {
Some(field_id) => Ok(field_id), let field_id = match fields_ids_map.id(key.as_str()) {
None => Err(PestError::new_from_span( Some(field_id) => field_id,
None => return Err(PestError::new_from_span(
ErrorVariant::CustomError { ErrorVariant::CustomError {
message: format!( message: format!(
"attribute `{}` not found, available attributes are: {}", "attribute `{}` not found, available attributes are: {}",
key.as_str(), key.as_str(),
fields_ids_map.iter().map(|(_, n)| n).collect::<Vec<_>>().join(", ") fields_ids_map.iter().map(|(_, n)| n).collect::<Vec<_>>().join(", "),
), ),
}, },
key.as_span(), key.as_span(),
)), )),
};
if !faceted_fields.contains(&field_id) {
return Err(PestError::new_from_span(
ErrorVariant::CustomError {
message: format!(
"attribute `{}` is not faceted, available faceted attributes are: {}",
key.as_str(),
faceted_fields.iter().flat_map(|id| {
fields_ids_map.name(*id)
}).collect::<Vec<_>>().join(", "),
),
},
key.as_span(),
));
} }
Ok(field_id)
} }
fn pest_parse<T>(pair: Pair<Rule>) -> (Result<T, pest::error::Error<Rule>>, String) fn pest_parse<T>(pair: Pair<Rule>) -> (Result<T, pest::error::Error<Rule>>, String)
@ -84,12 +102,10 @@ where T: FromStr,
{ {
let result = match pair.as_str().parse::<T>() { let result = match pair.as_str().parse::<T>() {
Ok(value) => Ok(value), Ok(value) => Ok(value),
Err(e) => { Err(e) => Err(PestError::<Rule>::new_from_span(
Err(PestError::<Rule>::new_from_span(
ErrorVariant::CustomError { message: e.to_string() }, ErrorVariant::CustomError { message: e.to_string() },
pair.as_span(), pair.as_span(),
)) )),
}
}; };
(result, pair.as_str().to_string()) (result, pair.as_str().to_string())
@ -106,8 +122,6 @@ impl FacetCondition {
A: AsRef<str>, A: AsRef<str>,
B: AsRef<str>, B: AsRef<str>,
{ {
let fields_ids_map = index.fields_ids_map(rtxn)?;
let faceted_fields = index.faceted_fields(rtxn)?;
let mut ands = None; let mut ands = None;
for either in array { for either in array {
@ -202,7 +216,6 @@ impl FacetCondition {
item: Pair<Rule>, item: Pair<Rule>,
) -> anyhow::Result<FacetCondition> ) -> anyhow::Result<FacetCondition>
{ {
let item_span = item.as_span();
let mut items = item.into_inner(); let mut items = item.into_inner();
let fid = field_id(fields_ids_map, faceted_fields, &mut items)?; let fid = field_id(fields_ids_map, faceted_fields, &mut items)?;
@ -236,7 +249,6 @@ impl FacetCondition {
item: Pair<Rule>, item: Pair<Rule>,
) -> anyhow::Result<FacetCondition> ) -> anyhow::Result<FacetCondition>
{ {
let item_span = item.as_span();
let mut items = item.into_inner(); let mut items = item.into_inner();
let fid = field_id(fields_ids_map, faceted_fields, &mut items)?; let fid = field_id(fields_ids_map, faceted_fields, &mut items)?;
@ -252,7 +264,6 @@ impl FacetCondition {
item: Pair<Rule>, item: Pair<Rule>,
) -> anyhow::Result<FacetCondition> ) -> anyhow::Result<FacetCondition>
{ {
let item_span = item.as_span();
let mut items = item.into_inner(); let mut items = item.into_inner();
let fid = field_id(fields_ids_map, faceted_fields, &mut items)?; let fid = field_id(fields_ids_map, faceted_fields, &mut items)?;
@ -268,7 +279,6 @@ impl FacetCondition {
item: Pair<Rule>, item: Pair<Rule>,
) -> anyhow::Result<FacetCondition> ) -> anyhow::Result<FacetCondition>
{ {
let item_span = item.as_span();
let mut items = item.into_inner(); let mut items = item.into_inner();
let fid = field_id(fields_ids_map, faceted_fields, &mut items)?; let fid = field_id(fields_ids_map, faceted_fields, &mut items)?;
@ -284,7 +294,6 @@ impl FacetCondition {
item: Pair<Rule>, item: Pair<Rule>,
) -> anyhow::Result<FacetCondition> ) -> anyhow::Result<FacetCondition>
{ {
let item_span = item.as_span();
let mut items = item.into_inner(); let mut items = item.into_inner();
let fid = field_id(fields_ids_map, faceted_fields, &mut items)?; let fid = field_id(fields_ids_map, faceted_fields, &mut items)?;

View File

@ -3,12 +3,12 @@ use std::ops::Bound::Unbounded;
use std::{cmp, fmt}; use std::{cmp, fmt};
use anyhow::Context; use anyhow::Context;
use heed::BytesDecode; use heed::{Database, BytesDecode};
use heed::types::{ByteSlice, Unit};
use roaring::RoaringBitmap; use roaring::RoaringBitmap;
use crate::facet::{FacetType, FacetValue}; use crate::facet::FacetType;
use crate::heed_codec::facet::{FacetValueStringCodec, FacetLevelValueF64Codec}; use crate::heed_codec::facet::FacetValueStringCodec;
use crate::heed_codec::facet::{FieldDocIdFacetStringCodec, FieldDocIdFacetF64Codec};
use crate::search::facet::{FacetIter, FacetRange}; use crate::search::facet::{FacetIter, FacetRange};
use crate::{Index, FieldId, DocumentId}; use crate::{Index, FieldId, DocumentId};
@ -60,86 +60,81 @@ impl<'a> FacetDistribution<'a> {
/// There is a small amount of candidates OR we ask for facet string values so we /// There is a small amount of candidates OR we ask for facet string values so we
/// decide to iterate over the facet values of each one of them, one by one. /// decide to iterate over the facet values of each one of them, one by one.
fn facet_values_from_documents( fn facet_distribution_from_documents(
&self, &self,
field_id: FieldId, field_id: FieldId,
facet_type: FacetType, facet_type: FacetType,
candidates: &RoaringBitmap, candidates: &RoaringBitmap,
) -> heed::Result<BTreeMap<FacetValue, u64>> distribution: &mut BTreeMap<String, u64>,
) -> heed::Result<()>
{ {
fn fetch_facet_values<'t, KC, K: 't>( fn fetch_facet_values<'t, KC, K: 't>(
index: &Index,
rtxn: &'t heed::RoTxn, rtxn: &'t heed::RoTxn,
db: Database<KC, Unit>,
field_id: FieldId, field_id: FieldId,
candidates: &RoaringBitmap, candidates: &RoaringBitmap,
) -> heed::Result<BTreeMap<FacetValue, u64>> distribution: &mut BTreeMap<String, u64>,
) -> heed::Result<()>
where where
K: fmt::Display,
KC: BytesDecode<'t, DItem = (FieldId, DocumentId, K)>, KC: BytesDecode<'t, DItem = (FieldId, DocumentId, K)>,
K: Into<FacetValue>,
{ {
let mut facet_values = BTreeMap::new();
let mut key_buffer = vec![field_id]; let mut key_buffer = vec![field_id];
for docid in candidates.into_iter().take(CANDIDATES_THRESHOLD as usize) { for docid in candidates.into_iter().take(CANDIDATES_THRESHOLD as usize) {
key_buffer.truncate(1); key_buffer.truncate(1);
key_buffer.extend_from_slice(&docid.to_be_bytes()); key_buffer.extend_from_slice(&docid.to_be_bytes());
let iter = index.field_id_docid_facet_values let iter = db
.remap_key_type::<ByteSlice>()
.prefix_iter(rtxn, &key_buffer)? .prefix_iter(rtxn, &key_buffer)?
.remap_key_type::<KC>(); .remap_key_type::<KC>();
for result in iter { for result in iter {
let ((_, _, value), ()) = result?; let ((_, _, value), ()) = result?;
*facet_values.entry(value.into()).or_insert(0) += 1; *distribution.entry(value.to_string()).or_insert(0) += 1;
} }
} }
Ok(facet_values) Ok(())
} }
let index = self.index;
let rtxn = self.rtxn;
match facet_type { match facet_type {
FacetType::String => {
fetch_facet_values::<FieldDocIdFacetStringCodec, _>(index, rtxn, field_id, candidates)
},
FacetType::Number => { FacetType::Number => {
fetch_facet_values::<FieldDocIdFacetF64Codec, _>(index, rtxn, field_id, candidates) let db = self.index.field_id_docid_facet_f64s;
fetch_facet_values(self.rtxn, db, field_id, candidates, distribution)
}, },
FacetType::String => {
let db = self.index.field_id_docid_facet_strings;
fetch_facet_values(self.rtxn, db, field_id, candidates, distribution)
}
} }
} }
/// There is too much documents, we use the facet levels to move throught /// There is too much documents, we use the facet levels to move throught
/// the facet values, to find the candidates and values associated. /// the facet values, to find the candidates and values associated.
fn facet_values_from_facet_levels( fn facet_numbers_distribution_from_facet_levels(
&self, &self,
field_id: FieldId, field_id: FieldId,
facet_type: FacetType,
candidates: &RoaringBitmap, candidates: &RoaringBitmap,
) -> heed::Result<BTreeMap<FacetValue, u64>> distribution: &mut BTreeMap<String, u64>,
) -> heed::Result<()>
{ {
let iter = match facet_type {
FacetType::String => unreachable!(),
FacetType::Number => {
let iter = FacetIter::new_non_reducing( let iter = FacetIter::new_non_reducing(
self.rtxn, self.index, field_id, candidates.clone(), self.rtxn, self.index, field_id, candidates.clone(),
)?; )?;
iter.map(|r| r.map(|(v, docids)| (FacetValue::from(v), docids)))
},
};
let mut facet_values = BTreeMap::new();
for result in iter { for result in iter {
let (value, mut docids) = result?; let (value, mut docids) = result?;
docids.intersect_with(candidates); docids.intersect_with(candidates);
if !docids.is_empty() { if !docids.is_empty() {
facet_values.insert(value, docids.len()); distribution.insert(value.to_string(), docids.len());
} }
if facet_values.len() == self.max_values_by_facet { if distribution.len() == self.max_values_by_facet {
break; break;
} }
} }
Ok(facet_values) Ok(())
} }
/// Placeholder search, a.k.a. no candidates were specified. We iterate throught the /// Placeholder search, a.k.a. no candidates were specified. We iterate throught the
@ -147,80 +142,73 @@ impl<'a> FacetDistribution<'a> {
fn facet_values_from_raw_facet_database( fn facet_values_from_raw_facet_database(
&self, &self,
field_id: FieldId, field_id: FieldId,
facet_type: FacetType, ) -> heed::Result<BTreeMap<String, u64>>
) -> heed::Result<BTreeMap<FacetValue, u64>>
{ {
let db = self.index.facet_field_id_value_docids; let mut distribution = BTreeMap::new();
let level = 0;
let iter = match facet_type {
FacetType::String => {
let iter = db
.prefix_iter(self.rtxn, &[field_id])?
.remap_key_type::<FacetValueStringCodec>()
.map(|r| r.map(|((_, v), docids)| (FacetValue::from(v), docids)));
Box::new(iter) as Box::<dyn Iterator<Item=_>>
},
FacetType::Number => {
let db = db.remap_key_type::<FacetLevelValueF64Codec>();
let range = FacetRange::new(
self.rtxn, db, field_id, level, Unbounded, Unbounded,
)?;
Box::new(range.map(|r| r.map(|((_, _, v, _), docids)| (FacetValue::from(v), docids))))
},
};
let mut facet_values = BTreeMap::new(); let db = self.index.facet_id_f64_docids;
for result in iter { let range = FacetRange::new(self.rtxn, db, field_id, 0, Unbounded, Unbounded)?;
let (value, docids) = result?;
facet_values.insert(value, docids.len()); for result in range {
if facet_values.len() == self.max_values_by_facet { let ((_, _, value, _), docids) = result?;
distribution.insert(value.to_string(), docids.len());
if distribution.len() == self.max_values_by_facet {
break; break;
} }
} }
Ok(facet_values) let iter = self.index
.facet_id_string_docids
.remap_key_type::<ByteSlice>()
.prefix_iter(self.rtxn, &[field_id])?
.remap_key_type::<FacetValueStringCodec>();
for result in iter {
let ((_, value), docids) = result?;
distribution.insert(value.to_string(), docids.len());
if distribution.len() == self.max_values_by_facet {
break;
}
} }
fn facet_values( Ok(distribution)
&self, }
field_id: FieldId,
facet_type: FacetType, fn facet_values(&self, field_id: FieldId) -> heed::Result<BTreeMap<String, u64>> {
) -> heed::Result<BTreeMap<FacetValue, u64>> use FacetType::{Number, String};
{
if let Some(candidates) = self.candidates.as_ref() { if let Some(candidates) = self.candidates.as_ref() {
// Classic search, candidates were specified, we must return facet values only related // Classic search, candidates were specified, we must return facet values only related
// to those candidates. We also enter here for facet strings for performance reasons. // to those candidates. We also enter here for facet strings for performance reasons.
if candidates.len() <= CANDIDATES_THRESHOLD || facet_type == FacetType::String { let mut distribution = BTreeMap::new();
self.facet_values_from_documents(field_id, facet_type, candidates) if candidates.len() <= CANDIDATES_THRESHOLD {
self.facet_distribution_from_documents(field_id, Number, candidates, &mut distribution)?;
self.facet_distribution_from_documents(field_id, String, candidates, &mut distribution)?;
} else { } else {
self.facet_values_from_facet_levels(field_id, facet_type, candidates) self.facet_numbers_distribution_from_facet_levels(field_id, candidates, &mut distribution)?;
self.facet_distribution_from_documents(field_id, String, candidates, &mut distribution)?;
} }
Ok(distribution)
} else { } else {
self.facet_values_from_raw_facet_database(field_id, facet_type) self.facet_values_from_raw_facet_database(field_id)
} }
} }
pub fn execute(&self) -> anyhow::Result<BTreeMap<String, BTreeMap<FacetValue, u64>>> { pub fn execute(&self) -> anyhow::Result<BTreeMap<String, BTreeMap<String, u64>>> {
let fields_ids_map = self.index.fields_ids_map(self.rtxn)?; let fields_ids_map = self.index.fields_ids_map(self.rtxn)?;
let faceted_fields = self.index.faceted_fields(self.rtxn)?; let faceted_fields = self.index.faceted_fields(self.rtxn)?;
let fields_ids: Vec<_> = match &self.facets {
Some(names) => names
.iter()
.filter_map(|n| faceted_fields.get(n).map(|t| (n.to_string(), *t)))
.collect(),
None => faceted_fields.into_iter().collect(),
};
let mut facets_values = BTreeMap::new(); let mut distribution = BTreeMap::new();
for (name, ftype) in fields_ids { for name in faceted_fields {
let fid = fields_ids_map.id(&name).with_context(|| { let fid = fields_ids_map.id(&name).with_context(|| {
format!("missing field name {:?} from the fields id map", name) format!("missing field name {:?} from the fields id map", name)
})?; })?;
let values = self.facet_values(fid, ftype)?; let values = self.facet_values(fid)?;
facets_values.insert(name, values); distribution.insert(name, values);
} }
Ok(facets_values) Ok(distribution)
} }
} }

View File

@ -141,19 +141,16 @@ impl<'a> Search<'a> {
let field_ids_map = self.index.fields_ids_map(self.rtxn)?; let field_ids_map = self.index.fields_ids_map(self.rtxn)?;
let id = field_ids_map.id(name).expect("distinct not present in field map"); let id = field_ids_map.id(name).expect("distinct not present in field map");
let faceted_fields = self.index.faceted_fields(self.rtxn)?; let faceted_fields = self.index.faceted_fields(self.rtxn)?;
match faceted_fields.get(name) { if faceted_fields.contains(name) {
Some(facet_type) => {
let distinct = FacetDistinct::new(id, self.index, self.rtxn); let distinct = FacetDistinct::new(id, self.index, self.rtxn);
self.perform_sort(distinct, matching_words, criteria) self.perform_sort(distinct, matching_words, criteria)
} } else {
None => {
let distinct = MapDistinct::new(id, self.index, self.rtxn); let distinct = MapDistinct::new(id, self.index, self.rtxn);
self.perform_sort(distinct, matching_words, criteria) self.perform_sort(distinct, matching_words, criteria)
} }
} }
} }
} }
}
fn perform_sort( fn perform_sort(
&self, &self,