add the possibility to sort by descending order on geoPoint

This commit is contained in:
Tamo 2021-09-13 14:27:14 +02:00 committed by Irevoire
parent 91ce4d1721
commit c695a1ffd2
No known key found for this signature in database
GPG Key ID: 7A6A970C96104F1B
2 changed files with 42 additions and 12 deletions

View File

@ -10,6 +10,7 @@ use crate::{GeoPoint, Index, Result};
pub struct Geo<'t> { pub struct Geo<'t> {
index: &'t Index, index: &'t Index,
rtxn: &'t heed::RoTxn<'t>, rtxn: &'t heed::RoTxn<'t>,
ascending: bool,
parent: Box<dyn Criterion + 't>, parent: Box<dyn Criterion + 't>,
candidates: Box<dyn Iterator<Item = RoaringBitmap>>, candidates: Box<dyn Iterator<Item = RoaringBitmap>>,
allowed_candidates: RoaringBitmap, allowed_candidates: RoaringBitmap,
@ -19,11 +20,30 @@ pub struct Geo<'t> {
} }
impl<'t> Geo<'t> { impl<'t> Geo<'t> {
pub fn new( pub fn asc(
index: &'t Index, index: &'t Index,
rtxn: &'t heed::RoTxn<'t>, rtxn: &'t heed::RoTxn<'t>,
parent: Box<dyn Criterion + 't>, parent: Box<dyn Criterion + 't>,
point: [f64; 2], point: [f64; 2],
) -> Result<Self> {
Self::new(index, rtxn, parent, point, true)
}
pub fn desc(
index: &'t Index,
rtxn: &'t heed::RoTxn<'t>,
parent: Box<dyn Criterion + 't>,
point: [f64; 2],
) -> Result<Self> {
Self::new(index, rtxn, parent, point, false)
}
fn new(
index: &'t Index,
rtxn: &'t heed::RoTxn<'t>,
parent: Box<dyn Criterion + 't>,
point: [f64; 2],
ascending: bool,
) -> Result<Self> { ) -> Result<Self> {
let candidates = Box::new(iter::empty()); let candidates = Box::new(iter::empty());
let allowed_candidates = index.geo_faceted_documents_ids(rtxn)?; let allowed_candidates = index.geo_faceted_documents_ids(rtxn)?;
@ -33,6 +53,7 @@ impl<'t> Geo<'t> {
Ok(Self { Ok(Self {
index, index,
rtxn, rtxn,
ascending,
parent, parent,
candidates, candidates,
allowed_candidates, allowed_candidates,
@ -89,9 +110,12 @@ impl Criterion for Geo<'_> {
} }
self.allowed_candidates = &candidates - params.excluded_candidates; self.allowed_candidates = &candidates - params.excluded_candidates;
self.candidates = match rtree { self.candidates = match rtree {
Some(rtree) => { Some(rtree) => geo_point(
geo_point(rtree, self.allowed_candidates.clone(), self.point) rtree,
} self.allowed_candidates.clone(),
self.point,
self.ascending,
),
None => Box::new(std::iter::empty()), None => Box::new(std::iter::empty()),
}; };
} }
@ -106,6 +130,7 @@ fn geo_point(
rtree: &RTree<GeoPoint>, rtree: &RTree<GeoPoint>,
mut candidates: RoaringBitmap, mut candidates: RoaringBitmap,
point: [f64; 2], point: [f64; 2],
ascending: bool,
) -> Box<dyn Iterator<Item = RoaringBitmap>> { ) -> Box<dyn Iterator<Item = RoaringBitmap>> {
let mut results = Vec::new(); let mut results = Vec::new();
for point in rtree.nearest_neighbor_iter(&point) { for point in rtree.nearest_neighbor_iter(&point) {
@ -117,5 +142,9 @@ fn geo_point(
} }
} }
Box::new(results.into_iter()) if ascending {
Box::new(results.into_iter())
} else {
Box::new(results.into_iter().rev())
}
} }

View File

@ -15,7 +15,7 @@ use super::query_tree::{Operation, PrimitiveQueryPart, Query, QueryKind};
use crate::criterion::{AscDesc as AscDescName, Member}; use crate::criterion::{AscDesc as AscDescName, Member};
use crate::search::criteria::geo::Geo; use crate::search::criteria::geo::Geo;
use crate::search::{word_derivations, WordDerivationsCache}; use crate::search::{word_derivations, WordDerivationsCache};
use crate::{DocumentId, FieldId, Index, Result, TreeLevel, UserError}; use crate::{DocumentId, FieldId, Index, Result, TreeLevel};
mod asc_desc; mod asc_desc;
mod attribute; mod attribute;
@ -304,17 +304,18 @@ impl<'t> CriteriaBuilder<'t> {
criterion, criterion,
field.to_string(), field.to_string(),
)?), )?),
AscDescName::Asc(Member::Geo(point)) => Box::new(Geo::new( AscDescName::Asc(Member::Geo(point)) => Box::new(Geo::asc(
&self.index,
&self.rtxn,
criterion,
point.clone(),
)?),
AscDescName::Desc(Member::Geo(point)) => Box::new(Geo::desc(
&self.index, &self.index,
&self.rtxn, &self.rtxn,
criterion, criterion,
point.clone(), point.clone(),
)?), )?),
AscDescName::Desc(Member::Geo(_point)) => {
return Err(UserError::InvalidSortName {
name: "sorting in descending order is not supported for the geosearch".to_string(),
})?
}
}; };
} }
criterion criterion