use std::collections::{HashSet, BTreeMap}; use std::ops::Bound::Unbounded; use std::{cmp, fmt}; use anyhow::Context; use roaring::RoaringBitmap; use crate::facet::{FacetType, FacetValue}; use crate::heed_codec::facet::{FacetValueStringCodec, FacetLevelValueF64Codec, FacetLevelValueI64Codec}; use crate::heed_codec::facet::{FieldDocIdFacetStringCodec, FieldDocIdFacetF64Codec, FieldDocIdFacetI64Codec}; use crate::search::facet::{FacetIter, FacetRange}; use crate::{Index, FieldId}; /// The default number of values by facets that will /// be fetched from the key-value store. const DEFAULT_VALUES_BY_FACET: usize = 100; /// The hard limit in the number of values by facets that will be fetched from /// the key-value store. Searching for more values could slow down the engine. const MAX_VALUES_BY_FACET: usize = 1000; /// Threshold on the number of candidates that will make /// the system to choose between one algorithm or another. const CANDIDATES_THRESHOLD: u64 = 1000; pub struct FacetDistribution<'a> { facets: Option>, candidates: Option, max_values_by_facet: usize, rtxn: &'a heed::RoTxn<'a>, index: &'a Index, } impl<'a> FacetDistribution<'a> { pub fn new(rtxn: &'a heed::RoTxn, index: &'a Index) -> FacetDistribution<'a> { FacetDistribution { facets: None, candidates: None, max_values_by_facet: DEFAULT_VALUES_BY_FACET, rtxn, index, } } pub fn facets, A: AsRef>(&mut self, names: I) -> &mut Self { self.facets = Some(names.into_iter().map(|s| s.as_ref().to_string()).collect()); self } pub fn candidates(&mut self, candidates: RoaringBitmap) -> &mut Self { self.candidates = Some(candidates); self } pub fn max_values_by_facet(&mut self, max: usize) -> &mut Self { self.max_values_by_facet = cmp::min(max, MAX_VALUES_BY_FACET); self } fn facet_values( &self, field_id: FieldId, facet_type: FacetType, ) -> heed::Result> { if let Some(candidates) = self.candidates.as_ref() { // Classic search, candidates were specified, we must return // facet values only related to those candidates. if candidates.len() <= CANDIDATES_THRESHOLD || facet_type == FacetType::String { // 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. let mut key_buffer = vec![field_id]; match facet_type { FacetType::String => { let mut facet_values = BTreeMap::new(); for docid in candidates.into_iter().take(CANDIDATES_THRESHOLD as usize) { key_buffer.truncate(1); key_buffer.extend_from_slice(&docid.to_be_bytes()); let iter = self.index.field_id_docid_facet_values .prefix_iter(self.rtxn, &key_buffer)? .remap_key_type::(); for result in iter { let ((_, _, value), ()) = result?; *facet_values.entry(FacetValue::from(value)).or_insert(0) += 1; } } Ok(facet_values) }, FacetType::Float => { let mut facet_values = BTreeMap::new(); for docid in candidates { key_buffer.truncate(1); key_buffer.extend_from_slice(&docid.to_be_bytes()); let iter = self.index.field_id_docid_facet_values .prefix_iter(self.rtxn, &key_buffer)? .remap_key_type::(); for result in iter { let ((_, _, value), ()) = result?; *facet_values.entry(FacetValue::from(value)).or_insert(0) += 1; } } Ok(facet_values) }, FacetType::Integer => { let mut facet_values = BTreeMap::new(); for docid in candidates { key_buffer.truncate(1); key_buffer.extend_from_slice(&docid.to_be_bytes()); let iter = self.index.field_id_docid_facet_values .prefix_iter(self.rtxn, &key_buffer)? .remap_key_type::(); for result in iter { let ((_, _, value), ()) = result?; *facet_values.entry(FacetValue::from(value)).or_insert(0) += 1; } } Ok(facet_values) }, } } else { // There is too much documents, we use the facet levels to move throught // the facet values, to find the candidates and values associated. let iter = match facet_type { FacetType::String => unreachable!(), FacetType::Float => { let iter = FacetIter::::new_non_reducing( self.rtxn, self.index, field_id, candidates.clone(), )?; let iter = iter.map(|r| r.map(|(v, docids)| (FacetValue::from(v), docids))); Box::new(iter) as Box::> }, FacetType::Integer => { let iter = FacetIter::::new_non_reducing( self.rtxn, self.index, field_id, candidates.clone(), )?; Box::new(iter.map(|r| r.map(|(v, docids)| (FacetValue::from(v), docids)))) }, }; let mut facet_values = BTreeMap::new(); for result in iter { let (value, mut docids) = result?; docids.intersect_with(candidates); if !docids.is_empty() { facet_values.insert(value, docids.len()); } if facet_values.len() == self.max_values_by_facet { break; } } Ok(facet_values) } } else { // Placeholder search, a.k.a. no candidates were specified. We iterate throught the // facet values one by one and iterate on the facet level 0 for numbers. let db = self.index.facet_field_id_value_docids; let iter = match facet_type { FacetType::String => { let iter = db .prefix_iter(self.rtxn, &[field_id])? .remap_key_type::() .map(|r| r.map(|((_, v), docids)| (FacetValue::from(v), docids))); Box::new(iter) as Box::> }, FacetType::Float => { let db = db.remap_key_type::(); let range = FacetRange::::new( self.rtxn, db, field_id, 0, Unbounded, Unbounded, )?; Box::new(range.map(|r| r.map(|((_, _, v, _), docids)| (FacetValue::from(v), docids)))) }, FacetType::Integer => { let db = db.remap_key_type::(); let range = FacetRange::::new( self.rtxn, db, field_id, 0, Unbounded, Unbounded, )?; Box::new(range.map(|r| r.map(|((_, _, v, _), docids)| (FacetValue::from(v), docids)))) }, }; let mut facet_values = BTreeMap::new(); for result in iter { let (value, docids) = result?; facet_values.insert(value, docids.len()); if facet_values.len() == self.max_values_by_facet { break; } } Ok(facet_values) } } pub fn execute(&self) -> anyhow::Result>> { let fields_ids_map = self.index.fields_ids_map(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(); for (name, ftype) in fields_ids { let fid = fields_ids_map.id(&name).with_context(|| { format!("missing field name {:?} from the fields id map", name) })?; let values = self.facet_values(fid, ftype)?; facets_values.insert(name, values); } Ok(facets_values) } } impl fmt::Debug for FacetDistribution<'_> { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { let FacetDistribution { facets, candidates, max_values_by_facet, rtxn: _, index: _, } = self; f.debug_struct("FacetDistribution") .field("facets", facets) .field("candidates", candidates) .field("max_values_by_facet", max_values_by_facet) .finish() } }