2023-04-13 13:45:34 +02:00
|
|
|
/*!
|
2023-05-02 11:25:10 +02:00
|
|
|
This module tests the `geo_sort` ranking rule
|
2023-04-13 13:45:34 +02:00
|
|
|
*/
|
|
|
|
|
|
|
|
use big_s::S;
|
|
|
|
use heed::RoTxn;
|
|
|
|
use maplit::hashset;
|
|
|
|
|
|
|
|
use crate::index::tests::TempIndex;
|
2023-06-14 22:07:06 +02:00
|
|
|
use crate::score_details::ScoreDetails;
|
2023-04-13 13:45:34 +02:00
|
|
|
use crate::search::new::tests::collect_field_values;
|
|
|
|
use crate::{AscDesc, Criterion, GeoSortStrategy, Member, Search, SearchResult};
|
|
|
|
|
|
|
|
fn create_index() -> TempIndex {
|
|
|
|
let index = TempIndex::new();
|
|
|
|
|
|
|
|
index
|
|
|
|
.update_settings(|s| {
|
|
|
|
s.set_primary_key("id".to_owned());
|
|
|
|
s.set_sortable_fields(hashset! { S("_geo") });
|
|
|
|
s.set_criteria(vec![Criterion::Words, Criterion::Sort]);
|
|
|
|
})
|
|
|
|
.unwrap();
|
|
|
|
index
|
|
|
|
}
|
|
|
|
|
|
|
|
#[track_caller]
|
|
|
|
fn execute_iterative_and_rtree_returns_the_same<'a>(
|
|
|
|
rtxn: &RoTxn<'a>,
|
|
|
|
index: &TempIndex,
|
|
|
|
search: &mut Search<'a>,
|
2023-06-14 22:07:06 +02:00
|
|
|
) -> (Vec<usize>, Vec<Vec<ScoreDetails>>) {
|
2023-04-13 13:45:34 +02:00
|
|
|
search.geo_sort_strategy(GeoSortStrategy::AlwaysIterative(2));
|
2023-06-14 22:07:06 +02:00
|
|
|
let SearchResult { documents_ids, document_scores: iterative_scores_bucketed, .. } =
|
|
|
|
search.execute().unwrap();
|
2023-04-29 11:40:00 +02:00
|
|
|
let iterative_ids_bucketed = collect_field_values(index, rtxn, "id", &documents_ids);
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
search.geo_sort_strategy(GeoSortStrategy::AlwaysIterative(1000));
|
2023-06-14 22:07:06 +02:00
|
|
|
let SearchResult { documents_ids, document_scores: iterative_scores, .. } =
|
|
|
|
search.execute().unwrap();
|
2023-04-29 11:40:00 +02:00
|
|
|
let iterative_ids = collect_field_values(index, rtxn, "id", &documents_ids);
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
assert_eq!(iterative_ids_bucketed, iterative_ids, "iterative bucket");
|
2023-06-14 22:07:06 +02:00
|
|
|
assert_eq!(iterative_scores_bucketed, iterative_scores, "iterative bucket score");
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
search.geo_sort_strategy(GeoSortStrategy::AlwaysRtree(2));
|
2023-06-14 22:07:06 +02:00
|
|
|
let SearchResult { documents_ids, document_scores: rtree_scores_bucketed, .. } =
|
|
|
|
search.execute().unwrap();
|
2023-04-29 11:40:00 +02:00
|
|
|
let rtree_ids_bucketed = collect_field_values(index, rtxn, "id", &documents_ids);
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
search.geo_sort_strategy(GeoSortStrategy::AlwaysRtree(1000));
|
2023-06-14 22:07:06 +02:00
|
|
|
let SearchResult { documents_ids, document_scores: rtree_scores, .. } =
|
|
|
|
search.execute().unwrap();
|
2023-04-29 11:40:00 +02:00
|
|
|
let rtree_ids = collect_field_values(index, rtxn, "id", &documents_ids);
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
assert_eq!(rtree_ids_bucketed, rtree_ids, "rtree bucket");
|
2023-06-14 22:07:06 +02:00
|
|
|
assert_eq!(rtree_scores_bucketed, rtree_scores, "rtree bucket score");
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
assert_eq!(iterative_ids, rtree_ids, "iterative vs rtree");
|
2023-06-14 22:07:06 +02:00
|
|
|
assert_eq!(iterative_scores, rtree_scores, "iterative vs rtree scores");
|
2023-04-13 13:45:34 +02:00
|
|
|
|
2023-06-14 22:07:06 +02:00
|
|
|
(iterative_ids.into_iter().map(|id| id.parse().unwrap()).collect(), iterative_scores)
|
2023-04-13 13:45:34 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn test_geo_sort() {
|
|
|
|
let index = create_index();
|
|
|
|
|
|
|
|
index
|
|
|
|
.add_documents(documents!([
|
2023-05-02 11:25:10 +02:00
|
|
|
{ "id": 2, "_geo": { "lat": 2, "lng": -1 } },
|
|
|
|
{ "id": 3, "_geo": { "lat": -2, "lng": -2 } },
|
|
|
|
{ "id": 5, "_geo": { "lat": 6, "lng": -5 } },
|
|
|
|
{ "id": 4, "_geo": { "lat": 3, "lng": 5 } },
|
|
|
|
{ "id": 0, "_geo": { "lat": 0, "lng": 0 } },
|
|
|
|
{ "id": 1, "_geo": { "lat": 1, "lng": 1 } },
|
|
|
|
{ "id": 6 }, { "id": 8 }, { "id": 7 }, { "id": 10 }, { "id": 9 },
|
2023-04-13 13:45:34 +02:00
|
|
|
]))
|
|
|
|
.unwrap();
|
|
|
|
|
2023-05-02 11:25:10 +02:00
|
|
|
let rtxn = index.read_txn().unwrap();
|
2023-04-13 13:45:34 +02:00
|
|
|
|
2023-05-02 11:25:10 +02:00
|
|
|
let mut s = Search::new(&rtxn, &index);
|
2023-06-14 22:07:06 +02:00
|
|
|
s.scoring_strategy(crate::score_details::ScoringStrategy::Detailed);
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
s.sort_criteria(vec![AscDesc::Asc(Member::Geo([0., 0.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-05-02 11:25:10 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[0, 1, 2, 3, 4, 5, 6, 8, 7, 10, 9]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
s.sort_criteria(vec![AscDesc::Desc(Member::Geo([0., 0.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-05-02 11:25:10 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[5, 4, 3, 2, 1, 0, 6, 8, 7, 10, 9]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn test_geo_sort_around_the_edge_of_the_flat_earth() {
|
|
|
|
let index = create_index();
|
|
|
|
|
|
|
|
index
|
|
|
|
.add_documents(documents!([
|
|
|
|
{ "id": 0, "_geo": { "lat": 0, "lng": 0 } },
|
|
|
|
{ "id": 1, "_geo": { "lat": 88, "lng": 0 } },
|
|
|
|
{ "id": 2, "_geo": { "lat": -89, "lng": 0 } },
|
|
|
|
|
|
|
|
{ "id": 3, "_geo": { "lat": 0, "lng": 178 } },
|
|
|
|
{ "id": 4, "_geo": { "lat": 0, "lng": -179 } },
|
|
|
|
]))
|
|
|
|
.unwrap();
|
|
|
|
|
|
|
|
let rtxn = index.read_txn().unwrap();
|
|
|
|
|
|
|
|
let mut s = Search::new(&rtxn, &index);
|
2023-06-14 22:07:06 +02:00
|
|
|
s.scoring_strategy(crate::score_details::ScoringStrategy::Detailed);
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
// --- asc
|
|
|
|
s.sort_criteria(vec![AscDesc::Asc(Member::Geo([0., 0.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[0, 1, 2, 3, 4]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
// ensuring the lat doesn't wrap around
|
|
|
|
s.sort_criteria(vec![AscDesc::Asc(Member::Geo([85., 0.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[1, 0, 3, 4, 2]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
s.sort_criteria(vec![AscDesc::Asc(Member::Geo([-85., 0.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[2, 0, 3, 4, 1]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
// ensuring the lng does wrap around
|
|
|
|
s.sort_criteria(vec![AscDesc::Asc(Member::Geo([0., 175.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[3, 4, 2, 1, 0]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
s.sort_criteria(vec![AscDesc::Asc(Member::Geo([0., -175.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[4, 3, 2, 1, 0]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
// --- desc
|
|
|
|
s.sort_criteria(vec![AscDesc::Desc(Member::Geo([0., 0.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[4, 3, 2, 1, 0]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
// ensuring the lat doesn't wrap around
|
|
|
|
s.sort_criteria(vec![AscDesc::Desc(Member::Geo([85., 0.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[2, 4, 3, 0, 1]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
s.sort_criteria(vec![AscDesc::Desc(Member::Geo([-85., 0.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[1, 4, 3, 0, 2]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
// ensuring the lng does wrap around
|
|
|
|
s.sort_criteria(vec![AscDesc::Desc(Member::Geo([0., 175.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[0, 1, 2, 4, 3]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
s.sort_criteria(vec![AscDesc::Desc(Member::Geo([0., -175.]))]);
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[0, 1, 2, 3, 4]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn geo_sort_mixed_with_words() {
|
|
|
|
let index = create_index();
|
|
|
|
|
|
|
|
index
|
|
|
|
.add_documents(documents!([
|
|
|
|
{ "id": 0, "doggo": "jean", "_geo": { "lat": 0, "lng": 0 } },
|
|
|
|
{ "id": 1, "doggo": "intel", "_geo": { "lat": 88, "lng": 0 } },
|
|
|
|
{ "id": 2, "doggo": "jean bob", "_geo": { "lat": -89, "lng": 0 } },
|
|
|
|
{ "id": 3, "doggo": "jean michel", "_geo": { "lat": 0, "lng": 178 } },
|
|
|
|
{ "id": 4, "doggo": "bob marley", "_geo": { "lat": 0, "lng": -179 } },
|
|
|
|
]))
|
|
|
|
.unwrap();
|
|
|
|
|
|
|
|
let rtxn = index.read_txn().unwrap();
|
|
|
|
|
|
|
|
let mut s = Search::new(&rtxn, &index);
|
2023-06-14 22:07:06 +02:00
|
|
|
s.scoring_strategy(crate::score_details::ScoringStrategy::Detailed);
|
2023-04-13 13:45:34 +02:00
|
|
|
s.sort_criteria(vec![AscDesc::Asc(Member::Geo([0., 0.]))]);
|
|
|
|
|
|
|
|
s.query("jean");
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[0, 2, 3]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
s.query("bob");
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[2, 4]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"), @r###"
|
|
|
|
[
|
|
|
|
[
|
|
|
|
Words(
|
|
|
|
Words {
|
|
|
|
matching_words: 1,
|
|
|
|
max_matching_words: 1,
|
|
|
|
},
|
|
|
|
),
|
|
|
|
GeoSort(
|
|
|
|
GeoSort {
|
|
|
|
target_point: [
|
|
|
|
0.0,
|
|
|
|
0.0,
|
|
|
|
],
|
|
|
|
ascending: true,
|
|
|
|
value: Some(
|
|
|
|
[
|
|
|
|
-89.0,
|
|
|
|
0.0,
|
|
|
|
],
|
|
|
|
),
|
|
|
|
},
|
|
|
|
),
|
|
|
|
],
|
|
|
|
[
|
|
|
|
Words(
|
|
|
|
Words {
|
|
|
|
matching_words: 1,
|
|
|
|
max_matching_words: 1,
|
|
|
|
},
|
|
|
|
),
|
|
|
|
GeoSort(
|
|
|
|
GeoSort {
|
|
|
|
target_point: [
|
|
|
|
0.0,
|
|
|
|
0.0,
|
|
|
|
],
|
|
|
|
ascending: true,
|
|
|
|
value: Some(
|
|
|
|
[
|
|
|
|
0.0,
|
|
|
|
-179.0,
|
|
|
|
],
|
|
|
|
),
|
|
|
|
},
|
|
|
|
),
|
|
|
|
],
|
|
|
|
]
|
|
|
|
"###);
|
2023-04-13 13:45:34 +02:00
|
|
|
|
|
|
|
s.query("intel");
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[1]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"), @r###"
|
|
|
|
[
|
|
|
|
[
|
|
|
|
Words(
|
|
|
|
Words {
|
|
|
|
matching_words: 1,
|
|
|
|
max_matching_words: 1,
|
|
|
|
},
|
|
|
|
),
|
|
|
|
GeoSort(
|
|
|
|
GeoSort {
|
|
|
|
target_point: [
|
|
|
|
0.0,
|
|
|
|
0.0,
|
|
|
|
],
|
|
|
|
ascending: true,
|
|
|
|
value: Some(
|
|
|
|
[
|
|
|
|
88.0,
|
|
|
|
0.0,
|
|
|
|
],
|
|
|
|
),
|
|
|
|
},
|
|
|
|
),
|
|
|
|
],
|
|
|
|
]
|
|
|
|
"###);
|
2023-04-13 13:45:34 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn geo_sort_without_any_geo_faceted_documents() {
|
|
|
|
let index = create_index();
|
|
|
|
|
|
|
|
index
|
|
|
|
.add_documents(documents!([
|
|
|
|
{ "id": 0, "doggo": "jean" },
|
|
|
|
{ "id": 1, "doggo": "intel" },
|
|
|
|
{ "id": 2, "doggo": "jean bob" },
|
|
|
|
{ "id": 3, "doggo": "jean michel" },
|
|
|
|
{ "id": 4, "doggo": "bob marley" },
|
|
|
|
]))
|
|
|
|
.unwrap();
|
|
|
|
|
|
|
|
let rtxn = index.read_txn().unwrap();
|
|
|
|
|
|
|
|
let mut s = Search::new(&rtxn, &index);
|
2023-06-14 22:07:06 +02:00
|
|
|
s.scoring_strategy(crate::score_details::ScoringStrategy::Detailed);
|
2023-04-13 13:45:34 +02:00
|
|
|
s.sort_criteria(vec![AscDesc::Asc(Member::Geo([0., 0.]))]);
|
|
|
|
|
|
|
|
s.query("jean");
|
2023-06-14 22:07:06 +02:00
|
|
|
let (ids, scores) = execute_iterative_and_rtree_returns_the_same(&rtxn, &index, &mut s);
|
2023-04-13 13:45:34 +02:00
|
|
|
insta::assert_snapshot!(format!("{ids:?}"), @"[0, 2, 3]");
|
2023-06-14 22:07:06 +02:00
|
|
|
insta::assert_snapshot!(format!("{scores:#?}"));
|
2023-04-13 13:45:34 +02:00
|
|
|
}
|