common_macro/
range_fn.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 proc_macro::TokenStream;
16use quote::quote;
17use syn::spanned::Spanned;
18use syn::{
19    Attribute, Ident, ItemFn, Path, Signature, Type, TypeReference, Visibility, parse_macro_input,
20};
21
22use crate::utils::extract_input_types;
23
24macro_rules! ok {
25    ($item:expr) => {
26        match $item {
27            Ok(item) => item,
28            Err(e) => return e.into_compile_error().into(),
29        }
30    };
31}
32
33pub(crate) fn process_range_fn(args: TokenStream, input: TokenStream) -> TokenStream {
34    let mut name: Option<Ident> = None;
35    let mut display_name: Option<Ident> = None;
36    let mut ret: Option<Ident> = None;
37    let mut evaluator: Option<Path> = None;
38
39    let parser = syn::meta::parser(|meta| {
40        if meta.path.is_ident("name") {
41            name = Some(meta.value()?.parse()?);
42            Ok(())
43        } else if meta.path.is_ident("display_name") {
44            display_name = Some(meta.value()?.parse()?);
45            Ok(())
46        } else if meta.path.is_ident("ret") {
47            ret = Some(meta.value()?.parse()?);
48            Ok(())
49        } else if meta.path.is_ident("evaluator") {
50            evaluator = Some(meta.value()?.parse()?);
51            Ok(())
52        } else {
53            Err(meta.error("unsupported property"))
54        }
55    });
56
57    // extract arg map
58    parse_macro_input!(args with parser);
59
60    // decompose the fn block
61    let compute_fn = parse_macro_input!(input as ItemFn);
62    let ItemFn {
63        attrs,
64        vis,
65        sig,
66        block,
67    } = compute_fn;
68
69    // extract fn arg list
70    let Signature {
71        inputs,
72        ident: fn_name,
73        ..
74    } = &sig;
75    let arg_types = ok!(extract_input_types(inputs));
76
77    // with format like Float64Array
78    let array_types = arg_types
79        .iter()
80        .map(|ty| {
81            if let Type::Reference(TypeReference { elem, .. }) = ty {
82                elem.as_ref().clone()
83            } else {
84                ty.clone()
85            }
86        })
87        .collect::<Vec<_>>();
88
89    let mut result = TokenStream::new();
90
91    // build the struct and its impl block
92    // only do this when `display_name` is specified
93    if let Some(display_name) = display_name {
94        let struct_code = build_struct(
95            attrs,
96            vis,
97            name.clone().expect("name required"),
98            display_name,
99            array_types,
100            ret.clone().expect("ret required"),
101        );
102        result.extend(struct_code);
103    }
104
105    let calc_fn_code = build_calc_fn(
106        name.expect("name required"),
107        arg_types,
108        fn_name.clone(),
109        ret.expect("ret required"),
110        evaluator.clone(),
111    );
112    // Preserve this fn, but remove its `pub` modifier. Specialized evaluators
113    // do not call it in production, while tests keep it as the slice oracle.
114    let input_fn_code: TokenStream = if evaluator.is_some() {
115        quote! {
116            #[cfg(test)]
117            #sig { #block }
118        }
119    } else {
120        quote! {
121            #sig { #block }
122        }
123    }
124    .into();
125
126    result.extend(calc_fn_code);
127    result.extend(input_fn_code);
128    result
129}
130
131fn build_struct(
132    attrs: Vec<Attribute>,
133    vis: Visibility,
134    name: Ident,
135    display_name_ident: Ident,
136    array_types: Vec<Type>,
137    return_array_type: Ident,
138) -> TokenStream {
139    let display_name = display_name_ident.to_string();
140    quote! {
141        #(#attrs)*
142        #[derive(Debug)]
143        #vis struct #name {}
144
145        impl #name {
146            pub const fn name() -> &'static str {
147                #display_name
148            }
149
150            pub fn scalar_udf() -> ScalarUDF {
151                datafusion_expr::create_udf(
152                    Self::name(),
153                    Self::input_type(),
154                    Self::return_type(),
155                    Volatility::Volatile,
156                    Arc::new(Self::calc) as _,
157                )
158            }
159
160            fn input_type() -> Vec<DataType> {
161                vec![#( RangeArray::convert_data_type(#array_types::new_null(0).data_type().clone()), )*]
162            }
163
164            fn return_type() -> DataType {
165                #return_array_type::new_null(0).data_type().clone()
166            }
167        }
168    }
169    .into()
170}
171
172fn build_calc_fn(
173    name: Ident,
174    param_types: Vec<Type>,
175    fn_name: Ident,
176    ret_type: Ident,
177    evaluator: Option<Path>,
178) -> TokenStream {
179    if let Some(evaluator) = evaluator {
180        return quote! {
181            impl #name {
182                fn calc(input: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
183                    #evaluator(input, Self::name())
184                }
185            }
186        }
187        .into();
188    }
189
190    let param_names = param_types
191        .iter()
192        .enumerate()
193        .map(|(i, ty)| Ident::new(&format!("param_{}", i), ty.span()))
194        .collect::<Vec<_>>();
195    let unref_param_types = param_types
196        .iter()
197        .map(|ty| {
198            if let Type::Reference(TypeReference { elem, .. }) = ty {
199                elem.as_ref().clone()
200            } else {
201                ty.clone()
202            }
203        })
204        .collect::<Vec<_>>();
205    let num_params = param_types.len();
206    let param_numbers = (0..num_params).collect::<Vec<_>>();
207    let range_array_names = param_names
208        .iter()
209        .map(|name| Ident::new(&format!("{}_range_array", name), name.span()))
210        .collect::<Vec<_>>();
211    let first_range_array_name = range_array_names.first().unwrap().clone();
212    let first_param_name = param_names.first().unwrap().clone();
213
214    quote! {
215        impl #name {
216            fn calc(input: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
217                assert_eq!(input.len(), #num_params);
218
219                #( let #range_array_names = RangeArray::try_new(extract_array(&input[#param_numbers])?.to_data().into())?; )*
220
221                // check arrays len
222                {
223                    let len_first = #first_range_array_name.len();
224                    #(
225                        if len_first != #range_array_names.len() {
226                            return Err(DataFusionError::Execution(format!("RangeArray have different lengths in PromQL function {}: array1={}, array2={}", #name::name(), len_first, #range_array_names.len())));
227                        }
228                    )*
229                }
230
231                let mut result_array = Vec::new();
232                for index in 0..#first_range_array_name.len(){
233                    #( let #param_names = #range_array_names.get(index).unwrap().as_any().downcast_ref::<#unref_param_types>().unwrap().clone(); )*
234
235                    // check element len
236                    {
237                        let len_first = #first_param_name.len();
238                        #(
239                            if len_first != #param_names.len() {
240                                return Err(DataFusionError::Execution(format!("RangeArray's element {} have different lengths in PromQL function {}: array1={}, array2={}", index, #name::name(), len_first, #param_names.len())));
241                            }
242                        )*
243                    }
244
245                    let result = #fn_name(#( &#param_names, )*);
246                    result_array.push(result);
247                }
248
249                let result = ColumnarValue::Array(Arc::new(#ret_type::from_iter(result_array)));
250                Ok(result)
251            }
252        }
253    }
254    .into()
255}
common_macro/range_fn.rs

common_macro/
range_fn.rs
