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 rate;
18
19use std::fmt;
20
21pub use clamp::{ClampFunction, ClampMaxFunction, ClampMinFunction};
22use datafusion::arrow::datatypes::DataType;
23use datafusion::error::DataFusionError;
24use datafusion::logical_expr::ColumnarValue;
25use datafusion_common::internal_err;
26use datafusion_expr::{ScalarFunctionArgs, Signature, Volatility};
27
28use crate::function::Function;
29use crate::function_registry::FunctionRegistry;
30use crate::scalars::math::modulo::ModuloFunction;
31use crate::scalars::math::rate::RateFunction;
32
33pub(crate) struct MathFunction;
34
35impl MathFunction {
36    pub fn register(registry: &FunctionRegistry) {
37        registry.register_scalar(ModuloFunction::default());
38        registry.register_scalar(RateFunction::default());
39        registry.register_scalar(RangeFunction::default());
40        registry.register_scalar(ClampFunction::default());
41        registry.register_scalar(ClampMinFunction::default());
42        registry.register_scalar(ClampMaxFunction::default());
43    }
44}
45
46/// `RangeFunction` will never be used as a normal function,
47/// just for datafusion to generate logical plan for RangeSelect
48#[derive(Clone, Debug)]
49struct RangeFunction {
50    signature: Signature,
51}
52
53impl fmt::Display for RangeFunction {
54    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
55        write!(f, "RANGE_FN")
56    }
57}
58
59impl Default for RangeFunction {
60    fn default() -> Self {
61        Self {
62            signature: Signature::variadic_any(Volatility::Immutable),
63        }
64    }
65}
66
67impl Function for RangeFunction {
68    fn name(&self) -> &str {
69        "range_fn"
70    }
71
72    // The first argument to range_fn is the expression to be evaluated
73    fn return_type(&self, input_types: &[DataType]) -> datafusion_common::Result<DataType> {
74        input_types
75            .first()
76            .cloned()
77            .ok_or(DataFusionError::Internal(
78                "No expr found in range_fn".into(),
79            ))
80    }
81
82    /// `range_fn` will never been used. As long as a legal signature is returned, the specific content of the signature does not matter.
83    /// In fact, the arguments loaded by `range_fn` are very complicated, and it is difficult to use `Signature` to describe
84    fn signature(&self) -> &Signature {
85        &self.signature
86    }
87
88    fn invoke_with_args(&self, _: ScalarFunctionArgs) -> datafusion_common::Result<ColumnarValue> {
89        internal_err!("not expected to invoke 'range_fn' directly")
90    }
91}
common_function/scalars/math.rs

common_function/scalars/
math.rs
