Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

wip: st_rotate #954

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 108 additions & 0 deletions rust/geoarrow/src/algorithm/broadcasting/iterator.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
//! Idiomatic iterators for [`Array`](crate::Array)

use arrow_buffer::NullBuffer;

use crate::array::PointArray;
use crate::scalar::Point;
use crate::trait_::ArrayAccessor;
use crate::ArrayBase;

/// An iterator that returns Some(T) or None, that can be used on any [`ArrayAccessor`]
///
/// # Performance
///
/// [`ArrayIter`] provides an idiomatic way to iterate over an array, however, this
/// comes at the cost of performance. In particular the interleaved handling of
/// the null mask is often sub-optimal.
///
/// If performing an infallible operation, it is typically faster to perform the operation
/// on every index of the array, and handle the null mask separately. For [`PrimitiveArray`]
/// this functionality is provided by [`compute::unary`]
///
/// If performing a fallible operation, it isn't possible to perform the operation independently
/// of the null mask, as this might result in a spurious failure on a null index. However,
/// there are more efficient ways to iterate over just the non-null indices, this functionality
/// is provided by [`compute::try_unary`]
///
/// [`PrimitiveArray`]: crate::PrimitiveArray
/// [`compute::unary`]: https://docs.rs/arrow/latest/arrow/compute/fn.unary.html
/// [`compute::try_unary`]: https://docs.rs/arrow/latest/arrow/compute/fn.try_unary.html
#[derive(Debug)]
pub struct PointArrayIter<'a> {
array: &'a PointArray,
logical_nulls: Option<NullBuffer>,
current: usize,
current_end: usize,
}

impl<'a> PointArrayIter<'a> {
/// create a new iterator
pub fn new(array: &'a PointArray) -> Self {
let len = array.len();
let logical_nulls = array.nulls().cloned();
PointArrayIter {
array,
logical_nulls,
current: 0,
current_end: len,
}
}

#[inline]
fn is_null(&self, idx: usize) -> bool {
self.logical_nulls
.as_ref()
.map(|x| x.is_null(idx))
.unwrap_or_default()
}
}

impl<'a> Iterator for PointArrayIter<'a> {
type Item = Option<Point<'a>>;

#[inline]
fn next(&mut self) -> Option<Self::Item> {
if self.current == self.current_end {
None
} else if self.is_null(self.current) {
self.current += 1;
Some(None)
} else {
let old = self.current;
self.current += 1;
// Safety:
// we just checked bounds in `self.current_end == self.current`
// this is safe on the premise that this struct is initialized with
// current = array.len()
// and that current_end is ever only decremented
unsafe { Some(Some(self.array.value_unchecked(old))) }
}
}

fn size_hint(&self) -> (usize, Option<usize>) {
(
self.array.len() - self.current,
Some(self.array.len() - self.current),
)
}
}

impl DoubleEndedIterator for PointArrayIter<'_> {
fn next_back(&mut self) -> Option<Self::Item> {
if self.current_end == self.current {
None
} else {
self.current_end -= 1;
Some(if self.is_null(self.current_end) {
None
} else {
// Safety:
// we just checked bounds in `self.current_end == self.current`
// this is safe on the premise that this struct is initialized with
// current = array.len()
// and that current_end is ever only decremented
unsafe { Some(self.array.value_unchecked(self.current_end)) }
})
}
}
}
5 changes: 3 additions & 2 deletions rust/geoarrow/src/algorithm/broadcasting/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,8 @@
// mod multilinestring;
// mod multipoint;
// mod multipolygon;
// mod point;
mod iterator;
mod point;
// mod polygon;
mod primitive;
mod vec;
Expand All @@ -28,7 +29,7 @@ mod vec;
// pub use multilinestring::BroadcastableMultiLineString;
// pub use multipoint::BroadcastableMultiPoint;
// pub use multipolygon::BroadcastableMultiPolygon;
// pub use point::BroadcastablePoint;
pub use point::BroadcastablePoint;
// pub use polygon::BroadcastablePolygon;
pub use primitive::BroadcastablePrimitive;
pub use vec::BroadcastableVec;
34 changes: 25 additions & 9 deletions rust/geoarrow/src/algorithm/broadcasting/point.rs
Original file line number Diff line number Diff line change
@@ -1,23 +1,39 @@
use geo_traits::to_geo::ToGeoPoint;
use geo_traits::PointTrait;

use crate::algorithm::broadcasting::iterator::PointArrayIter;
use crate::array::PointArray;
use crate::scalar::Point;
use crate::trait_::NativeScalar;

/// An enum over a [`Point`] scalar and [`PointArray`] array.
///
/// [`IntoIterator`] is implemented for this, where it will iterate over the `Array` variant
/// normally but will iterate over the `Scalar` variant forever.
#[derive(Debug)]
pub enum BroadcastablePoint<'a> {
Scalar(Point<'a>),
pub enum BroadcastablePoint {
Scalar(geo::Point),
Array(PointArray),
}

impl<G: PointTrait<T = f64>> From<G> for BroadcastablePoint {
fn from(value: G) -> Self {
Self::Scalar(value.to_point())
}
}

impl From<PointArray> for BroadcastablePoint {
fn from(value: PointArray) -> Self {
Self::Array(value)
}
}

pub enum BroadcastPointIter<'a> {
Scalar(Point<'a>),
Scalar(geo::Point),
Array(PointArrayIter<'a>),
}

impl<'a> IntoIterator for &'a BroadcastablePoint<'a> {
type Item = Option<Point<'a>>;
impl<'a> IntoIterator for &'a BroadcastablePoint {
type Item = Option<geo::Point>;
type IntoIter = BroadcastPointIter<'a>;

fn into_iter(self) -> Self::IntoIter {
Expand All @@ -28,12 +44,12 @@ impl<'a> IntoIterator for &'a BroadcastablePoint<'a> {
}
}

impl<'a> Iterator for BroadcastPointIter<'a> {
type Item = Option<Point<'a>>;
impl Iterator for BroadcastPointIter<'_> {
type Item = Option<geo::Point>;

fn next(&mut self) -> Option<Self::Item> {
match self {
BroadcastPointIter::Array(arr) => arr.next(),
BroadcastPointIter::Array(arr) => arr.next().map(|x| x.map(|y| y.to_geo())),
BroadcastPointIter::Scalar(val) => Some(Some(val.to_owned())),
}
}
Expand Down
Loading
Loading