feat(function): support length function for Array & Array<T>

common/functions/src/scalars/commons/common.rs

@@ -0,0 +1,24 @@
+// Copyright 2022 Datafuse Labs.
+//
+// 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 super::length::LengthFunction;
+use crate::scalars::FunctionFactory;
+
+pub struct CommonFunction;
+
+impl CommonFunction {
+    pub fn register(factory: &mut FunctionFactory) {
+        factory.register("length", LengthFunction::desc());
+    }
+}

common/functions/src/scalars/commons/length.rs

@@ -0,0 +1,54 @@
+// Copyright 2022 Datafuse Labs.
+//
+// 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 common_datavalues::prelude::*;
+use common_datavalues::DataTypeImpl;
+use common_exception::ErrorCode;
+use common_exception::Result;
+
+use crate::scalars::ArrayLengthFunction;
+use crate::scalars::Function;
+use crate::scalars::FunctionDescription;
+use crate::scalars::FunctionFeatures;
+use crate::scalars::StringLengthFunction;
+use crate::scalars::VariantArrayLengthFunction;
+
+pub struct LengthFunction {
+    _display_name: String,
+}
+
+impl LengthFunction {
+    pub fn try_create(display_name: &str, args: &[&DataTypeImpl]) -> Result<Box<dyn Function>> {
+        let data_type = args[0];
+
+        if data_type.data_type_id().is_string() {
+            StringLengthFunction::try_create(display_name, args)
+        } else if data_type.data_type_id().is_array() {
+            ArrayLengthFunction::try_create(display_name, args)
+        } else if data_type.data_type_id().is_variant_or_array() {
+            VariantArrayLengthFunction::try_create(display_name, args)
+        } else {
+            Err(ErrorCode::IllegalDataType(format!(
+                "Invalid argument types for function '{}': ({:?})",
+                display_name.to_uppercase(),
+                data_type.data_type_id(),
+            )))
+        }
+    }
+
+    pub fn desc() -> FunctionDescription {
+        FunctionDescription::creator(Box::new(Self::try_create))
+            .features(FunctionFeatures::default().deterministic().num_arguments(1))
+    }
+}

common/functions/src/scalars/commons/mod.rs

@@ -0,0 +1,19 @@
+// Copyright 2022 Datafuse Labs.
+//
+// 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.
+
+mod common;
+mod length;
+
+pub use common::CommonFunction;
+pub use length::LengthFunction;

common/functions/src/scalars/function_factory.rs

@@ -20,6 +20,7 @@ use common_exception::ErrorCode;
 use common_exception::Result;
 use once_cell::sync::Lazy;
 
+use super::commons::CommonFunction;
 use super::function::Function;
 use super::ArithmeticFunction;
 use super::ComparisonFunction;
@@ -69,6 +70,7 @@ static FUNCTION_FACTORY: Lazy<Arc<FunctionFactory>> = Lazy::new(|| {
     let mut function_factory = FunctionFactory::create();
 
     ArithmeticFunction::register(&mut function_factory);
+    CommonFunction::register(&mut function_factory);
     ToCastFunction::register(&mut function_factory);
     TupleClassFunction::register(&mut function_factory);
     ComparisonFunction::register(&mut function_factory);

common/functions/src/scalars/mod.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 mod arithmetics;
+mod commons;
 mod comparisons;
 mod conditionals;
 mod contexts;
@@ -34,6 +35,7 @@ mod tuples;
 mod uuids;
 
 pub use arithmetics::*;
+pub use commons::*;
 pub use comparisons::*;
 pub use conditionals::*;
 pub use contexts::*;

common/functions/src/scalars/semi_structureds/array_length.rs

@@ -0,0 +1,84 @@
+// Copyright 2022 Datafuse Labs.
+//
+// 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::fmt;
+use std::sync::Arc;
+
+use common_datavalues::prelude::*;
+use common_datavalues::DataTypeImpl;
+use common_exception::ErrorCode;
+use common_exception::Result;
+
+use crate::scalars::Function;
+use crate::scalars::FunctionContext;
+use crate::scalars::FunctionDescription;
+use crate::scalars::FunctionFeatures;
+
+#[derive(Clone)]
+pub struct ArrayLengthFunction {
+    display_name: String,
+}
+
+impl ArrayLengthFunction {
+    pub fn try_create(display_name: &str, args: &[&DataTypeImpl]) -> Result<Box<dyn Function>> {
+        let data_type = args[0];
+
+        if !data_type.data_type_id().is_array() {
+            return Err(ErrorCode::IllegalDataType(format!(
+                "Invalid argument types for function '{}': ({:?})",
+                display_name.to_uppercase(),
+                data_type.data_type_id(),
+            )));
+        }
+
+        Ok(Box::new(ArrayLengthFunction {
+            display_name: display_name.to_string(),
+        }))
+    }
+
+    pub fn desc() -> FunctionDescription {
+        FunctionDescription::creator(Box::new(Self::try_create))
+            .features(FunctionFeatures::default().deterministic().num_arguments(1))
+    }
+}
+
+impl Function for ArrayLengthFunction {
+    fn name(&self) -> &str {
+        &*self.display_name
+    }
+
+    fn return_type(&self) -> DataTypeImpl {
+        u64::to_data_type()
+    }
+
+    fn eval(
+        &self,
+        _func_ctx: FunctionContext,
+        columns: &ColumnsWithField,
+        input_rows: usize,
+    ) -> Result<ColumnRef> {
+        let array_column: &ArrayColumn = Series::check_get(columns[0].column())?;
+        let vec = (0..input_rows)
+            .into_iter()
+            .map(|index| array_column.size_at_index(index) as u64)
+            .collect();
+        Ok(Arc::new(PrimitiveColumn::new_from_vec(vec)))
+    }
+}
+
+impl fmt::Display for ArrayLengthFunction {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}()", self.display_name)
+    }
+}

common/functions/src/scalars/semi_structureds/length.rs

@@ -0,0 +1,93 @@
+// Copyright 2022 Datafuse Labs.
+//
+// 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::fmt;
+use std::sync::Arc;
+
+use common_datavalues::prelude::*;
+use common_datavalues::DataTypeImpl;
+use common_exception::ErrorCode;
+use common_exception::Result;
+
+use crate::scalars::Function;
+use crate::scalars::FunctionContext;
+use crate::scalars::FunctionDescription;
+use crate::scalars::FunctionFeatures;
+
+#[derive(Clone)]
+pub struct VariantArrayLengthFunction {
+    display_name: String,
+}
+
+impl VariantArrayLengthFunction {
+    pub fn try_create(display_name: &str, args: &[&DataTypeImpl]) -> Result<Box<dyn Function>> {
+        let data_type = args[0];
+
+        if !data_type.data_type_id().is_variant_or_array() {
+            return Err(ErrorCode::IllegalDataType(format!(
+                "Invalid argument types for function '{}': ({:?})",
+                display_name.to_uppercase(),
+                data_type.data_type_id(),
+            )));
+        }
+
+        Ok(Box::new(VariantArrayLengthFunction {
+            display_name: display_name.to_string(),
+        }))
+    }
+
+    pub fn desc() -> FunctionDescription {
+        FunctionDescription::creator(Box::new(Self::try_create))
+            .features(FunctionFeatures::default().deterministic().num_arguments(1))
+    }
+}
+
+impl Function for VariantArrayLengthFunction {
+    fn name(&self) -> &str {
+        &*self.display_name
+    }
+
+    fn return_type(&self) -> DataTypeImpl {
+        u64::to_data_type()
+    }
+
+    fn eval(
+        &self,
+        _func_ctx: FunctionContext,
+        columns: &ColumnsWithField,
+        input_rows: usize,
+    ) -> Result<ColumnRef> {
+        let array_column: &VariantColumn = Series::check_get(columns[0].column())?;
+        let mut res: Vec<u64> = Vec::with_capacity(input_rows);
+
+        for array in array_column.iter() {
+            if !array.is_array() {
+                return Err(ErrorCode::IllegalDataType(format!(
+                    "Invalid argument types for function '{}': Variant is not VariantArray",
+                    self.display_name.to_uppercase(),
+                )));
+            }
+            let len = array.to_owned_scalar().as_array().unwrap().len() as u64;
+            res.push(len);
+        }
+
+        Ok(Arc::new(PrimitiveColumn::new_from_vec(res)))
+    }
+}
+
+impl fmt::Display for VariantArrayLengthFunction {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}()", self.display_name)
+    }
+}

common/functions/src/scalars/semi_structureds/mod.rs

@@ -12,17 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod array_length;
 mod check_json;
 mod get;
 mod json_extract_path_text;
+mod length;
 mod parse_json;
 mod semi_structured;
 
+pub use array_length::ArrayLengthFunction;
 pub use check_json::CheckJsonFunction;
 pub use get::GetFunction;
 pub use get::GetIgnoreCaseFunction;
 pub use get::GetPathFunction;
 pub use json_extract_path_text::JsonExtractPathTextFunction;
+pub use length::VariantArrayLengthFunction;
 pub use parse_json::ParseJsonFunction;
 pub use parse_json::TryParseJsonFunction;
 pub use semi_structured::SemiStructuredFunction;

common/functions/src/scalars/strings/length.rs

@@ -16,9 +16,9 @@ use super::NumberOperator;
 use super::String2NumberFunction;
 
 #[derive(Clone, Default)]
-pub struct Length {}
+pub struct StringLength {}
 
-impl NumberOperator<u64> for Length {
+impl NumberOperator<u64> for StringLength {
     const IS_DETERMINISTIC: bool = true;
     const MAYBE_MONOTONIC: bool = false;
 
@@ -27,4 +27,4 @@ impl NumberOperator<u64> for Length {
     }
 }
 
-pub type LengthFunction = String2NumberFunction<Length, u64>;
+pub type StringLengthFunction = String2NumberFunction<StringLength, u64>;

common/functions/src/scalars/strings/mod.rs

@@ -72,7 +72,7 @@ pub use format::FormatFunction;
 pub use insert::InsertFunction;
 pub use leftright::LeftFunction;
 pub use leftright::RightFunction;
-pub use length::LengthFunction;
+pub use length::StringLengthFunction;
 pub use locate::InstrFunction;
 pub use locate::LocateFunction;
 pub use locate::PositionFunction;

common/functions/src/scalars/strings/string.rs

@@ -33,7 +33,6 @@ use crate::scalars::InstrFunction;
 use crate::scalars::LTrimFunction;
 use crate::scalars::LeftFunction;
 use crate::scalars::LeftPadFunction;
-use crate::scalars::LengthFunction;
 use crate::scalars::LocateFunction;
 use crate::scalars::LowerFunction;
 use crate::scalars::OctFunction;
@@ -83,7 +82,6 @@ impl StringFunction {
         factory.register("char_length", CharLengthFunction::desc());
         factory.register("character_length", CharLengthFunction::desc());
         factory.register("ord", OrdFunction::desc());
-        factory.register("length", LengthFunction::desc());
         factory.register("regexp_instr", RegexpInStrFunction::desc());
         factory.register("regexp_like", RegexpLikeFunction::desc());
         factory.register("regexp_replace", RegexpReplaceFunction::desc());