move the geosearch criteria out of asc_desc.rs

This commit is contained in:
Tamo
2021-09-01 15:14:23 +02:00
parent dc84ecc40b
commit a8a1f5bd55
3 changed files with 160 additions and 70 deletions

View File

@ -4,14 +4,12 @@ use itertools::Itertools;
use log::debug;
use ordered_float::OrderedFloat;
use roaring::RoaringBitmap;
use rstar::RTree;
use super::{Criterion, CriterionParameters, CriterionResult};
use crate::criterion::Member;
use crate::search::criteria::{resolve_query_tree, CriteriaBuilder};
use crate::search::facet::{FacetNumberIter, FacetStringIter};
use crate::search::query_tree::Operation;
use crate::{FieldId, GeoPoint, Index, Result};
use crate::{FieldId, Index, Result};
/// Threshold on the number of candidates that will make
/// the system to choose between one algorithm or another.
@ -20,11 +18,10 @@ const CANDIDATES_THRESHOLD: u64 = 1000;
pub struct AscDesc<'t> {
index: &'t Index,
rtxn: &'t heed::RoTxn<'t>,
member: Member,
field_name: String,
field_id: Option<FieldId>,
is_ascending: bool,
query_tree: Option<Operation>,
rtree: Option<RTree<GeoPoint>>,
candidates: Box<dyn Iterator<Item = heed::Result<RoaringBitmap>> + 't>,
allowed_candidates: RoaringBitmap,
bucket_candidates: RoaringBitmap,
@ -37,29 +34,29 @@ impl<'t> AscDesc<'t> {
index: &'t Index,
rtxn: &'t heed::RoTxn,
parent: Box<dyn Criterion + 't>,
member: Member,
field_name: String,
) -> Result<Self> {
Self::new(index, rtxn, parent, member, true)
Self::new(index, rtxn, parent, field_name, true)
}
pub fn desc(
index: &'t Index,
rtxn: &'t heed::RoTxn,
parent: Box<dyn Criterion + 't>,
member: Member,
field_name: String,
) -> Result<Self> {
Self::new(index, rtxn, parent, member, false)
Self::new(index, rtxn, parent, field_name, false)
}
fn new(
index: &'t Index,
rtxn: &'t heed::RoTxn,
parent: Box<dyn Criterion + 't>,
member: Member,
field_name: String,
is_ascending: bool,
) -> Result<Self> {
let fields_ids_map = index.fields_ids_map(rtxn)?;
let field_id = member.field().and_then(|field| fields_ids_map.id(&field));
let field_id = fields_ids_map.id(&field_name);
let faceted_candidates = match field_id {
Some(field_id) => {
let number_faceted = index.number_faceted_documents_ids(rtxn, field_id)?;
@ -68,16 +65,14 @@ impl<'t> AscDesc<'t> {
}
None => RoaringBitmap::default(),
};
let rtree = index.geo_rtree(rtxn)?;
Ok(AscDesc {
index,
rtxn,
member,
field_name,
field_id,
is_ascending,
query_tree: None,
rtree,
candidates: Box::new(std::iter::empty()),
allowed_candidates: RoaringBitmap::new(),
faceted_candidates,
@ -97,7 +92,7 @@ impl<'t> Criterion for AscDesc<'t> {
debug!(
"Facet {}({}) iteration",
if self.is_ascending { "Asc" } else { "Desc" },
self.member
self.field_name
);
match self.candidates.next().transpose()? {
@ -140,31 +135,15 @@ impl<'t> Criterion for AscDesc<'t> {
}
self.allowed_candidates = &candidates - params.excluded_candidates;
match &self.member {
Member::Field(field_name) => {
self.candidates = match self.field_id {
Some(field_id) => facet_ordered(
self.index,
self.rtxn,
field_id,
self.is_ascending,
candidates & &self.faceted_candidates,
)?,
None => Box::new(std::iter::empty()),
}
}
Member::Geo(point) => {
self.candidates = match &self.rtree {
Some(rtree) => {
// TODO: TAMO how to remove that?
let rtree = Box::new(rtree.clone());
let rtree = Box::leak(rtree);
geo_point(rtree, candidates, point.clone())?
}
None => Box::new(std::iter::empty()),
}
}
self.candidates = match self.field_id {
Some(field_id) => facet_ordered(
self.index,
self.rtxn,
field_id,
self.is_ascending,
candidates & &self.faceted_candidates,
)?,
None => Box::new(std::iter::empty()),
};
}
None => return Ok(None),
@ -184,22 +163,6 @@ impl<'t> Criterion for AscDesc<'t> {
}
}
fn geo_point<'t>(
rtree: &'t RTree<GeoPoint>,
candidates: RoaringBitmap,
point: [f64; 2],
) -> Result<Box<dyn Iterator<Item = heed::Result<RoaringBitmap>> + 't>> {
Ok(Box::new(
rtree
.nearest_neighbor_iter_with_distance_2(&point)
.filter_map(move |(point, _distance)| {
candidates.contains(point.data).then(|| point.data)
})
.map(|id| std::iter::once(id).collect::<RoaringBitmap>())
.map(Ok),
))
}
/// Returns an iterator over groups of the given candidates in ascending or descending order.
///
/// It will either use an iterative or a recursive method on the whole facet database depending