refactor: address review comments

This commit is contained in:
Ning Sun
2024-09-02 11:43:33 -07:00
parent e56709b545
commit 0bff41e038
3 changed files with 84 additions and 42 deletions

View File

@@ -44,21 +44,32 @@ impl Function for GeohashFunction {
} }
fn signature(&self) -> Signature { fn signature(&self) -> Signature {
Signature::one_of( let mut signatures = Vec::new();
vec![ for coord_type in &[
TypeSignature::Exact(vec![ ConcreteDataType::float32_datatype(),
ConcreteDataType::float32_datatype(), ConcreteDataType::float64_datatype(),
ConcreteDataType::float32_datatype(), ] {
ConcreteDataType::int64_datatype(), for resolution_type in &[
]), ConcreteDataType::int8_datatype(),
TypeSignature::Exact(vec![ ConcreteDataType::int16_datatype(),
ConcreteDataType::float64_datatype(), ConcreteDataType::int32_datatype(),
ConcreteDataType::float64_datatype(), ConcreteDataType::int64_datatype(),
ConcreteDataType::int64_datatype(), ConcreteDataType::uint8_datatype(),
]), ConcreteDataType::uint16_datatype(),
], ConcreteDataType::uint32_datatype(),
Volatility::Stable, ConcreteDataType::uint64_datatype(),
) ] {
signatures.push(TypeSignature::Exact(vec![
// latitude
coord_type.clone(),
// longitude
coord_type.clone(),
// resolution
resolution_type.clone(),
]));
}
}
Signature::one_of(signatures, Volatility::Stable)
} }
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> { fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
@@ -80,14 +91,24 @@ impl Function for GeohashFunction {
let mut results = StringVectorBuilder::with_capacity(size); let mut results = StringVectorBuilder::with_capacity(size);
for i in 0..size { for i in 0..size {
let lat = lat_vec.get(i).as_f64(); let lat = lat_vec.get(i).as_f64_lossy();
let lon = lon_vec.get(i).as_f64(); let lon = lon_vec.get(i).as_f64_lossy();
let resolution = resolution_vec.get(i); let r = match resolution_vec.get(i) {
Value::Int8(v) => v as usize,
Value::Int16(v) => v as usize,
Value::Int32(v) => v as usize,
Value::Int64(v) => v as usize,
Value::UInt8(v) => v as usize,
Value::UInt16(v) => v as usize,
Value::UInt32(v) => v as usize,
Value::UInt64(v) => v as usize,
_ => unreachable!(),
};
let result = match (lat, lon, resolution) { let result = match (lat, lon) {
(Some(lat), Some(lon), Value::Int64(r)) => { (Some(lat), Some(lon)) => {
let coord = Coord { x: lon, y: lat }; let coord = Coord { x: lon, y: lat };
let encoded = geohash::encode(coord, r as usize) let encoded = geohash::encode(coord, r)
.map_err(|e| { .map_err(|e| {
BoxedError::new(PlainError::new( BoxedError::new(PlainError::new(
format!("Geohash error: {}", e), format!("Geohash error: {}", e),

View File

@@ -44,21 +44,32 @@ impl Function for H3Function {
} }
fn signature(&self) -> Signature { fn signature(&self) -> Signature {
Signature::one_of( let mut signatures = Vec::new();
vec![ for coord_type in &[
TypeSignature::Exact(vec![ ConcreteDataType::float32_datatype(),
ConcreteDataType::float32_datatype(), ConcreteDataType::float64_datatype(),
ConcreteDataType::float32_datatype(), ] {
ConcreteDataType::int64_datatype(), for resolution_type in &[
]), ConcreteDataType::int8_datatype(),
TypeSignature::Exact(vec![ ConcreteDataType::int16_datatype(),
ConcreteDataType::float64_datatype(), ConcreteDataType::int32_datatype(),
ConcreteDataType::float64_datatype(), ConcreteDataType::int64_datatype(),
ConcreteDataType::int64_datatype(), ConcreteDataType::uint8_datatype(),
]), ConcreteDataType::uint16_datatype(),
], ConcreteDataType::uint32_datatype(),
Volatility::Stable, ConcreteDataType::uint64_datatype(),
) ] {
signatures.push(TypeSignature::Exact(vec![
// latitude
coord_type.clone(),
// longitude
coord_type.clone(),
// resolution
resolution_type.clone(),
]));
}
}
Signature::one_of(signatures, Volatility::Stable)
} }
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> { fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
@@ -80,12 +91,22 @@ impl Function for H3Function {
let mut results = StringVectorBuilder::with_capacity(size); let mut results = StringVectorBuilder::with_capacity(size);
for i in 0..size { for i in 0..size {
let lat = lat_vec.get(i).as_f64(); let lat = lat_vec.get(i).as_f64_lossy();
let lon = lon_vec.get(i).as_f64(); let lon = lon_vec.get(i).as_f64_lossy();
let resolution = resolution_vec.get(i); let r = match resolution_vec.get(i) {
Value::Int8(v) => v as u8,
Value::Int16(v) => v as u8,
Value::Int32(v) => v as u8,
Value::Int64(v) => v as u8,
Value::UInt8(v) => v,
Value::UInt16(v) => v as u8,
Value::UInt32(v) => v as u8,
Value::UInt64(v) => v as u8,
_ => unreachable!(),
};
let result = match (lat, lon, resolution) { let result = match (lat, lon) {
(Some(lat), Some(lon), Value::Int64(r)) => { (Some(lat), Some(lon)) => {
let coord = LatLng::new(lat, lon) let coord = LatLng::new(lat, lon)
.map_err(|e| { .map_err(|e| {
BoxedError::new(PlainError::new( BoxedError::new(PlainError::new(

View File

@@ -269,7 +269,7 @@ impl Value {
} }
/// Cast Value to f32. Return None if it's not castable; /// Cast Value to f32. Return None if it's not castable;
pub fn as_f64(&self) -> Option<f64> { pub fn as_f64_lossy(&self) -> Option<f64> {
match self { match self {
Value::Float32(v) => Some(v.0 as _), Value::Float32(v) => Some(v.0 as _),
Value::Float64(v) => Some(v.0), Value::Float64(v) => Some(v.0),