diff --git a/Cargo.lock b/Cargo.lock index 406c3956c..50c778435 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1811,12 +1811,13 @@ dependencies = [ [[package]] name = "geo-index" -version = "0.1.1" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "818dd7464a1edadbe7932a09b8e3672216597557100f23d1315f351e46c2c20e" +checksum = "8e6b2a121e60180a118037426b1e93d7f5904ebd95cf6e841a575195a65a31ce" dependencies = [ "bytemuck", "float_next_after", + "geo-traits", "num-traits", "thiserror 1.0.69", "tinyvec", diff --git a/rust/geoarrow/Cargo.toml b/rust/geoarrow/Cargo.toml index a748b3fc6..274b4e93d 100644 --- a/rust/geoarrow/Cargo.toml +++ b/rust/geoarrow/Cargo.toml @@ -64,7 +64,7 @@ flatgeobuf = { version = "4.6", optional = true, default-features = false } futures = { version = "0.3", optional = true } gdal = { version = "0.17", optional = true } geo = "0.29.3" -geo-index = "0.1.1" +geo-index = "0.2" geo-traits = "0.2" geos = { version = "10", features = ["v3_10_0"], optional = true } geozero = { version = "0.14", features = ["with-wkb"] } diff --git a/rust/geoarrow/src/algorithm/geo_index/rtree.rs b/rust/geoarrow/src/algorithm/geo_index/rtree.rs index fe3f42fa9..b4d35c72d 100644 --- a/rust/geoarrow/src/algorithm/geo_index/rtree.rs +++ b/rust/geoarrow/src/algorithm/geo_index/rtree.rs @@ -10,7 +10,7 @@ use crate::error::Result; use crate::trait_::ArrayAccessor; use crate::NativeArray; use geo_index::rtree::sort::HilbertSort; -use geo_index::rtree::{OwnedRTree, RTreeBuilder}; +use geo_index::rtree::{RTree as OwnedRTree, RTreeBuilder}; pub trait RTree { type Output; @@ -19,7 +19,7 @@ pub trait RTree { self.create_rtree_with_node_size(16) } - fn create_rtree_with_node_size(&self, node_size: usize) -> Self::Output; + fn create_rtree_with_node_size(&self, node_size: u16) -> Self::Output; } macro_rules! impl_rtree { @@ -27,9 +27,10 @@ macro_rules! impl_rtree { impl RTree for $struct_name { type Output = OwnedRTree; - fn create_rtree_with_node_size(&self, node_size: usize) -> Self::Output { + fn create_rtree_with_node_size(&self, node_size: u16) -> Self::Output { assert_eq!(self.null_count(), 0); - let mut builder = RTreeBuilder::new_with_node_size(self.len(), node_size); + let mut builder = + RTreeBuilder::new_with_node_size(self.len().try_into().unwrap(), node_size); self.iter().flatten().for_each(|geom| { let ([min_x, min_y], [max_x, max_y]) = $bounding_rect_fn(&geom); @@ -56,7 +57,7 @@ impl_rtree!(GeometryArray, bounding_rect_geometry); impl RTree for &dyn NativeArray { type Output = OwnedRTree; - fn create_rtree_with_node_size(&self, node_size: usize) -> Self::Output { + fn create_rtree_with_node_size(&self, node_size: u16) -> Self::Output { use NativeType::*; macro_rules! impl_method { @@ -82,7 +83,7 @@ impl RTree for &dyn NativeArray { impl RTree for ChunkedGeometryArray { type Output = Vec>; - fn create_rtree_with_node_size(&self, node_size: usize) -> Self::Output { + fn create_rtree_with_node_size(&self, node_size: u16) -> Self::Output { self.map(|chunk| chunk.as_ref().create_rtree_with_node_size(node_size)) } } @@ -90,7 +91,7 @@ impl RTree for ChunkedGeometryArray { impl RTree for &dyn ChunkedNativeArray { type Output = Result>>; - fn create_rtree_with_node_size(&self, node_size: usize) -> Self::Output { + fn create_rtree_with_node_size(&self, node_size: u16) -> Self::Output { use NativeType::*; macro_rules! impl_method { diff --git a/rust/geoarrow/src/indexed/array.rs b/rust/geoarrow/src/indexed/array.rs index aac1f8b18..625140f8c 100644 --- a/rust/geoarrow/src/indexed/array.rs +++ b/rust/geoarrow/src/indexed/array.rs @@ -9,7 +9,7 @@ use crate::NativeArray; use arrow_array::builder::BooleanBuilder; use arrow_array::BooleanArray; use arrow_buffer::{BooleanBufferBuilder, NullBuffer}; -use geo_index::rtree::{OwnedRTree, RTreeIndex}; +use geo_index::rtree::{RTree as OwnedRTree, RTreeIndex, RTreeMetadata}; use geo_traits::{CoordTrait, RectTrait}; // TODO: also store Option @@ -43,14 +43,14 @@ impl IndexedGeometryArray { self.len() == 0 } - pub fn search(&self, min_x: f64, min_y: f64, max_x: f64, max_y: f64) -> Vec { + pub fn search(&self, min_x: f64, min_y: f64, max_x: f64, max_y: f64) -> Vec { self.index.search(min_x, min_y, max_x, max_y) } pub fn intersection_candidates_with_other<'a, G2: NativeArray>( &'a self, other: &'a IndexedGeometryArray, - ) -> impl Iterator + 'a { + ) -> impl Iterator + 'a { self.index .intersection_candidates_with_other_tree(&other.index) } @@ -75,6 +75,7 @@ impl<'a, G: NativeArray + ArrayAccessor<'a>> IndexedGeometryArray { rhs_rect.max().x(), rhs_rect.max().y(), ) { + let candidate_idx = candidate_idx as usize; buffer.set_bit(candidate_idx, op(self.array.value(candidate_idx))); } @@ -115,6 +116,9 @@ impl<'a, G: NativeArray + ArrayAccessor<'a>> IndexedGeometryArray { continue; } + let left_candidate_idx = left_candidate_idx as usize; + let right_candidate_idx = right_candidate_idx as usize; + let left = self.array.value(left_candidate_idx); let right = other.array.value(right_candidate_idx); @@ -141,27 +145,15 @@ pub type IndexedRectArray = IndexedGeometryArray; pub type IndexedUnknownGeometryArray = IndexedGeometryArray>; impl RTreeIndex for IndexedGeometryArray { + fn metadata(&self) -> &RTreeMetadata { + self.index.metadata() + } + fn boxes(&self) -> &[f64] { self.index.boxes() } - fn indices(&self) -> std::borrow::Cow<'_, geo_index::indices::Indices> { + fn indices(&self) -> geo_index::indices::Indices { self.index.indices() } - - fn num_items(&self) -> usize { - self.index.num_items() - } - - fn num_nodes(&self) -> usize { - self.index.num_nodes() - } - - fn node_size(&self) -> usize { - self.index.node_size() - } - - fn level_bounds(&self) -> &[usize] { - self.index.level_bounds() - } }