common_query/
signature.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
15//! Signature module contains foundational types that are used to represent signatures, types,
16//! and return types of functions.
17//! Copied and modified from datafusion.
18pub use datafusion_expr::Volatility;
19use datafusion_expr::{Signature as DfSignature, TypeSignature as DfTypeSignature};
20use datatypes::arrow::datatypes::DataType as ArrowDataType;
21use datatypes::data_type::DataType;
22use datatypes::prelude::ConcreteDataType;
23
24/// A function's type signature, which defines the function's supported argument types.
25#[derive(Debug, Clone, PartialEq)]
26pub enum TypeSignature {
27    /// arbitrary number of arguments of an common type out of a list of valid types
28    // A function such as `concat` is `Variadic(vec![ConcreteDataType::String, ConcreteDataType::String])`
29    Variadic(Vec<ConcreteDataType>),
30    /// One or more arguments with arbitrary types
31    VariadicAny,
32    /// fixed number of arguments of an arbitrary but equal type out of a list of valid types
33    // A function of one argument of f64 is `Uniform(1, vec![ConcreteDataType::Float64])`
34    // A function of one argument of f64 or f32 is `Uniform(1, vec![ConcreteDataType::Float32, ConcreteDataType::Float64])`
35    Uniform(usize, Vec<ConcreteDataType>),
36    /// exact number of arguments of an exact type
37    Exact(Vec<ConcreteDataType>),
38    /// fixed number of arguments of arbitrary types
39    Any(usize),
40    /// One of a list of signatures
41    OneOf(Vec<TypeSignature>),
42    /// Zero argument
43    /// This is the new signature for functions with zero arguments
44    /// TODO(discord9): make all other usize nonzero usize
45    NullAry,
46}
47
48///The Signature of a function defines its supported input types as well as its volatility.
49#[derive(Debug, Clone, PartialEq)]
50pub struct Signature {
51    /// type_signature - The types that the function accepts. See [TypeSignature] for more information.
52    pub type_signature: TypeSignature,
53    /// volatility - The volatility of the function. See [Volatility] for more information.
54    pub volatility: Volatility,
55}
56
57#[inline]
58fn concrete_types_to_arrow_types(ts: Vec<ConcreteDataType>) -> Vec<ArrowDataType> {
59    ts.iter().map(ConcreteDataType::as_arrow_type).collect()
60}
61
62impl Signature {
63    /// new - Creates a new Signature from any type signature and the volatility.
64    pub fn new(type_signature: TypeSignature, volatility: Volatility) -> Self {
65        Signature {
66            type_signature,
67            volatility,
68        }
69    }
70
71    /// variadic - Creates a variadic signature that represents an arbitrary number of arguments all from a type in common_types.
72    pub fn variadic(common_types: Vec<ConcreteDataType>, volatility: Volatility) -> Self {
73        Self {
74            type_signature: TypeSignature::Variadic(common_types),
75            volatility,
76        }
77    }
78
79    /// variadic_any - Creates a variadic signature that represents an arbitrary number of arguments of any type.
80    pub fn variadic_any(volatility: Volatility) -> Self {
81        Self {
82            type_signature: TypeSignature::VariadicAny,
83            volatility,
84        }
85    }
86
87    /// uniform - Creates a function with a fixed number of arguments of the same type, which must be from valid_types.
88    pub fn uniform(
89        arg_count: usize,
90        valid_types: Vec<ConcreteDataType>,
91        volatility: Volatility,
92    ) -> Self {
93        Self {
94            type_signature: TypeSignature::Uniform(arg_count, valid_types),
95            volatility,
96        }
97    }
98    /// exact - Creates a signature which must match the types in exact_types in order.
99    pub fn exact(exact_types: Vec<ConcreteDataType>, volatility: Volatility) -> Self {
100        Signature {
101            type_signature: TypeSignature::Exact(exact_types),
102            volatility,
103        }
104    }
105    /// any - Creates a signature which can a be made of any type but of a specified number
106    pub fn any(arg_count: usize, volatility: Volatility) -> Self {
107        Signature {
108            type_signature: TypeSignature::Any(arg_count),
109            volatility,
110        }
111    }
112    /// one_of Creates a signature which can match any of the [TypeSignature]s which are passed in.
113    pub fn one_of(type_signatures: Vec<TypeSignature>, volatility: Volatility) -> Self {
114        Signature {
115            type_signature: TypeSignature::OneOf(type_signatures),
116            volatility,
117        }
118    }
119
120    pub fn nullary(volatility: Volatility) -> Self {
121        Signature {
122            type_signature: TypeSignature::NullAry,
123            volatility,
124        }
125    }
126}
127
128/// Conversations between datafusion signature and our signature
129impl From<TypeSignature> for DfTypeSignature {
130    fn from(type_signature: TypeSignature) -> DfTypeSignature {
131        match type_signature {
132            TypeSignature::Variadic(types) => {
133                DfTypeSignature::Variadic(concrete_types_to_arrow_types(types))
134            }
135            TypeSignature::Uniform(n, types) => {
136                if n == 0 {
137                    return DfTypeSignature::Nullary;
138                }
139                DfTypeSignature::Uniform(n, concrete_types_to_arrow_types(types))
140            }
141            TypeSignature::Exact(types) => {
142                DfTypeSignature::Exact(concrete_types_to_arrow_types(types))
143            }
144            TypeSignature::Any(n) => {
145                if n == 0 {
146                    return DfTypeSignature::Nullary;
147                }
148                DfTypeSignature::Any(n)
149            }
150            TypeSignature::OneOf(ts) => {
151                DfTypeSignature::OneOf(ts.into_iter().map(Into::into).collect())
152            }
153            TypeSignature::VariadicAny => DfTypeSignature::VariadicAny,
154            TypeSignature::NullAry => DfTypeSignature::Nullary,
155        }
156    }
157}
158
159impl From<&DfTypeSignature> for TypeSignature {
160    fn from(type_signature: &DfTypeSignature) -> TypeSignature {
161        match type_signature {
162            DfTypeSignature::Variadic(types) => TypeSignature::Variadic(
163                types
164                    .iter()
165                    .map(ConcreteDataType::from_arrow_type)
166                    .collect(),
167            ),
168            DfTypeSignature::Uniform(n, types) => {
169                if *n == 0 {
170                    return TypeSignature::NullAry;
171                }
172                TypeSignature::Uniform(
173                    *n,
174                    types
175                        .iter()
176                        .map(ConcreteDataType::from_arrow_type)
177                        .collect(),
178                )
179            }
180            DfTypeSignature::Exact(types) => TypeSignature::Exact(
181                types
182                    .iter()
183                    .map(ConcreteDataType::from_arrow_type)
184                    .collect(),
185            ),
186            DfTypeSignature::Any(n) => {
187                if *n == 0 {
188                    return TypeSignature::NullAry;
189                }
190                TypeSignature::Any(*n)
191            }
192            DfTypeSignature::OneOf(ts) => TypeSignature::OneOf(ts.iter().map(Into::into).collect()),
193            DfTypeSignature::VariadicAny => TypeSignature::VariadicAny,
194            DfTypeSignature::Nullary => TypeSignature::NullAry,
195            // Other type signatures are currently mapped to VariadicAny as a fallback.
196            // These cases are not used in the current UDF implementation.
197            _ => TypeSignature::VariadicAny,
198        }
199    }
200}
201
202impl From<Signature> for DfSignature {
203    fn from(sig: Signature) -> DfSignature {
204        DfSignature::new(sig.type_signature.into(), sig.volatility)
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use datatypes::arrow::datatypes::DataType;
211
212    use super::*;
213
214    #[test]
215    fn test_into_df_signature() {
216        let types = vec![
217            ConcreteDataType::int8_datatype(),
218            ConcreteDataType::float32_datatype(),
219            ConcreteDataType::float64_datatype(),
220        ];
221        let sig = Signature::exact(types.clone(), Volatility::Immutable);
222
223        assert_eq!(Volatility::Immutable, sig.volatility);
224        assert!(matches!(&sig.type_signature, TypeSignature::Exact(x) if x.clone() == types));
225
226        let df_sig = DfSignature::from(sig);
227        assert_eq!(Volatility::Immutable, df_sig.volatility);
228        let types = vec![DataType::Int8, DataType::Float32, DataType::Float64];
229        assert!(matches!(df_sig.type_signature, DfTypeSignature::Exact(x) if x == types));
230    }
231}
common_query/signature.rs

common_query/
signature.rs
