diff --git a/meilisearch/tests/search/geo.rs b/meilisearch/tests/search/geo.rs new file mode 100644 index 000000000..4253bc210 --- /dev/null +++ b/meilisearch/tests/search/geo.rs @@ -0,0 +1,62 @@ +use once_cell::sync::Lazy; +use serde_json::{json, Value}; + +use crate::common::Server; + +pub(self) static DOCUMENTS: Lazy = Lazy::new(|| { + json!([ + { + "id": 1, + "name": "Taco Truck", + "address": "444 Salsa Street, Burritoville", + "type": "Mexican", + "rating": 9, + "_geo": { + "lat": 34.0522, + "lng": -118.2437 + } + }, + { + "id": 2, + "name": "La Bella Italia", + "address": "456 Elm Street, Townsville", + "type": "Italian", + "rating": 9, + "_geo": { + "lat": "45.4777599", + "lng": "9.1967508" + } + }, + { + "id": 3, + "name": "Crêpe Truck", + "address": "2 Billig Avenue, Rouenville", + "type": "French", + "rating": 10 + } + ]) +}); + +#[actix_rt::test] +async fn geo_sort_with_geo_strings() { + let server = Server::new().await; + let index = server.index("test"); + + let documents = DOCUMENTS.clone(); + index.update_settings_filterable_attributes(json!(["_geo"])).await; + index.update_settings_sortable_attributes(json!(["_geo"])).await; + index.add_documents(documents, None).await; + index.wait_task(2).await; + + index + .search( + json!({ + "filter": "_geoRadius(45.472735, 9.184019, 10000)", + "sort": ["_geoPoint(0.0, 0.0):asc"] + }), + |response, code| { + assert_eq!(code, 200, "{}", response); + }, + ) + .await; +} diff --git a/meilisearch/tests/search/mod.rs b/meilisearch/tests/search/mod.rs index 9c80aed31..d560ba734 100644 --- a/meilisearch/tests/search/mod.rs +++ b/meilisearch/tests/search/mod.rs @@ -4,6 +4,7 @@ mod errors; mod facet_search; mod formatted; +mod geo; mod multi; mod pagination; mod restrict_searchable; diff --git a/milli/src/search/new/distinct.rs b/milli/src/search/new/distinct.rs index fff96bd5d..e90ffe878 100644 --- a/milli/src/search/new/distinct.rs +++ b/milli/src/search/new/distinct.rs @@ -100,7 +100,7 @@ fn facet_number_values<'a>( } /// Return an iterator over each string value in the given field of the given document. -fn facet_string_values<'a>( +pub fn facet_string_values<'a>( docid: u32, field_id: u16, index: &Index, diff --git a/milli/src/search/new/geo_sort.rs b/milli/src/search/new/geo_sort.rs index dddb7f426..bd9546048 100644 --- a/milli/src/search/new/geo_sort.rs +++ b/milli/src/search/new/geo_sort.rs @@ -6,6 +6,7 @@ use heed::{RoPrefix, RoTxn}; use roaring::RoaringBitmap; use rstar::RTree; +use super::facet_string_values; use super::ranking_rules::{RankingRule, RankingRuleOutput, RankingRuleQueryTrait}; use crate::heed_codec::facet::{FieldDocIdFacetCodec, OrderedF64Codec}; use crate::score_details::{self, ScoreDetails}; @@ -157,23 +158,7 @@ impl GeoSort { let mut documents = self .geo_candidates .iter() - .map(|id| -> Result<_> { - Ok(( - id, - [ - facet_number_values(id, lat, ctx.index, ctx.txn)? - .next() - .expect("A geo faceted document doesn't contain any lat")? - .0 - .2, - facet_number_values(id, lng, ctx.index, ctx.txn)? - .next() - .expect("A geo faceted document doesn't contain any lng")? - .0 - .2, - ], - )) - }) + .map(|id| -> Result<_> { Ok((id, geo_value(id, lat, lng, ctx.index, ctx.txn)?)) }) .collect::>>()?; // computing the distance between two points is expensive thus we cache the result documents @@ -185,6 +170,37 @@ impl GeoSort { } } +/// 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). +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 { + 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::().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 { fn id(&self) -> String { "geo_sort".to_owned() diff --git a/milli/src/search/new/mod.rs b/milli/src/search/new/mod.rs index 26f992be2..c27e02514 100644 --- a/milli/src/search/new/mod.rs +++ b/milli/src/search/new/mod.rs @@ -42,6 +42,7 @@ use roaring::RoaringBitmap; use sort::Sort; use space::Neighbor; +use self::distinct::facet_string_values; use self::geo_sort::GeoSort; pub use self::geo_sort::Strategy as GeoSortStrategy; use self::graph_based_ranking_rule::Words;