From e570c23153f4b4ce91dfd0fe80ed03802a396563 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Lecrenier?= Date: Wed, 31 Aug 2022 09:36:19 +0200 Subject: [PATCH] Reintroduce asc/desc functionality --- milli/src/search/criteria/asc_desc.rs | 33 ++++++++++--------- milli/src/search/facet/facet_distribution.rs | 4 +-- .../search/facet/facet_distribution_iter.rs | 6 ++-- milli/src/search/facet/facet_range_search.rs | 12 +++---- .../src/search/facet/facet_sort_ascending.rs | 21 +++++------- .../src/search/facet/facet_sort_descending.rs | 31 ++++++++--------- milli/src/search/facet/filter.rs | 2 +- milli/src/search/facet/incremental_update.rs | 4 +-- milli/src/search/facet/mod.rs | 10 +++--- 9 files changed, 60 insertions(+), 63 deletions(-) diff --git a/milli/src/search/criteria/asc_desc.rs b/milli/src/search/criteria/asc_desc.rs index bd08c54a5..a5ea9b058 100644 --- a/milli/src/search/criteria/asc_desc.rs +++ b/milli/src/search/criteria/asc_desc.rs @@ -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>) } 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::>(), + field_id, + candidates.clone(), + )?; + let string_iter = make_iter( + rtxn, + index.facet_id_f64_docids.remap_key_type::>(), + field_id, + candidates, + )?; - // Ok(Box::new(number_iter.chain(string_iter))) + Ok(Box::new(number_iter.chain(string_iter))) } } diff --git a/milli/src/search/facet/facet_distribution.rs b/milli/src/search/facet/facet_distribution.rs index 670719a9b..c7619c609 100644 --- a/milli/src/search/facet/facet_distribution.rs +++ b/milli/src/search/facet/facet_distribution.rs @@ -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::>(), + self.index.facet_id_f64_docids.remap_key_type::>(), 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::>(), + self.index.facet_id_string_docids.remap_key_type::>(), field_id, candidates, |facet_key, nbr_docids| { diff --git a/milli/src/search/facet/facet_distribution_iter.rs b/milli/src/search/facet/facet_distribution_iter.rs index 9f1031a85..f347b9d7e 100644 --- a/milli/src/search/facet/facet_distribution_iter.rs +++ b/milli/src/search/facet/facet_distribution_iter.rs @@ -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, FacetGroupValueCodec>, + db: heed::Database, 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::>(), field_id)?; + get_highest_level(rtxn, db.remap_key_type::>(), field_id)?; if let Some(first_bound) = get_first_facet_value::(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, FacetGroupValueCodec>, + db: heed::Database, FacetGroupValueCodec>, field_id: u16, callback: CB, } diff --git a/milli/src/search/facet/facet_range_search.rs b/milli/src/search/facet/facet_range_search.rs index 75db9fda2..b05a3c275 100644 --- a/milli/src/search/facet/facet_range_search.rs +++ b/milli/src/search/facet/facet_range_search.rs @@ -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, FacetGroupValueCodec>, + db: heed::Database, FacetGroupValueCodec>, field_id: u16, left: &'t Bound<>::EItem>, right: &'t Bound<>::EItem>, @@ -50,11 +50,11 @@ where }; let db = db.remap_key_type::>(); 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::(rtxn, &db, field_id)? { - let last_bound = get_last_facet_value::(rtxn, &db, field_id)?.unwrap(); + if let Some(first_bound) = get_first_facet_value::(rtxn, db, field_id)? { + let last_bound = get_last_facet_value::(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, FacetGroupValueCodec>, + db: heed::Database, FacetGroupValueCodec>, field_id: u16, left: Bound<&'b [u8]>, right: Bound<&'b [u8]>, diff --git a/milli/src/search/facet/facet_sort_ascending.rs b/milli/src/search/facet/facet_sort_ascending.rs index 73491d4ae..e4b77c691 100644 --- a/milli/src/search/facet/facet_sort_ascending.rs +++ b/milli/src/search/facet/facet_sort_ascending.rs @@ -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, FacetGroupValueCodec>, + db: heed::Database, FacetGroupValueCodec>, field_id: u16, candidates: RoaringBitmap, -) -> Result> + 't>> { - let highest_level = - get_highest_level(rtxn, &db.remap_key_type::>(), field_id)?; - if let Some(first_bound) = get_first_facet_value::( - rtxn, - &db.remap_key_type::>(), - field_id, - )? { +) -> Result> + 't>> { + let highest_level = get_highest_level(rtxn, db, field_id)?; + if let Some(first_bound) = get_first_facet_value::(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, FacetGroupValueCodec>, + db: heed::Database, 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; fn next(&mut self) -> Option { '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 }; diff --git a/milli/src/search/facet/facet_sort_descending.rs b/milli/src/search/facet/facet_sort_descending.rs index 81b0eb09d..fc62b894f 100644 --- a/milli/src/search/facet/facet_sort_descending.rs +++ b/milli/src/search/facet/facet_sort_descending.rs @@ -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, FacetGroupValueCodec>, + db: heed::Database, FacetGroupValueCodec>, field_id: u16, candidates: RoaringBitmap, -) -> Result> + 't>> { +) -> Result> + 't>> { let highest_level = get_highest_level(rtxn, db, field_id)?; if let Some(first_bound) = get_first_facet_value::(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, FacetGroupValueCodec>, + db: heed::Database, 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; fn next(&mut self) -> Option { '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::>().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; diff --git a/milli/src/search/facet/filter.rs b/milli/src/search/facet/filter.rs index 79d7f5e0f..6ec626a5c 100644 --- a/milli/src/search/facet/filter.rs +++ b/milli/src/search/facet/filter.rs @@ -274,7 +274,7 @@ impl<'a> Filter<'a> { (_, _) => (), } let x = facet_range_search::find_docids_of_facet_within_bounds::( - 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; diff --git a/milli/src/search/facet/incremental_update.rs b/milli/src/search/facet/incremental_update.rs index f01b19dab..fd4e1eeb5 100644 --- a/milli/src/search/facet/incremental_update.rs +++ b/milli/src/search/facet/incremental_update.rs @@ -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(); diff --git a/milli/src/search/facet/mod.rs b/milli/src/search/facet/mod.rs index 023d433ad..8405c0141 100644 --- a/milli/src/search/facet/mod.rs +++ b/milli/src/search/facet/mod.rs @@ -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, FacetGroupValueCodec>, + db: heed::Database, FacetGroupValueCodec>, field_id: u16, ) -> heed::Result> where @@ -40,7 +40,7 @@ where } pub(crate) fn get_last_facet_value<'t, BoundCodec>( txn: &'t RoTxn, - db: &'t heed::Database, FacetGroupValueCodec>, + db: heed::Database, FacetGroupValueCodec>, field_id: u16, ) -> heed::Result> where @@ -63,7 +63,7 @@ where } pub(crate) fn get_highest_level<'t>( txn: &'t RoTxn<'t>, - db: &'t heed::Database, FacetGroupValueCodec>, + db: heed::Database, FacetGroupValueCodec>, field_id: u16, ) -> heed::Result { let field_id_prefix = &field_id.to_be_bytes();