Bump geo-index to 0.2

geoarrow · Feb 13, 2025 · f8f6abc · f8f6abc
1 parent 7832713
commit f8f6abc
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 30 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git 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
@@ -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,17 +19,18 @@ 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 {
     ($struct_name:ty, $bounding_rect_fn:ident) => {
         impl RTree for $struct_name {
             type Output = OwnedRTree<f64>;
 
-            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<f64>;
 
-    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,15 +83,15 @@ impl RTree for &dyn NativeArray {
 impl<G: NativeArray> RTree for ChunkedGeometryArray<G> {
     type Output = Vec<OwnedRTree<f64>>;
 
-    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))
     }
 }
 
 impl RTree for &dyn ChunkedNativeArray {
     type Output = Result<Vec<OwnedRTree<f64>>>;
 
-    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
@@ -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<ValidOffsets>
@@ -43,14 +43,14 @@ impl<G: NativeArray> IndexedGeometryArray<G> {
         self.len() == 0
     }
 
-    pub fn search(&self, min_x: f64, min_y: f64, max_x: f64, max_y: f64) -> Vec<usize> {
+    pub fn search(&self, min_x: f64, min_y: f64, max_x: f64, max_y: f64) -> Vec<u32> {
         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<G2>,
-    ) -> impl Iterator<Item = (usize, usize)> + 'a {
+    ) -> impl Iterator<Item = (u32, u32)> + 'a {
         self.index
             .intersection_candidates_with_other_tree(&other.index)
     }
@@ -75,6 +75,7 @@ impl<'a, G: NativeArray + ArrayAccessor<'a>> IndexedGeometryArray<G> {
             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<G> {
                 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<RectArray>;
 pub type IndexedUnknownGeometryArray = IndexedGeometryArray<Arc<dyn NativeArray>>;
 
 impl<G: NativeArray> RTreeIndex<f64> for IndexedGeometryArray<G> {
+    fn metadata(&self) -> &RTreeMetadata<f64> {
+        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()
-    }
 }