common_function/scalars/
math.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15pub mod clamp;
16mod modulo;
17mod pow;
18mod rate;
19
20use std::fmt;
21use std::sync::Arc;
22
23pub use clamp::{ClampFunction, ClampMaxFunction, ClampMinFunction};
24use common_query::error::{GeneralDataFusionSnafu, Result};
25use common_query::prelude::Signature;
26use datafusion::error::DataFusionError;
27use datafusion::logical_expr::Volatility;
28use datatypes::prelude::ConcreteDataType;
29use datatypes::vectors::VectorRef;
30pub use pow::PowFunction;
31pub use rate::RateFunction;
32use snafu::ResultExt;
33
34use crate::function::{Function, FunctionContext};
35use crate::function_registry::FunctionRegistry;
36use crate::scalars::math::modulo::ModuloFunction;
37
38pub(crate) struct MathFunction;
39
40impl MathFunction {
41    pub fn register(registry: &FunctionRegistry) {
42        registry.register(Arc::new(ModuloFunction));
43        registry.register(Arc::new(PowFunction));
44        registry.register(Arc::new(RateFunction));
45        registry.register(Arc::new(RangeFunction));
46        registry.register(Arc::new(ClampFunction));
47        registry.register(Arc::new(ClampMinFunction));
48        registry.register(Arc::new(ClampMaxFunction));
49    }
50}
51
52/// `RangeFunction` will never be used as a normal function,
53/// just for datafusion to generate logical plan for RangeSelect
54#[derive(Clone, Debug, Default)]
55struct RangeFunction;
56
57impl fmt::Display for RangeFunction {
58    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
59        write!(f, "RANGE_FN")
60    }
61}
62
63impl Function for RangeFunction {
64    fn name(&self) -> &str {
65        "range_fn"
66    }
67
68    // The first argument to range_fn is the expression to be evaluated
69    fn return_type(&self, input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
70        input_types
71            .first()
72            .cloned()
73            .ok_or(DataFusionError::Internal(
74                "No expr found in range_fn".into(),
75            ))
76            .context(GeneralDataFusionSnafu)
77    }
78
79    /// `range_fn` will never been used. As long as a legal signature is returned, the specific content of the signature does not matter.
80    /// In fact, the arguments loaded by `range_fn` are very complicated, and it is difficult to use `Signature` to describe
81    fn signature(&self) -> Signature {
82        Signature::variadic_any(Volatility::Immutable)
83    }
84
85    fn eval(&self, _func_ctx: &FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {
86        Err(DataFusionError::Internal(
87            "range_fn just a empty function used in range select, It should not be eval!".into(),
88        ))
89        .context(GeneralDataFusionSnafu)
90    }
91}