Reintroduce asc/desc functionality

This commit is contained in:
Loïc Lecrenier 2022-08-31 09:36:19 +02:00 committed by Loïc Lecrenier
parent bd2c0e1ab6
commit e570c23153
9 changed files with 60 additions and 63 deletions

View File

@ -6,7 +6,10 @@ use ordered_float::OrderedFloat;
use roaring::RoaringBitmap;
use super::{Criterion, CriterionParameters, CriterionResult};
use crate::heed_codec::facet::new::{FacetKeyCodec, MyByteSlice};
use crate::search::criteria::{resolve_query_tree, CriteriaBuilder};
use crate::search::facet::facet_sort_ascending::ascending_facet_sort;
use crate::search::facet::facet_sort_descending::descending_facet_sort;
// use crate::search::facet::FacetStringIter;
use crate::search::query_tree::Operation;
use crate::{FieldId, Index, Result};
@ -186,24 +189,22 @@ fn facet_ordered<'t>(
iterative_facet_string_ordered_iter(index, rtxn, field_id, is_ascending, candidates)?;
Ok(Box::new(number_iter.chain(string_iter).map(Ok)) as Box<dyn Iterator<Item = _>>)
} else {
todo!()
// let facet_number_fn = if is_ascending {
// FacetNumberIter::new_reducing
// } else {
// FacetNumberIter::new_reverse_reducing
// };
// let number_iter = facet_number_fn(rtxn, index, field_id, candidates.clone())?
// .map(|res| res.map(|(_, docids)| docids));
let make_iter = if is_ascending { ascending_facet_sort } else { descending_facet_sort };
// let facet_string_fn = if is_ascending {
// FacetStringIter::new_reducing
// } else {
// FacetStringIter::new_reverse_reducing
// };
// let string_iter = facet_string_fn(rtxn, index, field_id, candidates)?
// .map(|res| res.map(|(_, _, docids)| docids));
let number_iter = make_iter(
rtxn,
index.facet_id_f64_docids.remap_key_type::<FacetKeyCodec<MyByteSlice>>(),
field_id,
candidates.clone(),
)?;
let string_iter = make_iter(
rtxn,
index.facet_id_f64_docids.remap_key_type::<FacetKeyCodec<MyByteSlice>>(),
field_id,
candidates,
)?;
// Ok(Box::new(number_iter.chain(string_iter)))
Ok(Box::new(number_iter.chain(string_iter)))
}
}

View File

@ -138,7 +138,7 @@ impl<'a> FacetDistribution<'a> {
) -> heed::Result<()> {
facet_distribution_iter::iterate_over_facet_distribution(
self.rtxn,
&self.index.facet_id_f64_docids.remap_key_type::<FacetKeyCodec<MyByteSlice>>(),
self.index.facet_id_f64_docids.remap_key_type::<FacetKeyCodec<MyByteSlice>>(),
field_id,
candidates,
|facet_key, nbr_docids| {
@ -161,7 +161,7 @@ impl<'a> FacetDistribution<'a> {
) -> heed::Result<()> {
facet_distribution_iter::iterate_over_facet_distribution(
self.rtxn,
&self.index.facet_id_string_docids.remap_key_type::<FacetKeyCodec<MyByteSlice>>(),
self.index.facet_id_string_docids.remap_key_type::<FacetKeyCodec<MyByteSlice>>(),
field_id,
candidates,
|facet_key, nbr_docids| {

View File

@ -7,7 +7,7 @@ use super::{get_first_facet_value, get_highest_level};
pub fn iterate_over_facet_distribution<'t, CB>(
rtxn: &'t heed::RoTxn<'t>,
db: &'t heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
field_id: u16,
candidates: &RoaringBitmap,
callback: CB,
@ -17,7 +17,7 @@ where
{
let mut fd = FacetDistribution { rtxn, db, field_id, callback };
let highest_level =
get_highest_level(rtxn, &db.remap_key_type::<FacetKeyCodec<MyByteSlice>>(), field_id)?;
get_highest_level(rtxn, db.remap_key_type::<FacetKeyCodec<MyByteSlice>>(), field_id)?;
if let Some(first_bound) = get_first_facet_value::<MyByteSlice>(rtxn, db, field_id)? {
fd.iterate(candidates, highest_level, first_bound, usize::MAX)?;
@ -32,7 +32,7 @@ where
CB: FnMut(&'t [u8], u64) -> ControlFlow<()>,
{
rtxn: &'t heed::RoTxn<'t>,
db: &'t heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
field_id: u16,
callback: CB,
}

View File

@ -15,7 +15,7 @@ use super::get_last_facet_value;
pub fn find_docids_of_facet_within_bounds<'t, BoundCodec>(
rtxn: &'t heed::RoTxn<'t>,
db: &'t heed::Database<FacetKeyCodec<BoundCodec>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<BoundCodec>, FacetGroupValueCodec>,
field_id: u16,
left: &'t Bound<<BoundCodec as BytesEncode<'t>>::EItem>,
right: &'t Bound<<BoundCodec as BytesEncode<'t>>::EItem>,
@ -50,11 +50,11 @@ where
};
let db = db.remap_key_type::<FacetKeyCodec<MyByteSlice>>();
let mut docids = RoaringBitmap::new();
let mut f = FacetRangeSearch { rtxn, db: &db, field_id, left, right, docids: &mut docids };
let highest_level = get_highest_level(rtxn, &db, field_id)?;
let mut f = FacetRangeSearch { rtxn, db, field_id, left, right, docids: &mut docids };
let highest_level = get_highest_level(rtxn, db, field_id)?;
if let Some(first_bound) = get_first_facet_value::<MyByteSlice>(rtxn, &db, field_id)? {
let last_bound = get_last_facet_value::<MyByteSlice>(rtxn, &db, field_id)?.unwrap();
if let Some(first_bound) = get_first_facet_value::<MyByteSlice>(rtxn, db, field_id)? {
let last_bound = get_last_facet_value::<MyByteSlice>(rtxn, db, field_id)?.unwrap();
f.run(highest_level, first_bound, Bound::Included(last_bound), usize::MAX)?;
Ok(docids)
} else {
@ -65,7 +65,7 @@ where
/// Fetch the document ids that have a facet with a value between the two given bounds
struct FacetRangeSearch<'t, 'b, 'bitmap> {
rtxn: &'t heed::RoTxn<'t>,
db: &'t heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
field_id: u16,
left: Bound<&'b [u8]>,
right: Bound<&'b [u8]>,

View File

@ -1,24 +1,19 @@
use crate::heed_codec::facet::new::{
FacetGroupValue, FacetGroupValueCodec, FacetKey, FacetKeyCodec, MyByteSlice,
};
use crate::Result;
use heed::Result;
use roaring::RoaringBitmap;
use super::{get_first_facet_value, get_highest_level};
pub fn ascending_facet_sort<'t>(
rtxn: &'t heed::RoTxn<'t>,
db: &'t heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
field_id: u16,
candidates: RoaringBitmap,
) -> Result<Box<dyn Iterator<Item = Result<(&'t [u8], RoaringBitmap)>> + 't>> {
let highest_level =
get_highest_level(rtxn, &db.remap_key_type::<FacetKeyCodec<MyByteSlice>>(), field_id)?;
if let Some(first_bound) = get_first_facet_value::<MyByteSlice>(
rtxn,
&db.remap_key_type::<FacetKeyCodec<MyByteSlice>>(),
field_id,
)? {
) -> Result<Box<dyn Iterator<Item = Result<RoaringBitmap>> + 't>> {
let highest_level = get_highest_level(rtxn, db, field_id)?;
if let Some(first_bound) = get_first_facet_value::<MyByteSlice>(rtxn, db, field_id)? {
let first_key = FacetKey { field_id, level: highest_level, left_bound: first_bound };
let iter = db.range(rtxn, &(first_key..)).unwrap().take(usize::MAX);
@ -30,7 +25,7 @@ pub fn ascending_facet_sort<'t>(
struct AscendingFacetSort<'t, 'e> {
rtxn: &'t heed::RoTxn<'e>,
db: &'t heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
field_id: u16,
stack: Vec<(
RoaringBitmap,
@ -39,7 +34,7 @@ struct AscendingFacetSort<'t, 'e> {
}
impl<'t, 'e> Iterator for AscendingFacetSort<'t, 'e> {
type Item = Result<(&'t [u8], RoaringBitmap)>;
type Item = Result<RoaringBitmap>;
fn next(&mut self) -> Option<Self::Item> {
'outer: loop {
@ -67,7 +62,7 @@ impl<'t, 'e> Iterator for AscendingFacetSort<'t, 'e> {
*documents_ids -= &bitmap;
if level == 0 {
return Some(Ok((left_bound, bitmap)));
return Some(Ok(bitmap));
}
let starting_key_below =
FacetKey { field_id: self.field_id, level: level - 1, left_bound };

View File

@ -3,17 +3,17 @@ use std::ops::Bound;
use crate::heed_codec::facet::new::{
FacetGroupValue, FacetGroupValueCodec, FacetKey, FacetKeyCodec, MyByteSlice,
};
use crate::Result;
use heed::Result;
use roaring::RoaringBitmap;
use super::{get_first_facet_value, get_highest_level, get_last_facet_value};
fn descending_facet_sort<'t>(
pub fn descending_facet_sort<'t>(
rtxn: &'t heed::RoTxn<'t>,
db: &'t heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
field_id: u16,
candidates: RoaringBitmap,
) -> Result<Box<dyn Iterator<Item = Result<(&'t [u8], RoaringBitmap)>> + 't>> {
) -> Result<Box<dyn Iterator<Item = Result<RoaringBitmap>> + 't>> {
let highest_level = get_highest_level(rtxn, db, field_id)?;
if let Some(first_bound) = get_first_facet_value::<MyByteSlice>(rtxn, db, field_id)? {
let first_key = FacetKey { field_id, level: highest_level, left_bound: first_bound };
@ -33,7 +33,7 @@ fn descending_facet_sort<'t>(
struct DescendingFacetSort<'t> {
rtxn: &'t heed::RoTxn<'t>,
db: &'t heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
field_id: u16,
stack: Vec<(
RoaringBitmap,
@ -43,7 +43,7 @@ struct DescendingFacetSort<'t> {
}
impl<'t> Iterator for DescendingFacetSort<'t> {
type Item = Result<(&'t [u8], RoaringBitmap)>;
type Item = Result<RoaringBitmap>;
fn next(&mut self) -> Option<Self::Item> {
'outer: loop {
@ -70,7 +70,7 @@ impl<'t> Iterator for DescendingFacetSort<'t> {
*documents_ids -= &bitmap;
if level == 0 {
return Some(Ok((left_bound, bitmap)));
return Some(Ok(bitmap));
}
let starting_key_below = FacetKey { field_id, level: level - 1, left_bound };
@ -89,14 +89,15 @@ impl<'t> Iterator for DescendingFacetSort<'t> {
};
let prev_right_bound = *right_bound;
*right_bound = Bound::Excluded(left_bound);
let iter = match self.db.rev_range(
&self.rtxn,
&(Bound::Included(starting_key_below), end_key_kelow),
) {
Ok(iter) => iter,
Err(e) => return Some(Err(e.into())),
}
.take(group_size as usize);
let iter =
match self.db.remap_key_type::<FacetKeyCodec<MyByteSlice>>().rev_range(
&self.rtxn,
&(Bound::Included(starting_key_below), end_key_kelow),
) {
Ok(iter) => iter,
Err(e) => return Some(Err(e.into())),
}
.take(group_size as usize);
self.stack.push((bitmap, iter, prev_right_bound));
continue 'outer;

View File

@ -274,7 +274,7 @@ impl<'a> Filter<'a> {
(_, _) => (),
}
let x = facet_range_search::find_docids_of_facet_within_bounds::<OrderedF64Codec>(
rtxn, &db, field_id, &left, &right,
rtxn, db, field_id, &left, &right,
)?;
// TODO: the facet range search should take a mutable roaring bitmap as argument
*output = x;

View File

@ -264,7 +264,7 @@ impl<'i> IncrementalFacetUpdate<'i> {
}
let group_size = self.group_size;
let highest_level = get_highest_level(&txn, &self.db, field_id)?;
let highest_level = get_highest_level(&txn, *self.db, field_id)?;
let result =
self.insert_in_level(txn, field_id, highest_level as u8, new_key, new_values)?;
@ -413,7 +413,7 @@ impl<'i> IncrementalFacetUpdate<'i> {
if self.db.get(txn, &FacetKey { field_id, level: 0, left_bound: key })?.is_none() {
return Ok(());
}
let highest_level = get_highest_level(&txn, &self.db, field_id)?;
let highest_level = get_highest_level(&txn, *self.db, field_id)?;
// let key_bytes = BoundCodec::bytes_encode(&key).unwrap();

View File

@ -11,14 +11,14 @@ pub use self::filter::Filter;
mod facet_distribution;
mod facet_distribution_iter;
mod facet_range_search;
mod facet_sort_ascending;
mod facet_sort_descending;
pub mod facet_sort_ascending;
pub mod facet_sort_descending;
mod filter;
mod incremental_update;
pub(crate) fn get_first_facet_value<'t, BoundCodec>(
txn: &'t RoTxn,
db: &'t heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
field_id: u16,
) -> heed::Result<Option<BoundCodec::DItem>>
where
@ -40,7 +40,7 @@ where
}
pub(crate) fn get_last_facet_value<'t, BoundCodec>(
txn: &'t RoTxn,
db: &'t heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
field_id: u16,
) -> heed::Result<Option<BoundCodec::DItem>>
where
@ -63,7 +63,7 @@ where
}
pub(crate) fn get_highest_level<'t>(
txn: &'t RoTxn<'t>,
db: &'t heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
db: heed::Database<FacetKeyCodec<MyByteSlice>, FacetGroupValueCodec>,
field_id: u16,
) -> heed::Result<u8> {
let field_id_prefix = &field_id.to_be_bytes();