h3ron_polars/spatial_index/
mod.rs

1//! Spatial search indexes
2//!
3//! For some background on spatial search algorithms see [A dive into spatial search algorithms](https://blog.mapbox.com/a-dive-into-spatial-search-algorithms-ebd0c5e39d2a).
4//!
5//! Available implementations:
6//! * `KDTreeIndex`: Fast to create and query, works only on centroids.
7//! * `PackedHilbertRTreeIndex`: Still fast to create and query, works on envelopes. Creation is a bit slower than `KDTreeIndex`
8//! * `RTreeIndex`: Also envelope based.
9//!
10//! All of the above spatial indexes provide a second stage which can perform fine-grained filtering
11//! by doing geometry intersections on the returned index-matches using the exact geometry of the indexed
12//! entity and the queried `Polygon` or `MultiPolygon`. See the [SpatialIndexGeomOp] trait.
13//! This of course comes with an additional runtime cost.
14//!
15//! For a more detailed comparison of the runtime characteristics see the included `spatialindex` benchmark.
16
17#[cfg(feature = "si_kdtree")]
18pub mod kdtree;
19
20#[cfg(feature = "si_rtree")]
21pub mod rtree;
22
23#[cfg(feature = "si_packed_hilbert_rtree")]
24pub mod packed_hilbert_rtree;
25
26#[cfg(test)]
27pub(crate) mod tests;
28
29use crate::{Error, IndexChunked, IndexValue};
30use geo::bounding_rect::BoundingRect;
31use geo::{Contains, Intersects};
32use geo_types::{Coord, MultiPolygon, Polygon, Rect};
33use h3ron::to_geo::ToLine;
34use h3ron::{H3Cell, H3DirectedEdge, ToCoordinate, ToPolygon};
35use polars::export::arrow::array::BooleanArray;
36use polars::export::arrow::bitmap::{Bitmap, MutableBitmap};
37use polars::prelude::BooleanChunked;
38use polars_core::prelude::{FromData, TakeRandom, UInt64Chunked};
39
40#[cfg(feature = "si_kdtree")]
41pub use crate::spatial_index::kdtree::*;
42
43#[cfg(feature = "si_rtree")]
44pub use crate::spatial_index::rtree::*;
45
46#[cfg(feature = "si_packed_hilbert_rtree")]
47pub use crate::spatial_index::packed_hilbert_rtree::*;
48
49/// marker trait to restrict on what kind of geometries a spatial index
50/// operates.
51pub trait SIKind {}
52
53/// Marks spatial indexes operating on [Coord] points
54pub struct CoordinateSIKind {}
55impl SIKind for CoordinateSIKind {}
56
57/// Marks spatial indexes operating on [Rect] envelopes
58pub struct RectSIKind {}
59impl SIKind for RectSIKind {}
60
61pub trait SpatialIndex<IX: IndexValue, Kind: SIKind> {
62    fn h3indexchunked(&self) -> IndexChunked<IX>;
63
64    /// internal
65    fn envelopes_intersect_impl(&self, rect: &Rect) -> MutableBitmap;
66
67    /// The envelope of the indexed elements has some overlap with the given `rect`
68    fn envelopes_intersect(&self, rect: &Rect) -> BooleanChunked {
69        finish_mask(
70            self.envelopes_intersect_impl(rect).into(),
71            &self.h3indexchunked(),
72        )
73    }
74
75    /// The envelope of the indexed elements is with `distance` of the given [Coord] `coord`.
76    fn envelopes_within_distance(&self, coord: Coord, distance: f64) -> BooleanChunked;
77}
78
79pub trait SpatialIndexGeomOp<IX: IndexValue, Kind: SIKind> {
80    /// The geometry of the indexed elements is with in the given [Rect]
81    fn geometries_intersect(&self, rect: &Rect) -> BooleanChunked;
82
83    /// The geometry of the indexed elements is with in the given [Polygon]
84    fn geometries_intersect_polygon(&self, polygon: &Polygon) -> BooleanChunked;
85
86    /// The geometry of the indexed elements is with in the given [MultiPolygon]
87    fn geometries_intersect_multipolygon(&self, multipolygon: &MultiPolygon) -> BooleanChunked;
88}
89
90impl<T, IX: IndexValue> SpatialIndexGeomOp<IX, CoordinateSIKind> for T
91where
92    T: SpatialIndex<IX, CoordinateSIKind>,
93    IX: CoordinateIndexable,
94{
95    fn geometries_intersect(&self, rect: &Rect) -> BooleanChunked {
96        self.envelopes_intersect(rect) // index only works with points, so this is the same
97    }
98
99    fn geometries_intersect_polygon(&self, polygon: &Polygon) -> BooleanChunked {
100        geometries_intersect_polygon(self, polygon, validate_coordinate_containment)
101    }
102
103    fn geometries_intersect_multipolygon(&self, multipolygon: &MultiPolygon) -> BooleanChunked {
104        geometries_intersect_multipolygon(self, multipolygon, validate_coordinate_containment)
105    }
106}
107
108impl<T, IX: IndexValue> SpatialIndexGeomOp<IX, RectSIKind> for T
109where
110    T: SpatialIndex<IX, RectSIKind>,
111    IX: RectIndexable,
112{
113    fn geometries_intersect(&self, rect: &Rect) -> BooleanChunked {
114        let mask = self.envelopes_intersect_impl(rect);
115        let ic = self.h3indexchunked();
116        finish_mask(
117            validate_geometry_intersection(mask, &ic, &rect.to_polygon()).into(),
118            &ic,
119        )
120    }
121
122    fn geometries_intersect_polygon(&self, polygon: &Polygon) -> BooleanChunked {
123        geometries_intersect_polygon(self, polygon, validate_geometry_intersection)
124    }
125
126    fn geometries_intersect_multipolygon(&self, multipolygon: &MultiPolygon) -> BooleanChunked {
127        geometries_intersect_multipolygon(self, multipolygon, validate_geometry_intersection)
128    }
129}
130
131pub trait CoordinateIndexable {
132    /// coordinate to use for spatial indexing
133    fn spatial_index_coordinate(&self) -> Result<Coord, Error>;
134}
135
136impl CoordinateIndexable for H3Cell {
137    fn spatial_index_coordinate(&self) -> Result<Coord, Error> {
138        self.to_coordinate().map_err(Error::from)
139    }
140}
141
142impl CoordinateIndexable for H3DirectedEdge {
143    fn spatial_index_coordinate(&self) -> Result<Coord, Error> {
144        let cells = self.cells()?;
145        let c1 = cells.destination.to_coordinate()?;
146        let c2 = cells.origin.to_coordinate()?;
147        Ok(((c1.x + c2.x) / 2.0, (c1.y + c2.y) / 2.0).into())
148    }
149}
150
151pub trait RectIndexable {
152    fn spatial_index_rect(&self) -> Result<Option<Rect>, Error>;
153    fn intersects_with_polygon(&self, poly: &Polygon) -> Result<bool, Error>;
154}
155
156impl RectIndexable for H3Cell {
157    fn spatial_index_rect(&self) -> Result<Option<Rect>, Error> {
158        Ok(self.to_polygon()?.bounding_rect())
159    }
160
161    fn intersects_with_polygon(&self, poly: &Polygon) -> Result<bool, Error> {
162        Ok(poly.intersects(&self.to_polygon()?))
163    }
164}
165
166impl RectIndexable for H3DirectedEdge {
167    fn spatial_index_rect(&self) -> Result<Option<Rect>, Error> {
168        Ok(Some(self.to_line()?.bounding_rect()))
169    }
170
171    fn intersects_with_polygon(&self, poly: &Polygon) -> Result<bool, Error> {
172        Ok(poly.intersects(&self.to_line()?))
173    }
174}
175
176pub(crate) fn negative_mask(ca: &UInt64Chunked) -> MutableBitmap {
177    let mut mask = MutableBitmap::new();
178    mask.extend_constant(ca.len(), false);
179    mask
180}
181
182pub(crate) fn finish_mask<IX: IndexValue>(mask: Bitmap, ic: &IndexChunked<IX>) -> BooleanChunked {
183    let validites = ic.validity_bitmap();
184    let bool_arr = BooleanArray::from_data_default(mask, Some(validites));
185    BooleanChunked::from(bool_arr)
186}
187
188fn geometries_intersect_polygon<IX: IndexValue, SI, Kind, Validator>(
189    spatial_index: &SI,
190    polygon: &Polygon,
191    validator: Validator,
192) -> BooleanChunked
193where
194    SI: SpatialIndex<IX, Kind>,
195    Kind: SIKind,
196    Validator: Fn(MutableBitmap, &IndexChunked<IX>, &Polygon) -> MutableBitmap,
197{
198    let mask = if let Some(rect) = polygon.bounding_rect() {
199        let mask = spatial_index.envelopes_intersect_impl(&rect);
200        validator(mask, &spatial_index.h3indexchunked(), polygon)
201    } else {
202        negative_mask(spatial_index.h3indexchunked().chunked_array)
203    };
204    finish_mask(mask.into(), &spatial_index.h3indexchunked())
205}
206
207fn geometries_intersect_multipolygon<IX: IndexValue, SI, Kind, Validator>(
208    spatial_index: &SI,
209    multipolygon: &MultiPolygon,
210    validator: Validator,
211) -> BooleanChunked
212where
213    SI: SpatialIndex<IX, Kind>,
214    Kind: SIKind,
215    Validator: Fn(MutableBitmap, &IndexChunked<IX>, &Polygon) -> MutableBitmap,
216{
217    let mask = multipolygon
218        .0
219        .iter()
220        .filter_map(|polygon| {
221            if let Some(rect) = polygon.bounding_rect() {
222                let mask = spatial_index.envelopes_intersect_impl(&rect);
223                Some(validator(mask, &spatial_index.h3indexchunked(), polygon))
224            } else {
225                None
226            }
227        })
228        .fold(
229            negative_mask(spatial_index.h3indexchunked().chunked_array),
230            |acc_mask, mask| acc_mask | &(mask.into()),
231        );
232    finish_mask(mask.into(), &spatial_index.h3indexchunked())
233}
234
235pub(crate) fn validate_geometry_intersection<IX>(
236    mut mask: MutableBitmap,
237    indexchunked: &IndexChunked<IX>,
238    polygon: &Polygon,
239) -> MutableBitmap
240where
241    IX: RectIndexable + IndexValue,
242{
243    for i in 0..mask.len() {
244        if mask.get(i) {
245            if let Some(index) = indexchunked.get(i) {
246                if let Ok(true) = index.intersects_with_polygon(polygon) {
247                    mask.set(i, true)
248                }
249            }
250        }
251    }
252    mask
253}
254
255pub(crate) fn validate_coordinate_containment<IX>(
256    mut mask: MutableBitmap,
257    indexchunked: &IndexChunked<IX>,
258    polygon: &Polygon,
259) -> MutableBitmap
260where
261    IX: CoordinateIndexable + IndexValue,
262{
263    for i in 0..mask.len() {
264        if mask.get(i) {
265            if let Some(index) = indexchunked.get(i) {
266                if let Ok(coord) = index.spatial_index_coordinate() {
267                    if polygon.contains(&coord) {
268                        mask.set(i, true)
269                    }
270                }
271            }
272        }
273    }
274    mask
275}
h3ron_polars/spatial_index/mod.rs

h3ron_polars/spatial_index/
mod.rs
