common_function/
helper.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
15use api::v1::meta::ResolveStrategy;
16use common_query::error::{
17    InvalidFuncArgsSnafu, InvalidInputTypeSnafu, Result, UnsupportedInputDataTypeSnafu,
18};
19use datafusion_expr::{Signature, TypeSignature, Volatility};
20use datatypes::arrow::datatypes::DataType;
21use datatypes::prelude::ConcreteDataType;
22use datatypes::types::cast::cast;
23use datatypes::value::ValueRef;
24use snafu::{OptionExt, ResultExt};
25
26/// Integer types accepted by geospatial function signatures.
27pub(crate) const INTEGER_TYPES: &[DataType] = &[
28    DataType::Int8,
29    DataType::Int16,
30    DataType::Int32,
31    DataType::Int64,
32    DataType::UInt8,
33    DataType::UInt16,
34    DataType::UInt32,
35    DataType::UInt64,
36];
37
38/// Legacy primitive numeric signature types; Decimal values are coerced to `Float64`.
39pub(crate) const NUMERICS: &[DataType] = &[
40    DataType::Int8,
41    DataType::Int16,
42    DataType::Int32,
43    DataType::Int64,
44    DataType::UInt8,
45    DataType::UInt16,
46    DataType::UInt32,
47    DataType::UInt64,
48    DataType::Float16,
49    DataType::Float32,
50    DataType::Float64,
51];
52
53/// Create a function signature with oneof signatures of interleaving two arguments.
54pub(crate) fn one_of_sigs2(args1: Vec<DataType>, args2: Vec<DataType>) -> Signature {
55    let mut sigs = Vec::with_capacity(args1.len() * args2.len());
56
57    for arg1 in &args1 {
58        for arg2 in &args2 {
59            sigs.push(TypeSignature::Exact(vec![arg1.clone(), arg2.clone()]));
60        }
61    }
62
63    Signature::one_of(sigs, Volatility::Immutable)
64}
65
66/// Cast a [`ValueRef`] to u64, returns `None` if fails
67pub fn cast_u64(value: &ValueRef) -> Result<Option<u64>> {
68    cast(value.clone().into(), &ConcreteDataType::uint64_datatype())
69        .context(InvalidInputTypeSnafu {
70            err_msg: format!(
71                "Failed to cast input into uint64, actual type: {:#?}",
72                value.data_type(),
73            ),
74        })
75        .map(|v| v.as_u64())
76}
77
78/// Cast a [`ValueRef`] to u32, returns `None` if fails
79pub fn cast_u32(value: &ValueRef) -> Result<Option<u32>> {
80    cast(value.clone().into(), &ConcreteDataType::uint32_datatype())
81        .context(InvalidInputTypeSnafu {
82            err_msg: format!(
83                "Failed to cast input into uint32, actual type: {:#?}",
84                value.data_type(),
85            ),
86        })
87        .map(|v| v.as_u64().map(|v| v as u32))
88}
89
90/// Parse a resolve strategy from a string.
91pub fn parse_resolve_strategy(strategy: &str) -> Result<ResolveStrategy> {
92    ResolveStrategy::from_str_name(strategy).context(InvalidFuncArgsSnafu {
93        err_msg: format!("Invalid resolve strategy: {}", strategy),
94    })
95}
96
97/// Default parallelism for reconcile operations.
98pub fn default_parallelism() -> u32 {
99    64
100}
101
102/// Default resolve strategy for reconcile operations.
103pub fn default_resolve_strategy() -> ResolveStrategy {
104    ResolveStrategy::UseLatest
105}
106
107/// Get the string value from the params.
108///
109/// # Errors
110/// Returns an error if the input type is not a string.
111pub fn get_string_from_params<'a>(
112    params: &'a [ValueRef<'a>],
113    index: usize,
114    fn_name: &'a str,
115) -> Result<&'a str> {
116    let ValueRef::String(s) = &params[index] else {
117        return UnsupportedInputDataTypeSnafu {
118            function: fn_name,
119            datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
120        }
121        .fail();
122    };
123    Ok(s)
124}
125
126macro_rules! with_match_timestamp_types {
127    ($data_type:expr, | $_t:tt $T:ident | $body:tt) => {{
128        macro_rules! __with_ty__ {
129            ( $_t $T:ident ) => {
130                $body
131            };
132        }
133
134        use datafusion_common::DataFusionError;
135        use datafusion_common::arrow::datatypes::{
136            TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
137            TimestampSecondType,
138        };
139
140        match $data_type {
141            DataType::Timestamp(TimeUnit::Second, _) => Ok(__with_ty__! { TimestampSecondType }),
142            DataType::Timestamp(TimeUnit::Millisecond, _) => {
143                Ok(__with_ty__! { TimestampMillisecondType })
144            }
145            DataType::Timestamp(TimeUnit::Microsecond, _) => {
146                Ok(__with_ty__! { TimestampMicrosecondType })
147            }
148            DataType::Timestamp(TimeUnit::Nanosecond, _) => {
149                Ok(__with_ty__! { TimestampNanosecondType })
150            }
151            _ => Err(DataFusionError::Execution(format!(
152                "not expected data type: '{}'",
153                $data_type
154            ))),
155        }
156    }};
157}
158
159pub(crate) use with_match_timestamp_types;
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164
165    #[test]
166    fn test_parse_resolve_strategy() {
167        assert_eq!(
168            parse_resolve_strategy("UseLatest").unwrap(),
169            ResolveStrategy::UseLatest
170        );
171    }
172}
common_function/helper.rs

common_function/
helper.rs
