Put rstar creation inside trait (#206)

src/algorithm/mod.rs

@@ -9,3 +9,4 @@ pub mod geos;
 pub mod native;
 #[cfg(feature = "proj")]
 pub mod proj;
+pub mod rstar;

src/algorithm/rstar.rs

@@ -0,0 +1,57 @@
+use crate::array::*;
+use arrow2::types::Offset;
+use rstar::primitives::CachedEnvelope;
+
+/// Construct an R-Tree from a geometry array.
+pub trait RTree<'a> {
+    /// The object type to store in the RTree.
+    type RTreeObject: rstar::RTreeObject;
+
+    /// Build an [`RTree`] spatial index containing this array's geometries.
+    fn rstar_tree(&'a self) -> rstar::RTree<Self::RTreeObject>;
+}
+
+impl<'a> RTree<'a> for PointArray {
+    type RTreeObject = crate::scalar::Point<'a>;
+
+    fn rstar_tree(&'a self) -> rstar::RTree<Self::RTreeObject> {
+        // Note: for points we don't memoize with CachedEnvelope
+        rstar::RTree::bulk_load(self.iter().flatten().collect())
+    }
+}
+
+impl<'a> RTree<'a> for RectArray {
+    type RTreeObject = crate::scalar::Rect<'a>;
+
+    fn rstar_tree(&'a self) -> rstar::RTree<Self::RTreeObject> {
+        // Note: for rects we don't memoize with CachedEnvelope
+        rstar::RTree::bulk_load(self.iter().flatten().collect())
+    }
+}
+
+macro_rules! iter_cached_impl {
+    ($type:ty, $scalar_type:ty) => {
+        impl<'a, O: Offset> RTree<'a> for $type {
+            type RTreeObject = CachedEnvelope<$scalar_type>;
+
+            fn rstar_tree(&'a self) -> rstar::RTree<Self::RTreeObject> {
+                rstar::RTree::bulk_load(self.iter().flatten().map(CachedEnvelope::new).collect())
+            }
+        }
+    };
+}
+
+iter_cached_impl!(LineStringArray<O>, crate::scalar::LineString<'a, O>);
+iter_cached_impl!(PolygonArray<O>, crate::scalar::Polygon<'a, O>);
+iter_cached_impl!(MultiPointArray<O>, crate::scalar::MultiPoint<'a, O>);
+iter_cached_impl!(
+    MultiLineStringArray<O>,
+    crate::scalar::MultiLineString<'a, O>
+);
+iter_cached_impl!(MultiPolygonArray<O>, crate::scalar::MultiPolygon<'a, O>);
+iter_cached_impl!(WKBArray<O>, crate::scalar::WKB<'a, O>);
+iter_cached_impl!(MixedGeometryArray<O>, crate::scalar::Geometry<'a, O>);
+iter_cached_impl!(
+    GeometryCollectionArray<O>,
+    crate::scalar::GeometryCollection<'a, O>
+);

src/array/binary/array.rs

@@ -8,8 +8,6 @@ use arrow2::bitmap::utils::{BitmapIter, ZipValidity};
 use arrow2::bitmap::Bitmap;
 use arrow2::datatypes::DataType;
 use arrow2::types::Offset;
-use rstar::primitives::CachedEnvelope;
-use rstar::RTree;
 
 /// An immutable array of WKB geometries using GeoArrow's in-memory representation.
 ///
@@ -43,7 +41,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for WKBArray<O> {
     type Scalar = WKB<'a, O>;
     type ScalarGeo = geo::Geometry;
     type ArrowArray = BinaryArray<O>;
-    type RTreeObject = CachedEnvelope<Self::Scalar>;
 
     fn value(&'a self, i: usize) -> Self::Scalar {
         WKB::new_borrowed(&self.0, i)
@@ -87,11 +84,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for WKBArray<O> {
         self
     }
 
-    /// Build a spatial index containing this array's geometries
-    fn rstar_tree(&'a self) -> RTree<Self::RTreeObject> {
-        RTree::bulk_load(self.iter().flatten().map(CachedEnvelope::new).collect())
-    }
-
     /// Returns the number of geometries in this array
     #[inline]
     fn len(&self) -> usize {

src/array/coord/combined/array.rs

@@ -8,7 +8,6 @@ use crate::GeometryArrayTrait;
 use arrow2::array::{Array, FixedSizeListArray, StructArray};
 use arrow2::datatypes::DataType;
 use itertools::Itertools;
-use rstar::RTree;
 
 /// An Arrow representation of an array of coordinates.
 ///
@@ -44,7 +43,6 @@ impl<'a> GeometryArrayTrait<'a> for CoordBuffer {
     type ArrowArray = Box<dyn Array>;
     type Scalar = Coord<'a>;
     type ScalarGeo = geo::Coord;
-    type RTreeObject = Self::Scalar;
 
     fn value(&'a self, i: usize) -> Self::Scalar {
         match self {
@@ -109,10 +107,6 @@ impl<'a> GeometryArrayTrait<'a> for CoordBuffer {
         }
     }
 
-    fn rstar_tree(&'a self) -> RTree<Self::Scalar> {
-        panic!("not implemented for coords");
-    }
-
     fn len(&self) -> usize {
         match self {
             CoordBuffer::Interleaved(c) => c.len(),

src/array/coord/interleaved/array.rs

@@ -5,7 +5,6 @@ use crate::GeometryArrayTrait;
 use arrow2::array::{Array, FixedSizeListArray, PrimitiveArray};
 use arrow2::buffer::Buffer;
 use arrow2::datatypes::{DataType, Field};
-use rstar::RTree;
 
 /// A an array of XY coordinates stored interleaved in a single buffer.
 #[derive(Debug, Clone, PartialEq)]
@@ -57,7 +56,6 @@ impl<'a> GeometryArrayTrait<'a> for InterleavedCoordBuffer {
     type ArrowArray = FixedSizeListArray;
     type Scalar = InterleavedCoord<'a>;
     type ScalarGeo = geo::Coord;
-    type RTreeObject = Self::Scalar;
 
     fn value(&'a self, i: usize) -> Self::Scalar {
         InterleavedCoord {
@@ -94,10 +92,6 @@ impl<'a> GeometryArrayTrait<'a> for InterleavedCoordBuffer {
         panic!("into_coord_type only implemented on CoordBuffer");
     }
 
-    fn rstar_tree(&'a self) -> RTree<Self::Scalar> {
-        panic!("not implemented for coords");
-    }
-
     fn len(&self) -> usize {
         self.coords.len() / 2
     }

src/array/coord/separated/array.rs

@@ -1,7 +1,6 @@
 use arrow2::array::{Array, PrimitiveArray, StructArray};
 use arrow2::buffer::Buffer;
 use arrow2::datatypes::{DataType, Field};
-use rstar::RTree;
 
 use crate::array::CoordType;
 use crate::error::{GeoArrowError, Result};
@@ -64,7 +63,6 @@ impl<'a> GeometryArrayTrait<'a> for SeparatedCoordBuffer {
     type ArrowArray = StructArray;
     type Scalar = SeparatedCoord<'a>;
     type ScalarGeo = geo::Coord;
-    type RTreeObject = Self::Scalar;
 
     fn value(&'a self, i: usize) -> Self::Scalar {
         SeparatedCoord {
@@ -102,10 +100,6 @@ impl<'a> GeometryArrayTrait<'a> for SeparatedCoordBuffer {
         panic!("into_coord_type only implemented on CoordBuffer");
     }
 
-    fn rstar_tree(&'a self) -> RTree<Self::Scalar> {
-        panic!("not implemented for coords");
-    }
-
     fn len(&self) -> usize {
         self.x.len()
     }

src/array/geometry/array.rs

@@ -2,7 +2,6 @@ use arrow2::array::Array;
 use arrow2::bitmap::Bitmap;
 use arrow2::datatypes::DataType;
 use arrow2::types::Offset;
-use rstar::RTree;
 
 use crate::algorithm::native::type_id::TypeIds;
 use crate::array::{
@@ -32,7 +31,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for GeometryArray<O> {
     type Scalar = crate::scalar::Geometry<'a, O>;
     type ScalarGeo = geo::Geometry;
     type ArrowArray = Box<dyn Array>;
-    type RTreeObject = Self::Scalar;
 
     fn value(&'a self, i: usize) -> Self::Scalar {
         match self {
@@ -134,13 +132,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for GeometryArray<O> {
         }
     }
 
-    fn rstar_tree(&'a self) -> RTree<Self::Scalar> {
-        let elements: Vec<_> = (0..self.len())
-            .filter_map(|geom_idx| self.get(geom_idx))
-            .collect();
-        RTree::bulk_load(elements)
-    }
-
     /// The length of the [`GeometryArray`]. Every array has a length corresponding to the number
     /// of geometries it contains.
     fn len(&self) -> usize {

src/array/geometrycollection/array.rs

@@ -4,8 +4,6 @@ use arrow2::bitmap::Bitmap;
 use arrow2::datatypes::DataType;
 use arrow2::offset::OffsetsBuffer;
 use arrow2::types::Offset;
-use rstar::primitives::CachedEnvelope;
-use rstar::RTree;
 
 use crate::array::{CoordBuffer, CoordType, MixedGeometryArray};
 use crate::scalar::GeometryCollection;
@@ -50,7 +48,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for GeometryCollectionArray<O> {
     type Scalar = GeometryCollection<'a, O>;
     type ScalarGeo = geo::GeometryCollection;
     type ArrowArray = ListArray<O>;
-    type RTreeObject = CachedEnvelope<Self::Scalar>;
 
     fn value(&'a self, i: usize) -> Self::Scalar {
         GeometryCollection {
@@ -95,11 +92,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for GeometryCollectionArray<O> {
         todo!()
     }
 
-    /// Build a spatial index containing this array's geometries
-    fn rstar_tree(&'a self) -> RTree<Self::RTreeObject> {
-        RTree::bulk_load(self.iter().flatten().map(CachedEnvelope::new).collect())
-    }
-
     /// Returns the number of geometries in this array
     #[inline]
     fn len(&self) -> usize {

src/array/linestring/array.rs

@@ -9,8 +9,6 @@ use arrow2::bitmap::Bitmap;
 use arrow2::datatypes::{DataType, Field};
 use arrow2::offset::OffsetsBuffer;
 use arrow2::types::Offset;
-use rstar::primitives::CachedEnvelope;
-use rstar::RTree;
 
 use super::MutableLineStringArray;
 
@@ -113,7 +111,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for LineStringArray<O> {
     type Scalar = LineString<'a, O>;
     type ScalarGeo = geo::LineString;
     type ArrowArray = ListArray<O>;
-    type RTreeObject = CachedEnvelope<Self::Scalar>;
 
     /// Gets the value at slot `i`
     fn value(&'a self, i: usize) -> Self::Scalar {
@@ -160,11 +157,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for LineStringArray<O> {
         )
     }
 
-    /// Build a spatial index containing this array's geometries
-    fn rstar_tree(&'a self) -> RTree<Self::RTreeObject> {
-        RTree::bulk_load(self.iter().flatten().map(CachedEnvelope::new).collect())
-    }
-
     /// Returns the number of geometries in this array
     #[inline]
     fn len(&self) -> usize {

src/array/mixed/array.rs

@@ -3,8 +3,6 @@ use arrow2::bitmap::Bitmap;
 use arrow2::buffer::Buffer;
 use arrow2::datatypes::{DataType, Field, UnionMode};
 use arrow2::types::Offset;
-use rstar::primitives::CachedEnvelope;
-use rstar::RTree;
 
 use crate::array::mixed::mutable::MutableMixedGeometryArray;
 use crate::array::{
@@ -170,7 +168,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for MixedGeometryArray<O> {
     type Scalar = Geometry<'a, O>;
     type ScalarGeo = geo::Geometry;
     type ArrowArray = UnionArray;
-    type RTreeObject = CachedEnvelope<Self::Scalar>;
 
     /// Gets the value at slot `i`
     fn value(&'a self, i: usize) -> Self::Scalar {
@@ -280,11 +277,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for MixedGeometryArray<O> {
         todo!();
     }
 
-    /// Build a spatial index containing this array's geometries
-    fn rstar_tree(&'a self) -> RTree<Self::RTreeObject> {
-        RTree::bulk_load(self.iter().flatten().map(CachedEnvelope::new).collect())
-    }
-
     /// Returns the number of geometries in this array
     #[inline]
     fn len(&self) -> usize {

src/array/multilinestring/array.rs

@@ -9,8 +9,6 @@ use arrow2::bitmap::Bitmap;
 use arrow2::datatypes::{DataType, Field};
 use arrow2::offset::{Offsets, OffsetsBuffer};
 use arrow2::types::Offset;
-use rstar::primitives::CachedEnvelope;
-use rstar::RTree;
 
 use super::MutableMultiLineStringArray;
 
@@ -148,7 +146,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for MultiLineStringArray<O> {
     type Scalar = MultiLineString<'a, O>;
     type ScalarGeo = geo::MultiLineString;
     type ArrowArray = ListArray<O>;
-    type RTreeObject = CachedEnvelope<Self::Scalar>;
 
     fn value(&'a self, i: usize) -> Self::Scalar {
         MultiLineString::new_borrowed(&self.coords, &self.geom_offsets, &self.ring_offsets, i)
@@ -198,11 +195,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for MultiLineStringArray<O> {
         )
     }
 
-    /// Build a spatial index containing this array's geometries
-    fn rstar_tree(&'a self) -> RTree<Self::RTreeObject> {
-        RTree::bulk_load(self.iter().flatten().map(CachedEnvelope::new).collect())
-    }
-
     /// Returns the number of geometries in this array
     #[inline]
     fn len(&self) -> usize {

src/array/multipoint/array.rs

@@ -10,8 +10,6 @@ use arrow2::bitmap::Bitmap;
 use arrow2::datatypes::{DataType, Field};
 use arrow2::offset::{Offsets, OffsetsBuffer};
 use arrow2::types::Offset;
-use rstar::primitives::CachedEnvelope;
-use rstar::RTree;
 
 /// An immutable array of MultiPoint geometries using GeoArrow's in-memory representation.
 ///
@@ -112,7 +110,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for MultiPointArray<O> {
     type Scalar = MultiPoint<'a, O>;
     type ScalarGeo = geo::MultiPoint;
     type ArrowArray = ListArray<O>;
-    type RTreeObject = CachedEnvelope<Self::Scalar>;
 
     fn value(&'a self, i: usize) -> Self::Scalar {
         MultiPoint::new_borrowed(&self.coords, &self.geom_offsets, i)
@@ -158,10 +155,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for MultiPointArray<O> {
         )
     }
 
-    fn rstar_tree(&'a self) -> RTree<Self::RTreeObject> {
-        RTree::bulk_load(self.iter().flatten().map(CachedEnvelope::new).collect())
-    }
-
     /// Returns the number of geometries in this array
     #[inline]
     fn len(&self) -> usize {

src/array/multipolygon/array.rs

@@ -9,8 +9,6 @@ use arrow2::bitmap::Bitmap;
 use arrow2::datatypes::{DataType, Field};
 use arrow2::offset::{Offsets, OffsetsBuffer};
 use arrow2::types::Offset;
-use rstar::primitives::CachedEnvelope;
-use rstar::RTree;
 
 use super::MutableMultiPolygonArray;
 
@@ -173,7 +171,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for MultiPolygonArray<O> {
     type Scalar = MultiPolygon<'a, O>;
     type ScalarGeo = geo::MultiPolygon;
     type ArrowArray = ListArray<O>;
-    type RTreeObject = CachedEnvelope<Self::Scalar>;
 
     fn value(&'a self, i: usize) -> Self::Scalar {
         MultiPolygon::new_borrowed(
@@ -238,11 +235,6 @@ impl<'a, O: Offset> GeometryArrayTrait<'a> for MultiPolygonArray<O> {
         )
     }
 
-    /// Build a spatial index containing this array's geometries
-    fn rstar_tree(&'a self) -> RTree<Self::RTreeObject> {
-        RTree::bulk_load(self.iter().flatten().map(CachedEnvelope::new).collect())
-    }
-
     /// Returns the number of geometries in this array
     #[inline]
     fn len(&self) -> usize {