gistools/data_structures/
point_index.rs

1use crate::{
2    data_store::vector::{Vector, VectorStore},
3    geometry::{LonLat, S1ChordAngle, S2Cap, S2CellId, S2Point, convert},
4    parsers::FeatureReader,
5};
6use alloc::{vec, vec::Vec};
7use core::marker::PhantomData;
8use s2json::{
9    Face, GetM, GetXY, GetZ, JSONCollection, MValue, Projection, VectorFeature, VectorFeatureType,
10    VectorGeometry, VectorPoint,
11};
12use serde::{Deserialize, Serialize, de::DeserializeOwned};
13
14/// Locally allocated Point Index
15pub type LocalPointIndex<M = MValue> = PointIndex<M, Vector<S2CellId, IndexPoint<M>>>;
16
17/// An Indexed Point
18#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
19pub struct IndexPoint<M> {
20    /// S2Point (3D vector)
21    pub s2point: S2Point,
22    /// Data associated with the point
23    pub data: M,
24}
25impl<M> IndexPoint<M> {
26    /// Create a new Indexed Point
27    pub fn new(s2point: S2Point, data: M) -> Self {
28        IndexPoint { s2point, data }
29    }
30}
31impl<M> GetXY for IndexPoint<M> {
32    fn x(&self) -> f64 {
33        self.s2point.x()
34    }
35    fn y(&self) -> f64 {
36        self.s2point.y()
37    }
38}
39impl<M> GetZ for IndexPoint<M> {
40    fn z(&self) -> Option<f64> {
41        self.s2point.z()
42    }
43}
44impl<M> GetM<M> for IndexPoint<M> {
45    fn m(&self) -> Option<&M> {
46        Some(&self.data)
47    }
48}
49
50/// # Point Index
51///
52/// ## Description
53/// An index of cells with radius queries
54/// Assumes the data is compatible with S2JSON MValues with serde_json serialization
55#[derive(Debug, Default)]
56pub struct PointIndex<
57    M: Clone + Default + Serialize + DeserializeOwned = MValue,
58    S: VectorStore<S2CellId, IndexPoint<M>> = Vector<S2CellId, IndexPoint<M>>,
59> {
60    store: S,
61    sorted: bool,
62    projection: Projection,
63    _marker: PhantomData<M>,
64}
65impl<M: Clone + Default + Serialize + DeserializeOwned, S: VectorStore<S2CellId, IndexPoint<M>>>
66    PointIndex<M, S>
67{
68    /// Create a new PointIndex
69    pub fn new(store: Option<S>, projection: Option<Projection>) -> Self {
70        let store = store.unwrap_or_else(|| S::new(None));
71        let projection = projection.unwrap_or(Projection::S2);
72        PointIndex { store, sorted: true, projection, _marker: PhantomData }
73    }
74
75    /// Returns the number of points in the index
76    pub fn len(&self) -> u64 {
77        self.store.len()
78    }
79
80    /// Returns true if the index is empty
81    pub fn is_empty(&self) -> bool {
82        self.store.is_empty()
83    }
84
85    /// Get a reference to the point store
86    pub fn get_index(&self, index: u64) -> Option<&(S2CellId, IndexPoint<M>)> {
87        self.store.get_index(index)
88    }
89
90    /// Get a mutable reference to the point store
91    pub fn get_index_mut(&mut self, index: u64) -> Option<&mut (S2CellId, IndexPoint<M>)> {
92        self.store.get_index_mut(index)
93    }
94
95    /// Insert a cell with the point and its corresponding data to the index
96    pub fn insert(&mut self, cell: S2CellId, point: S2Point, data: Option<M>) {
97        self.store.push(cell, IndexPoint::new(point, data.unwrap_or_default()));
98        self.sorted = false;
99    }
100
101    /// Add a lon-lat pair as with any shape
102    pub fn insert_point<T: GetXY>(&mut self, point: T, data: Option<M>) {
103        self.insert_vector_point(VectorPoint::new(point.x(), point.y(), None, data));
104    }
105
106    /// Add a lon-lat pair to the cluster
107    pub fn insert_lon_lat(&mut self, mut vp: LonLat<M>) {
108        self.insert_vector_point(vp.take());
109    }
110
111    /// Add a lon-lat pair via a VectorPoint to the cluster
112    pub fn insert_vector_point(&mut self, vp: VectorPoint<M>) {
113        self.insert_feature(JSONCollection::<(), M, M>::VectorFeature(VectorFeature {
114            _type: VectorFeatureType::VectorFeature,
115            geometry: VectorGeometry::new_point(vp, None),
116            ..Default::default()
117        }));
118    }
119
120    /// Insert an STPoint to the index
121    pub fn insert_face_st(&mut self, face: Face, s: f64, t: f64, data: M) {
122        self.insert_feature(JSONCollection::<(), M, M>::VectorFeature(VectorFeature {
123            _type: VectorFeatureType::S2Feature,
124            face,
125            geometry: VectorGeometry::new_point(VectorPoint::new(s, t, None, Some(data)), None),
126            ..Default::default()
127        }));
128    }
129
130    /// Add all points from a reader. It will try to use the M-value first, but if it doesn't exist
131    /// it will use the feature properties data
132    pub fn insert_reader<T: Clone, F: FeatureReader<T, M, M>>(&mut self, reader: &F) {
133        for feature in reader.iter() {
134            self.insert_feature(JSONCollection::<T, M, M>::VectorFeature(feature))
135        }
136    }
137
138    /// Add a vector feature. It will try to use the M-value first, but if it doesn't exist
139    /// it will use the feature properties data
140    pub fn insert_feature<T: Clone>(&mut self, data: JSONCollection<T, M, M>) {
141        let features = convert(self.projection, &data, Some(true), Some(true));
142        for feature in features {
143            match feature.geometry {
144                VectorGeometry::Point(geometry) => {
145                    let coordinates = geometry.coordinates;
146                    self._insert_face_st(
147                        feature.face.into(),
148                        coordinates.x,
149                        coordinates.y,
150                        coordinates.m.or(Some(feature.properties)),
151                    );
152                }
153                VectorGeometry::MultiPoint(geometry) => {
154                    for point in geometry.coordinates {
155                        self._insert_face_st(
156                            feature.face.into(),
157                            point.x,
158                            point.y,
159                            point.m.or(Some(feature.properties.clone())),
160                        );
161                    }
162                }
163                _ => {}
164            }
165        }
166    }
167
168    /// Sort the index
169    pub fn sort(&mut self) {
170        if self.sorted {
171            return;
172        }
173        self.store.sort();
174        self.sorted = true;
175    }
176
177    /// Find the starting index of a search
178    pub fn lower_bound(&self, id: S2CellId) -> u64 {
179        assert!(self.sorted);
180        // lower bound search
181        let mut lo: u64 = 0;
182        let mut hi: u64 = self.store.len();
183        let mut mid: u64;
184
185        while lo < hi {
186            mid = (lo + hi) / 2;
187            let (mid_cell, _) = self.store.get_index(mid).unwrap();
188            if *mid_cell < id {
189                lo = mid + 1;
190            } else {
191                hi = mid;
192            }
193        }
194
195        lo
196    }
197
198    /// Search for points given a range of low and high ids
199    pub fn search_range(
200        &self,
201        mut low: S2CellId,
202        high: Option<S2CellId>,
203        max_results: Option<usize>,
204    ) -> Vec<&(S2CellId, IndexPoint<M>)> {
205        assert!(self.sorted);
206        let max_results = max_results.unwrap_or(usize::MAX);
207        let mut res = vec![];
208        let high = high.unwrap_or_else(|| {
209            let (lo, hi) = low.range();
210            low = lo;
211            hi
212        });
213        let mut lo_idx = self.lower_bound(low);
214
215        let len = self.store.len();
216        loop {
217            if lo_idx >= len {
218                break;
219            }
220            let curr_lo = self.store.get_index(lo_idx).unwrap();
221            if curr_lo.0 > high {
222                break;
223            }
224            res.push(curr_lo);
225            if res.len() >= max_results {
226                break;
227            }
228            lo_idx += 1;
229        }
230
231        res
232    }
233
234    /// Search for points given a range of low and high ids
235    pub fn search_range_mut(
236        &mut self,
237        mut low: S2CellId,
238        high: Option<S2CellId>,
239        max_results: Option<usize>,
240    ) -> Vec<&mut (S2CellId, IndexPoint<M>)> {
241        // setup
242        let max_results = max_results.unwrap_or(usize::MAX);
243        self.sort();
244        let high = high.unwrap_or_else(|| {
245            let (lo, hi) = low.range();
246            low = lo;
247            hi
248        });
249        let mut lo_idx = self.lower_bound(low);
250
251        let mut res = vec![];
252        let len = self.store.len();
253        while lo_idx < len {
254            let entry_ptr: *mut (S2CellId, IndexPoint<M>) =
255                self.store.get_index_mut(lo_idx).unwrap() as *mut _; // Convert to raw ptr
256            if unsafe { (*entry_ptr).0 } > high {
257                break;
258            }
259            // store
260            unsafe {
261                res.push(&mut *entry_ptr);
262            }
263            // if were at max results, break
264            if res.len() >= max_results {
265                break;
266            }
267            lo_idx += 1;
268        }
269
270        res
271    }
272
273    /// TODO: WG Proj case: is the radius correct?
274    /// Search for points within a given radius of a target point
275    pub fn search_radius<M2: Clone + Default>(
276        &self,
277        target: &LonLat<M2>,
278        radius: S1ChordAngle,
279        max_results: Option<usize>,
280    ) -> Vec<&(S2CellId, IndexPoint<M>)> {
281        self.s2_search_radius(self.adjust_lon_lat(target), radius, max_results)
282    }
283
284    /// Search for points within a given radius of a target point in S2 space
285    pub fn s2_search_radius(
286        &self,
287        target: S2Point,
288        radius: S1ChordAngle,
289        max_results: Option<usize>,
290    ) -> Vec<&(S2CellId, IndexPoint<M>)> {
291        assert!(self.sorted);
292        let max_results = max_results.unwrap_or(usize::MAX);
293        if radius < 0. {
294            return vec![];
295        }
296        let cap = S2Cap::new(target, radius, ());
297        let mut res = vec![];
298        for cell in cap.get_intersecting_cells() {
299            // iterate each covering s2cell min-max range on store. check distance from found
300            // store Cells to target and if within radius add to results
301            let (min, max) = cell.range();
302            for point in self.search_range(min, Some(max), Some(max_results)) {
303                if S1ChordAngle::from_s2_points(&target, &(point.1.s2point)) < radius {
304                    res.push(point);
305                }
306                if res.len() >= max_results {
307                    break;
308                }
309            }
310        }
311
312        res
313    }
314
315    /// TODO: WG Proj case: is the radius correct?
316    /// Search for points within a given radius of a target point
317    pub fn search_radius_mut<M2: Clone + Default>(
318        &mut self,
319        target: &LonLat<M2>,
320        radius: S1ChordAngle,
321        max_results: Option<usize>,
322    ) -> Vec<&mut (S2CellId, IndexPoint<M>)> {
323        self.s2_search_radius_mut(self.adjust_lon_lat(target), radius, max_results)
324    }
325
326    /// Search for points within a given radius of a target point in S2 space
327    pub fn s2_search_radius_mut(
328        &mut self,
329        target: S2Point,
330        radius: S1ChordAngle,
331        max_results: Option<usize>,
332    ) -> Vec<&mut (S2CellId, IndexPoint<M>)> {
333        let max_results = max_results.unwrap_or(usize::MAX);
334        self.sort();
335        if radius < 0. {
336            return vec![];
337        }
338        // first pass is to store the raw pointers
339        let cap = S2Cap::new(target, radius, ());
340        let mut raw_ptrs = vec![];
341        for cell in cap.get_intersecting_cells() {
342            // iterate each covering s2cell min-max range on store. check distance from found
343            // store Cells to target and if within radius add to results
344            let (min, max) = cell.range();
345            let search_results = self.search_range_mut(min, Some(max), Some(max_results));
346            for point in search_results {
347                let s2point = &point.1.s2point;
348                let angle = S1ChordAngle::from_s2_points(&target, s2point);
349                let ptr: *mut (S2CellId, IndexPoint<M>) = point as *mut _; // Convert to raw pointer
350                if angle < radius {
351                    raw_ptrs.push(ptr);
352                }
353                if raw_ptrs.len() >= max_results {
354                    break;
355                }
356            }
357        }
358
359        // Second pass: Convert raw pointers back to &mut references safely
360        let mut res = vec![];
361        unsafe {
362            for ptr in raw_ptrs {
363                res.push(&mut *ptr);
364            }
365        }
366        res
367    }
368
369    fn adjust_lon_lat<M2: Clone + Default>(&self, lonlat: &LonLat<M2>) -> S2Point {
370        if self.projection == Projection::WG {
371            let mut projected = lonlat.0.clone();
372            projected.project(None);
373            return S2Point::from_face_st(0, projected.x, projected.y);
374        }
375        lonlat.into()
376    }
377
378    /// Insert an STPoint to the index
379    fn _insert_face_st(&mut self, face: u8, s: f64, t: f64, data: Option<M>) {
380        let point = S2Point::from_face_st(face, s, t);
381        self.insert((&point).into(), point, data);
382    }
383
384    /// Get a reference iterator over the point store
385    pub fn iter<'a>(&'a self) -> impl Iterator<Item = &'a (S2CellId, IndexPoint<M>)>
386    where
387        IndexPoint<M>: 'a,
388    {
389        self.store.iter()
390    }
391
392    /// Get a mutable iterator over the point store
393    pub fn iter_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut (S2CellId, IndexPoint<M>)>
394    where
395        IndexPoint<M>: 'a,
396    {
397        self.store.iter_mut()
398    }
399}
gistools/data_structures/point_index.rs

gistools/data_structures/
point_index.rs
