h3ron_polars/spatial_index/
kdtree.rs

1use crate::spatial_index::{
2    finish_mask, negative_mask, CoordinateIndexable, CoordinateSIKind, SpatialIndex,
3};
4use crate::{AsH3IndexChunked, IndexChunked, IndexValue};
5use geo_types::{Coord, Rect};
6use kdbush::{KDBush, PointReader};
7use polars::export::arrow::bitmap::MutableBitmap;
8use polars::prelude::BooleanChunked;
9use polars_core::prelude::UInt64Chunked;
10use std::marker::PhantomData;
11
12struct Points(Vec<(usize, Coord)>);
13
14impl PointReader for Points {
15    fn size_hint(&self) -> usize {
16        self.0.len()
17    }
18
19    fn visit_all<F>(&self, mut visitor: F)
20    where
21        F: FnMut(usize, f64, f64),
22    {
23        for (pos, coord) in self.0.iter() {
24            visitor(*pos, coord.x, coord.y)
25        }
26    }
27}
28
29pub trait BuildKDTreeIndex<'a, IX>
30where
31    IX: IndexValue + CoordinateIndexable,
32{
33    /// Build a [KDTreeIndex] using the default parameters.
34    ///
35    /// # Example
36    ///
37    /// ```
38    /// use geo_types::Rect;
39    /// use polars::prelude::UInt64Chunked;
40    /// use polars_core::prelude::TakeRandom;
41    /// use h3ron::{H3Cell, Index};
42    /// use h3ron_polars::{AsH3CellChunked, NamedFromIndexes};
43    /// use h3ron_polars::spatial_index::{BuildKDTreeIndex, SpatialIndex, SpatialIndexGeomOp};
44    ///
45    /// let ca = UInt64Chunked::new_from_indexes(
46    ///     "",
47    ///     vec![
48    ///         H3Cell::from_coordinate((45.5, 45.5).into(), 7).unwrap(),
49    ///         H3Cell::from_coordinate((-60.5, -60.5).into(), 7).unwrap(),
50    ///         H3Cell::from_coordinate((120.5, 70.5).into(), 7).unwrap(),
51    ///     ],
52    /// );
53    ///
54    /// let kdbush = ca.h3cell().kdtree_index();
55    /// let mask = kdbush.geometries_intersect(&Rect::new((40.0, 40.0), (50.0, 50.0)));
56    ///
57    /// assert_eq!(mask.len(), 3);
58    /// assert_eq!(mask.get(0), Some(true));
59    /// assert_eq!(mask.get(1), Some(false));
60    /// assert_eq!(mask.get(2), Some(false));
61    /// ```
62    fn kdtree_index(&self) -> KDTreeIndex<IX> {
63        self.kdtree_index_with_node_size(64)
64    }
65
66    /// Build a [KDTreeIndex] using custom parameters.
67    ///
68    /// `node_size` - Size of the KD-tree node, 64 by default. Higher means faster indexing but slower search, and vise versa
69    fn kdtree_index_with_node_size(&self, node_size: u8) -> KDTreeIndex<IX>;
70}
71
72impl<'a, IX: IndexValue> BuildKDTreeIndex<'a, IX> for IndexChunked<'a, IX>
73where
74    IX: CoordinateIndexable,
75{
76    fn kdtree_index_with_node_size(&self, node_size: u8) -> KDTreeIndex<IX> {
77        // KDBush requires at least one entry to be successful build, so we need to inspect
78        // the data first.
79        let entries: Vec<_> = self
80            .iter_indexes_nonvalidated()
81            .enumerate()
82            .filter_map(|(pos, maybe_index)| match maybe_index {
83                Some(index) => index.spatial_index_coordinate().ok().map(|c| (pos, c)),
84                _ => None,
85            })
86            .collect();
87
88        let kdbush = if entries.is_empty() {
89            None
90        } else {
91            Some(KDBush::create(Points(entries), node_size))
92        };
93
94        KDTreeIndex {
95            index_phantom: PhantomData::<IX>,
96
97            // clones of arrow-backed arrays are cheap, so we clone this for the benefit of not
98            // requiring a lifetime dependency
99            chunked_array: self.chunked_array.clone(),
100
101            kdbush,
102        }
103    }
104}
105
106/// A very fast static spatial index for 2D points based on a flat [KD-tree](https://en.wikipedia.org/wiki/K-d_tree).
107///
108/// Operates only the centroid coordinate, not on envelopes / bounding boxes.
109///
110/// # Example
111///
112/// ```
113/// use polars::prelude::UInt64Chunked;
114/// use h3ron::H3Cell;
115/// use h3ron_polars::{AsH3CellChunked, NamedFromIndexes};
116/// use h3ron_polars::spatial_index::BuildKDTreeIndex;
117///
118/// let uc = UInt64Chunked::new_from_indexes(
119///     "",
120///     vec![
121///         H3Cell::from_coordinate((45.5, 45.5).into(), 7).unwrap(),
122///         H3Cell::from_coordinate((-60.5, -60.5).into(), 7).unwrap(),
123///         H3Cell::from_coordinate((120.5, 70.5).into(), 7).unwrap(),
124///     ],
125/// );
126///
127/// let idx = uc.h3cell().kdtree_index();
128/// ```
129pub struct KDTreeIndex<IX: IndexValue> {
130    index_phantom: PhantomData<IX>,
131
132    chunked_array: UInt64Chunked,
133
134    pub kdbush: Option<KDBush>,
135}
136
137impl<IX: IndexValue> SpatialIndex<IX, CoordinateSIKind> for KDTreeIndex<IX>
138where
139    IX: CoordinateIndexable,
140{
141    fn h3indexchunked(&self) -> IndexChunked<IX> {
142        self.chunked_array.h3indexchunked()
143    }
144
145    fn envelopes_intersect_impl(&self, rect: &Rect) -> MutableBitmap {
146        let mut mask = negative_mask(&self.chunked_array);
147        if let Some(kdbush) = self.kdbush.as_ref() {
148            kdbush.range(
149                rect.min().x,
150                rect.min().y,
151                rect.max().x,
152                rect.max().y,
153                |id| mask.set(id, true),
154            );
155        }
156        mask
157    }
158
159    fn envelopes_within_distance(&self, coord: Coord, distance: f64) -> BooleanChunked {
160        let mut mask = negative_mask(&self.chunked_array);
161        if let Some(kdbush) = self.kdbush.as_ref() {
162            kdbush.within(coord.x, coord.y, distance, |id| mask.set(id, true));
163        }
164        finish_mask(mask.into(), &self.h3indexchunked())
165    }
166}
167
168#[cfg(test)]
169mod test {
170    use crate::spatial_index::{BuildKDTreeIndex, KDTreeIndex};
171    use crate::IndexChunked;
172    use h3ron::H3Cell;
173
174    fn build_index(cc: &IndexChunked<H3Cell>) -> KDTreeIndex<H3Cell> {
175        cc.kdtree_index()
176    }
177    crate::spatial_index::tests::impl_std_tests!(build_index);
178}
h3ron_polars/spatial_index/kdtree.rs

h3ron_polars/spatial_index/
kdtree.rs
