From 50b521c526d8f6f02ee4467c587b14d956a337f2 Mon Sep 17 00:00:00 2001 From: xiaoniaoyouhuajiang <2583473505@qq.com> Date: Thu, 27 Feb 2025 23:54:48 +0800 Subject: [PATCH] feat: add `vec_dim` function (#5587) * feat:add `vec_dim` function * delete unused imports * Modified to be implemented correctly * fix comment * add order for sqlness test --- src/common/function/src/scalars/vector.rs | 2 + .../function/src/scalars/vector/vector_dim.rs | 172 ++++++++++++++++++ .../common/function/vector/vector.result | 42 +++++ .../common/function/vector/vector.sql | 20 ++ 4 files changed, 236 insertions(+) create mode 100644 src/common/function/src/scalars/vector/vector_dim.rs diff --git a/src/common/function/src/scalars/vector.rs b/src/common/function/src/scalars/vector.rs index 90aed7cbd7a..381c757d9b9 100644 --- a/src/common/function/src/scalars/vector.rs +++ b/src/common/function/src/scalars/vector.rs @@ -22,6 +22,7 @@ mod scalar_add; mod scalar_mul; pub(crate) mod sum; mod vector_add; +mod vector_dim; mod vector_div; mod vector_mul; mod vector_norm; @@ -54,6 +55,7 @@ impl VectorFunction { registry.register(Arc::new(vector_mul::VectorMulFunction)); registry.register(Arc::new(vector_div::VectorDivFunction)); registry.register(Arc::new(vector_norm::VectorNormFunction)); + registry.register(Arc::new(vector_dim::VectorDimFunction)); registry.register(Arc::new(elem_sum::ElemSumFunction)); registry.register(Arc::new(elem_product::ElemProductFunction)); } diff --git a/src/common/function/src/scalars/vector/vector_dim.rs b/src/common/function/src/scalars/vector/vector_dim.rs new file mode 100644 index 00000000000..6a7c0781004 --- /dev/null +++ b/src/common/function/src/scalars/vector/vector_dim.rs @@ -0,0 +1,172 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::borrow::Cow; +use std::fmt::Display; + +use common_query::error::InvalidFuncArgsSnafu; +use common_query::prelude::{Signature, TypeSignature, Volatility}; +use datatypes::prelude::ConcreteDataType; +use datatypes::scalars::ScalarVectorBuilder; +use datatypes::vectors::{MutableVector, UInt64VectorBuilder, VectorRef}; +use snafu::ensure; + +use crate::function::{Function, FunctionContext}; +use crate::scalars::vector::impl_conv::{as_veclit, as_veclit_if_const}; + +const NAME: &str = "vec_dim"; + +/// Returns the dimension of the vector. +/// +/// # Example +/// +/// ```sql +/// SELECT vec_dim('[7.0, 8.0, 9.0, 10.0]'); +/// +/// +---------------------------------------------------------------+ +/// | vec_dim(Utf8("[7.0, 8.0, 9.0, 10.0]")) | +/// +---------------------------------------------------------------+ +/// | 4 | +/// +---------------------------------------------------------------+ +/// +#[derive(Debug, Clone, Default)] +pub struct VectorDimFunction; + +impl Function for VectorDimFunction { + fn name(&self) -> &str { + NAME + } + + fn return_type( + &self, + _input_types: &[ConcreteDataType], + ) -> common_query::error::Result { + Ok(ConcreteDataType::uint64_datatype()) + } + + fn signature(&self) -> Signature { + Signature::one_of( + vec![ + TypeSignature::Exact(vec![ConcreteDataType::string_datatype()]), + TypeSignature::Exact(vec![ConcreteDataType::binary_datatype()]), + ], + Volatility::Immutable, + ) + } + + fn eval( + &self, + _func_ctx: FunctionContext, + columns: &[VectorRef], + ) -> common_query::error::Result { + ensure!( + columns.len() == 1, + InvalidFuncArgsSnafu { + err_msg: format!( + "The length of the args is not correct, expect exactly one, have: {}", + columns.len() + ) + } + ); + let arg0 = &columns[0]; + + let len = arg0.len(); + let mut result = UInt64VectorBuilder::with_capacity(len); + if len == 0 { + return Ok(result.to_vector()); + } + + let arg0_const = as_veclit_if_const(arg0)?; + + for i in 0..len { + let arg0 = match arg0_const.as_ref() { + Some(arg0) => Some(Cow::Borrowed(arg0.as_ref())), + None => as_veclit(arg0.get_ref(i))?, + }; + let Some(arg0) = arg0 else { + result.push_null(); + continue; + }; + result.push(Some(arg0.len() as u64)); + } + + Ok(result.to_vector()) + } +} + +impl Display for VectorDimFunction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", NAME.to_ascii_uppercase()) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_query::error::Error; + use datatypes::vectors::StringVector; + + use super::*; + + #[test] + fn test_vec_dim() { + let func = VectorDimFunction; + + let input0 = Arc::new(StringVector::from(vec![ + Some("[0.0,2.0,3.0]".to_string()), + Some("[1.0,2.0,3.0,4.0]".to_string()), + None, + Some("[5.0]".to_string()), + ])); + + let result = func.eval(FunctionContext::default(), &[input0]).unwrap(); + + let result = result.as_ref(); + assert_eq!(result.len(), 4); + assert_eq!(result.get_ref(0).as_u64().unwrap(), Some(3)); + assert_eq!(result.get_ref(1).as_u64().unwrap(), Some(4)); + assert!(result.get_ref(2).is_null()); + assert_eq!(result.get_ref(3).as_u64().unwrap(), Some(1)); + } + + #[test] + fn test_dim_error() { + let func = VectorDimFunction; + + let input0 = Arc::new(StringVector::from(vec![ + Some("[1.0,2.0,3.0]".to_string()), + Some("[4.0,5.0,6.0]".to_string()), + None, + Some("[2.0,3.0,3.0]".to_string()), + ])); + let input1 = Arc::new(StringVector::from(vec![ + Some("[1.0,1.0,1.0]".to_string()), + Some("[6.0,5.0,4.0]".to_string()), + Some("[3.0,2.0,2.0]".to_string()), + ])); + + let result = func.eval(FunctionContext::default(), &[input0, input1]); + + match result { + Err(Error::InvalidFuncArgs { err_msg, .. }) => { + assert_eq!( + err_msg, + "The length of the args is not correct, expect exactly one, have: 2" + ) + } + _ => unreachable!(), + } + } +} diff --git a/tests/cases/standalone/common/function/vector/vector.result b/tests/cases/standalone/common/function/vector/vector.result index 1b81fa98b08..7f40c73636b 100644 --- a/tests/cases/standalone/common/function/vector/vector.result +++ b/tests/cases/standalone/common/function/vector/vector.result @@ -284,3 +284,45 @@ FROM ( | [-4,-20,-54] | +-------------------------------+ +SELECT vec_dim('[7.0, 8.0, 9.0, 10.0]'); + ++----------------------------------------+ +| vec_dim(Utf8("[7.0, 8.0, 9.0, 10.0]")) | ++----------------------------------------+ +| 4 | ++----------------------------------------+ + +SELECT v, vec_dim(v) +FROM ( + SELECT '[1.0, 2.0, 3.0]' AS v + UNION ALL + SELECT '[-1.0]' AS v + UNION ALL + SELECT '[4.0, 5.0, 6.0]' AS v + ) Order By vec_dim(v) ASC; + ++-----------------+------------+ +| v | vec_dim(v) | ++-----------------+------------+ +| [-1.0] | 1 | +| [1.0, 2.0, 3.0] | 3 | +| [4.0, 5.0, 6.0] | 3 | ++-----------------+------------+ + +SELECT v, vec_dim(v) +FROM ( + SELECT '[1.0, 2.0, 3.0]' AS v + UNION ALL + SELECT '[-1.0]' AS v + UNION ALL + SELECT '[7.0, 8.0, 9.0, 10.0]' AS v + ) Order By vec_dim(v) ASC; + ++-----------------------+------------+ +| v | vec_dim(v) | ++-----------------------+------------+ +| [-1.0] | 1 | +| [1.0, 2.0, 3.0] | 3 | +| [7.0, 8.0, 9.0, 10.0] | 4 | ++-----------------------+------------+ + diff --git a/tests/cases/standalone/common/function/vector/vector.sql b/tests/cases/standalone/common/function/vector/vector.sql index 49c8e88f283..b53b6af453f 100644 --- a/tests/cases/standalone/common/function/vector/vector.sql +++ b/tests/cases/standalone/common/function/vector/vector.sql @@ -79,3 +79,23 @@ FROM ( UNION ALL SELECT '[4.0, 5.0, 6.0]' AS v ); + +SELECT vec_dim('[7.0, 8.0, 9.0, 10.0]'); + +SELECT v, vec_dim(v) +FROM ( + SELECT '[1.0, 2.0, 3.0]' AS v + UNION ALL + SELECT '[-1.0]' AS v + UNION ALL + SELECT '[4.0, 5.0, 6.0]' AS v + ) Order By vec_dim(v) ASC; + +SELECT v, vec_dim(v) +FROM ( + SELECT '[1.0, 2.0, 3.0]' AS v + UNION ALL + SELECT '[-1.0]' AS v + UNION ALL + SELECT '[7.0, 8.0, 9.0, 10.0]' AS v + ) Order By vec_dim(v) ASC;