use std::collections::HashMap;
use std::sync::Arc;
use super::MutableMultiPointArray;
use crate::algorithm::native::eq::offset_buffer_eq;
use crate::array::mutable_offset::OffsetsBuilder;
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, LineStringArray, PointArray, WKBArray};
use crate::datatypes::GeoDataType;
use crate::error::{GeoArrowError, Result};
use crate::geo_traits::MultiPointTrait;
use crate::scalar::MultiPoint;
use crate::trait_::{GeoArrayAccessor, IntoArrow};
use crate::util::{owned_slice_offsets, owned_slice_validity};
use crate::GeometryArrayTrait;
use arrow_array::{Array, GenericListArray, LargeListArray, ListArray, OffsetSizeTrait};
use arrow_buffer::bit_iterator::BitIterator;
use arrow_buffer::{NullBuffer, OffsetBuffer};
use arrow_schema::{DataType, Field};
#[derive(Debug, Clone)]
pub struct MultiPointArray<O: OffsetSizeTrait> {
data_type: GeoDataType,
pub coords: CoordBuffer,
pub geom_offsets: OffsetBuffer<O>,
pub validity: Option<NullBuffer>,
}
pub(super) fn check<O: OffsetSizeTrait>(
coords: &CoordBuffer,
validity_len: Option<usize>,
geom_offsets: &OffsetBuffer<O>,
) -> Result<()> {
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 geom_offsets.last().to_usize().unwrap() != coords.len() {
return Err(GeoArrowError::General(
"largest geometry offset must match coords length".to_string(),
));
}
Ok(())
}
impl<O: OffsetSizeTrait> MultiPointArray<O> {
pub fn new(
coords: CoordBuffer,
geom_offsets: OffsetBuffer<O>,
validity: Option<NullBuffer>,
) -> Self {
Self::try_new(coords, geom_offsets, validity).unwrap()
}
pub fn try_new(
coords: CoordBuffer,
geom_offsets: OffsetBuffer<O>,
validity: Option<NullBuffer>,
) -> Result<Self> {
check(&coords, validity.as_ref().map(|v| v.len()), &geom_offsets)?;
let coord_type = coords.coord_type();
let data_type = match O::IS_LARGE {
true => GeoDataType::LargeMultiPoint(coord_type),
false => GeoDataType::MultiPoint(coord_type),
};
Ok(Self {
data_type,
coords,
geom_offsets,
validity,
})
}
fn vertices_field(&self) -> Arc<Field> {
Field::new("points", self.coords.storage_type(), true).into()
}
fn outer_type(&self) -> DataType {
match O::IS_LARGE {
true => DataType::LargeList(self.vertices_field()),
false => DataType::List(self.vertices_field()),
}
}
}
impl<'a, O: OffsetSizeTrait> GeometryArrayTrait<'a> for MultiPointArray<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.multipoint"
}
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.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.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),
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_coord_idx, _) = self.geom_offsets.start_end(offset);
let (_, end_coord_idx) = self.geom_offsets.start_end(offset + length - 1);
let geom_offsets = owned_slice_offsets(&self.geom_offsets, offset, length);
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, validity)
}
}
impl<'a, O: OffsetSizeTrait> GeoArrayAccessor<'a> for MultiPointArray<O> {
type Item = MultiPoint<'a, O>;
type ItemGeo = geo::MultiPoint;
unsafe fn value_unchecked(&'a self, index: usize) -> Self::Item {
MultiPoint::new_borrowed(&self.coords, &self.geom_offsets, index)
}
}
impl<O: OffsetSizeTrait> IntoArrow for MultiPointArray<O> {
type ArrowArray = GenericListArray<O>;
fn into_arrow(self) -> Self::ArrowArray {
let vertices_field = self.vertices_field();
let validity = self.validity;
let coord_array = self.coords.into_arrow();
GenericListArray::new(vertices_field, self.geom_offsets, coord_array, validity)
}
}
impl<O: OffsetSizeTrait> MultiPointArray<O> {
pub fn iter_geo_values(&self) -> impl Iterator<Item = geo::MultiPoint> + '_ {
(0..self.len()).map(|i| self.value_as_geo(i))
}
pub fn iter_geo(
&self,
) -> ZipValidity<geo::MultiPoint, impl Iterator<Item = geo::MultiPoint> + '_, 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 MultiPointArray<O> {
type Error = GeoArrowError;
fn try_from(value: &GenericListArray<O>) -> Result<Self> {
let coords: CoordBuffer = value.values().as_ref().try_into()?;
let geom_offsets = value.offsets();
let validity = value.nulls();
Ok(Self::new(coords, geom_offsets.clone(), validity.cloned()))
}
}
impl TryFrom<&dyn Array> for MultiPointArray<i32> {
type Error = GeoArrowError;
fn try_from(value: &dyn Array) -> Result<Self> {
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: MultiPointArray<i64> = downcasted.try_into()?;
geom_array.try_into()
}
_ => Err(GeoArrowError::General(format!(
"Unexpected type: {:?}",
value.data_type()
))),
}
}
}
impl TryFrom<&dyn Array> for MultiPointArray<i64> {
type Error = GeoArrowError;
fn try_from(value: &dyn Array) -> Result<Self> {
match value.data_type() {
DataType::List(_) => {
let downcasted = value.as_any().downcast_ref::<ListArray>().unwrap();
let geom_array: MultiPointArray<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: MultiPointTrait<T = f64>> From<Vec<Option<G>>> for MultiPointArray<O> {
fn from(other: Vec<Option<G>>) -> Self {
let mut_arr: MutableMultiPointArray<O> = other.into();
mut_arr.into()
}
}
impl<O: OffsetSizeTrait, G: MultiPointTrait<T = f64>> From<Vec<G>> for MultiPointArray<O> {
fn from(other: Vec<G>) -> Self {
let mut_arr: MutableMultiPointArray<O> = other.into();
mut_arr.into()
}
}
impl<O: OffsetSizeTrait, G: MultiPointTrait<T = f64>> From<bumpalo::collections::Vec<'_, Option<G>>>
for MultiPointArray<O>
{
fn from(other: bumpalo::collections::Vec<'_, Option<G>>) -> Self {
let mut_arr: MutableMultiPointArray<O> = other.into();
mut_arr.into()
}
}
impl<O: OffsetSizeTrait, G: MultiPointTrait<T = f64>> From<bumpalo::collections::Vec<'_, G>>
for MultiPointArray<O>
{
fn from(other: bumpalo::collections::Vec<'_, G>) -> Self {
let mut_arr: MutableMultiPointArray<O> = other.into();
mut_arr.into()
}
}
impl<O: OffsetSizeTrait> TryFrom<WKBArray<O>> for MultiPointArray<O> {
type Error = GeoArrowError;
fn try_from(value: WKBArray<O>) -> Result<Self> {
let mut_arr: MutableMultiPointArray<O> = value.try_into()?;
Ok(mut_arr.into())
}
}
impl<O: OffsetSizeTrait> From<MultiPointArray<O>> for LineStringArray<O> {
fn from(value: MultiPointArray<O>) -> Self {
Self::new(value.coords, value.geom_offsets, value.validity)
}
}
impl<O: OffsetSizeTrait> TryFrom<PointArray> for MultiPointArray<O> {
type Error = GeoArrowError;
fn try_from(value: PointArray) -> Result<Self> {
let geom_length = value.len();
let coords = value.coords;
let validity = value.validity;
let mut geom_offsets = OffsetsBuilder::with_capacity(geom_length);
for _ in 0..coords.len() {
geom_offsets.try_push_usize(1)?;
}
Ok(Self::new(coords, geom_offsets.into(), validity))
}
}
impl From<MultiPointArray<i32>> for MultiPointArray<i64> {
fn from(value: MultiPointArray<i32>) -> Self {
Self::new(
value.coords,
offsets_buffer_i32_to_i64(&value.geom_offsets),
value.validity,
)
}
}
impl TryFrom<MultiPointArray<i64>> for MultiPointArray<i32> {
type Error = GeoArrowError;
fn try_from(value: MultiPointArray<i64>) -> Result<Self> {
Ok(Self::new(
value.coords,
offsets_buffer_i64_to_i32(&value.geom_offsets)?,
value.validity,
))
}
}
impl<O: OffsetSizeTrait> Default for MultiPointArray<O> {
fn default() -> Self {
MutableMultiPointArray::default().into()
}
}
impl<O: OffsetSizeTrait> PartialEq for MultiPointArray<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 self.coords != other.coords {
return false;
}
true
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::test::geoarrow_data::{
example_multipoint_interleaved, example_multipoint_separated, example_multipoint_wkb,
};
use crate::test::multipoint::{mp0, mp1};
#[test]
fn geo_roundtrip_accurate() {
let arr: MultiPointArray<i64> = vec![mp0(), mp1()].into();
assert_eq!(arr.value_as_geo(0), mp0());
assert_eq!(arr.value_as_geo(1), mp1());
}
#[test]
fn geo_roundtrip_accurate_option_vec() {
let arr: MultiPointArray<i64> = vec![Some(mp0()), Some(mp1()), None].into();
assert_eq!(arr.get_as_geo(0), Some(mp0()));
assert_eq!(arr.get_as_geo(1), Some(mp1()));
assert_eq!(arr.get_as_geo(2), None);
}
#[test]
fn slice() {
let arr: MultiPointArray<i64> = vec![mp0(), mp1()].into();
let sliced = arr.slice(1, 1);
assert_eq!(sliced.len(), 1);
assert_eq!(sliced.get_as_geo(0), Some(mp1()));
}
#[test]
fn owned_slice() {
let arr: MultiPointArray<i64> = vec![mp0(), mp1()].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(mp1()));
}
#[test]
fn parse_wkb_geoarrow_interleaved_example() {
let geom_arr = example_multipoint_interleaved();
let wkb_arr = example_multipoint_wkb();
let parsed_geom_arr: MultiPointArray<i64> = wkb_arr.try_into().unwrap();
assert_eq!(geom_arr, parsed_geom_arr);
}
#[test]
fn parse_wkb_geoarrow_separated_example() {
let geom_arr = example_multipoint_separated().into_coord_type(CoordType::Interleaved);
let wkb_arr = example_multipoint_wkb();
let parsed_geom_arr: MultiPointArray<i64> = wkb_arr.try_into().unwrap();
assert_eq!(geom_arr, parsed_geom_arr);
}
}