Skip to main content

datafusion_functions_nested/
lib.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#![doc(
19    html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg",
20    html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg"
21)]
22#![cfg_attr(docsrs, feature(doc_cfg))]
23// Make sure fast / cheap clones on Arc are explicit:
24// https://github.com/apache/datafusion/issues/11143
25#![deny(clippy::clone_on_ref_ptr)]
26#![cfg_attr(test, allow(clippy::needless_pass_by_value))]
27
28//! Nested type Functions for [DataFusion].
29//!
30//! This crate contains a collection of nested type functions implemented using the
31//! extension API.
32//!
33//! [DataFusion]: https://crates.io/crates/datafusion
34//!
35//! You can register the functions in this crate using the [`register_all`] function.
36
37#[macro_use]
38pub mod macros;
39
40#[macro_use]
41pub mod macros_lambda;
42
43pub mod array_add;
44pub mod array_any_match;
45pub mod array_avg;
46pub mod array_compact;
47pub mod array_filter;
48pub mod array_first;
49pub mod array_has;
50pub mod array_normalize;
51pub mod array_product;
52pub mod array_scale;
53pub mod array_subtract;
54pub mod array_sum;
55pub mod array_transform;
56pub mod arrays_zip;
57pub mod cardinality;
58pub mod concat;
59pub mod cosine_distance;
60pub mod dimension;
61pub mod distance;
62pub mod empty;
63pub mod except;
64pub mod expr_ext;
65pub mod extract;
66pub mod flatten;
67pub mod inner_product;
68pub(crate) mod lambda_utils;
69pub mod length;
70pub mod make_array;
71pub mod map;
72pub mod map_entries;
73pub mod map_extract;
74pub mod map_keys;
75pub mod map_values;
76pub mod min_max;
77pub mod planner;
78pub mod position;
79pub mod range;
80pub mod remove;
81pub mod repeat;
82pub mod replace;
83pub mod resize;
84pub mod reverse;
85pub mod set_ops;
86pub mod sort;
87pub mod string;
88pub mod utils;
89
90use datafusion_common::Result;
91use datafusion_execution::FunctionRegistry;
92use datafusion_expr::{HigherOrderUDF, ScalarUDF};
93use log::debug;
94use std::sync::Arc;
95
96/// Fluent-style API for creating `Expr`s
97pub mod expr_fn {
98    pub use super::array_add::array_add;
99    pub use super::array_any_match::array_any_match;
100    pub use super::array_avg::array_avg;
101    pub use super::array_compact::array_compact;
102    pub use super::array_filter::array_filter;
103    pub use super::array_first::array_first;
104    pub use super::array_has::array_has;
105    pub use super::array_has::array_has_all;
106    pub use super::array_has::array_has_any;
107    pub use super::array_normalize::array_normalize;
108    pub use super::array_product::array_product;
109    pub use super::array_scale::array_scale;
110    pub use super::array_subtract::array_subtract;
111    pub use super::array_sum::array_sum;
112    pub use super::array_transform::array_transform;
113    pub use super::arrays_zip::arrays_zip;
114    pub use super::cardinality::cardinality;
115    pub use super::concat::array_append;
116    pub use super::concat::array_concat;
117    pub use super::concat::array_prepend;
118    pub use super::cosine_distance::cosine_distance;
119    pub use super::dimension::array_dims;
120    pub use super::dimension::array_ndims;
121    pub use super::distance::array_distance;
122    pub use super::empty::array_empty;
123    pub use super::except::array_except;
124    pub use super::extract::array_any_value;
125    pub use super::extract::array_element;
126    pub use super::extract::array_pop_back;
127    pub use super::extract::array_pop_front;
128    pub use super::extract::array_slice;
129    pub use super::flatten::flatten;
130    pub use super::inner_product::inner_product;
131    pub use super::length::array_length;
132    pub use super::make_array::make_array;
133    pub use super::map_entries::map_entries;
134    pub use super::map_extract::map_extract;
135    pub use super::map_keys::map_keys;
136    pub use super::map_values::map_values;
137    pub use super::min_max::array_max;
138    pub use super::min_max::array_min;
139    pub use super::position::array_position;
140    pub use super::position::array_positions;
141    pub use super::range::gen_series;
142    pub use super::range::range;
143    pub use super::remove::array_remove;
144    pub use super::remove::array_remove_all;
145    pub use super::remove::array_remove_n;
146    pub use super::repeat::array_repeat;
147    pub use super::replace::array_replace;
148    pub use super::replace::array_replace_all;
149    pub use super::replace::array_replace_n;
150    pub use super::resize::array_resize;
151    pub use super::reverse::array_reverse;
152    pub use super::set_ops::array_distinct;
153    pub use super::set_ops::array_intersect;
154    pub use super::set_ops::array_union;
155    pub use super::sort::array_sort;
156    pub use super::string::array_to_string;
157    pub use super::string::string_to_array;
158}
159
160/// Return all default nested type functions
161pub fn all_default_nested_functions() -> Vec<Arc<ScalarUDF>> {
162    vec![
163        array_compact::array_compact_udf(),
164        string::array_to_string_udf(),
165        string::string_to_array_udf(),
166        range::range_udf(),
167        range::gen_series_udf(),
168        dimension::array_dims_udf(),
169        cardinality::cardinality_udf(),
170        dimension::array_ndims_udf(),
171        concat::array_append_udf(),
172        concat::array_prepend_udf(),
173        concat::array_concat_udf(),
174        except::array_except_udf(),
175        extract::array_element_udf(),
176        extract::array_pop_back_udf(),
177        extract::array_pop_front_udf(),
178        extract::array_slice_udf(),
179        extract::array_any_value_udf(),
180        make_array::make_array_udf(),
181        array_has::array_has_udf(),
182        array_has::array_has_all_udf(),
183        array_has::array_has_any_udf(),
184        empty::array_empty_udf(),
185        length::array_length_udf(),
186        array_normalize::array_normalize_udf(),
187        array_add::array_add_udf(),
188        array_avg::array_avg_udf(),
189        array_product::array_product_udf(),
190        array_scale::array_scale_udf(),
191        array_subtract::array_subtract_udf(),
192        array_sum::array_sum_udf(),
193        cosine_distance::cosine_distance_udf(),
194        inner_product::inner_product_udf(),
195        distance::array_distance_udf(),
196        flatten::flatten_udf(),
197        min_max::array_max_udf(),
198        min_max::array_min_udf(),
199        sort::array_sort_udf(),
200        repeat::array_repeat_udf(),
201        resize::array_resize_udf(),
202        reverse::array_reverse_udf(),
203        set_ops::array_distinct_udf(),
204        set_ops::array_intersect_udf(),
205        set_ops::array_union_udf(),
206        arrays_zip::arrays_zip_udf(),
207        position::array_position_udf(),
208        position::array_positions_udf(),
209        remove::array_remove_udf(),
210        remove::array_remove_all_udf(),
211        remove::array_remove_n_udf(),
212        replace::array_replace_n_udf(),
213        replace::array_replace_all_udf(),
214        replace::array_replace_udf(),
215        map::map_udf(),
216        map_entries::map_entries_udf(),
217        map_extract::map_extract_udf(),
218        map_keys::map_keys_udf(),
219        map_values::map_values_udf(),
220    ]
221}
222
223pub fn all_default_higher_order_functions() -> Vec<Arc<HigherOrderUDF>> {
224    vec![
225        array_any_match::array_any_match_higher_order_function(),
226        array_filter::array_filter_higher_order_function(),
227        array_first::array_first_higher_order_function(),
228        array_transform::array_transform_higher_order_function(),
229    ]
230}
231
232/// Registers all enabled packages with a [`FunctionRegistry`]
233pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> {
234    let functions: Vec<Arc<ScalarUDF>> = all_default_nested_functions();
235    functions.into_iter().try_for_each(|udf| {
236        let existing_udf = registry.register_udf(udf)?;
237        if let Some(existing_udf) = existing_udf {
238            debug!("Overwrite existing UDF: {}", existing_udf.name());
239        }
240        Ok(()) as Result<()>
241    })?;
242
243    let functions: Vec<Arc<HigherOrderUDF>> = all_default_higher_order_functions();
244    functions.into_iter().try_for_each(|function| {
245        let existing_function = registry.register_higher_order_function(function)?;
246        if let Some(existing_function) = existing_function {
247            debug!(
248                "Overwrite existing higher-order function: {}",
249                existing_function.name()
250            );
251        }
252        Ok(()) as Result<()>
253    })?;
254
255    Ok(())
256}
257
258#[cfg(test)]
259mod tests {
260    use crate::{all_default_higher_order_functions, all_default_nested_functions};
261    use datafusion_common::Result;
262    use std::collections::HashSet;
263
264    #[test]
265    fn test_no_duplicate_name() -> Result<()> {
266        let scalars = all_default_nested_functions();
267        let scalars = scalars.iter().map(|s| (s.name(), s.aliases()));
268
269        let lambdas = all_default_higher_order_functions();
270        let lambdas = lambdas.iter().map(|l| (l.name(), l.aliases()));
271
272        let mut names = HashSet::new();
273
274        for (name, aliases) in scalars.chain(lambdas) {
275            assert!(
276                names.insert(name.to_string().to_lowercase()),
277                "duplicate function name: {name}",
278            );
279            for alias in aliases {
280                assert!(
281                    names.insert(alias.to_string().to_lowercase()),
282                    "duplicate function name: {alias}"
283                );
284            }
285        }
286        Ok(())
287    }
288}