From 7aed777bc416b735bc5f60e50ca092a2e692b839 Mon Sep 17 00:00:00 2001 From: liangxingjian <965662709@qq.com> Date: Wed, 12 Oct 2022 17:06:42 +0800 Subject: [PATCH] feat:add iter and ref of geo types --- src/datatypes/src/scalars.rs | 19 +++++-- src/datatypes/src/value.rs | 50 +++++++++++++++-- src/datatypes/src/vectors/geometry.rs | 61 +++++++++++---------- src/datatypes/src/vectors/geometry/point.rs | 11 +--- 4 files changed, 94 insertions(+), 47 deletions(-) diff --git a/src/datatypes/src/scalars.rs b/src/datatypes/src/scalars.rs index 2badff7229..2018cefac7 100644 --- a/src/datatypes/src/scalars.rs +++ b/src/datatypes/src/scalars.rs @@ -1,9 +1,10 @@ use std::any::Any; +use std::default; use common_time::{Date, DateTime, Timestamp}; use crate::prelude::*; -use crate::value::{GeometryValue, ListValue, ListValueRef}; +use crate::value::{GeometryValue, GeometryValueRef, ListValue, ListValueRef}; use crate::vectors::all::GeometryVector; use crate::vectors::*; @@ -344,24 +345,30 @@ impl<'a> ScalarRef<'a> for ListValueRef<'a> { impl Scalar for GeometryValue { type VectorType = GeometryVector; - type RefType<'a> = GeometryValue; + type RefType<'a> = GeometryValueRef<'a>; fn as_scalar_ref(&self) -> Self::RefType<'_> { - todo!() + GeometryValueRef::Ref { val: self } } fn upcast_gat<'short, 'long: 'short>(long: Self::RefType<'long>) -> Self::RefType<'short> { - todo!() + long } } -impl<'a> ScalarRef<'a> for GeometryValue { +impl<'a> ScalarRef<'a> for GeometryValueRef<'a> { type VectorType = GeometryVector; type ScalarType = GeometryValue; fn to_owned_scalar(&self) -> Self::ScalarType { - todo!() + match self { + GeometryValueRef::Indexed { vector, idx } => match vector.get(*idx) { + Value::Geometry(value) => value, + _ => unreachable!(), + }, + GeometryValueRef::Ref { val } => (*val).clone(), + } } } diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 9527858d97..608bc43994 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -10,6 +10,7 @@ use serde::{Deserialize, Serialize}; use crate::error::{self, Result}; use crate::prelude::*; +use crate::vectors::all::GeometryVector; use crate::vectors::ListVector; pub type OrderedF32 = OrderedFloat; @@ -73,7 +74,7 @@ impl Value { Value::Date(_) => ConcreteDataType::date_datatype(), Value::DateTime(_) => ConcreteDataType::date_datatype(), Value::Timestamp(v) => ConcreteDataType::timestamp_datatype(v.unit()), - Value::Geometry(_) => todo!(), + Value::Geometry(_) => ConcreteDataType::geometry_datatype(), } } @@ -115,7 +116,7 @@ impl Value { Value::DateTime(v) => ValueRef::DateTime(*v), Value::List(v) => ValueRef::List(ListValueRef::Ref { val: v }), Value::Timestamp(v) => ValueRef::Timestamp(*v), - Value::Geometry(_) => todo!(), + Value::Geometry(v) => ValueRef::Geometry(GeometryValueRef::Ref { val: v }), } } } @@ -253,7 +254,7 @@ impl TryFrom for serde_json::Value { Value::DateTime(v) => serde_json::Value::Number(v.val().into()), Value::List(v) => serde_json::to_value(v)?, Value::Timestamp(v) => serde_json::to_value(v.value())?, - Value::Geometry(_) => todo!(), + Value::Geometry(_) => unimplemented!(), }; Ok(json_value) @@ -346,7 +347,7 @@ pub enum ValueRef<'a> { DateTime(DateTime), Timestamp(Timestamp), List(ListValueRef<'a>), - Geometry(GeometryValue), + Geometry(GeometryValueRef<'a>), } macro_rules! impl_as_for_value_ref { @@ -537,6 +538,47 @@ impl<'a> PartialOrd for ListValueRef<'a> { } } +#[derive(Debug, Clone, Copy)] +pub enum GeometryValueRef<'a> { + Indexed { + vector: &'a GeometryVector, + idx: usize, + }, + Ref { + val: &'a GeometryValue, + }, +} + +impl<'a> GeometryValueRef<'a> { + /// Convert self to [Value]. This method would clone the underlying data. + fn to_value(self) -> Value { + match self { + GeometryValueRef::Indexed { vector, idx } => vector.get(idx), + GeometryValueRef::Ref { val } => Value::Geometry(val.clone()), + } + } +} + +impl<'a> PartialEq for GeometryValueRef<'a> { + fn eq(&self, other: &Self) -> bool { + self.to_value().eq(&other.to_value()) + } +} + +impl<'a> Eq for GeometryValueRef<'a> {} + +impl<'a> Ord for GeometryValueRef<'a> { + fn cmp(&self, other: &Self) -> Ordering { + unreachable!() + } +} + +impl<'a> PartialOrd for GeometryValueRef<'a> { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/datatypes/src/vectors/geometry.rs b/src/datatypes/src/vectors/geometry.rs index 098f4b1472..4e10fd0196 100644 --- a/src/datatypes/src/vectors/geometry.rs +++ b/src/datatypes/src/vectors/geometry.rs @@ -5,7 +5,8 @@ use snafu::{ensure, OptionExt}; use self::point::{PointVector, PointVectorBuilder}; use super::{MutableVector, Validity, Value, Vector}; -use crate::value::ValueRef; +use crate::prelude::ScalarRef; +use crate::value::{GeometryValueRef, ValueRef}; use crate::vectors::{impl_try_from_arrow_array_for_vector, impl_validity_for_vector}; use crate::{ data_type::ConcreteDataType, @@ -86,27 +87,34 @@ impl Vector for GeometryVector { fn get_ref(&self, index: usize) -> crate::value::ValueRef { match self { - GeometryVector::PointVector(vec) => vec.get_ref(index), + GeometryVector::PointVector(vector) => { + if vector.is_null(index) { + return ValueRef::Null; + } + } } + + ValueRef::Geometry(GeometryValueRef::Indexed { + vector: self, + idx: index, + }) } } impl ScalarVector for GeometryVector { type OwnedItem = GeometryValue; - type RefItem<'a> = GeometryValue; + type RefItem<'a> = GeometryValueRef<'a>; type Iter<'a> = GeometryVectorIter<'a>; type Builder = GeometryVectorBuilder; fn get_data(&self, idx: usize) -> Option> { - match self { - GeometryVector::PointVector(vec) => match self.get(idx) { - Value::Null => None, - Value::Geometry(geo_value) => Some(geo_value), - _ => unreachable!(), - }, + match self.get_ref(idx) { + ValueRef::Null => None, + ValueRef::Geometry(geo_ref) => Some(geo_ref), + _ => unreachable!(), } } @@ -124,7 +132,7 @@ pub struct GeometryVectorIter<'a> { } impl<'a> Iterator for GeometryVectorIter<'a> { - type Item = Option; + type Item = Option>; fn next(&mut self) -> Option { let pos = self.pos; self.pos = self.pos + 1; @@ -133,13 +141,10 @@ impl<'a> Iterator for GeometryVectorIter<'a> { 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!(), - }, - } + Some(Some(GeometryValueRef::Indexed { + vector: self.vector, + idx: pos, + })) } } @@ -171,17 +176,14 @@ impl MutableVector for GeometryVectorBuilder { } fn push_value_ref(&mut self, value: crate::value::ValueRef) -> crate::Result<()> { - match value { - ValueRef::Geometry(value) => { - self.push(Some(value)); - Ok(()) - } - ValueRef::Null => { - self.push(None); - Ok(()) - } - _ => unimplemented!(), + match self { + GeometryVectorBuilder::PointVectorBuilder(builder) => match value { + ValueRef::Null => builder.push(None), + ValueRef::Geometry(geo_ref) => builder.push(Some(geo_ref.to_owned_scalar())), + _ => unreachable!(), + }, } + Ok(()) } fn extend_slice_of( @@ -217,8 +219,9 @@ impl ScalarVectorBuilder for GeometryVectorBuilder { } fn push(&mut self, value: Option<::RefItem<'_>>) { - match self { - GeometryVectorBuilder::PointVectorBuilder(builder) => builder.push(value), + match value { + Some(geo_ref) => self.push_value_ref(ValueRef::Geometry(geo_ref)).unwrap(), + None => self.push_value_ref(ValueRef::Null).unwrap(), } } diff --git a/src/datatypes/src/vectors/geometry/point.rs b/src/datatypes/src/vectors/geometry/point.rs index 67ee3a3447..7fb4e78378 100644 --- a/src/datatypes/src/vectors/geometry/point.rs +++ b/src/datatypes/src/vectors/geometry/point.rs @@ -8,7 +8,7 @@ use arrow::datatypes::DataType::{self, Float64, List}; use arrow::datatypes::Field; use geo::Point; -use crate::value::{GeometryValue, OrderedF64, Value, ValueRef}; +use crate::value::{GeometryValue, GeometryValueRef, OrderedF64, Value, ValueRef}; use crate::vectors::impl_validity_for_vector; use crate::{ prelude::{ScalarVector, ScalarVectorBuilder, Validity, Vector}, @@ -66,13 +66,8 @@ impl PointVector { 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 fn is_null(&self, index: usize) -> bool { + self.get(index).is_null() } }