Return the facets of a placeholder facet-search sorted by count

This commit is contained in:
Clément Renault 2024-03-12 18:04:38 +01:00
parent d3a95ea2f6
commit 9f7a4fbfeb
No known key found for this signature in database
GPG Key ID: F250A4C4E3AE5F5F
2 changed files with 65 additions and 18 deletions

View File

@ -715,9 +715,8 @@ pub fn perform_facet_search(
let rtxn = index.read_txn()?; let rtxn = index.read_txn()?;
let (search, _, _, _) = prepare_search(index, &rtxn, &search_query, features, None)?; let (search, _, _, _) = prepare_search(index, &rtxn, &search_query, features, None)?;
let sort_by = index.sort_facet_values_by(&rtxn)?.get(&facet_name);
let mut facet_search = let mut facet_search =
SearchForFacetValues::new(facet_name, search, sort_by, search_query.hybrid.is_some()); SearchForFacetValues::new(facet_name, search, search_query.hybrid.is_some());
if let Some(facet_query) = &facet_query { if let Some(facet_query) = &facet_query {
facet_search.query(facet_query); facet_search.query(facet_query);
} }

View File

@ -1,3 +1,5 @@
use std::cmp::{Ordering, Reverse};
use std::collections::BinaryHeap;
use std::fmt; use std::fmt;
use std::ops::ControlFlow; use std::ops::ControlFlow;
@ -307,7 +309,6 @@ pub struct SearchForFacetValues<'a> {
facet: String, facet: String,
search_query: Search<'a>, search_query: Search<'a>,
max_values: usize, max_values: usize,
sort_by: OrderBy,
is_hybrid: bool, is_hybrid: bool,
} }
@ -315,7 +316,6 @@ impl<'a> SearchForFacetValues<'a> {
pub fn new( pub fn new(
facet: String, facet: String,
search_query: Search<'a>, search_query: Search<'a>,
sort_by: OrderBy,
is_hybrid: bool, is_hybrid: bool,
) -> SearchForFacetValues<'a> { ) -> SearchForFacetValues<'a> {
SearchForFacetValues { SearchForFacetValues {
@ -323,7 +323,6 @@ impl<'a> SearchForFacetValues<'a> {
facet, facet,
search_query, search_query,
max_values: DEFAULT_MAX_NUMBER_OF_VALUES_PER_FACET, max_values: DEFAULT_MAX_NUMBER_OF_VALUES_PER_FACET,
sort_by,
is_hybrid, is_hybrid,
} }
} }
@ -384,6 +383,8 @@ impl<'a> SearchForFacetValues<'a> {
.search_query .search_query
.execute_for_candidates(self.is_hybrid || self.search_query.vector.is_some())?; .execute_for_candidates(self.is_hybrid || self.search_query.vector.is_some())?;
let sort_by = index.sort_facet_values_by(rtxn)?.get(&self.facet);
match self.query.as_ref() { match self.query.as_ref() {
Some(query) => { Some(query) => {
let options = NormalizerOption { lossy: true, ..Default::default() }; let options = NormalizerOption { lossy: true, ..Default::default() };
@ -465,8 +466,11 @@ impl<'a> SearchForFacetValues<'a> {
} }
} }
None => { None => {
let mut results = vec![];
let prefix = FacetGroupKey { field_id: fid, level: 0, left_bound: "" }; let prefix = FacetGroupKey { field_id: fid, level: 0, left_bound: "" };
match sort_by {
OrderBy::Lexicographic => {
let mut results = vec![];
for result in index.facet_id_string_docids.prefix_iter(rtxn, &prefix)? { for result in index.facet_id_string_docids.prefix_iter(rtxn, &prefix)? {
let (FacetGroupKey { left_bound, .. }, FacetGroupValue { bitmap, .. }) = let (FacetGroupKey { left_bound, .. }, FacetGroupValue { bitmap, .. }) =
result?; result?;
@ -483,6 +487,36 @@ impl<'a> SearchForFacetValues<'a> {
} }
Ok(results) Ok(results)
} }
OrderBy::Count => {
let mut top_counts = BinaryHeap::new();
for result in index.facet_id_string_docids.prefix_iter(rtxn, &prefix)? {
let (FacetGroupKey { left_bound, .. }, FacetGroupValue { bitmap, .. }) =
result?;
let count = search_candidates.intersection_len(&bitmap);
if count != 0 {
let value = self
.one_original_value_of(fid, left_bound, bitmap.min().unwrap())?
.unwrap_or_else(|| left_bound.to_string());
if top_counts.len() >= self.max_values {
top_counts.pop();
}
// It is a max heap and we need to move the smallest counts at the
// top to be able to pop them when we reach the max_values limit.
top_counts.push(Reverse(FacetValueHit { value, count }));
}
}
// Convert the heap into a vec of hits by removing the Reverse wrapper.
// Hits are already in the right order as they were reversed and there
// are output in ascending order.
Ok(top_counts
.into_sorted_vec()
.into_iter()
.map(|Reverse(hit)| hit)
.collect())
}
}
}
} }
} }
@ -539,6 +573,20 @@ pub struct FacetValueHit {
pub count: u64, pub count: u64,
} }
impl PartialOrd for FacetValueHit {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for FacetValueHit {
fn cmp(&self, other: &Self) -> Ordering {
self.count.cmp(&other.count).then_with(|| self.value.cmp(&other.value))
}
}
impl Eq for FacetValueHit {}
#[cfg(test)] #[cfg(test)]
mod test { mod test {
#[allow(unused_imports)] #[allow(unused_imports)]