use std::collections::HashMap;
use std::sync::Arc;
use crate::algorithm::native::eq::offset_buffer_eq;
use crate::array::util::{offsets_buffer_i32_to_i64, offsets_buffer_i64_to_i32, OffsetBufferUtils};
use crate::array::zip_validity::ZipValidity;
use crate::array::{CoordBuffer, CoordType, MultiLineStringArray, RectArray, WKBArray};
use crate::datatypes::GeoDataType;
use crate::error::GeoArrowError;
use crate::geo_traits::PolygonTrait;
use crate::scalar::Polygon;
use crate::trait_::{GeoArrayAccessor, IntoArrow};
use crate::util::{owned_slice_offsets, owned_slice_validity};
use crate::GeometryArrayTrait;
use arrow_array::{Array, OffsetSizeTrait};
use arrow_array::{GenericListArray, LargeListArray, ListArray};
use arrow_buffer::bit_iterator::BitIterator;
use arrow_buffer::{NullBuffer, OffsetBuffer};
use arrow_schema::{DataType, Field};
use super::MutablePolygonArray;
#[derive(Debug, Clone)]
pub struct PolygonArray<O: OffsetSizeTrait> {
data_type: GeoDataType,
pub coords: CoordBuffer,
pub geom_offsets: OffsetBuffer<O>,
pub ring_offsets: OffsetBuffer<O>,
pub validity: Option<NullBuffer>,
}
pub(super) fn check<O: OffsetSizeTrait>(
coords: &CoordBuffer,
geom_offsets: &OffsetBuffer<O>,
ring_offsets: &OffsetBuffer<O>,
validity_len: Option<usize>,
) -> Result<(), GeoArrowError> {
if validity_len.map_or(false, |len| len != geom_offsets.len_proxy()) {
return Err(GeoArrowError::General(
"validity mask length must match the number of values".to_string(),
));
}
if ring_offsets.last().to_usize().unwrap() != coords.len() {
return Err(GeoArrowError::General(
"largest ring offset must match coords length".to_string(),
));
}
if geom_offsets.last().to_usize().unwrap() != ring_offsets.len_proxy() {
return Err(GeoArrowError::General(
"largest geometry offset must match ring offsets length".to_string(),
));
}
Ok(())
}
impl<O: OffsetSizeTrait> PolygonArray<O> {
pub fn new(
coords: CoordBuffer,
geom_offsets: OffsetBuffer<O>,
ring_offsets: OffsetBuffer<O>,
validity: Option<NullBuffer>,
) -> Self {
Self::try_new(coords, geom_offsets, ring_offsets, validity).unwrap()
}
pub fn try_new(
coords: CoordBuffer,
geom_offsets: OffsetBuffer<O>,
ring_offsets: OffsetBuffer<O>,
validity: Option<NullBuffer>,
) -> Result<Self, GeoArrowError> {
check(
&coords,
&geom_offsets,
&ring_offsets,
validity.as_ref().map(|v| v.len()),
)?;
let coord_type = coords.coord_type();
let data_type = match O::IS_LARGE {
true => GeoDataType::LargePolygon(coord_type),
false => GeoDataType::Polygon(coord_type),
};
Ok(Self {
data_type,
coords,
geom_offsets,
ring_offsets,
validity,
})
}
fn vertices_field(&self) -> Arc<Field> {
Field::new("vertices", self.coords.storage_type(), false).into()
}
fn rings_field(&self) -> Arc<Field> {
let name = "rings";
match O::IS_LARGE {
true => Field::new_large_list(name, self.vertices_field(), true).into(),
false => Field::new_list(name, self.vertices_field(), true).into(),
}
}
fn outer_type(&self) -> DataType {
match O::IS_LARGE {
true => DataType::LargeList(self.rings_field()),
false => DataType::List(self.rings_field()),
}
}
}
impl<'a, O: OffsetSizeTrait> GeometryArrayTrait<'a> for PolygonArray<O> {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn data_type(&self) -> &GeoDataType {
&self.data_type
}
fn storage_type(&self) -> DataType {
self.outer_type()
}
fn extension_field(&self) -> Arc<Field> {
let mut metadata = HashMap::new();
metadata.insert(
"ARROW:extension:name".to_string(),
self.extension_name().to_string(),
);
Arc::new(Field::new("geometry", self.storage_type(), true).with_metadata(metadata))
}
fn extension_name(&self) -> &str {
"geoarrow.polygon"
}
fn into_array_ref(self) -> Arc<dyn Array> {
Arc::new(self.into_arrow())
}
fn with_coords(self, coords: CoordBuffer) -> Self {
assert_eq!(coords.len(), self.coords.len());
Self::new(coords, self.geom_offsets, self.ring_offsets, self.validity)
}
fn coord_type(&self) -> CoordType {
self.coords.coord_type()
}
fn into_coord_type(self, coord_type: CoordType) -> Self {
Self::new(
self.coords.into_coord_type(coord_type),
self.geom_offsets,
self.ring_offsets,
self.validity,
)
}
#[inline]
fn len(&self) -> usize {
self.geom_offsets.len() - 1
}
#[inline]
fn validity(&self) -> Option<&NullBuffer> {
self.validity.as_ref()
}
#[inline]
fn slice(&self, offset: usize, length: usize) -> Self {
assert!(
offset + length <= self.len(),
"offset + length may not exceed length of array"
);
Self {
data_type: self.data_type.clone(),
coords: self.coords.clone(),
geom_offsets: self.geom_offsets.slice(offset, length),
ring_offsets: self.ring_offsets.clone(),
validity: self.validity.as_ref().map(|v| v.slice(offset, length)),
}
}
fn owned_slice(&self, offset: usize, length: usize) -> Self {
assert!(
offset + length <= self.len(),
"offset + length may not exceed length of array"
);
assert!(length >= 1, "length must be at least 1");
let (start_ring_idx, _) = self.geom_offsets.start_end(offset);
let (_, end_ring_idx) = self.geom_offsets.start_end(offset + length - 1);
let (start_coord_idx, _) = self.ring_offsets.start_end(start_ring_idx);
let (_, end_coord_idx) = self.ring_offsets.start_end(end_ring_idx - 1);
let geom_offsets = owned_slice_offsets(&self.geom_offsets, offset, length);
let ring_offsets = owned_slice_offsets(
&self.ring_offsets,
start_ring_idx,
end_ring_idx - start_ring_idx,
);
let coords = self
.coords
.owned_slice(start_coord_idx, end_coord_idx - start_coord_idx);
let validity = owned_slice_validity(self.nulls(), offset, length);
Self::new(coords, geom_offsets, ring_offsets, validity)
}
}
impl<'a, O: OffsetSizeTrait> GeoArrayAccessor<'a> for PolygonArray<O> {
type Item = Polygon<'a, O>;
type ItemGeo = geo::Polygon;
unsafe fn value_unchecked(&'a self, index: usize) -> Self::Item {
Polygon::new_borrowed(&self.coords, &self.geom_offsets, &self.ring_offsets, index)
}
}
impl<O: OffsetSizeTrait> IntoArrow for PolygonArray<O> {
type ArrowArray = GenericListArray<O>;
fn into_arrow(self) -> Self::ArrowArray {
let vertices_field = self.vertices_field();
let rings_field = self.rings_field();
let validity = self.validity;
let coord_array = self.coords.into_arrow();
let ring_array = Arc::new(GenericListArray::new(
vertices_field,
self.ring_offsets,
coord_array,
None,
));
GenericListArray::new(rings_field, self.geom_offsets, ring_array, validity)
}
}
impl<O: OffsetSizeTrait> PolygonArray<O> {
pub fn iter_geo_values(&self) -> impl Iterator<Item = geo::Polygon> + '_ {
(0..self.len()).map(|i| self.value_as_geo(i))
}
pub fn iter_geo(
&self,
) -> ZipValidity<geo::Polygon, impl Iterator<Item = geo::Polygon> + '_, BitIterator> {
ZipValidity::new_with_validity(self.iter_geo_values(), self.nulls())
}
#[cfg(feature = "geos")]
pub fn value_as_geos(&self, i: usize) -> geos::Geometry {
self.value(i).try_into().unwrap()
}
#[cfg(feature = "geos")]
pub fn get_as_geos(&self, i: usize) -> Option<geos::Geometry> {
if self.is_null(i) {
return None;
}
Some(self.value_as_geos(i))
}
#[cfg(feature = "geos")]
pub fn iter_geos_values(&self) -> impl Iterator<Item = geos::Geometry> + '_ {
(0..self.len()).map(|i| self.value_as_geos(i))
}
#[cfg(feature = "geos")]
pub fn iter_geos(
&self,
) -> ZipValidity<geos::Geometry, impl Iterator<Item = geos::Geometry> + '_, BitIterator> {
ZipValidity::new_with_validity(self.iter_geos_values(), self.nulls())
}
}
impl<O: OffsetSizeTrait> TryFrom<&GenericListArray<O>> for PolygonArray<O> {
type Error = GeoArrowError;
fn try_from(geom_array: &GenericListArray<O>) -> Result<Self, Self::Error> {
let geom_offsets = geom_array.offsets();
let validity = geom_array.nulls();
let rings_dyn_array = geom_array.values();
let rings_array = rings_dyn_array
.as_any()
.downcast_ref::<GenericListArray<O>>()
.unwrap();
let ring_offsets = rings_array.offsets();
let coords: CoordBuffer = rings_array.values().as_ref().try_into()?;
Ok(Self::new(
coords,
geom_offsets.clone(),
ring_offsets.clone(),
validity.cloned(),
))
}
}
impl TryFrom<&dyn Array> for PolygonArray<i32> {
type Error = GeoArrowError;
fn try_from(value: &dyn Array) -> Result<Self, Self::Error> {
match value.data_type() {
DataType::List(_) => {
let downcasted = value.as_any().downcast_ref::<ListArray>().unwrap();
downcasted.try_into()
}
DataType::LargeList(_) => {
let downcasted = value.as_any().downcast_ref::<LargeListArray>().unwrap();
let geom_array: PolygonArray<i64> = downcasted.try_into()?;
geom_array.try_into()
}
_ => Err(GeoArrowError::General(format!(
"Unexpected type: {:?}",
value.data_type()
))),
}
}
}
impl TryFrom<&dyn Array> for PolygonArray<i64> {
type Error = GeoArrowError;
fn try_from(value: &dyn Array) -> Result<Self, Self::Error> {
match value.data_type() {
DataType::List(_) => {
let downcasted = value.as_any().downcast_ref::<ListArray>().unwrap();
let geom_array: PolygonArray<i32> = downcasted.try_into()?;
Ok(geom_array.into())
}
DataType::LargeList(_) => {
let downcasted = value.as_any().downcast_ref::<LargeListArray>().unwrap();
downcasted.try_into()
}
_ => Err(GeoArrowError::General(format!(
"Unexpected type: {:?}",
value.data_type()
))),
}
}
}
impl<O: OffsetSizeTrait, G: PolygonTrait<T = f64>> From<Vec<Option<G>>> for PolygonArray<O> {
fn from(other: Vec<Option<G>>) -> Self {
let mut_arr: MutablePolygonArray<O> = other.into();
mut_arr.into()
}
}
impl<O: OffsetSizeTrait, G: PolygonTrait<T = f64>> From<Vec<G>> for PolygonArray<O> {
fn from(other: Vec<G>) -> Self {
let mut_arr: MutablePolygonArray<O> = other.into();
mut_arr.into()
}
}
impl<O: OffsetSizeTrait, G: PolygonTrait<T = f64>> From<bumpalo::collections::Vec<'_, G>>
for PolygonArray<O>
{
fn from(value: bumpalo::collections::Vec<G>) -> Self {
let mut_arr: MutablePolygonArray<O> = value.into();
mut_arr.into()
}
}
impl<O: OffsetSizeTrait, G: PolygonTrait<T = f64>> From<bumpalo::collections::Vec<'_, Option<G>>>
for PolygonArray<O>
{
fn from(value: bumpalo::collections::Vec<Option<G>>) -> Self {
let mut_arr: MutablePolygonArray<O> = value.into();
mut_arr.into()
}
}
impl<O: OffsetSizeTrait> TryFrom<WKBArray<O>> for PolygonArray<O> {
type Error = GeoArrowError;
fn try_from(value: WKBArray<O>) -> Result<Self, Self::Error> {
let mut_arr: MutablePolygonArray<O> = value.try_into()?;
Ok(mut_arr.into())
}
}
impl<O: OffsetSizeTrait> From<PolygonArray<O>> for MultiLineStringArray<O> {
fn from(value: PolygonArray<O>) -> Self {
Self::new(
value.coords,
value.geom_offsets,
value.ring_offsets,
value.validity,
)
}
}
impl From<PolygonArray<i32>> for PolygonArray<i64> {
fn from(value: PolygonArray<i32>) -> Self {
Self::new(
value.coords,
offsets_buffer_i32_to_i64(&value.geom_offsets),
offsets_buffer_i32_to_i64(&value.ring_offsets),
value.validity,
)
}
}
impl TryFrom<PolygonArray<i64>> for PolygonArray<i32> {
type Error = GeoArrowError;
fn try_from(value: PolygonArray<i64>) -> Result<Self, Self::Error> {
Ok(Self::new(
value.coords,
offsets_buffer_i64_to_i32(&value.geom_offsets)?,
offsets_buffer_i64_to_i32(&value.ring_offsets)?,
value.validity,
))
}
}
impl<O: OffsetSizeTrait> From<RectArray> for PolygonArray<O> {
fn from(value: RectArray) -> Self {
let geom_capacity = value.len();
let ring_capacity = geom_capacity;
let coord_capacity = (value.len() - value.null_count()) * 5;
let mut output_array =
MutablePolygonArray::with_capacities(coord_capacity, ring_capacity, geom_capacity);
value.iter_geo().for_each(|maybe_g| {
output_array
.push_polygon(maybe_g.map(|geom| geom.to_polygon()).as_ref())
.unwrap()
});
output_array.into()
}
}
impl<O: OffsetSizeTrait> Default for PolygonArray<O> {
fn default() -> Self {
MutablePolygonArray::default().into()
}
}
impl<O: OffsetSizeTrait> PartialEq for PolygonArray<O> {
fn eq(&self, other: &Self) -> bool {
if self.validity != other.validity {
return false;
}
if !offset_buffer_eq(&self.geom_offsets, &other.geom_offsets) {
return false;
}
if !offset_buffer_eq(&self.ring_offsets, &other.ring_offsets) {
return false;
}
if self.coords != other.coords {
return false;
}
true
}
}
#[cfg(test)]
mod test {
use crate::test::geoarrow_data::{
example_polygon_interleaved, example_polygon_separated, example_polygon_wkb,
};
use crate::test::polygon::{p0, p1};
use super::*;
#[test]
fn geo_roundtrip_accurate() {
let arr: PolygonArray<i64> = vec![p0(), p1()].into();
assert_eq!(arr.value_as_geo(0), p0());
assert_eq!(arr.value_as_geo(1), p1());
}
#[test]
fn geo_roundtrip_accurate_option_vec() {
let arr: PolygonArray<i64> = vec![Some(p0()), Some(p1()), None].into();
assert_eq!(arr.get_as_geo(0), Some(p0()));
assert_eq!(arr.get_as_geo(1), Some(p1()));
assert_eq!(arr.get_as_geo(2), None);
}
#[test]
fn slice() {
let arr: PolygonArray<i64> = vec![p0(), p1()].into();
let sliced = arr.slice(1, 1);
assert_eq!(sliced.len(), 1);
assert_eq!(sliced.get_as_geo(0), Some(p1()));
}
#[test]
fn owned_slice() {
let arr: PolygonArray<i64> = vec![p0(), p1()].into();
let sliced = arr.owned_slice(1, 1);
assert_eq!(arr.len(), 2);
assert_eq!(sliced.len(), 1);
assert_eq!(sliced.get_as_geo(0), Some(p1()));
}
#[test]
fn parse_wkb_geoarrow_interleaved_example() {
let geom_arr = example_polygon_interleaved();
let wkb_arr = example_polygon_wkb();
let parsed_geom_arr: PolygonArray<i64> = wkb_arr.try_into().unwrap();
assert_eq!(geom_arr, parsed_geom_arr);
}
#[test]
fn parse_wkb_geoarrow_separated_example() {
let geom_arr = example_polygon_separated().into_coord_type(CoordType::Interleaved);
let wkb_arr = example_polygon_wkb();
let parsed_geom_arr: PolygonArray<i64> = wkb_arr.try_into().unwrap();
assert_eq!(geom_arr, parsed_geom_arr);
}
}