h3ron_polars/spatial_index/
rtree.rs

1use crate::spatial_index::{finish_mask, negative_mask, RectIndexable, RectSIKind, SpatialIndex};
2use crate::{AsH3IndexChunked, IndexChunked, IndexValue};
3use geo_types::{Coord, Rect};
4use polars::export::arrow::bitmap::MutableBitmap;
5use polars::prelude::UInt64Chunked;
6use polars_core::datatypes::BooleanChunked;
7use rstar::primitives::{GeomWithData, Rectangle};
8use rstar::{RTree, AABB};
9use std::marker::PhantomData;
10
11// todo: Use the Line type supported by rtree instead of Rectangle to index H3DirectedEdges
12
13type RTreeCoord = [f64; 2];
14type RTreeBBox = Rectangle<RTreeCoord>;
15type LocatedArrayPosition = GeomWithData<RTreeBBox, usize>;
16
17/// [R-Tree](https://en.wikipedia.org/wiki/R-tree) spatial index
18pub struct RTreeIndex<IX: IndexValue> {
19    index_phantom: PhantomData<IX>,
20    chunked_array: UInt64Chunked,
21    pub rtree: RTree<LocatedArrayPosition>,
22}
23
24#[inline]
25fn to_coord(coord: Coord) -> RTreeCoord {
26    [coord.x, coord.y]
27}
28
29#[inline]
30fn to_bbox(rect: &Rect) -> RTreeBBox {
31    RTreeBBox::from_corners(to_coord(rect.min()), to_coord(rect.max()))
32}
33
34pub trait BuildRTreeIndex<'a, IX>
35where
36    IX: IndexValue + RectIndexable,
37{
38    fn rtree_index(&self) -> RTreeIndex<IX>;
39}
40
41impl<'a, IX: IndexValue> BuildRTreeIndex<'a, IX> for IndexChunked<'a, IX>
42where
43    IX: RectIndexable,
44{
45    /// Build a [R-Tree](https://en.wikipedia.org/wiki/R-tree) spatial index
46    ///
47    /// # Example
48    ///
49    /// ```
50    /// use geo_types::Rect;
51    /// use polars::prelude::UInt64Chunked;
52    /// use polars_core::prelude::TakeRandom;
53    /// use h3ron::{H3Cell, Index};
54    /// use h3ron_polars::{AsH3CellChunked, NamedFromIndexes};
55    /// use h3ron_polars::spatial_index::{BuildRTreeIndex, SpatialIndex, SpatialIndexGeomOp};
56    ///
57    /// let ca = UInt64Chunked::new_from_indexes(
58    ///     "",
59    ///     vec![
60    ///         H3Cell::from_coordinate((45.5, 45.5).into(), 7).unwrap(),
61    ///         H3Cell::from_coordinate((-60.5, -60.5).into(), 7).unwrap(),
62    ///         H3Cell::from_coordinate((120.5, 70.5).into(), 7).unwrap(),
63    ///     ],
64    /// );
65    ///
66    /// let rtree = ca.h3cell().rtree_index();
67    /// let mask = rtree.geometries_intersect(&Rect::new((40.0, 40.0), (50.0, 50.0)));
68    ///
69    /// assert_eq!(mask.len(), 3);
70    /// assert_eq!(mask.get(0), Some(true));
71    /// assert_eq!(mask.get(1), Some(false));
72    /// assert_eq!(mask.get(2), Some(false));
73    /// ```
74    fn rtree_index(&self) -> RTreeIndex<IX> {
75        let entries: Vec<_> = self
76            .iter_indexes_nonvalidated()
77            .enumerate()
78            .filter_map(|(pos, maybe_index)| match maybe_index {
79                Some(index) => index.spatial_index_rect().ok().and_then(|maybe_rect| {
80                    maybe_rect.map(|rect| LocatedArrayPosition::new(to_bbox(&rect), pos))
81                }),
82                _ => None,
83            })
84            .collect();
85
86        RTreeIndex {
87            index_phantom: PhantomData::<IX>,
88            chunked_array: self.chunked_array.clone(),
89            rtree: RTree::bulk_load(entries),
90        }
91    }
92}
93
94impl<IX: IndexValue> SpatialIndex<IX, RectSIKind> for RTreeIndex<IX>
95where
96    IX: RectIndexable,
97{
98    fn h3indexchunked(&self) -> IndexChunked<IX> {
99        self.chunked_array.h3indexchunked()
100    }
101
102    fn envelopes_intersect_impl(&self, rect: &Rect) -> MutableBitmap {
103        let mut mask = negative_mask(&self.chunked_array);
104        let envelope = AABB::from_corners(to_coord(rect.min()), to_coord(rect.max()));
105        let locator = self.rtree.locate_in_envelope_intersecting(&envelope);
106        for located_array_position in locator {
107            mask.set(located_array_position.data, true);
108        }
109        mask
110    }
111
112    fn envelopes_within_distance(&self, coord: Coord, distance: f64) -> BooleanChunked {
113        let mut mask = negative_mask(&self.chunked_array);
114
115        let locator = self.rtree.locate_within_distance(to_coord(coord), distance);
116        for located_array_position in locator {
117            mask.set(located_array_position.data, true);
118        }
119
120        finish_mask(mask.into(), &self.h3indexchunked())
121    }
122}
123
124#[cfg(test)]
125mod test {
126    use crate::spatial_index::{BuildRTreeIndex, RTreeIndex};
127    use crate::IndexChunked;
128    use h3ron::H3Cell;
129
130    fn build_index(cc: &IndexChunked<H3Cell>) -> RTreeIndex<H3Cell> {
131        cc.rtree_index()
132    }
133    crate::spatial_index::tests::impl_std_tests!(build_index);
134}
h3ron_polars/spatial_index/rtree.rs

h3ron_polars/spatial_index/
rtree.rs
