use FactoryCreatorWithTypes for math functions

databendlabs · Apr 6, 2022 · 89db29f · 89db29f
1 parent 2fc53b2
commit 89db29f
Show file tree

Hide file tree

Showing 31 changed files with 503 additions and 341 deletions.
diff --git a/common/datablocks/src/kernels/data_block_filter.rs b/common/datablocks/src/kernels/data_block_filter.rs
@@ -67,7 +67,7 @@ impl DataBlock {
             let inner = col.inner();
             if inner.is_const() {
                 let col: &ConstColumn = unsafe { Series::static_cast(inner) };
-                let inner_boolean = Self::cast_to_nonull_boolean(col.inner())?;
+                let inner_boolean = Self::cast_to_nonull_boolean(inner)?;
                 return Ok(Arc::new(ConstColumn::new(inner_boolean, col.len())));
             }
 

diff --git a/common/functions/src/scalars/maths/abs.rs b/common/functions/src/scalars/maths/abs.rs
@@ -20,7 +20,7 @@ use common_exception::Result;
 
 use crate::scalars::function::Function;
 use crate::scalars::function_common::assert_numeric;
-use crate::scalars::function_factory::FunctionDescription;
+use crate::scalars::function_factory::TypedFunctionDescription;
 use crate::scalars::scalar_unary_op;
 use crate::scalars::EvalContext;
 use crate::scalars::FunctionFeatures;
@@ -29,17 +29,33 @@ use crate::scalars::Monotonicity;
 #[derive(Clone)]
 pub struct AbsFunction {
     _display_name: String,
+    result_type: DataTypePtr,
 }
 
 impl AbsFunction {
-    pub fn try_create(display_name: &str) -> Result<Box<dyn Function>> {
+    pub fn try_create(display_name: &str, args: &[&DataTypePtr]) -> Result<Box<dyn Function>> {
+        assert_numeric(args[0])?;
+        let result_type = match args[0].data_type_id() {
+            TypeID::Int8 => u8::to_data_type(),
+            TypeID::Int16 => u16::to_data_type(),
+            TypeID::Int32 => u32::to_data_type(),
+            TypeID::Int64 => u64::to_data_type(),
+            TypeID::UInt8 => u8::to_data_type(),
+            TypeID::UInt16 => u16::to_data_type(),
+            TypeID::UInt32 => u32::to_data_type(),
+            TypeID::UInt64 => u64::to_data_type(),
+            TypeID::Float32 => f32::to_data_type(),
+            TypeID::Float64 => f64::to_data_type(),
+            _ => unreachable!(),
+        };
         Ok(Box::new(AbsFunction {
             _display_name: display_name.to_string(),
+            result_type,
         }))
     }
 
-    pub fn desc() -> FunctionDescription {
-        FunctionDescription::creator(Box::new(Self::try_create)).features(
+    pub fn desc() -> TypedFunctionDescription {
+        TypedFunctionDescription::creator(Box::new(Self::try_create)).features(
             FunctionFeatures::default()
                 .deterministic()
                 .monotonicity()
@@ -75,22 +91,8 @@ impl Function for AbsFunction {
         "abs"
     }
 
-    fn return_type(&self, args: &[&DataTypePtr]) -> Result<DataTypePtr> {
-        assert_numeric(args[0])?;
-        let data_type = match args[0].data_type_id() {
-            TypeID::Int8 => u8::to_data_type(),
-            TypeID::Int16 => u16::to_data_type(),
-            TypeID::Int32 => u32::to_data_type(),
-            TypeID::Int64 => u64::to_data_type(),
-            TypeID::UInt8 => u8::to_data_type(),
-            TypeID::UInt16 => u16::to_data_type(),
-            TypeID::UInt32 => u32::to_data_type(),
-            TypeID::UInt64 => u64::to_data_type(),
-            TypeID::Float32 => f32::to_data_type(),
-            TypeID::Float64 => f64::to_data_type(),
-            _ => unreachable!(),
-        };
-        Ok(data_type)
+    fn return_type(&self, _args: &[&DataTypePtr]) -> Result<DataTypePtr> {
+        Ok(self.result_type.clone())
     }
 
     fn eval(&self, columns: &ColumnsWithField, _input_rows: usize) -> Result<ColumnRef> {

diff --git a/common/functions/src/scalars/maths/angle.rs b/common/functions/src/scalars/maths/angle.rs
@@ -21,7 +21,7 @@ use common_exception::Result;
 use num::cast::AsPrimitive;
 
 use crate::scalars::function_common::assert_numeric;
-use crate::scalars::function_factory::FunctionDescription;
+use crate::scalars::function_factory::TypedFunctionDescription;
 use crate::scalars::scalar_unary_op;
 use crate::scalars::EvalContext;
 use crate::scalars::Function;
@@ -40,15 +40,16 @@ pub trait AngleConvertFunction {
 impl<T> AngleFunction<T>
 where T: AngleConvertFunction + Clone + Sync + Send + 'static
 {
-    pub fn try_create(display_name: &str) -> Result<Box<dyn Function>> {
+    pub fn try_create(display_name: &str, args: &[&DataTypePtr]) -> Result<Box<dyn Function>> {
+        assert_numeric(args[0])?;
         Ok(Box::new(AngleFunction::<T> {
             _display_name: display_name.to_string(),
             t: PhantomData,
         }))
     }
 
-    pub fn desc() -> FunctionDescription {
-        FunctionDescription::creator(Box::new(Self::try_create))
+    pub fn desc() -> TypedFunctionDescription {
+        TypedFunctionDescription::creator(Box::new(Self::try_create))
             .features(FunctionFeatures::default().deterministic().num_arguments(1))
     }
 }
@@ -60,8 +61,7 @@ where T: AngleConvertFunction + Clone + Sync + Send + 'static
         "AngleFunction"
     }
 
-    fn return_type(&self, args: &[&DataTypePtr]) -> Result<DataTypePtr> {
-        assert_numeric(args[0])?;
+    fn return_type(&self, _args: &[&DataTypePtr]) -> Result<DataTypePtr> {
         Ok(Float64Type::arc())
     }
 

diff --git a/common/functions/src/scalars/maths/ceil.rs b/common/functions/src/scalars/maths/ceil.rs
@@ -21,7 +21,7 @@ use common_exception::Result;
 use num::cast::AsPrimitive;
 
 use crate::scalars::function_common::assert_numeric;
-use crate::scalars::function_factory::FunctionDescription;
+use crate::scalars::function_factory::TypedFunctionDescription;
 use crate::scalars::scalar_unary_op;
 use crate::scalars::EvalContext;
 use crate::scalars::Function;
@@ -34,14 +34,15 @@ pub struct CeilFunction {
 }
 
 impl CeilFunction {
-    pub fn try_create(display_name: &str) -> Result<Box<dyn Function>> {
+    pub fn try_create(display_name: &str, args: &[&DataTypePtr]) -> Result<Box<dyn Function>> {
+        assert_numeric(args[0])?;
         Ok(Box::new(CeilFunction {
             display_name: display_name.to_string(),
         }))
     }
 
-    pub fn desc() -> FunctionDescription {
-        FunctionDescription::creator(Box::new(Self::try_create)).features(
+    pub fn desc() -> TypedFunctionDescription {
+        TypedFunctionDescription::creator(Box::new(Self::try_create)).features(
             FunctionFeatures::default()
                 .deterministic()
                 .monotonicity()
@@ -60,8 +61,7 @@ impl Function for CeilFunction {
         &*self.display_name
     }
 
-    fn return_type(&self, args: &[&DataTypePtr]) -> Result<DataTypePtr> {
-        assert_numeric(args[0])?;
+    fn return_type(&self, _args: &[&DataTypePtr]) -> Result<DataTypePtr> {
         Ok(Float64Type::arc())
     }
 

diff --git a/common/functions/src/scalars/maths/exp.rs b/common/functions/src/scalars/maths/exp.rs
@@ -20,7 +20,7 @@ use common_exception::Result;
 use num::cast::AsPrimitive;
 
 use crate::scalars::function_common::assert_numeric;
-use crate::scalars::function_factory::FunctionDescription;
+use crate::scalars::function_factory::TypedFunctionDescription;
 use crate::scalars::scalar_unary_op;
 use crate::scalars::EvalContext;
 use crate::scalars::Function;
@@ -32,14 +32,15 @@ pub struct ExpFunction {
 }
 
 impl ExpFunction {
-    pub fn try_create(_display_name: &str) -> Result<Box<dyn Function>> {
+    pub fn try_create(_display_name: &str, args: &[&DataTypePtr]) -> Result<Box<dyn Function>> {
+        assert_numeric(args[0])?;
         Ok(Box::new(ExpFunction {
             _display_name: _display_name.to_string(),
         }))
     }
 
-    pub fn desc() -> FunctionDescription {
-        FunctionDescription::creator(Box::new(Self::try_create))
+    pub fn desc() -> TypedFunctionDescription {
+        TypedFunctionDescription::creator(Box::new(Self::try_create))
             .features(FunctionFeatures::default().deterministic().num_arguments(1))
     }
 }
@@ -54,8 +55,7 @@ impl Function for ExpFunction {
         &*self._display_name
     }
 
-    fn return_type(&self, args: &[&DataTypePtr]) -> Result<DataTypePtr> {
-        assert_numeric(args[0])?;
+    fn return_type(&self, _args: &[&DataTypePtr]) -> Result<DataTypePtr> {
         Ok(Float64Type::arc())
     }
 

diff --git a/common/functions/src/scalars/maths/floor.rs b/common/functions/src/scalars/maths/floor.rs
@@ -21,7 +21,7 @@ use common_exception::Result;
 use num::cast::AsPrimitive;
 
 use crate::scalars::function_common::assert_numeric;
-use crate::scalars::function_factory::FunctionDescription;
+use crate::scalars::function_factory::TypedFunctionDescription;
 use crate::scalars::scalar_unary_op;
 use crate::scalars::EvalContext;
 use crate::scalars::Function;
@@ -34,14 +34,15 @@ pub struct FloorFunction {
 }
 
 impl FloorFunction {
-    pub fn try_create(display_name: &str) -> Result<Box<dyn Function>> {
+    pub fn try_create(display_name: &str, args: &[&DataTypePtr]) -> Result<Box<dyn Function>> {
+        assert_numeric(args[0])?;
         Ok(Box::new(FloorFunction {
             display_name: display_name.to_string(),
         }))
     }
 
-    pub fn desc() -> FunctionDescription {
-        FunctionDescription::creator(Box::new(Self::try_create)).features(
+    pub fn desc() -> TypedFunctionDescription {
+        TypedFunctionDescription::creator(Box::new(Self::try_create)).features(
             FunctionFeatures::default()
                 .deterministic()
                 .monotonicity()
@@ -60,8 +61,7 @@ impl Function for FloorFunction {
         &*self.display_name
     }
 
-    fn return_type(&self, args: &[&DataTypePtr]) -> Result<DataTypePtr> {
-        assert_numeric(args[0])?;
+    fn return_type(&self, _args: &[&DataTypePtr]) -> Result<DataTypePtr> {
         Ok(Float64Type::arc())
     }
 

diff --git a/common/functions/src/scalars/maths/log.rs b/common/functions/src/scalars/maths/log.rs
@@ -27,8 +27,8 @@ use crate::scalars::scalar_binary_op;
 use crate::scalars::scalar_unary_op;
 use crate::scalars::EvalContext;
 use crate::scalars::Function;
-use crate::scalars::FunctionDescription;
 use crate::scalars::FunctionFeatures;
+use crate::scalars::TypedFunctionDescription;
 
 /// Const f64 is now allowed.
 /// feature(adt_const_params) is not stable & complete
@@ -69,15 +69,18 @@ pub struct GenericLogFunction<T> {
 }
 
 impl<T: Base> GenericLogFunction<T> {
-    pub fn try_create(display_name: &str) -> Result<Box<dyn Function>> {
+    pub fn try_create(display_name: &str, args: &[&DataTypePtr]) -> Result<Box<dyn Function>> {
+        for arg in args {
+            assert_numeric(*arg)?;
+        }
         Ok(Box::new(Self {
             display_name: display_name.to_string(),
             t: PhantomData,
         }))
     }
 
-    pub fn desc() -> FunctionDescription {
-        FunctionDescription::creator(Box::new(Self::try_create)).features(
+    pub fn desc() -> TypedFunctionDescription {
+        TypedFunctionDescription::creator(Box::new(Self::try_create)).features(
             FunctionFeatures::default()
                 .deterministic()
                 .variadic_arguments(1, 2),
@@ -103,11 +106,8 @@ impl<T: Base> Function for GenericLogFunction<T> {
         &*self.display_name
     }
 
-    fn return_type(&self, args: &[&DataTypePtr]) -> Result<DataTypePtr> {
-        for arg in args {
-            assert_numeric(*arg)?;
-        }
-        Ok(f64::to_data_type())
+    fn return_type(&self, _args: &[&DataTypePtr]) -> Result<DataTypePtr> {
+        Ok(Float64Type::arc())
     }
 
     fn eval(&self, columns: &ColumnsWithField, _input_rows: usize) -> Result<ColumnRef> {

diff --git a/common/functions/src/scalars/maths/math.rs b/common/functions/src/scalars/maths/math.rs
@@ -48,36 +48,36 @@ pub struct MathsFunction;
 
 impl MathsFunction {
     pub fn register(factory: &mut FunctionFactory) {
-        factory.register("abs", AbsFunction::desc());
-        factory.register("sign", SignFunction::desc());
-        factory.register("pi", PiFunction::desc());
+        factory.register_typed("abs", AbsFunction::desc());
+        factory.register_typed("sign", SignFunction::desc());
+        factory.register_typed("pi", PiFunction::desc());
         factory.register_typed("crc32", CRC32Function::desc());
-        factory.register("exp", ExpFunction::desc());
-        factory.register("sqrt", SqrtFunction::desc());
-        factory.register("ceil", CeilFunction::desc());
-        factory.register("ceiling", CeilFunction::desc());
-        factory.register("floor", FloorFunction::desc());
+        factory.register_typed("exp", ExpFunction::desc());
+        factory.register_typed("sqrt", SqrtFunction::desc());
+        factory.register_typed("ceil", CeilFunction::desc());
+        factory.register_typed("ceiling", CeilFunction::desc());
+        factory.register_typed("floor", FloorFunction::desc());
 
-        factory.register("log", LogFunction::desc());
-        factory.register("log10", Log10Function::desc());
-        factory.register("log2", Log2Function::desc());
-        factory.register("ln", LnFunction::desc());
-        factory.register("pow", PowFunction::desc());
-        factory.register("power", PowFunction::desc());
-        factory.register("rand", RandomFunction::desc());
-        factory.register("round", RoundNumberFunction::desc());
-        factory.register("truncate", TruncNumberFunction::desc());
+        factory.register_typed("log", LogFunction::desc());
+        factory.register_typed("log10", Log10Function::desc());
+        factory.register_typed("log2", Log2Function::desc());
+        factory.register_typed("ln", LnFunction::desc());
+        factory.register_typed("pow", PowFunction::desc());
+        factory.register_typed("power", PowFunction::desc());
+        factory.register_typed("rand", RandomFunction::desc());
+        factory.register_typed("round", RoundNumberFunction::desc());
+        factory.register_typed("truncate", TruncNumberFunction::desc());
 
-        factory.register("sin", TrigonometricSinFunction::desc());
-        factory.register("cos", TrigonometricCosFunction::desc());
-        factory.register("tan", TrigonometricTanFunction::desc());
-        factory.register("cot", TrigonometricCotFunction::desc());
-        factory.register("asin", TrigonometricAsinFunction::desc());
-        factory.register("acos", TrigonometricAcosFunction::desc());
-        factory.register("atan", TrigonometricAtanFunction::desc());
-        factory.register("atan2", TrigonometricAtan2Function::desc());
+        factory.register_typed("sin", TrigonometricSinFunction::desc());
+        factory.register_typed("cos", TrigonometricCosFunction::desc());
+        factory.register_typed("tan", TrigonometricTanFunction::desc());
+        factory.register_typed("cot", TrigonometricCotFunction::desc());
+        factory.register_typed("asin", TrigonometricAsinFunction::desc());
+        factory.register_typed("acos", TrigonometricAcosFunction::desc());
+        factory.register_typed("atan", TrigonometricAtanFunction::desc());
+        factory.register_typed("atan2", TrigonometricAtan2Function::desc());
 
-        factory.register("degrees", DegressFunction::desc());
-        factory.register("radians", RadiansFunction::desc());
+        factory.register_typed("degrees", DegressFunction::desc());
+        factory.register_typed("radians", RadiansFunction::desc());
     }
 }
diff --git a/common/functions/src/scalars/maths/pi.rs b/common/functions/src/scalars/maths/pi.rs
@@ -18,7 +18,7 @@ use std::fmt;
 use common_datavalues::prelude::*;
 use common_exception::Result;
 
-use crate::scalars::function_factory::FunctionDescription;
+use crate::scalars::function_factory::TypedFunctionDescription;
 use crate::scalars::Function;
 use crate::scalars::FunctionFeatures;
 
@@ -28,14 +28,14 @@ pub struct PiFunction {
 }
 
 impl PiFunction {
-    pub fn try_create(display_name: &str) -> Result<Box<dyn Function>> {
+    pub fn try_create(display_name: &str, _args: &[&DataTypePtr]) -> Result<Box<dyn Function>> {
         Ok(Box::new(PiFunction {
             display_name: display_name.to_string(),
         }))
     }
 
-    pub fn desc() -> FunctionDescription {
-        FunctionDescription::creator(Box::new(Self::try_create))
+    pub fn desc() -> TypedFunctionDescription {
+        TypedFunctionDescription::creator(Box::new(Self::try_create))
             .features(FunctionFeatures::default().deterministic())
     }
 }