use std::collections::HashMap;
use std::sync::Arc;
use crate::array::util::{offsets_buffer_i32_to_i64, offsets_buffer_i64_to_i32};
use crate::array::zip_validity::ZipValidity;
use crate::array::{CoordType, MutableWKBArray};
use crate::datatypes::GeoDataType;
use crate::error::GeoArrowError;
use crate::scalar::WKB;
use crate::trait_::{GeoArrayAccessor, IntoArrow};
use crate::GeometryArrayTrait;
use arrow_array::OffsetSizeTrait;
use arrow_array::{Array, BinaryArray, GenericBinaryArray, LargeBinaryArray};
use arrow_buffer::bit_iterator::BitIterator;
use arrow_buffer::NullBuffer;
use arrow_schema::{DataType, Field};
#[derive(Debug, Clone, PartialEq)]
pub struct WKBArray<O: OffsetSizeTrait>(GenericBinaryArray<O>, GeoDataType);
impl<O: OffsetSizeTrait> WKBArray<O> {
pub fn new(arr: GenericBinaryArray<O>) -> Self {
let data_type = match O::IS_LARGE {
true => GeoDataType::LargeWKB,
false => GeoDataType::WKB,
};
Self(arr, data_type)
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl<'a, O: OffsetSizeTrait> GeometryArrayTrait<'a> for WKBArray<O> {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn data_type(&self) -> &GeoDataType {
&self.1
}
fn storage_type(&self) -> DataType {
self.0.data_type().clone()
}
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.wkb"
}
fn into_array_ref(self) -> Arc<dyn Array> {
Arc::new(self.into_arrow())
}
fn with_coords(self, _coords: crate::array::CoordBuffer) -> Self {
unimplemented!()
}
fn coord_type(&self) -> CoordType {
CoordType::Interleaved
}
fn into_coord_type(self, _coord_type: CoordType) -> Self {
self
}
#[inline]
fn len(&self) -> usize {
self.0.len()
}
fn validity(&self) -> Option<&NullBuffer> {
self.0.nulls()
}
#[inline]
fn slice(&self, offset: usize, length: usize) -> Self {
assert!(
offset + length <= self.len(),
"offset + length may not exceed length of array"
);
Self(self.0.slice(offset, length), self.1.clone())
}
fn owned_slice(&self, _offset: usize, _length: usize) -> Self {
todo!()
}
}
impl<'a, O: OffsetSizeTrait> GeoArrayAccessor<'a> for WKBArray<O> {
type Item = WKB<'a, O>;
type ItemGeo = geo::Geometry;
unsafe fn value_unchecked(&'a self, index: usize) -> Self::Item {
WKB::new_borrowed(&self.0, index)
}
}
impl<O: OffsetSizeTrait> IntoArrow for WKBArray<O> {
type ArrowArray = GenericBinaryArray<O>;
fn into_arrow(self) -> Self::ArrowArray {
GenericBinaryArray::new(
self.0.offsets().clone(),
self.0.values().clone(),
self.0.nulls().cloned(),
)
}
}
impl<O: OffsetSizeTrait> WKBArray<O> {
#[cfg(feature = "geos")]
pub fn value_as_geos(&self, i: usize) -> geos::Geometry {
let buf = self.0.value(i);
geos::Geometry::new_from_wkb(buf).expect("Unable to parse WKB")
}
#[cfg(feature = "geos")]
pub fn get_as_geos(&self, i: usize) -> Option<geos::Geometry> {
if self.is_null(i) {
return None;
}
let buf = self.0.value(i);
Some(geos::Geometry::new_from_wkb(buf).expect("Unable to parse WKB"))
}
pub fn iter_geo_values(&self) -> impl Iterator<Item = geo::Geometry> + '_ {
(0..self.len()).map(|i| self.value_as_geo(i))
}
pub fn iter_geo(
&self,
) -> ZipValidity<geo::Geometry, impl Iterator<Item = geo::Geometry> + '_, BitIterator> {
ZipValidity::new_with_validity(self.iter_geo_values(), self.nulls())
}
#[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> From<GenericBinaryArray<O>> for WKBArray<O> {
fn from(value: GenericBinaryArray<O>) -> Self {
Self::new(value)
}
}
impl TryFrom<&dyn Array> for WKBArray<i32> {
type Error = GeoArrowError;
fn try_from(value: &dyn Array) -> Result<Self, Self::Error> {
match value.data_type() {
DataType::Binary => {
let downcasted = value.as_any().downcast_ref::<BinaryArray>().unwrap();
Ok(downcasted.clone().into())
}
DataType::LargeBinary => {
let downcasted = value.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
let geom_array: WKBArray<i64> = downcasted.clone().into();
geom_array.try_into()
}
_ => Err(GeoArrowError::General(format!(
"Unexpected type: {:?}",
value.data_type()
))),
}
}
}
impl TryFrom<&dyn Array> for WKBArray<i64> {
type Error = GeoArrowError;
fn try_from(value: &dyn Array) -> Result<Self, Self::Error> {
match value.data_type() {
DataType::Binary => {
let downcasted = value.as_any().downcast_ref::<BinaryArray>().unwrap();
let geom_array: WKBArray<i32> = downcasted.clone().into();
Ok(geom_array.into())
}
DataType::LargeBinary => {
let downcasted = value.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
Ok(downcasted.clone().into())
}
_ => Err(GeoArrowError::General(format!(
"Unexpected type: {:?}",
value.data_type()
))),
}
}
}
impl From<WKBArray<i32>> for WKBArray<i64> {
fn from(value: WKBArray<i32>) -> Self {
let binary_array = value.0;
let (offsets, values, nulls) = binary_array.into_parts();
Self::new(LargeBinaryArray::new(
offsets_buffer_i32_to_i64(&offsets),
values,
nulls,
))
}
}
impl TryFrom<WKBArray<i64>> for WKBArray<i32> {
type Error = GeoArrowError;
fn try_from(value: WKBArray<i64>) -> Result<Self, Self::Error> {
let binary_array = value.0;
let (offsets, values, nulls) = binary_array.into_parts();
Ok(Self::new(BinaryArray::new(
offsets_buffer_i64_to_i32(&offsets)?,
values,
nulls,
)))
}
}
impl<O: OffsetSizeTrait> From<Vec<Option<geo::Geometry>>> for WKBArray<O> {
fn from(other: Vec<Option<geo::Geometry>>) -> Self {
let mut_arr: MutableWKBArray<O> = other.into();
mut_arr.into()
}
}
impl<O: OffsetSizeTrait> From<bumpalo::collections::Vec<'_, Option<geo::Geometry>>>
for WKBArray<O>
{
fn from(other: bumpalo::collections::Vec<'_, Option<geo::Geometry>>) -> Self {
let mut_arr: MutableWKBArray<O> = other.into();
mut_arr.into()
}
}