datafusion_functions_extra/
macros.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18/// Copied from, `/datafusion/functions-array/src/macros.rs`.
19macro_rules! make_udaf_expr {
20    ($EXPR_FN:ident, $($arg:ident)*, $DOC:expr, $AGGREGATE_UDF_FN:ident) => {
21        // "fluent expr_fn" style function
22        #[doc = $DOC]
23        pub fn $EXPR_FN(
24            $($arg: datafusion::logical_expr::Expr,)*
25        ) -> datafusion::logical_expr::Expr {
26            datafusion::logical_expr::Expr::AggregateFunction(datafusion::logical_expr::expr::AggregateFunction::new_udf(
27                $AGGREGATE_UDF_FN(),
28                vec![$($arg),*],
29                false,
30                None,
31                None,
32                None,
33            ))
34        }
35    };
36}
37
38macro_rules! make_udaf_expr_and_func {
39    ($UDAF:ty, $EXPR_FN:ident, $($arg:ident)*, $DOC:expr, $AGGREGATE_UDF_FN:ident) => {
40        make_udaf_expr!($EXPR_FN, $($arg)*, $DOC, $AGGREGATE_UDF_FN);
41        create_func!($UDAF, $AGGREGATE_UDF_FN);
42    };
43    ($UDAF:ty, $EXPR_FN:ident, $DOC:expr, $AGGREGATE_UDF_FN:ident) => {
44        // "fluent expr_fn" style function
45        #[doc = $DOC]
46        pub fn $EXPR_FN(
47            args: Vec<datafusion::logical_expr::Expr>,
48        ) -> datafusion::logical_expr::Expr {
49            datafusion::logical_expr::Expr::AggregateFunction(datafusion::logical_expr::expr::AggregateFunction::new_udf(
50                $AGGREGATE_UDF_FN(),
51                args,
52                false,
53                None,
54                None,
55                None,
56            ))
57        }
58
59        create_func!($UDAF, $AGGREGATE_UDF_FN);
60    };
61}
62
63macro_rules! create_func {
64    ($UDAF:ty, $AGGREGATE_UDF_FN:ident) => {
65        create_func!($UDAF, $AGGREGATE_UDF_FN, <$UDAF>::default());
66    };
67    ($UDAF:ty, $AGGREGATE_UDF_FN:ident, $CREATE:expr) => {
68        paste::paste! {
69            /// Singleton instance of [$UDAF], ensures the UDAF is only created once
70            /// named STATIC_$(UDAF). For example `STATIC_FirstValue`
71            #[allow(non_upper_case_globals)]
72            static [< STATIC_ $UDAF >]: std::sync::OnceLock<std::sync::Arc<datafusion::logical_expr::AggregateUDF>> =
73                std::sync::OnceLock::new();
74
75            #[doc = concat!("AggregateFunction that returns a [`AggregateUDF`](datafusion_expr::AggregateUDF) for [`", stringify!($UDAF), "`]")]
76            pub fn $AGGREGATE_UDF_FN() -> std::sync::Arc<datafusion::logical_expr::AggregateUDF> {
77                [< STATIC_ $UDAF >]
78                    .get_or_init(|| {
79                        std::sync::Arc::new(datafusion::logical_expr::AggregateUDF::from($CREATE))
80                    })
81                    .clone()
82            }
83        }
84    }
85}