diff --git a/Cargo.lock b/Cargo.lock index 7c9c798..4a2c53e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -23,6 +23,15 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7079075b41f533b8c61d2a4d073c4676e1f8b249ff94a393b0595db304e0dd87" +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + [[package]] name = "autocfg" version = "1.1.0" @@ -215,6 +224,7 @@ dependencies = [ "bytemuck", "criterion", "float_next_after", + "geo-traits", "num-traits", "rayon", "rstar", @@ -222,6 +232,26 @@ dependencies = [ "tinyvec", ] +[[package]] +name = "geo-traits" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b018fc19fa58202b03f1c809aebe654f7d70fd3887dace34c3d05c11aeb474b5" +dependencies = [ + "geo-types", +] + +[[package]] +name = "geo-types" +version = "0.7.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f47c611187777bbca61ea7aba780213f5f3441fd36294ab333e96cfa791b65" +dependencies = [ + "approx", + "num-traits", + "serde", +] + [[package]] name = "half" version = "1.8.2" diff --git a/Cargo.toml b/Cargo.toml index b4828a7..2c1f27b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,7 @@ categories = ["data-structures", "algorithms", "science::geo"] [dependencies] bytemuck = "1" float_next_after = "1" +geo-traits = "0.2" num-traits = "0.2" rayon = { version = "1.8.0", optional = true } thiserror = "1" diff --git a/src/kdtree/builder.rs b/src/kdtree/builder.rs index be0d004..60f8b26 100644 --- a/src/kdtree/builder.rs +++ b/src/kdtree/builder.rs @@ -2,11 +2,14 @@ use std::cmp; use std::marker::PhantomData; use bytemuck::cast_slice_mut; +use geo_traits::{CoordTrait, PointTrait}; +use crate::error::Result; use crate::indices::MutableIndices; use crate::kdtree::constants::{KDBUSH_HEADER_SIZE, KDBUSH_MAGIC, KDBUSH_VERSION}; use crate::kdtree::OwnedKDTree; use crate::r#type::IndexableNum; +use crate::GeoIndexError; const DEFAULT_NODE_SIZE: u16 = 64; @@ -68,7 +71,10 @@ impl KDTreeBuilder { } } - /// Add a point to the index. + /// Add a point to the KDTree. + /// + /// This returns a positional index that provides a lookup back into the original data. + #[inline] pub fn add(&mut self, x: N, y: N) -> usize { let index = self.pos >> 1; let (coords, mut ids) = split_data_borrow( @@ -88,6 +94,29 @@ impl KDTreeBuilder { index } + /// Add a coord to the KDTree. + /// + /// This returns a positional index that provides a lookup back into the original data. + #[inline] + pub fn add_coord(&mut self, coord: &impl CoordTrait) -> usize { + self.add(coord.x(), coord.y()) + } + + /// Add a point to the KDTree. + /// + /// This returns a positional index that provides a lookup back into the original data. + /// + /// ## Errors + /// + /// - If the point is empty. + #[inline] + pub fn add_point(&mut self, point: &impl PointTrait) -> Result { + let coord = point.coord().ok_or(GeoIndexError::General( + "Unable to add empty point to KDTree".to_string(), + ))?; + Ok(self.add_coord(&coord)) + } + /// Consume this builder, perfoming the k-d sort and generating a KDTree ready for queries. pub fn finish(mut self) -> OwnedKDTree { assert_eq!( diff --git a/src/kdtree/trait.rs b/src/kdtree/trait.rs index 575c5db..f0f6fac 100644 --- a/src/kdtree/trait.rs +++ b/src/kdtree/trait.rs @@ -1,5 +1,6 @@ use std::borrow::Cow; +use geo_traits::{CoordTrait, RectTrait}; use tinyvec::TinyVec; use crate::indices::Indices; @@ -80,6 +81,18 @@ pub trait KDTreeIndex { result } + /// Search the index for items within a given bounding box. + /// + /// Returns indices of found items + fn range_rect(&self, rect: &impl RectTrait) -> Vec { + self.range( + rect.min().x(), + rect.min().y(), + rect.max().x(), + rect.max().y(), + ) + } + /// Search the index for items within a given radius. /// /// - qx: x value of query point @@ -144,6 +157,16 @@ pub trait KDTreeIndex { } result } + + /// Search the index for items within a given radius. + /// + /// - coord: coordinate of query point + /// - r: radius + /// + /// Returns indices of found items + fn within_coord(&self, coord: &impl CoordTrait, r: N) -> Vec { + self.within(coord.x(), coord.y(), r) + } } impl KDTreeIndex for KDTreeRef<'_, N> { diff --git a/src/rtree/builder.rs b/src/rtree/builder.rs index 450e77e..75148cf 100644 --- a/src/rtree/builder.rs +++ b/src/rtree/builder.rs @@ -1,4 +1,5 @@ use bytemuck::cast_slice_mut; +use geo_traits::{CoordTrait, RectTrait}; use crate::indices::MutableIndices; use crate::r#type::IndexableNum; @@ -73,7 +74,9 @@ impl RTreeBuilder { } } - /// Add a given rectangle to the index. + /// Add a given rectangle to the RTree. + /// + /// This returns a positional index that provides a lookup back into the original data. #[inline] pub fn add(&mut self, min_x: N, min_y: N, max_x: N, max_y: N) -> usize { let index = self.pos >> 2; @@ -110,6 +113,19 @@ impl RTreeBuilder { index } + /// Add a given rectangle to the RTree. + /// + /// This returns a positional index that provides a lookup back into the original data. + #[inline] + pub fn add_rect(&mut self, rect: &impl RectTrait) -> usize { + self.add( + rect.min().x(), + rect.min().y(), + rect.max().x(), + rect.max().y(), + ) + } + /// Consume this builder, perfoming the sort and generating an RTree ready for queries. pub fn finish>(mut self) -> OwnedRTree { assert_eq!( diff --git a/src/rtree/index.rs b/src/rtree/index.rs index e6e2a97..3622004 100644 --- a/src/rtree/index.rs +++ b/src/rtree/index.rs @@ -23,7 +23,7 @@ pub(crate) struct TreeMetadata { } impl TreeMetadata { - pub fn try_new(data: &[u8]) -> Result { + fn try_new(data: &[u8]) -> Result { let magic = data[0]; if magic != 0xfb { return Err(GeoIndexError::General( diff --git a/src/rtree/trait.rs b/src/rtree/trait.rs index e8ff067..6e7feab 100644 --- a/src/rtree/trait.rs +++ b/src/rtree/trait.rs @@ -1,3 +1,5 @@ +use geo_traits::{CoordTrait, RectTrait}; + use crate::error::Result; use crate::indices::Indices; use crate::r#type::IndexableNum; @@ -97,6 +99,18 @@ pub trait RTreeIndex: Sized { results } + /// Search an RTree given the provided bounding box. + /// + /// Results are the indexes of the inserted objects in insertion order. + fn search_rect(&self, rect: &impl RectTrait) -> Vec { + self.search( + rect.min().x(), + rect.min().y(), + rect.max().x(), + rect.max().y(), + ) + } + // #[allow(unused_mut, unused_labels, unused_variables)] // fn neighbors(&self, x: N, y: N, max_distance: Option) -> Vec { // let boxes = self.boxes();