Skip to content

Commit

Permalink
Binary operations simplification
Browse files Browse the repository at this point in the history
  • Loading branch information
@kylebarron
kylebarron committed Oct 6, 2024
1 parent 697bb36 commit 107501c
Show file tree
Hide file tree
Showing 5 changed files with 211 additions and 252 deletions.
330 changes: 82 additions & 248 deletions src/algorithm/geo/within.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
use crate::algorithm::native::binary::try_binary_boolean_native_geometry;
use crate::algorithm::native::Unary;
use crate::array::*;
use crate::scalar::*;
use crate::trait_::ArrayAccessor;
use crate::datatypes::{Dimension, NativeType};
use crate::error::{GeoArrowError, Result};
use crate::geo_traits::GeometryTrait;
use crate::io::geo::geometry_to_geo;
use crate::trait_::NativeGeometryAccessor;
use crate::trait_::NativeScalar;
use arrow_array::builder::BooleanBuilder;
use arrow_array::BooleanArray;
use geo::Within as _Within;

Expand Down Expand Up @@ -40,269 +44,99 @@ use geo::Within as _Within;
/// ```
///
/// [DE-9IM]: https://en.wikipedia.org/wiki/DE-9IM
pub trait Within<Other = Self> {
fn is_within(&self, b: &Other) -> BooleanArray;
pub trait Within<'a, Other> {
fn is_within(&'a self, b: &'a Other) -> Result<BooleanArray>;
}

// ┌────────────────────────────────┐
// │ Implementations for RHS arrays │
// └────────────────────────────────┘

// Note: this implementation is outside the macro because it is not generic over O
impl Within for PointArray<2> {
fn is_within(&self, rhs: &Self) -> BooleanArray {
assert_eq!(self.len(), rhs.len());

let mut output_array = BooleanBuilder::with_capacity(self.len());

self.iter_geo()
.zip(rhs.iter_geo())
.for_each(|(first, second)| match (first, second) {
(Some(first), Some(second)) => output_array.append_value(first.is_within(&second)),
_ => output_array.append_null(),
});

output_array.finish()
}
}

// Implementation that iterates over geo objects
macro_rules! iter_geo_impl {
($first:ty, $second:ty) => {
impl<'a> Within<$second> for $first {
fn is_within(&self, rhs: &$second) -> BooleanArray {
assert_eq!(self.len(), rhs.len());

let mut output_array = BooleanBuilder::with_capacity(self.len());

self.iter_geo()
.zip(rhs.iter_geo())
.for_each(|(first, second)| match (first, second) {
(Some(first), Some(second)) => {
output_array.append_value(first.is_within(&second))
}
_ => output_array.append_null(),
});

output_array.finish()
($array_type:ty) => {
impl<'a, R: NativeGeometryAccessor<'a, 2>> Within<'a, R> for $array_type {
fn is_within(&'a self, rhs: &'a R) -> Result<BooleanArray> {
try_binary_boolean_native_geometry(self, rhs, |l, r| {
Ok(l.to_geo().is_within(&r.to_geo()))
})
}
}
};
}

// Implementations on PointArray
iter_geo_impl!(PointArray<2>, LineStringArray<2>);
iter_geo_impl!(PointArray<2>, PolygonArray<2>);
iter_geo_impl!(PointArray<2>, MultiPointArray<2>);
iter_geo_impl!(PointArray<2>, MultiLineStringArray<2>);
iter_geo_impl!(PointArray<2>, MultiPolygonArray<2>);

// Implementations on LineStringArray
iter_geo_impl!(LineStringArray<2>, PointArray<2>);
iter_geo_impl!(LineStringArray<2>, LineStringArray<2>);
iter_geo_impl!(LineStringArray<2>, PolygonArray<2>);
iter_geo_impl!(LineStringArray<2>, MultiPointArray<2>);
iter_geo_impl!(LineStringArray<2>, MultiLineStringArray<2>);
iter_geo_impl!(LineStringArray<2>, MultiPolygonArray<2>);

// Implementations on PolygonArray
iter_geo_impl!(PolygonArray<2>, PointArray<2>);
iter_geo_impl!(PolygonArray<2>, LineStringArray<2>);
iter_geo_impl!(PolygonArray<2>, PolygonArray<2>);
iter_geo_impl!(PolygonArray<2>, MultiPointArray<2>);
iter_geo_impl!(PolygonArray<2>, MultiLineStringArray<2>);
iter_geo_impl!(PolygonArray<2>, MultiPolygonArray<2>);

// Implementations on MultiPointArray
iter_geo_impl!(MultiPointArray<2>, PointArray<2>);
iter_geo_impl!(MultiPointArray<2>, LineStringArray<2>);
iter_geo_impl!(MultiPointArray<2>, PolygonArray<2>);
iter_geo_impl!(MultiPointArray<2>, MultiPointArray<2>);
iter_geo_impl!(MultiPointArray<2>, MultiLineStringArray<2>);
iter_geo_impl!(MultiPointArray<2>, MultiPolygonArray<2>);

// Implementations on MultiLineStringArray
iter_geo_impl!(MultiLineStringArray<2>, PointArray<2>);
iter_geo_impl!(MultiLineStringArray<2>, LineStringArray<2>);
iter_geo_impl!(MultiLineStringArray<2>, PolygonArray<2>);
iter_geo_impl!(MultiLineStringArray<2>, MultiPointArray<2>);
iter_geo_impl!(MultiLineStringArray<2>, MultiLineStringArray<2>);
iter_geo_impl!(MultiLineStringArray<2>, MultiPolygonArray<2>);

// Implementations on MultiPolygonArray
iter_geo_impl!(MultiPolygonArray<2>, PointArray<2>);
iter_geo_impl!(MultiPolygonArray<2>, LineStringArray<2>);
iter_geo_impl!(MultiPolygonArray<2>, PolygonArray<2>);
iter_geo_impl!(MultiPolygonArray<2>, MultiPointArray<2>);
iter_geo_impl!(MultiPolygonArray<2>, MultiLineStringArray<2>);
iter_geo_impl!(MultiPolygonArray<2>, MultiPolygonArray<2>);

// ┌──────────────────────────────────────────┐
// │ Implementations for RHS geoarrow scalars │
// └──────────────────────────────────────────┘

// Note: this implementation is outside the macro because it is not generic over O
impl<'a> Within<Point<'a, 2>> for PointArray<2> {
fn is_within(&self, rhs: &Point<'a, 2>) -> BooleanArray {
let mut output_array = BooleanBuilder::with_capacity(self.len());

self.iter_geo().for_each(|maybe_point| {
let output = maybe_point.map(|point| point.is_within(&rhs.to_geo()));
output_array.append_option(output)
});

output_array.finish()
iter_geo_impl!(PointArray<2>);
iter_geo_impl!(LineStringArray<2>);
iter_geo_impl!(PolygonArray<2>);
iter_geo_impl!(MultiPointArray<2>);
iter_geo_impl!(MultiLineStringArray<2>);
iter_geo_impl!(MultiPolygonArray<2>);
iter_geo_impl!(MixedGeometryArray<2>);
iter_geo_impl!(GeometryCollectionArray<2>);
iter_geo_impl!(RectArray<2>);

impl<'a, R: NativeGeometryAccessor<'a, 2>> Within<'a, R> for &dyn NativeArray {
fn is_within(&'a self, rhs: &'a R) -> Result<BooleanArray> {
use Dimension::*;
use NativeType::*;

match self.data_type() {
Point(_, XY) => Within::is_within(self.as_point::<2>(), rhs),
LineString(_, XY) => Within::is_within(self.as_line_string::<2>(), rhs),
Polygon(_, XY) => Within::is_within(self.as_polygon::<2>(), rhs),
MultiPoint(_, XY) => Within::is_within(self.as_multi_point::<2>(), rhs),
MultiLineString(_, XY) => Within::is_within(self.as_multi_line_string::<2>(), rhs),
MultiPolygon(_, XY) => Within::is_within(self.as_multi_polygon::<2>(), rhs),
Mixed(_, XY) => Within::is_within(self.as_mixed::<2>(), rhs),
GeometryCollection(_, XY) => Within::is_within(self.as_geometry_collection::<2>(), rhs),
Rect(XY) => Within::is_within(self.as_rect::<2>(), rhs),
_ => Err(GeoArrowError::IncorrectType("".into())),
}
}
}

/// Implementation that iterates over geo objects
macro_rules! iter_geo_impl_geoarrow_scalar {
($first:ty, $second:ty) => {
impl<'a> Within<$second> for $first {
fn is_within(&self, rhs: &$second) -> BooleanArray {
let mut output_array = BooleanBuilder::with_capacity(self.len());
let rhs_geo = rhs.to_geo();

self.iter_geo().for_each(|maybe_geom| {
let output = maybe_geom.map(|geom| geom.is_within(&rhs_geo));
output_array.append_option(output)
});

output_array.finish()
}
}
};
pub trait WithinScalar<'a, G: GeometryTrait> {
fn is_within(&'a self, b: &'a G) -> Result<BooleanArray>;
}

// Implementations on PointArray
iter_geo_impl_geoarrow_scalar!(PointArray<2>, LineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(PointArray<2>, Polygon<'a, 2>);
iter_geo_impl_geoarrow_scalar!(PointArray<2>, MultiPoint<'a, 2>);
iter_geo_impl_geoarrow_scalar!(PointArray<2>, MultiLineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(PointArray<2>, MultiPolygon<'a, 2>);

// Implementations on LineStringArray
iter_geo_impl_geoarrow_scalar!(LineStringArray<2>, Point<'a, 2>);
iter_geo_impl_geoarrow_scalar!(LineStringArray<2>, LineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(LineStringArray<2>, Polygon<'a, 2>);
iter_geo_impl_geoarrow_scalar!(LineStringArray<2>, MultiPoint<'a, 2>);
iter_geo_impl_geoarrow_scalar!(LineStringArray<2>, MultiLineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(LineStringArray<2>, MultiPolygon<'a, 2>);

// Implementations on PolygonArray
iter_geo_impl_geoarrow_scalar!(PolygonArray<2>, Point<'a, 2>);
iter_geo_impl_geoarrow_scalar!(PolygonArray<2>, LineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(PolygonArray<2>, Polygon<'a, 2>);
iter_geo_impl_geoarrow_scalar!(PolygonArray<2>, MultiPoint<'a, 2>);
iter_geo_impl_geoarrow_scalar!(PolygonArray<2>, MultiLineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(PolygonArray<2>, MultiPolygon<'a, 2>);

// Implementations on MultiPointArray
iter_geo_impl_geoarrow_scalar!(MultiPointArray<2>, Point<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiPointArray<2>, LineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiPointArray<2>, Polygon<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiPointArray<2>, MultiPoint<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiPointArray<2>, MultiLineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiPointArray<2>, MultiPolygon<'a, 2>);

// Implementations on MultiLineStringArray
iter_geo_impl_geoarrow_scalar!(MultiLineStringArray<2>, Point<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiLineStringArray<2>, LineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiLineStringArray<2>, Polygon<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiLineStringArray<2>, MultiPoint<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiLineStringArray<2>, MultiLineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiLineStringArray<2>, MultiPolygon<'a, 2>);

// Implementations on MultiPolygonArray
iter_geo_impl_geoarrow_scalar!(MultiPolygonArray<2>, Point<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiPolygonArray<2>, LineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiPolygonArray<2>, Polygon<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiPolygonArray<2>, MultiPoint<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiPolygonArray<2>, MultiLineString<'a, 2>);
iter_geo_impl_geoarrow_scalar!(MultiPolygonArray<2>, MultiPolygon<'a, 2>);

// ┌─────────────────────────────────────┐
// │ Implementations for RHS geo scalars │
// └─────────────────────────────────────┘

macro_rules! non_generic_iter_geo_impl_geo_scalar {
($first:ty, $second:ty) => {
impl<'a> Within<$second> for $first {
fn is_within(&self, rhs: &$second) -> BooleanArray {
let mut output_array = BooleanBuilder::with_capacity(self.len());

self.iter_geo().for_each(|maybe_geom| {
let output = maybe_geom.map(|geom| geom.is_within(rhs));
output_array.append_option(output)
});

output_array.finish()
macro_rules! scalar_impl {
($array_type:ty) => {
impl<'a, G: GeometryTrait<T = f64>> WithinScalar<'a, G> for $array_type {
fn is_within(&'a self, rhs: &'a G) -> Result<BooleanArray> {
let right = geometry_to_geo(rhs);
self.try_unary_boolean(|left| {
Ok::<_, GeoArrowError>(left.to_geo().is_within(&right))
})
}
}
};
}

// Implementations on PointArray
non_generic_iter_geo_impl_geo_scalar!(PointArray<2>, geo::Point);
non_generic_iter_geo_impl_geo_scalar!(PointArray<2>, geo::LineString);
non_generic_iter_geo_impl_geo_scalar!(PointArray<2>, geo::Polygon);
non_generic_iter_geo_impl_geo_scalar!(PointArray<2>, geo::MultiPoint);
non_generic_iter_geo_impl_geo_scalar!(PointArray<2>, geo::MultiLineString);
non_generic_iter_geo_impl_geo_scalar!(PointArray<2>, geo::MultiPolygon);

macro_rules! iter_geo_impl_geo_scalar {
($first:ty, $second:ty) => {
impl<'a> Within<$second> for $first {
fn is_within(&self, rhs: &$second) -> BooleanArray {
let mut output_array = BooleanBuilder::with_capacity(self.len());

self.iter_geo().for_each(|maybe_geom| {
let output = maybe_geom.map(|geom| geom.is_within(rhs));
output_array.append_option(output)
});

output_array.finish()
scalar_impl!(PointArray<2>);
scalar_impl!(LineStringArray<2>);
scalar_impl!(PolygonArray<2>);
scalar_impl!(MultiPointArray<2>);
scalar_impl!(MultiLineStringArray<2>);
scalar_impl!(MultiPolygonArray<2>);
scalar_impl!(MixedGeometryArray<2>);
scalar_impl!(GeometryCollectionArray<2>);
scalar_impl!(RectArray<2>);

impl<'a, G: GeometryTrait<T = f64>> WithinScalar<'a, G> for &dyn NativeArray {
fn is_within(&'a self, rhs: &'a G) -> Result<BooleanArray> {
use Dimension::*;
use NativeType::*;

match self.data_type() {
Point(_, XY) => WithinScalar::is_within(self.as_point::<2>(), rhs),
LineString(_, XY) => WithinScalar::is_within(self.as_line_string::<2>(), rhs),
Polygon(_, XY) => WithinScalar::is_within(self.as_polygon::<2>(), rhs),
MultiPoint(_, XY) => WithinScalar::is_within(self.as_multi_point::<2>(), rhs),
MultiLineString(_, XY) => {
WithinScalar::is_within(self.as_multi_line_string::<2>(), rhs)
}
MultiPolygon(_, XY) => WithinScalar::is_within(self.as_multi_polygon::<2>(), rhs),
Mixed(_, XY) => WithinScalar::is_within(self.as_mixed::<2>(), rhs),
GeometryCollection(_, XY) => {
WithinScalar::is_within(self.as_geometry_collection::<2>(), rhs)
}
Rect(XY) => WithinScalar::is_within(self.as_rect::<2>(), rhs),
_ => Err(GeoArrowError::IncorrectType("".into())),
}
};
}
}

// Implementations on LineStringArray
iter_geo_impl_geo_scalar!(LineStringArray<2>, geo::Point);
iter_geo_impl_geo_scalar!(LineStringArray<2>, geo::LineString);
iter_geo_impl_geo_scalar!(LineStringArray<2>, geo::Polygon);
iter_geo_impl_geo_scalar!(LineStringArray<2>, geo::MultiPoint);
iter_geo_impl_geo_scalar!(LineStringArray<2>, geo::MultiLineString);
iter_geo_impl_geo_scalar!(LineStringArray<2>, geo::MultiPolygon);

// Implementations on PolygonArray
iter_geo_impl_geo_scalar!(PolygonArray<2>, geo::Point);
iter_geo_impl_geo_scalar!(PolygonArray<2>, geo::LineString);
iter_geo_impl_geo_scalar!(PolygonArray<2>, geo::Polygon);
iter_geo_impl_geo_scalar!(PolygonArray<2>, geo::MultiPoint);
iter_geo_impl_geo_scalar!(PolygonArray<2>, geo::MultiLineString);
iter_geo_impl_geo_scalar!(PolygonArray<2>, geo::MultiPolygon);

// Implementations on MultiPointArray
iter_geo_impl_geo_scalar!(MultiPointArray<2>, geo::Point);
iter_geo_impl_geo_scalar!(MultiPointArray<2>, geo::LineString);
iter_geo_impl_geo_scalar!(MultiPointArray<2>, geo::Polygon);
iter_geo_impl_geo_scalar!(MultiPointArray<2>, geo::MultiPoint);
iter_geo_impl_geo_scalar!(MultiPointArray<2>, geo::MultiLineString);
iter_geo_impl_geo_scalar!(MultiPointArray<2>, geo::MultiPolygon);

// Implementations on MultiLineStringArray
iter_geo_impl_geo_scalar!(MultiLineStringArray<2>, geo::Point);
iter_geo_impl_geo_scalar!(MultiLineStringArray<2>, geo::LineString);
iter_geo_impl_geo_scalar!(MultiLineStringArray<2>, geo::Polygon);
iter_geo_impl_geo_scalar!(MultiLineStringArray<2>, geo::MultiPoint);
iter_geo_impl_geo_scalar!(MultiLineStringArray<2>, geo::MultiLineString);
iter_geo_impl_geo_scalar!(MultiLineStringArray<2>, geo::MultiPolygon);

// Implementations on MultiPolygonArray
iter_geo_impl_geo_scalar!(MultiPolygonArray<2>, geo::Point);
iter_geo_impl_geo_scalar!(MultiPolygonArray<2>, geo::LineString);
iter_geo_impl_geo_scalar!(MultiPolygonArray<2>, geo::Polygon);
iter_geo_impl_geo_scalar!(MultiPolygonArray<2>, geo::MultiPoint);
iter_geo_impl_geo_scalar!(MultiPolygonArray<2>, geo::MultiLineString);
iter_geo_impl_geo_scalar!(MultiPolygonArray<2>, geo::MultiPolygon);
Loading

0 comments on commit 107501c

Please sign in to comment.