use std::any::Any;
use std::fmt::Debug;
use std::sync::Arc;
use arrow_array::{Array, ArrayRef};
use arrow_buffer::NullBuffer;
use arrow_schema::extension::ExtensionType;
use geo_traits::GeometryTrait;
use geoarrow_schema::error::GeoArrowResult;
use geoarrow_schema::{GeoArrowType, Metadata};
use crate::array::from_arrow_array;
pub trait IntoArrow {
type ArrowArray: Array;
type ExtensionType: ExtensionType;
fn into_arrow(self) -> Self::ArrowArray;
fn extension_type(&self) -> &Self::ExtensionType;
}
pub trait GeoArrowArray: Debug + Send + Sync {
fn as_any(&self) -> &dyn Any;
fn data_type(&self) -> GeoArrowType;
#[must_use]
fn into_array_ref(self) -> ArrayRef;
#[must_use]
fn to_array_ref(&self) -> ArrayRef;
fn len(&self) -> usize;
fn is_empty(&self) -> bool {
self.len() == 0
}
fn logical_nulls(&self) -> Option<NullBuffer>;
fn logical_null_count(&self) -> usize;
fn is_null(&self, i: usize) -> bool;
#[inline]
fn is_valid(&self, i: usize) -> bool {
!self.is_null(i)
}
#[must_use]
fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray>;
fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray>;
}
impl GeoArrowArray for Arc<dyn GeoArrowArray> {
fn as_any(&self) -> &dyn Any {
self.as_ref().as_any()
}
fn data_type(&self) -> GeoArrowType {
self.as_ref().data_type()
}
fn into_array_ref(self) -> ArrayRef {
self.as_ref().to_array_ref()
}
fn to_array_ref(&self) -> ArrayRef {
self.as_ref().to_array_ref()
}
fn len(&self) -> usize {
self.as_ref().len()
}
fn logical_nulls(&self) -> Option<NullBuffer> {
self.as_ref().logical_nulls()
}
fn logical_null_count(&self) -> usize {
self.as_ref().logical_null_count()
}
fn is_null(&self, i: usize) -> bool {
self.as_ref().is_null(i)
}
fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
self.as_ref().slice(offset, length)
}
fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
let field = self.data_type().with_metadata(metadata).to_field("", true);
let array = self.as_ref().to_array_ref();
from_arrow_array(array.as_ref(), &field).unwrap()
}
}
impl<T: GeoArrowArray> GeoArrowArray for &T {
fn as_any(&self) -> &dyn Any {
T::as_any(self)
}
fn data_type(&self) -> GeoArrowType {
T::data_type(self)
}
fn into_array_ref(self) -> ArrayRef {
T::to_array_ref(self)
}
fn to_array_ref(&self) -> ArrayRef {
T::to_array_ref(self)
}
fn len(&self) -> usize {
T::len(self)
}
fn logical_nulls(&self) -> Option<NullBuffer> {
T::logical_nulls(self)
}
fn logical_null_count(&self) -> usize {
T::logical_null_count(self)
}
fn is_null(&self, i: usize) -> bool {
T::is_null(self, i)
}
fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
T::slice(self, offset, length)
}
fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
let field = self.data_type().with_metadata(metadata).to_field("", true);
let array = T::to_array_ref(self);
from_arrow_array(array.as_ref(), &field).unwrap()
}
}
pub trait GeoArrowArrayAccessor<'a>: GeoArrowArray {
type Item: Send + Sync + GeometryTrait<T = f64>;
fn value(&'a self, index: usize) -> GeoArrowResult<Self::Item> {
assert!(index < self.len());
unsafe { self.value_unchecked(index) }
}
unsafe fn value_unchecked(&'a self, index: usize) -> GeoArrowResult<Self::Item>;
fn get(&'a self, index: usize) -> GeoArrowResult<Option<Self::Item>> {
if self.is_null(index) {
return Ok(None);
}
Ok(Some(self.value(index)?))
}
unsafe fn get_unchecked(&'a self, index: usize) -> Option<GeoArrowResult<Self::Item>> {
if self.is_null(index) {
return None;
}
Some(unsafe { self.value_unchecked(index) })
}
fn iter(&'a self) -> impl ExactSizeIterator<Item = Option<GeoArrowResult<Self::Item>>> + 'a {
(0..self.len()).map(|i| unsafe { self.get_unchecked(i) })
}
fn iter_values(&'a self) -> impl ExactSizeIterator<Item = GeoArrowResult<Self::Item>> + 'a {
(0..self.len()).map(|i| unsafe { self.value_unchecked(i) })
}
}
pub(crate) trait GeoArrowArrayBuilder: Debug + Send + Sync {
fn len(&self) -> usize;
fn is_empty(&self) -> bool {
self.len() == 0
}
fn push_null(&mut self);
#[allow(dead_code)]
fn push_geometry(
&mut self,
geometry: Option<&impl GeometryTrait<T = f64>>,
) -> GeoArrowResult<()>;
#[allow(dead_code)]
fn finish(self) -> Arc<dyn GeoArrowArray>;
}
pub trait GeoArrowArrayReader: Iterator<Item = GeoArrowResult<Arc<dyn GeoArrowArray>>> {
fn data_type(&self) -> GeoArrowType;
}
impl<R: GeoArrowArrayReader + ?Sized> GeoArrowArrayReader for Box<R> {
fn data_type(&self) -> GeoArrowType {
self.as_ref().data_type()
}
}
pub struct GeoArrowArrayIterator<I>
where
I: IntoIterator<Item = GeoArrowResult<Arc<dyn GeoArrowArray>>>,
{
inner: I::IntoIter,
inner_type: GeoArrowType,
}
impl<I> GeoArrowArrayIterator<I>
where
I: IntoIterator<Item = GeoArrowResult<Arc<dyn GeoArrowArray>>>,
{
pub fn new(iter: I, data_type: GeoArrowType) -> Self {
Self {
inner: iter.into_iter(),
inner_type: data_type,
}
}
}
impl<I> Iterator for GeoArrowArrayIterator<I>
where
I: IntoIterator<Item = GeoArrowResult<Arc<dyn GeoArrowArray>>>,
{
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> {
self.inner.next()
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<I> GeoArrowArrayReader for GeoArrowArrayIterator<I>
where
I: IntoIterator<Item = GeoArrowResult<Arc<dyn GeoArrowArray>>>,
{
fn data_type(&self) -> GeoArrowType {
self.inner_type.clone()
}
}
#[cfg(test)]
mod test {
use std::sync::Arc;
use arrow_array::Array;
use arrow_array::builder::{ArrayBuilder, FixedSizeListBuilder, Float64Builder, StructBuilder};
use arrow_schema::{DataType, Field};
use geoarrow_schema::{CoordType, Dimension, GeometryType, PointType};
use super::*;
use crate::builder::GeometryBuilder;
use crate::trait_::GeoArrowArray;
#[test]
fn infer_type_interleaved_point() {
let test_cases = [
(2, Dimension::XY),
(3, Dimension::XYZ),
(4, Dimension::XYZM),
];
for (list_size, dim) in test_cases.into_iter() {
let array = FixedSizeListBuilder::new(Float64Builder::new(), list_size).finish();
let t =
GeoArrowType::from_arrow_field(&Field::new("", array.data_type().clone(), true))
.unwrap();
assert_eq!(
t,
GeoArrowType::Point(
PointType::new(dim, Default::default()).with_coord_type(CoordType::Interleaved)
)
);
}
}
#[test]
fn infer_type_separated_point() {
let test_cases = [
(
vec![
Arc::new(Field::new("x", DataType::Float64, true)),
Arc::new(Field::new("y", DataType::Float64, true)),
],
vec![
Box::new(Float64Builder::new()) as Box<dyn ArrayBuilder>,
Box::new(Float64Builder::new()),
],
Dimension::XY,
),
(
vec![
Arc::new(Field::new("x", DataType::Float64, true)),
Arc::new(Field::new("y", DataType::Float64, true)),
Arc::new(Field::new("z", DataType::Float64, true)),
],
vec![
Box::new(Float64Builder::new()) as Box<dyn ArrayBuilder>,
Box::new(Float64Builder::new()),
Box::new(Float64Builder::new()),
],
Dimension::XYZ,
),
(
vec![
Arc::new(Field::new("x", DataType::Float64, true)),
Arc::new(Field::new("y", DataType::Float64, true)),
Arc::new(Field::new("z", DataType::Float64, true)),
Arc::new(Field::new("m", DataType::Float64, true)),
],
vec![
Box::new(Float64Builder::new()) as Box<dyn ArrayBuilder>,
Box::new(Float64Builder::new()),
Box::new(Float64Builder::new()),
Box::new(Float64Builder::new()),
],
Dimension::XYZM,
),
];
for (fields, builders, dim) in test_cases.into_iter() {
let array = StructBuilder::new(fields, builders).finish();
let t =
GeoArrowType::from_arrow_field(&Field::new("", array.data_type().clone(), true))
.unwrap();
assert_eq!(
t,
GeoArrowType::Point(
PointType::new(dim, Default::default()).with_coord_type(CoordType::Separated)
)
);
}
}
#[test]
fn native_type_round_trip() {
let point_array = crate::test::point::point_array(CoordType::Interleaved);
let field = point_array.data_type.to_field("geometry", true);
let data_type: GeoArrowType = (&field).try_into().unwrap();
assert_eq!(point_array.data_type(), data_type);
let ml_array = crate::test::multilinestring::ml_array(CoordType::Interleaved);
let field = ml_array.data_type.to_field("geometry", true);
let data_type: GeoArrowType = (&field).try_into().unwrap();
assert_eq!(ml_array.data_type(), data_type);
let mut builder = GeometryBuilder::new(
GeometryType::new(Default::default()).with_coord_type(CoordType::Interleaved),
);
builder
.push_geometry(Some(&crate::test::point::p0()))
.unwrap();
builder
.push_geometry(Some(&crate::test::point::p1()))
.unwrap();
builder
.push_geometry(Some(&crate::test::point::p2()))
.unwrap();
builder
.push_geometry(Some(&crate::test::multilinestring::ml0()))
.unwrap();
builder
.push_geometry(Some(&crate::test::multilinestring::ml1()))
.unwrap();
let geom_array = builder.finish();
let field = geom_array.data_type.to_field("geometry", true);
let data_type: GeoArrowType = (&field).try_into().unwrap();
assert_eq!(geom_array.data_type(), data_type);
}
}