Introduce ComputedFacets and compute_facet_distribution_stats

This commit is contained in:
Louis Dureuil 2024-09-12 17:49:03 +02:00
parent 7c084b1286
commit f6114a1ff2
No known key found for this signature in database

View File

@ -989,39 +989,13 @@ pub fn perform_search(
HitsInfo::OffsetLimit { limit, offset, estimated_total_hits: number_of_hits } HitsInfo::OffsetLimit { limit, offset, estimated_total_hits: number_of_hits }
}; };
let (facet_distribution, facet_stats) = match facets { let (facet_distribution, facet_stats) = facets
Some(ref fields) => { .map(move |facets| {
let mut facet_distribution = index.facets_distribution(&rtxn); compute_facet_distribution_stats(&facets, index, &rtxn, candidates, None, None)
})
let max_values_by_facet = index .transpose()?
.max_values_per_facet(&rtxn) .map(|ComputedFacets { distribution, stats }| (distribution, stats))
.map_err(milli::Error::from)? .unzip();
.map(|x| x as usize)
.unwrap_or(DEFAULT_VALUES_PER_FACET);
facet_distribution.max_values_per_facet(max_values_by_facet);
let sort_facet_values_by =
index.sort_facet_values_by(&rtxn).map_err(milli::Error::from)?;
if fields.iter().all(|f| f != "*") {
let fields: Vec<_> =
fields.iter().map(|n| (n, sort_facet_values_by.get(n))).collect();
facet_distribution.facets(fields);
}
let distribution = facet_distribution
.candidates(candidates)
.default_order_by(sort_facet_values_by.get("*"))
.execute()?;
let stats = facet_distribution.compute_stats()?;
(Some(distribution), Some(stats))
}
None => (None, None),
};
let facet_stats = facet_stats.map(|stats| {
stats.into_iter().map(|(k, (min, max))| (k, FacetStats { min, max })).collect()
});
let result = SearchResult { let result = SearchResult {
hits: documents, hits: documents,
@ -1037,6 +1011,55 @@ pub fn perform_search(
Ok(result) Ok(result)
} }
#[derive(Debug, Clone, Default, Serialize)]
pub struct ComputedFacets {
pub distribution: BTreeMap<String, IndexMap<String, u64>>,
pub stats: BTreeMap<String, FacetStats>,
}
fn compute_facet_distribution_stats<S: AsRef<str>>(
facets: &[S],
index: &Index,
rtxn: &RoTxn,
candidates: roaring::RoaringBitmap,
override_max_values_per_facet: Option<usize>,
override_sort_facet_values_by: Option<OrderBy>,
) -> Result<ComputedFacets, ResponseError> {
let mut facet_distribution = index.facets_distribution(rtxn);
let max_values_by_facet = match override_max_values_per_facet {
Some(max_values_by_facet) => max_values_by_facet,
None => index
.max_values_per_facet(rtxn)
.map_err(milli::Error::from)?
.map(|x| x as usize)
.unwrap_or(DEFAULT_VALUES_PER_FACET),
};
facet_distribution.max_values_per_facet(max_values_by_facet);
let sort_facet_values_by = index.sort_facet_values_by(rtxn).map_err(milli::Error::from)?;
let sort_facet_values_by = |n: &str| match override_sort_facet_values_by {
Some(order_by) => order_by,
None => sort_facet_values_by.get(n),
};
// add specific facet if there is no placeholder
if facets.iter().all(|f| f.as_ref() != "*") {
let fields: Vec<_> = facets.iter().map(|n| (n, sort_facet_values_by(n.as_ref()))).collect();
facet_distribution.facets(fields);
}
let distribution = facet_distribution
.candidates(candidates)
.default_order_by(sort_facet_values_by("*"))
.execute()?;
let stats = facet_distribution.compute_stats()?;
let stats = stats.into_iter().map(|(k, (min, max))| (k, FacetStats { min, max })).collect();
Ok(ComputedFacets { distribution, stats })
}
pub fn search_from_kind( pub fn search_from_kind(
search_kind: SearchKind, search_kind: SearchKind,
search: milli::Search<'_>, search: milli::Search<'_>,