From d44887bada7fcd5bc3df8b096286901ef37ad3a9 Mon Sep 17 00:00:00 2001 From: liangxingjian <965662709@qq.com> Date: Mon, 10 Oct 2022 17:52:17 +0800 Subject: [PATCH] feat:add some geo vec impl --- src/datatypes/src/data_type.rs | 4 ++ src/datatypes/src/types/geometry.rs | 4 +- src/datatypes/src/vectors/geometry.rs | 57 +++++++++++++++------ src/datatypes/src/vectors/geometry/point.rs | 40 ++++++++++++++- 4 files changed, 85 insertions(+), 20 deletions(-) diff --git a/src/datatypes/src/data_type.rs b/src/datatypes/src/data_type.rs index 8207a375b7..dfcaa497b8 100644 --- a/src/datatypes/src/data_type.rs +++ b/src/datatypes/src/data_type.rs @@ -188,6 +188,10 @@ impl ConcreteDataType { ConcreteDataType::Timestamp(TimestampType::new(TimeUnit::Millisecond)) } + pub fn geometry_datatype() -> Self { + ConcreteDataType::Geometry(GeometryType::Point) + } + /// Converts from arrow timestamp unit to // TODO(hl): maybe impl From for our timestamp ? pub fn from_arrow_time_unit(t: &arrow::datatypes::TimeUnit) -> Self { diff --git a/src/datatypes/src/types/geometry.rs b/src/datatypes/src/types/geometry.rs index 3fb23ee644..3d45c102a1 100644 --- a/src/datatypes/src/types/geometry.rs +++ b/src/datatypes/src/types/geometry.rs @@ -4,7 +4,9 @@ use crate::data_type::DataType; use crate::prelude::{DataTypeRef, LogicalTypeId, Value}; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub enum GeometryType {} +pub enum GeometryType { + Point, +} impl DataType for GeometryType { fn name(&self) -> &str { diff --git a/src/datatypes/src/vectors/geometry.rs b/src/datatypes/src/vectors/geometry.rs index 05b5a6e086..fba9ceb9ad 100644 --- a/src/datatypes/src/vectors/geometry.rs +++ b/src/datatypes/src/vectors/geometry.rs @@ -1,9 +1,12 @@ -use arrow::array::Array; +use std::sync::Arc; + +use arrow::array::{Array, MutableArray}; use snafu::ensure; -use self::point::PointVector; +use self::point::{PointVector, PointVectorBuilder}; use super::{MutableVector, Vector}; use crate::{ + data_type::ConcreteDataType, error, prelude::{ScalarVector, ScalarVectorBuilder}, serialize::Serializable, @@ -18,27 +21,33 @@ pub enum GeometryVector { impl Vector for GeometryVector { fn data_type(&self) -> crate::data_type::ConcreteDataType { - todo!() + ConcreteDataType::geometry_datatype() } fn vector_type_name(&self) -> String { - todo!() + "GeometryVector".to_string() } fn as_any(&self) -> &dyn std::any::Any { - todo!() + self } fn len(&self) -> usize { - todo!() + match self { + GeometryVector::PointVector(vec) => vec.array.len(), + } } fn to_arrow_array(&self) -> arrow::array::ArrayRef { - todo!() + match self { + GeometryVector::PointVector(vec) => Arc::new(vec.array.clone()), + } } fn to_boxed_arrow_array(&self) -> Box { - todo!() + match self { + GeometryVector::PointVector(vec) => Box::new(vec.array.clone()), + } } fn validity(&self) -> super::Validity { @@ -100,32 +109,40 @@ impl<'a> Iterator for GeometryVectorIter<'a> { } pub enum GeometryVectorBuilder { - + PointVectorBuilder(PointVectorBuilder), } impl MutableVector for GeometryVectorBuilder { fn data_type(&self) -> crate::data_type::ConcreteDataType { - todo!() + ConcreteDataType::geometry_datatype() } fn len(&self) -> usize { - todo!() + match self { + GeometryVectorBuilder::PointVectorBuilder(builder) => builder.array_x.len(), + } } fn as_any(&self) -> &dyn std::any::Any { - todo!() + self } fn as_mut_any(&mut self) -> &mut dyn std::any::Any { - todo!() + self } fn to_vector(&mut self) -> super::VectorRef { - todo!() + Arc::new(self.finish()) } fn push_value_ref(&mut self, value: crate::value::ValueRef) -> crate::Result<()> { - todo!() + match value { + crate::value::ValueRef::Geometry(value) => { + self.push(Some(value)); + Ok(()) + } + _ => todo!(), + } } fn extend_slice_of( @@ -146,11 +163,17 @@ impl ScalarVectorBuilder for GeometryVectorBuilder { } fn push(&mut self, value: Option<::RefItem<'_>>) { - todo!() + match self { + GeometryVectorBuilder::PointVectorBuilder(builder) => builder.push(value), + } } fn finish(&mut self) -> Self::VectorType { - todo!() + match self { + GeometryVectorBuilder::PointVectorBuilder(builder) => { + GeometryVector::PointVector(builder.finish()) + } + } } } diff --git a/src/datatypes/src/vectors/geometry/point.rs b/src/datatypes/src/vectors/geometry/point.rs index 3a83ce4631..8bf510e2cf 100644 --- a/src/datatypes/src/vectors/geometry/point.rs +++ b/src/datatypes/src/vectors/geometry/point.rs @@ -1,7 +1,12 @@ -use arrow::array::{Array, FixedSizeListArray, ListArray, MutableFixedSizeListArray, StructArray}; -use arrow::datatypes::DataType::List; +use arrow::array::{ + Array, FixedSizeListArray, Float64Vec, ListArray, MutableArray, MutableFixedSizeListArray, + StructArray, +}; +use arrow::datatypes::DataType::{self, Float64, List}; +use arrow::datatypes::Field; use geo::Point; +use crate::value::{GeometryValue, ValueRef}; use crate::{ prelude::{ScalarVector, ScalarVectorBuilder, Vector}, vectors::MutableVector, @@ -19,4 +24,35 @@ impl PointVector { pub struct PointVectorBuilder { //pub array: MutableFixedSizeListArray, + pub array_x: Float64Vec, + pub array_y: Float64Vec, +} + +impl PointVectorBuilder { + pub fn push(&mut self, value: Option<&GeometryValue>) { + match value { + Some(val) => match *val { + GeometryValue::Point(xy) => { + self.array_x.push(Some(*xy.x())); + self.array_y.push(Some(*xy.y())); + } + }, + None => { + self.array_x.push_null(); + self.array_y.push_null(); + } + } + } + pub fn finish(&mut self) -> PointVector { + let (x, y) = (self.array_x.as_arc(), self.array_y.as_arc()); + let fields = vec![ + Field::new("x", Float64, true), + Field::new("y", Float64, true), + ]; + + let array = StructArray::new(DataType::Struct(fields), vec![x, y], None); + //how to get validity of struct? + + PointVector { array } + } }