uqa_storage/
spatial_index.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! In-memory spatial index over geographic points.
8//!
9//! The current implementation uses a brute-force Haversine scan: simple,
10//! correct, and sufficient for the algebraic operators that consume it.
11
12use std::collections::BTreeMap;
13use std::sync::Arc;
14
15use uqa_core::{DocId, Payload, PostingEntry, PostingList};
16
17const EARTH_RADIUS_M: f64 = 6_371_000.0;
18
19/// Great-circle distance in meters between two `(longitude, latitude)`
20/// pairs in degrees.
21pub fn haversine_distance(lon1: f64, lat1: f64, lon2: f64, lat2: f64) -> f64 {
22    let phi1 = lat1.to_radians();
23    let phi2 = lat2.to_radians();
24    let d_phi = (lat2 - lat1).to_radians();
25    let d_lambda = (lon2 - lon1).to_radians();
26    let a = (d_phi / 2.0).sin().powi(2) + phi1.cos() * phi2.cos() * (d_lambda / 2.0).sin().powi(2);
27    let c = 2.0 * a.sqrt().atan2((1.0 - a).sqrt());
28    EARTH_RADIUS_M * c
29}
30
31pub trait SpatialIndex: Send + Sync {
32    fn add(&mut self, doc_id: DocId, lon: f64, lat: f64);
33    fn remove(&mut self, doc_id: DocId);
34    fn clear(&mut self);
35    fn search_within(&self, center_lon: f64, center_lat: f64, radius_m: f64) -> PostingList;
36    fn count(&self) -> usize;
37    fn snapshot(&self) -> Arc<dyn SpatialIndex>;
38}
39
40#[derive(Debug, Default, Clone)]
41pub struct MemorySpatialIndex {
42    field: String,
43    points: BTreeMap<DocId, (f64, f64)>,
44}
45
46impl MemorySpatialIndex {
47    pub fn new(field: impl Into<String>) -> Self {
48        Self {
49            field: field.into(),
50            points: BTreeMap::new(),
51        }
52    }
53
54    pub fn field(&self) -> &str {
55        &self.field
56    }
57}
58
59impl SpatialIndex for MemorySpatialIndex {
60    fn add(&mut self, doc_id: DocId, lon: f64, lat: f64) {
61        self.points.insert(doc_id, (lon, lat));
62    }
63
64    fn remove(&mut self, doc_id: DocId) {
65        self.points.remove(&doc_id);
66    }
67
68    fn clear(&mut self) {
69        self.points.clear();
70    }
71
72    fn search_within(&self, center_lon: f64, center_lat: f64, radius_m: f64) -> PostingList {
73        let mut entries: Vec<PostingEntry> = self
74            .points
75            .iter()
76            .filter_map(|(&doc_id, &(lon, lat))| {
77                let d = haversine_distance(center_lon, center_lat, lon, lat);
78                if d <= radius_m {
79                    let score = if radius_m > 0.0 {
80                        1.0 - (d / radius_m)
81                    } else {
82                        1.0
83                    };
84                    Some(PostingEntry::new(doc_id, Payload::with_score(score)))
85                } else {
86                    None
87                }
88            })
89            .collect();
90        entries.sort_by_key(|e| e.doc_id);
91        PostingList::from_sorted_unchecked(entries)
92    }
93
94    fn count(&self) -> usize {
95        self.points.len()
96    }
97
98    fn snapshot(&self) -> Arc<dyn SpatialIndex> {
99        Arc::new(self.clone())
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    fn approx(a: f64, b: f64, eps: f64) {
108        assert!((a - b).abs() < eps, "expected {a} ~ {b} within {eps}");
109    }
110
111    #[test]
112    fn haversine_zero_for_identical_point() {
113        approx(haversine_distance(0.0, 0.0, 0.0, 0.0), 0.0, 1e-6);
114    }
115
116    #[test]
117    fn haversine_known_distance() {
118        // London (-0.1276, 51.5074) to Paris (2.3522, 48.8566) ~= 343 km
119        let d = haversine_distance(-0.1276, 51.5074, 2.3522, 48.8566);
120        approx(d, 343_000.0, 5_000.0);
121    }
122
123    #[test]
124    fn search_within_filters_by_distance() {
125        let mut idx = MemorySpatialIndex::new("location");
126        idx.add(1, 0.0, 0.0); // origin
127        idx.add(2, 0.001, 0.0); // ~111 m east
128        idx.add(3, 1.0, 1.0); // ~157 km away
129        let pl = idx.search_within(0.0, 0.0, 500.0);
130        let docs: Vec<DocId> = pl.iter().map(|e| e.doc_id).collect();
131        assert_eq!(docs, vec![1, 2]);
132        // Score is 1 - distance/radius, so doc 1 (distance 0) gets the
133        // top score.
134        let s1 = pl.get_entry(1).unwrap().payload.score;
135        let s2 = pl.get_entry(2).unwrap().payload.score;
136        assert!(s1 > s2);
137    }
138}
uqa_storage/spatial_index.rs

uqa_storage/
spatial_index.rs
