Move geosort code out of search

crates/milli/src/search/mod.rs

@@ -9,6 +9,8 @@ use roaring::bitmap::RoaringBitmap;
 pub use self::facet::{FacetDistribution, Filter, OrderBy, DEFAULT_VALUES_PER_FACET};
 pub use self::new::matches::{FormatOptions, MatchBounds, MatcherBuilder, MatchingWords};
 use self::new::{execute_vector_search, PartialSearchResult, VectorStoreStats};
+use crate::documents::GeoSortParameter;
+use crate::documents::GeoSortStrategy;
 use crate::filterable_attributes_rules::{filtered_matching_patterns, matching_features};
 use crate::index::MatchingStrategy;
 use crate::score_details::{ScoreDetails, ScoringStrategy};
@@ -46,7 +48,7 @@ pub struct Search<'a> {
     sort_criteria: Option<Vec<AscDesc>>,
     distinct: Option<String>,
     searchable_attributes: Option<&'a [String]>,
-    geo_param: new::GeoSortParameter,
+    geo_param: GeoSortParameter,
     terms_matching_strategy: TermsMatchingStrategy,
     scoring_strategy: ScoringStrategy,
     words_limit: usize,
@@ -69,7 +71,7 @@ impl<'a> Search<'a> {
             sort_criteria: None,
             distinct: None,
             searchable_attributes: None,
-            geo_param: new::GeoSortParameter::default(),
+            geo_param: GeoSortParameter::default(),
             terms_matching_strategy: TermsMatchingStrategy::default(),
             scoring_strategy: Default::default(),
             exhaustive_number_hits: false,
@@ -145,7 +147,7 @@ impl<'a> Search<'a> {
     }
 
     #[cfg(test)]
-    pub fn geo_sort_strategy(&mut self, strategy: new::GeoSortStrategy) -> &mut Search<'a> {
+    pub fn geo_sort_strategy(&mut self, strategy: GeoSortStrategy) -> &mut Search<'a> {
         self.geo_param.strategy = strategy;
         self
     }

crates/milli/src/search/new/distinct.rs

@@ -118,7 +118,7 @@ pub fn facet_string_values<'a>(
 }
 
 #[allow(clippy::drop_non_drop)]
-fn facet_values_prefix_key(distinct: u16, id: u32) -> [u8; FID_SIZE + DOCID_SIZE] {
+pub(crate) fn facet_values_prefix_key(distinct: u16, id: u32) -> [u8; FID_SIZE + DOCID_SIZE] {
     concat_arrays::concat_arrays!(distinct.to_be_bytes(), id.to_be_bytes())
 }
 

crates/milli/src/search/new/geo_sort.rs

@@ -8,6 +8,7 @@ use rstar::RTree;
 use super::facet_string_values;
 use super::ranking_rules::{RankingRule, RankingRuleOutput, RankingRuleQueryTrait};
 use crate::documents::geo_sort::{fill_cache, next_bucket};
+use crate::documents::{GeoSortParameter, GeoSortStrategy};
 use crate::heed_codec::facet::{FieldDocIdFacetCodec, OrderedF64Codec};
 use crate::score_details::{self, ScoreDetails};
 use crate::{GeoPoint, Index, Result, SearchContext, SearchLogger};
@@ -20,75 +21,10 @@ fn facet_values_prefix_key(distinct: u16, id: u32) -> [u8; FID_SIZE + DOCID_SIZE
     concat_arrays::concat_arrays!(distinct.to_be_bytes(), id.to_be_bytes())
 }
 
-/// Return an iterator over each number value in the given field of the given document.
-fn facet_number_values<'a>(
-    docid: u32,
-    field_id: u16,
-    index: &Index,
-    txn: &'a RoTxn<'a>,
-) -> Result<RoPrefix<'a, FieldDocIdFacetCodec<OrderedF64Codec>, Unit>> {
-    let key = facet_values_prefix_key(field_id, docid);
-
-    let iter = index
-        .field_id_docid_facet_f64s
-        .remap_key_type::<Bytes>()
-        .prefix_iter(txn, &key)?
-        .remap_key_type();
-
-    Ok(iter)
-}
-
-#[derive(Debug, Clone, Copy)]
-pub struct Parameter {
-    // Define the strategy used by the geo sort
-    pub strategy: Strategy,
-    // Limit the number of docs in a single bucket to avoid unexpectedly large overhead
-    pub max_bucket_size: u64,
-    // Considering the errors of GPS and geographical calculations, distances less than distance_error_margin will be treated as equal
-    pub distance_error_margin: f64,
-}
-
-impl Default for Parameter {
-    fn default() -> Self {
-        Self { strategy: Strategy::default(), max_bucket_size: 1000, distance_error_margin: 1.0 }
-    }
-}
-/// Define the strategy used by the geo sort.
-/// The parameter represents the cache size, and, in the case of the Dynamic strategy,
-/// the point where we move from using the iterative strategy to the rtree.
-#[derive(Debug, Clone, Copy)]
-pub enum Strategy {
-    AlwaysIterative(usize),
-    AlwaysRtree(usize),
-    Dynamic(usize),
-}
-
-impl Default for Strategy {
-    fn default() -> Self {
-        Strategy::Dynamic(1000)
-    }
-}
-
-impl Strategy {
-    pub fn use_rtree(&self, candidates: usize) -> bool {
-        match self {
-            Strategy::AlwaysIterative(_) => false,
-            Strategy::AlwaysRtree(_) => true,
-            Strategy::Dynamic(i) => candidates >= *i,
-        }
-    }
-
-    pub fn cache_size(&self) -> usize {
-        match self {
-            Strategy::AlwaysIterative(i) | Strategy::AlwaysRtree(i) | Strategy::Dynamic(i) => *i,
-        }
-    }
-}
-
 pub struct GeoSort<Q: RankingRuleQueryTrait> {
     query: Option<Q>,
 
-    strategy: Strategy,
+    strategy: GeoSortStrategy,
     ascending: bool,
     point: [f64; 2],
     field_ids: Option<[u16; 2]>,
@@ -105,12 +41,12 @@ pub struct GeoSort<Q: RankingRuleQueryTrait> {
 
 impl<Q: RankingRuleQueryTrait> GeoSort<Q> {
     pub fn new(
-        parameter: Parameter,
+        parameter: GeoSortParameter,
         geo_faceted_docids: RoaringBitmap,
         point: [f64; 2],
         ascending: bool,
     ) -> Result<Self> {
-        let Parameter { strategy, max_bucket_size, distance_error_margin } = parameter;
+        let GeoSortParameter { strategy, max_bucket_size, distance_error_margin } = parameter;
         Ok(Self {
             query: None,
             strategy,
@@ -148,37 +84,6 @@ impl<Q: RankingRuleQueryTrait> GeoSort<Q> {
     }
 }
 
-/// Extracts the lat and long values from a single document.
-///
-/// If it is not able to find it in the facet number index it will extract it
-/// from the facet string index and parse it as f64 (as the geo extraction behaves).
-pub(crate) fn geo_value(
-    docid: u32,
-    field_lat: u16,
-    field_lng: u16,
-    index: &Index,
-    rtxn: &RoTxn<'_>,
-) -> Result<[f64; 2]> {
-    let extract_geo = |geo_field: u16| -> Result<f64> {
-        match facet_number_values(docid, geo_field, index, rtxn)?.next() {
-            Some(Ok(((_, _, geo), ()))) => Ok(geo),
-            Some(Err(e)) => Err(e.into()),
-            None => match facet_string_values(docid, geo_field, index, rtxn)?.next() {
-                Some(Ok((_, geo))) => {
-                    Ok(geo.parse::<f64>().expect("cannot parse geo field as f64"))
-                }
-                Some(Err(e)) => Err(e.into()),
-                None => panic!("A geo faceted document doesn't contain any lat or lng"),
-            },
-        }
-    };
-
-    let lat = extract_geo(field_lat)?;
-    let lng = extract_geo(field_lng)?;
-
-    Ok([lat, lng])
-}
-
 impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for GeoSort<Q> {
     fn id(&self) -> String {
         "geo_sort".to_owned()
@@ -224,15 +129,17 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for GeoSort<Q> {
             ctx.index,
             ctx.txn,
             universe,
-            self.strategy,
             self.ascending,
             self.point,
             &self.field_ids,
             &mut self.rtree,
             &mut self.cached_sorted_docids,
             &self.geo_candidates,
-            self.max_bucket_size,
-            self.distance_error_margin,
+            GeoSortParameter {
+                strategy: self.strategy,
+                max_bucket_size: self.max_bucket_size,
+                distance_error_margin: self.distance_error_margin,
+            },
         )
         .map(|o| {
             o.map(|(candidates, point)| RankingRuleOutput {
@@ -254,16 +161,3 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for GeoSort<Q> {
         self.cached_sorted_docids.clear();
     }
 }
-
-/// Compute the antipodal coordinate of `coord`
-pub(crate) fn opposite_of(mut coord: [f64; 2]) -> [f64; 2] {
-    coord[0] *= -1.;
-    // in the case of x,0 we want to return x,180
-    if coord[1] > 0. {
-        coord[1] -= 180.;
-    } else {
-        coord[1] += 180.;
-    }
-
-    coord
-}

crates/milli/src/search/new/mod.rs

@@ -46,14 +46,14 @@ use resolve_query_graph::{compute_query_graph_docids, PhraseDocIdsCache};
 use roaring::RoaringBitmap;
 use sort::Sort;
 
-use self::distinct::facet_string_values;
+pub(crate) use self::distinct::{facet_string_values, facet_values_prefix_key};
 use self::geo_sort::GeoSort;
-pub use self::geo_sort::{Parameter as GeoSortParameter, Strategy as GeoSortStrategy};
 use self::graph_based_ranking_rule::Words;
 use self::interner::Interned;
 use self::vector_sort::VectorSort;
 use crate::attribute_patterns::{match_pattern, PatternMatch};
 use crate::constants::RESERVED_GEO_FIELD_NAME;
+use crate::documents::GeoSortParameter;
 use crate::index::PrefixSearch;
 use crate::localized_attributes_rules::LocalizedFieldIds;
 use crate::score_details::{ScoreDetails, ScoringStrategy};
@@ -319,7 +319,7 @@ fn resolve_negative_phrases(
 fn get_ranking_rules_for_placeholder_search<'ctx>(
     ctx: &SearchContext<'ctx>,
     sort_criteria: &Option<Vec<AscDesc>>,
-    geo_param: geo_sort::Parameter,
+    geo_param: GeoSortParameter,
 ) -> Result<Vec<BoxRankingRule<'ctx, PlaceholderQuery>>> {
     let mut sort = false;
     let mut sorted_fields = HashSet::new();
@@ -371,7 +371,7 @@ fn get_ranking_rules_for_placeholder_search<'ctx>(
 fn get_ranking_rules_for_vector<'ctx>(
     ctx: &SearchContext<'ctx>,
     sort_criteria: &Option<Vec<AscDesc>>,
-    geo_param: geo_sort::Parameter,
+    geo_param: GeoSortParameter,
     limit_plus_offset: usize,
     target: &[f32],
     embedder_name: &str,
@@ -448,7 +448,7 @@ fn get_ranking_rules_for_vector<'ctx>(
 fn get_ranking_rules_for_query_graph_search<'ctx>(
     ctx: &SearchContext<'ctx>,
     sort_criteria: &Option<Vec<AscDesc>>,
-    geo_param: geo_sort::Parameter,
+    geo_param: GeoSortParameter,
     terms_matching_strategy: TermsMatchingStrategy,
 ) -> Result<Vec<BoxRankingRule<'ctx, QueryGraph>>> {
     // query graph search
@@ -559,7 +559,7 @@ fn resolve_sort_criteria<'ctx, Query: RankingRuleQueryTrait>(
     ranking_rules: &mut Vec<BoxRankingRule<'ctx, Query>>,
     sorted_fields: &mut HashSet<String>,
     geo_sorted: &mut bool,
-    geo_param: geo_sort::Parameter,
+    geo_param: GeoSortParameter,
 ) -> Result<()> {
     let sort_criteria = sort_criteria.clone().unwrap_or_default();
     ranking_rules.reserve(sort_criteria.len());
@@ -629,7 +629,7 @@ pub fn execute_vector_search(
     universe: RoaringBitmap,
     sort_criteria: &Option<Vec<AscDesc>>,
     distinct: &Option<String>,
-    geo_param: geo_sort::Parameter,
+    geo_param: GeoSortParameter,
     from: usize,
     length: usize,
     embedder_name: &str,
@@ -692,7 +692,7 @@ pub fn execute_search(
     mut universe: RoaringBitmap,
     sort_criteria: &Option<Vec<AscDesc>>,
     distinct: &Option<String>,
-    geo_param: geo_sort::Parameter,
+    geo_param: GeoSortParameter,
     from: usize,
     length: usize,
     words_limit: Option<usize>,