h3ron_polars/spatial_index/
packed_hilbert_rtree.rs

1use crate::algorithm::bounding_rect::BoundingRect;
2use crate::spatial_index::{finish_mask, negative_mask, RectIndexable, RectSIKind, SpatialIndex};
3use crate::{AsH3IndexChunked, Error, IndexChunked, IndexValue};
4use geo_types::{coord, Coord, Rect};
5use polars::export::arrow::bitmap::MutableBitmap;
6use polars::prelude::{BooleanChunked, UInt64Chunked};
7use static_aabb2d_index::{NeighborVisitor, StaticAABB2DIndex, StaticAABB2DIndexBuilder};
8use std::marker::PhantomData;
9
10/// Spatial index implementation using the [packed Hilbert R-tree](https://en.wikipedia.org/wiki/Hilbert_R-tree#Packed_Hilbert_R-trees) algorithm
11///
12/// Based on [flatbush](https://github.com/mourner/flatbush) and the rust port [static_aabb2d_index](https://github.com/jbuckmccready/static_aabb2d_index).
13///
14/// # Example
15///
16/// ```
17/// use polars::prelude::UInt64Chunked;
18/// use h3ron::H3Cell;
19/// use h3ron_polars::{AsH3CellChunked, NamedFromIndexes};
20/// use h3ron_polars::spatial_index::BuildPackedHilbertRTreeIndex;
21///
22/// let uc = UInt64Chunked::new_from_indexes(
23///     "",
24///     vec![
25///         H3Cell::from_coordinate((45.5, 45.5).into(), 7).unwrap(),
26///         H3Cell::from_coordinate((-60.5, -60.5).into(), 7).unwrap(),
27///         H3Cell::from_coordinate((120.5, 70.5).into(), 7).unwrap(),
28///     ],
29/// );
30///
31/// let idx = uc.h3cell().packed_hilbert_rtree_index();
32/// ```
33pub struct PackedHilbertRTreeIndex<IX: IndexValue> {
34    pub index: Option<StaticAABB2DIndex<f64>>,
35    index_phantom: PhantomData<IX>,
36    chunked_array: UInt64Chunked,
37
38    /// maps the positions of the index contents to the position in the `chunked_array`
39    positions_in_chunked_array: Box<[usize]>,
40}
41
42pub trait BuildPackedHilbertRTreeIndex<IX: IndexValue> {
43    /// Build a spatial index using the [packed Hilbert R-tree](https://en.wikipedia.org/wiki/Hilbert_R-tree#Packed_Hilbert_R-trees) algorithm
44    ///
45    /// Based on [flatbush](https://github.com/mourner/flatbush) and the rust port [static_aabb2d_index](https://github.com/jbuckmccready/static_aabb2d_index).
46    fn packed_hilbert_rtree_index(&self) -> Result<PackedHilbertRTreeIndex<IX>, Error>;
47}
48
49impl<'a, IX: IndexValue> BuildPackedHilbertRTreeIndex<IX> for IndexChunked<'a, IX>
50where
51    IX: RectIndexable,
52{
53    fn packed_hilbert_rtree_index(&self) -> Result<PackedHilbertRTreeIndex<IX>, Error> {
54        let (positions_in_chunked_array, rects) =
55            self.iter_indexes_nonvalidated().enumerate().fold(
56                (
57                    Vec::with_capacity(self.len()),
58                    Vec::with_capacity(self.len()),
59                ),
60                |(mut positions, mut rects), (pos, maybe_index)| {
61                    if let Some(index) = maybe_index {
62                        if let Ok(Some(rect)) = index.spatial_index_rect() {
63                            positions.push(pos);
64                            rects.push(rect)
65                        }
66                    }
67                    (positions, rects)
68                },
69            );
70
71        let index = if !positions_in_chunked_array.is_empty() {
72            let mut builder = StaticAABB2DIndexBuilder::new(positions_in_chunked_array.len());
73            for rect in rects {
74                // add takes in (min_x, min_y, max_x, max_y) of the bounding box
75                builder.add(rect.min().x, rect.min().y, rect.max().x, rect.max().y);
76            }
77            Some(
78                builder
79                    .build()
80                    .map_err(|e| Error::SpatialIndex(e.to_string()))?,
81            )
82        } else {
83            None
84        };
85        Ok(PackedHilbertRTreeIndex {
86            index,
87            index_phantom: PhantomData::<IX>,
88            chunked_array: self.chunked_array.clone(),
89            positions_in_chunked_array: positions_in_chunked_array.into_boxed_slice(),
90        })
91    }
92}
93
94impl<IX: IndexValue> SpatialIndex<IX, RectSIKind> for PackedHilbertRTreeIndex<IX> {
95    fn h3indexchunked(&self) -> IndexChunked<IX> {
96        self.chunked_array.h3indexchunked()
97    }
98
99    fn envelopes_intersect_impl(&self, rect: &Rect) -> MutableBitmap {
100        let mut mask = negative_mask(&self.chunked_array);
101        if let Some(index) = self.index.as_ref() {
102            for index_position in
103                index.query(rect.min().x, rect.min().y, rect.max().x, rect.max().y)
104            {
105                mask.set(self.positions_in_chunked_array[index_position], true);
106            }
107        }
108        mask
109    }
110
111    fn envelopes_within_distance(&self, coord: Coord, distance: f64) -> BooleanChunked {
112        let mut mask = negative_mask(&self.chunked_array);
113
114        if let Some(index) = self.index.as_ref() {
115            let mut visitor = Visitor {
116                found: vec![],
117                distance,
118            };
119            let _ = index.visit_neighbors(coord.x, coord.y, &mut visitor);
120            for index_position in visitor.found {
121                mask.set(self.positions_in_chunked_array[index_position], true);
122            }
123        }
124
125        finish_mask(mask.into(), &self.h3indexchunked())
126    }
127}
128
129impl<IX: IndexValue> BoundingRect for PackedHilbertRTreeIndex<IX> {
130    fn bounding_rect(&self) -> Result<Option<Rect>, Error> {
131        if let Some(index) = self.index.as_ref() {
132            if let Some(bounds) = index.bounds() {
133                Ok(Some(Rect::new(
134                    coord! {x: bounds.min_x, y: bounds.min_y},
135                    coord! {x: bounds.max_x, y: bounds.max_y},
136                )))
137            } else {
138                Ok(None)
139            }
140        } else {
141            Ok(None)
142        }
143    }
144}
145
146struct Visitor {
147    found: Vec<usize>,
148    distance: f64,
149}
150
151impl NeighborVisitor<f64, Result<(), ()>> for Visitor {
152    fn visit(&mut self, index_pos: usize, dist_squared: f64) -> Result<(), ()> {
153        if dist_squared <= self.distance {
154            self.found.push(index_pos);
155            Ok(())
156        } else {
157            Err(())
158        }
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use crate::spatial_index::packed_hilbert_rtree::BuildPackedHilbertRTreeIndex;
165    use crate::spatial_index::PackedHilbertRTreeIndex;
166    use crate::IndexChunked;
167    use h3ron::H3Cell;
168
169    fn build_index(cc: &IndexChunked<H3Cell>) -> PackedHilbertRTreeIndex<H3Cell> {
170        cc.packed_hilbert_rtree_index().unwrap()
171    }
172    crate::spatial_index::tests::impl_std_tests!(build_index);
173}
h3ron_polars/spatial_index/packed_hilbert_rtree.rs

h3ron_polars/spatial_index/
packed_hilbert_rtree.rs
