From ebcd18d3c4bc71a309ab08d0aef1d80e6c8d35c1 Mon Sep 17 00:00:00 2001 From: liangxingjian <965662709@qq.com> Date: Tue, 11 Oct 2022 19:07:15 +0800 Subject: [PATCH] feat:add some impl of geo type --- src/datatypes/src/scalars.rs | 4 +- src/datatypes/src/value.rs | 7 +- src/datatypes/src/vectors/geometry.rs | 84 +++++++++++++++++---- src/datatypes/src/vectors/geometry/point.rs | 64 ++++++++++++++-- 4 files changed, 134 insertions(+), 25 deletions(-) diff --git a/src/datatypes/src/scalars.rs b/src/datatypes/src/scalars.rs index 48fc3ffd04..2badff7229 100644 --- a/src/datatypes/src/scalars.rs +++ b/src/datatypes/src/scalars.rs @@ -344,7 +344,7 @@ impl<'a> ScalarRef<'a> for ListValueRef<'a> { impl Scalar for GeometryValue { type VectorType = GeometryVector; - type RefType<'a> = &'a GeometryValue; + type RefType<'a> = GeometryValue; fn as_scalar_ref(&self) -> Self::RefType<'_> { todo!() @@ -355,7 +355,7 @@ impl Scalar for GeometryValue { } } -impl<'a> ScalarRef<'a> for &'a GeometryValue { +impl<'a> ScalarRef<'a> for GeometryValue { type VectorType = GeometryVector; type ScalarType = GeometryValue; diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 78e551c385..9527858d97 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -4,6 +4,7 @@ use common_base::bytes::{Bytes, StringBytes}; use common_time::date::Date; use common_time::datetime::DateTime; use common_time::timestamp::Timestamp; +use geo::Point; pub use ordered_float::OrderedFloat; use serde::{Deserialize, Serialize}; @@ -308,9 +309,9 @@ impl Ord for ListValue { } } -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq, Copy)] pub enum GeometryValue { - Point(geo::Point), + Point(Point), } impl Default for GeometryValue { @@ -345,7 +346,7 @@ pub enum ValueRef<'a> { DateTime(DateTime), Timestamp(Timestamp), List(ListValueRef<'a>), - Geometry(&'a GeometryValue), + Geometry(GeometryValue), } macro_rules! impl_as_for_value_ref { diff --git a/src/datatypes/src/vectors/geometry.rs b/src/datatypes/src/vectors/geometry.rs index fba9ceb9ad..098f4b1472 100644 --- a/src/datatypes/src/vectors/geometry.rs +++ b/src/datatypes/src/vectors/geometry.rs @@ -1,10 +1,12 @@ use std::sync::Arc; use arrow::array::{Array, MutableArray}; -use snafu::ensure; +use snafu::{ensure, OptionExt}; use self::point::{PointVector, PointVectorBuilder}; -use super::{MutableVector, Vector}; +use super::{MutableVector, Validity, Value, Vector}; +use crate::value::ValueRef; +use crate::vectors::{impl_try_from_arrow_array_for_vector, impl_validity_for_vector}; use crate::{ data_type::ConcreteDataType, error, @@ -51,7 +53,9 @@ impl Vector for GeometryVector { } fn validity(&self) -> super::Validity { - todo!() + match self { + GeometryVector::PointVector(vec) => vec.validity(), + } } fn memory_size(&self) -> usize { @@ -67,44 +71,75 @@ impl Vector for GeometryVector { } fn slice(&self, offset: usize, length: usize) -> super::VectorRef { - todo!() + match self { + GeometryVector::PointVector(vec) => { + Arc::new(GeometryVector::PointVector(vec.slice(offset, length))) + } + } } fn get(&self, index: usize) -> crate::value::Value { - todo!() + match self { + GeometryVector::PointVector(vec) => vec.get(index), + } } fn get_ref(&self, index: usize) -> crate::value::ValueRef { - todo!() + match self { + GeometryVector::PointVector(vec) => vec.get_ref(index), + } } } impl ScalarVector for GeometryVector { type OwnedItem = GeometryValue; - type RefItem<'a> = &'a GeometryValue; + type RefItem<'a> = GeometryValue; type Iter<'a> = GeometryVectorIter<'a>; type Builder = GeometryVectorBuilder; fn get_data(&self, idx: usize) -> Option> { - todo!() + match self { + GeometryVector::PointVector(vec) => match self.get(idx) { + Value::Null => None, + Value::Geometry(geo_value) => Some(geo_value), + _ => unreachable!(), + }, + } } fn iter_data(&self) -> Self::Iter<'_> { - todo!() + GeometryVectorIter { + vector: self, + pos: 0, + } } } pub struct GeometryVectorIter<'a> { vector: &'a GeometryVector, + pos: usize, } impl<'a> Iterator for GeometryVectorIter<'a> { - type Item = Option<&'a GeometryValue>; + type Item = Option; fn next(&mut self) -> Option { - todo!() + let pos = self.pos; + self.pos = self.pos + 1; + + if self.vector.len() <= pos { + return None; + } + + match self.vector { + GeometryVector::PointVector(vec) => match vec.get(pos) { + Value::Null => Some(None), + Value::Geometry(geo_value) => Some(Some(geo_value)), + _ => unreachable!(), + }, + } } } @@ -137,11 +172,15 @@ impl MutableVector for GeometryVectorBuilder { fn push_value_ref(&mut self, value: crate::value::ValueRef) -> crate::Result<()> { match value { - crate::value::ValueRef::Geometry(value) => { + ValueRef::Geometry(value) => { self.push(Some(value)); Ok(()) } - _ => todo!(), + ValueRef::Null => { + self.push(None); + Ok(()) + } + _ => unimplemented!(), } } @@ -151,7 +190,22 @@ impl MutableVector for GeometryVectorBuilder { offset: usize, length: usize, ) -> crate::Result<()> { - todo!() + let concrete_vector = vector + .as_any() + .downcast_ref::() + .with_context(|| crate::error::CastTypeSnafu { + msg: format!( + "Failed to cast vector from {} to {}", + vector.vector_type_name(), + stringify!(GeometryVector) + ), + })?; + + for idx in offset..offset + length { + let value = concrete_vector.get_ref(idx); + self.push_value_ref(value)?; + } + Ok(()) } } @@ -159,7 +213,7 @@ impl ScalarVectorBuilder for GeometryVectorBuilder { type VectorType = GeometryVector; fn with_capacity(capacity: usize) -> Self { - todo!() + unimplemented!() } fn push(&mut self, value: Option<::RefItem<'_>>) { diff --git a/src/datatypes/src/vectors/geometry/point.rs b/src/datatypes/src/vectors/geometry/point.rs index 8bf510e2cf..67ee3a3447 100644 --- a/src/datatypes/src/vectors/geometry/point.rs +++ b/src/datatypes/src/vectors/geometry/point.rs @@ -1,14 +1,17 @@ +use std::sync::Arc; + use arrow::array::{ Array, FixedSizeListArray, Float64Vec, ListArray, MutableArray, MutableFixedSizeListArray, - StructArray, + PrimitiveArray, StructArray, }; use arrow::datatypes::DataType::{self, Float64, List}; use arrow::datatypes::Field; use geo::Point; -use crate::value::{GeometryValue, ValueRef}; +use crate::value::{GeometryValue, OrderedF64, Value, ValueRef}; +use crate::vectors::impl_validity_for_vector; use crate::{ - prelude::{ScalarVector, ScalarVectorBuilder, Vector}, + prelude::{ScalarVector, ScalarVectorBuilder, Validity, Vector}, vectors::MutableVector, }; #[derive(Debug, Clone, PartialEq)] @@ -20,6 +23,57 @@ impl PointVector { pub fn memory_size(&self) -> usize { 2 * self.array.len() * std::mem::size_of::() } + + pub fn validity(&self) -> Validity { + impl_validity_for_vector!(self.array) + } + + pub fn slice(&self, offset: usize, length: usize) -> Self { + return Self { + array: self.array.slice(offset, length), + }; + } + + pub fn get(&self, index: usize) -> Value { + let ref_x_array = self + .array + .values() + .get(0) + .unwrap() + .as_any() + .downcast_ref::>() + .unwrap(); + + let validity = ref_x_array.validity(); + + if let Some(bm) = validity { + if !bm.get_bit(index) { + return Value::Null; + } + } + + let ref_y_array = self + .array + .values() + .get(1) + .unwrap() + .as_any() + .downcast_ref::>() + .unwrap(); + + let (x, y) = (ref_x_array.value(index), ref_x_array.value(index)); + let geo_value = GeometryValue::Point(Point::::new(x.into(), y.into())); + Value::Geometry(geo_value) + } + + pub fn get_ref(&self, index: usize) -> ValueRef { + if let Value::Geometry(geo_value) = self.get(index) { + ValueRef::Geometry(geo_value) + } else { + //NUll case + ValueRef::Null + } + } } pub struct PointVectorBuilder { @@ -29,9 +83,9 @@ pub struct PointVectorBuilder { } impl PointVectorBuilder { - pub fn push(&mut self, value: Option<&GeometryValue>) { + pub fn push(&mut self, value: Option) { match value { - Some(val) => match *val { + Some(val) => match val { GeometryValue::Point(xy) => { self.array_x.push(Some(*xy.x())); self.array_y.push(Some(*xy.y()));