Skip to content

Commit

Permalink
Implement geos bool ops (#924)
Browse files Browse the repository at this point in the history
Closes #440
  • Loading branch information
kylebarron authored Dec 10, 2024
1 parent 5e53eec commit e4add1d
Show file tree
Hide file tree
Showing 15 changed files with 447 additions and 75 deletions.
149 changes: 149 additions & 0 deletions rust/geoarrow/src/algorithm/geos/bool_ops.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
use arrow_array::BooleanArray;
use geo_traits::GeometryTrait;
use geos::Geom;

use crate::algorithm::native::{Binary, Unary};
use crate::array::GeometryArray;
use crate::error::{GeoArrowError, Result};
use crate::io::geos::scalar::{to_geos_geometry, GEOSGeometry};
use crate::trait_::NativeScalar;

pub trait BooleanOps<Rhs> {
fn intersects(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn crosses(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn disjoint(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn touches(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn overlaps(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn within(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn equals(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn equals_exact(&self, rhs: &Rhs, precision: f64) -> Result<BooleanArray>;
fn covers(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn covered_by(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn contains(&self, rhs: &Rhs) -> Result<BooleanArray>;

fn difference(&self, rhs: &Rhs) -> Result<GeometryArray>;
fn sym_difference(&self, rhs: &Rhs) -> Result<GeometryArray>;
fn union(&self, rhs: &Rhs) -> Result<GeometryArray>;
fn intersection(&self, rhs: &Rhs) -> Result<GeometryArray>;
}

macro_rules! impl_method {
($method_name:ident) => {
fn $method_name(&self, rhs: &GeometryArray) -> Result<BooleanArray> {
self.try_binary_boolean(rhs, |left, right| {
Ok(left.to_geos()?.$method_name(&right.to_geos()?)?)
})
}
};
}

macro_rules! impl_method_geometry {
($method_name:ident) => {
fn $method_name(&self, rhs: &GeometryArray) -> Result<GeometryArray> {
self.try_binary_geometry(
rhs,
|left, right| {
let left = to_geos_geometry(&left)?;
let right = to_geos_geometry(&right)?;
let out = left.difference(&right)?;
Ok(GEOSGeometry::new(out))
},
false,
)
}
};
}

impl BooleanOps<GeometryArray> for GeometryArray {
impl_method!(intersects);
impl_method!(crosses);
impl_method!(disjoint);
impl_method!(touches);
impl_method!(overlaps);
impl_method!(within);
impl_method!(equals);
impl_method!(covers);
impl_method!(covered_by);
impl_method!(contains);

fn equals_exact(&self, rhs: &GeometryArray, precision: f64) -> Result<BooleanArray> {
self.try_binary_boolean(rhs, |left, right| {
Ok(left.to_geos()?.equals_exact(&right.to_geos()?, precision)?)
})
}

impl_method_geometry!(difference);
impl_method_geometry!(sym_difference);
impl_method_geometry!(union);
impl_method_geometry!(intersection);
}

pub trait BooleanOpsScalar<Rhs> {
fn intersects(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn crosses(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn disjoint(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn touches(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn overlaps(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn within(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn equals(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn equals_exact(&self, rhs: &Rhs, precision: f64) -> Result<BooleanArray>;
fn covers(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn covered_by(&self, rhs: &Rhs) -> Result<BooleanArray>;
fn contains(&self, rhs: &Rhs) -> Result<BooleanArray>;

fn difference(&self, rhs: &Rhs) -> Result<GeometryArray>;
fn sym_difference(&self, rhs: &Rhs) -> Result<GeometryArray>;
fn union(&self, rhs: &Rhs) -> Result<GeometryArray>;
fn intersection(&self, rhs: &Rhs) -> Result<GeometryArray>;
}

macro_rules! impl_method_scalar {
($method_name:ident) => {
fn $method_name(&self, rhs: &G) -> Result<BooleanArray> {
let rhs = to_geos_geometry(rhs)?;
self.try_unary_boolean::<_, GeoArrowError>(|geom| {
Ok(geom.to_geos()?.$method_name(&rhs)?)
})
}
};
}

macro_rules! impl_method_geometry_scalar {
($method_name:ident) => {
fn $method_name(&self, rhs: &G) -> Result<GeometryArray> {
let rhs = to_geos_geometry(rhs)?;
self.try_unary_geometry(
|geom| {
let geom = to_geos_geometry(&geom)?;
let out = geom.$method_name(&rhs)?;
Ok(GEOSGeometry::new(out))
},
false,
)
}
};
}
impl<G: GeometryTrait<T = f64>> BooleanOpsScalar<G> for GeometryArray {
impl_method_scalar!(intersects);
impl_method_scalar!(crosses);
impl_method_scalar!(disjoint);
impl_method_scalar!(touches);
impl_method_scalar!(overlaps);
impl_method_scalar!(within);
impl_method_scalar!(equals);
impl_method_scalar!(covers);
impl_method_scalar!(covered_by);
impl_method_scalar!(contains);

fn equals_exact(&self, rhs: &G, precision: f64) -> Result<BooleanArray> {
let rhs = to_geos_geometry(rhs)?;
self.try_unary_boolean::<_, GeoArrowError>(|geom| {
Ok(geom.to_geos()?.equals_exact(&rhs, precision)?)
})
}

impl_method_geometry_scalar!(difference);
impl_method_geometry_scalar!(sym_difference);
impl_method_geometry_scalar!(union);
impl_method_geometry_scalar!(intersection);
}
2 changes: 2 additions & 0 deletions rust/geoarrow/src/algorithm/geos/mod.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
//! Bindings to the [`geos`] crate for geometry operations.
mod area;
mod bool_ops;
mod buffer;
mod is_empty;
mod is_ring;
Expand All @@ -10,6 +11,7 @@ mod length;
mod util;

pub use area::Area;
pub use bool_ops::{BooleanOps, BooleanOpsScalar};
pub use buffer::Buffer;
pub use is_empty::IsEmpty;
pub use is_ring::IsRing;
Expand Down
88 changes: 79 additions & 9 deletions rust/geoarrow/src/algorithm/native/binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,13 @@ use arrow_array::{BooleanArray, PrimitiveArray};
use arrow_buffer::ArrowNativeType;
use arrow_buffer::{BooleanBufferBuilder, BufferBuilder, MutableBuffer, NullBuffer};
use arrow_data::ArrayData;
use geo_traits::GeometryTrait;

use crate::array::*;
use crate::error::{GeoArrowError, Result};
use crate::trait_::ArrayAccessor;

pub trait Binary<'a, Rhs: ArrayAccessor<'a> = Self>: ArrayAccessor<'a> {
pub trait Binary<'a, Rhs: ArrayAccessor<'a> = Self>: ArrayAccessor<'a> + NativeArray {
fn binary_boolean<F>(&'a self, rhs: &'a Rhs, op: F) -> Result<BooleanArray>
where
F: Fn(Self::Item, Rhs::Item) -> bool,
Expand Down Expand Up @@ -116,92 +117,161 @@ pub trait Binary<'a, Rhs: ArrayAccessor<'a> = Self>: ArrayAccessor<'a> {
Ok(PrimitiveArray::new(values, Some(nulls)))
}
}

fn try_binary_geometry<F, G>(
&'a self,
rhs: &'a Rhs,
op: F,
prefer_multi: bool,
) -> Result<GeometryArray>
where
G: GeometryTrait<T = f64>,
F: Fn(Self::Item, Rhs::Item) -> Result<G>,
{
if self.len() != rhs.len() {
return Err(GeoArrowError::General(
"Cannot perform binary operation on arrays of different length".to_string(),
));
}

let mut builder = GeometryBuilder::with_capacity_and_options(
Default::default(),
self.coord_type(),
self.metadata().clone(),
prefer_multi,
);

if self.is_empty() {
return Ok(builder.finish());
}

for (left, right) in self.iter().zip(rhs.iter()) {
if let (Some(left), Some(right)) = (left, right) {
builder.push_geometry(Some(&op(left, right)?))?;
} else {
builder.push_null();
}
}
Ok(builder.finish())
}
}

// Implementations on PointArray
impl Binary<'_, PointArray> for PointArray {}
impl Binary<'_, PointArray> for RectArray {}
impl Binary<'_, PointArray> for LineStringArray {}
impl Binary<'_, PointArray> for PolygonArray {}
impl Binary<'_, PointArray> for MultiPointArray {}
impl Binary<'_, PointArray> for MultiLineStringArray {}
impl Binary<'_, PointArray> for MultiPolygonArray {}
impl Binary<'_, PointArray> for MixedGeometryArray {}
impl Binary<'_, PointArray> for GeometryCollectionArray {}
impl Binary<'_, PointArray> for RectArray {}
impl Binary<'_, PointArray> for GeometryArray {}

// Implementations on LineStringArray
impl Binary<'_, LineStringArray> for PointArray {}
impl Binary<'_, LineStringArray> for RectArray {}
impl Binary<'_, LineStringArray> for LineStringArray {}
impl Binary<'_, LineStringArray> for PolygonArray {}
impl Binary<'_, LineStringArray> for MultiPointArray {}
impl Binary<'_, LineStringArray> for MultiLineStringArray {}
impl Binary<'_, LineStringArray> for MultiPolygonArray {}
impl Binary<'_, LineStringArray> for MixedGeometryArray {}
impl Binary<'_, LineStringArray> for GeometryCollectionArray {}
impl Binary<'_, LineStringArray> for RectArray {}
impl Binary<'_, LineStringArray> for GeometryArray {}

// Implementations on PolygonArray
impl Binary<'_, PolygonArray> for PointArray {}
impl Binary<'_, PolygonArray> for RectArray {}
impl Binary<'_, PolygonArray> for LineStringArray {}
impl Binary<'_, PolygonArray> for PolygonArray {}
impl Binary<'_, PolygonArray> for MultiPointArray {}
impl Binary<'_, PolygonArray> for MultiLineStringArray {}
impl Binary<'_, PolygonArray> for MultiPolygonArray {}
impl Binary<'_, PolygonArray> for MixedGeometryArray {}
impl Binary<'_, PolygonArray> for GeometryCollectionArray {}
impl Binary<'_, PolygonArray> for RectArray {}
impl Binary<'_, PolygonArray> for GeometryArray {}

// Implementations on MultiPointArray
impl Binary<'_, MultiPointArray> for PointArray {}
impl Binary<'_, MultiPointArray> for RectArray {}
impl Binary<'_, MultiPointArray> for LineStringArray {}
impl Binary<'_, MultiPointArray> for PolygonArray {}
impl Binary<'_, MultiPointArray> for MultiPointArray {}
impl Binary<'_, MultiPointArray> for MultiLineStringArray {}
impl Binary<'_, MultiPointArray> for MultiPolygonArray {}
impl Binary<'_, MultiPointArray> for MixedGeometryArray {}
impl Binary<'_, MultiPointArray> for GeometryCollectionArray {}
impl Binary<'_, MultiPointArray> for RectArray {}
impl Binary<'_, MultiPointArray> for GeometryArray {}

// Implementations on MultiLineStringArray
impl Binary<'_, MultiLineStringArray> for PointArray {}
impl Binary<'_, MultiLineStringArray> for RectArray {}
impl Binary<'_, MultiLineStringArray> for LineStringArray {}
impl Binary<'_, MultiLineStringArray> for PolygonArray {}
impl Binary<'_, MultiLineStringArray> for MultiPointArray {}
impl Binary<'_, MultiLineStringArray> for MultiLineStringArray {}
impl Binary<'_, MultiLineStringArray> for MultiPolygonArray {}
impl Binary<'_, MultiLineStringArray> for MixedGeometryArray {}
impl Binary<'_, MultiLineStringArray> for GeometryCollectionArray {}
impl Binary<'_, MultiLineStringArray> for RectArray {}
impl Binary<'_, MultiLineStringArray> for GeometryArray {}

// Implementations on MultiPolygonArray
impl Binary<'_, MultiPolygonArray> for PointArray {}
impl Binary<'_, MultiPolygonArray> for RectArray {}
impl Binary<'_, MultiPolygonArray> for LineStringArray {}
impl Binary<'_, MultiPolygonArray> for PolygonArray {}
impl Binary<'_, MultiPolygonArray> for MultiPointArray {}
impl Binary<'_, MultiPolygonArray> for MultiLineStringArray {}
impl Binary<'_, MultiPolygonArray> for MultiPolygonArray {}
impl Binary<'_, MultiPolygonArray> for MixedGeometryArray {}
impl Binary<'_, MultiPolygonArray> for GeometryCollectionArray {}
impl Binary<'_, MultiPolygonArray> for RectArray {}
impl Binary<'_, MultiPolygonArray> for GeometryArray {}

// Implementations on MixedGeometryArray
impl Binary<'_, MixedGeometryArray> for PointArray {}
impl Binary<'_, MixedGeometryArray> for RectArray {}
impl Binary<'_, MixedGeometryArray> for LineStringArray {}
impl Binary<'_, MixedGeometryArray> for PolygonArray {}
impl Binary<'_, MixedGeometryArray> for MultiPointArray {}
impl Binary<'_, MixedGeometryArray> for MultiLineStringArray {}
impl Binary<'_, MixedGeometryArray> for MultiPolygonArray {}
impl Binary<'_, MixedGeometryArray> for MixedGeometryArray {}
impl Binary<'_, MixedGeometryArray> for GeometryCollectionArray {}
impl Binary<'_, MixedGeometryArray> for RectArray {}
impl Binary<'_, MixedGeometryArray> for GeometryArray {}

// Implementations on GeometryCollectionArray
impl Binary<'_, GeometryCollectionArray> for PointArray {}
impl Binary<'_, GeometryCollectionArray> for RectArray {}
impl Binary<'_, GeometryCollectionArray> for LineStringArray {}
impl Binary<'_, GeometryCollectionArray> for PolygonArray {}
impl Binary<'_, GeometryCollectionArray> for MultiPointArray {}
impl Binary<'_, GeometryCollectionArray> for MultiLineStringArray {}
impl Binary<'_, GeometryCollectionArray> for MultiPolygonArray {}
impl Binary<'_, GeometryCollectionArray> for MixedGeometryArray {}
impl Binary<'_, GeometryCollectionArray> for GeometryCollectionArray {}
impl Binary<'_, GeometryCollectionArray> for RectArray {}
impl Binary<'_, GeometryCollectionArray> for GeometryArray {}

// Implementations on RectArray
impl Binary<'_, RectArray> for PointArray {}
impl Binary<'_, RectArray> for LineStringArray {}
impl Binary<'_, RectArray> for PolygonArray {}
impl Binary<'_, RectArray> for MultiPointArray {}
impl Binary<'_, RectArray> for MultiLineStringArray {}
impl Binary<'_, RectArray> for MultiPolygonArray {}
impl Binary<'_, RectArray> for MixedGeometryArray {}
impl Binary<'_, RectArray> for GeometryCollectionArray {}
impl Binary<'_, RectArray> for RectArray {}
impl Binary<'_, RectArray> for GeometryArray {}

// Implementations on GeometryArray
impl Binary<'_, GeometryArray> for PointArray {}
impl Binary<'_, GeometryArray> for LineStringArray {}
impl Binary<'_, GeometryArray> for PolygonArray {}
impl Binary<'_, GeometryArray> for MultiPointArray {}
impl Binary<'_, GeometryArray> for MultiLineStringArray {}
impl Binary<'_, GeometryArray> for MultiPolygonArray {}
impl Binary<'_, GeometryArray> for MixedGeometryArray {}
impl Binary<'_, GeometryArray> for GeometryCollectionArray {}
impl Binary<'_, GeometryArray> for RectArray {}
impl Binary<'_, GeometryArray> for GeometryArray {}
Loading

0 comments on commit e4add1d

Please sign in to comment.