datafusion_functions_json/
lib.rs

1use log::debug;
2use std::sync::Arc;
3
4use datafusion::common::Result;
5use datafusion::execution::FunctionRegistry;
6use datafusion::logical_expr::ScalarUDF;
7
8mod common;
9mod common_macros;
10mod common_union;
11mod json_as_text;
12mod json_contains;
13mod json_from_scalar;
14mod json_get;
15mod json_get_array;
16mod json_get_bool;
17mod json_get_float;
18mod json_get_int;
19mod json_get_json;
20mod json_get_str;
21mod json_length;
22mod json_object_keys;
23mod rewrite;
24
25pub use common_union::{JsonUnionEncoder, JsonUnionValue, JSON_UNION_DATA_TYPE};
26
27pub mod functions {
28    pub use crate::json_as_text::json_as_text;
29    pub use crate::json_contains::json_contains;
30    pub use crate::json_from_scalar::json_from_scalar;
31    pub use crate::json_get::json_get;
32    pub use crate::json_get_array::json_get_array;
33    pub use crate::json_get_bool::json_get_bool;
34    pub use crate::json_get_float::json_get_float;
35    pub use crate::json_get_int::json_get_int;
36    pub use crate::json_get_json::json_get_json;
37    pub use crate::json_get_str::json_get_str;
38    pub use crate::json_length::json_length;
39    pub use crate::json_object_keys::json_object_keys;
40}
41
42pub mod udfs {
43    pub use crate::json_as_text::json_as_text_udf;
44    pub use crate::json_contains::json_contains_udf;
45    pub use crate::json_from_scalar::json_from_scalar_udf;
46    pub use crate::json_get::json_get_udf;
47    pub use crate::json_get_array::json_get_array_udf;
48    pub use crate::json_get_bool::json_get_bool_udf;
49    pub use crate::json_get_float::json_get_float_udf;
50    pub use crate::json_get_int::json_get_int_udf;
51    pub use crate::json_get_json::json_get_json_udf;
52    pub use crate::json_get_str::json_get_str_udf;
53    pub use crate::json_length::json_length_udf;
54    pub use crate::json_object_keys::json_object_keys_udf;
55}
56
57/// Register all JSON UDFs, and [`rewrite::JsonFunctionRewriter`] with the provided [`FunctionRegistry`].
58///
59/// # Arguments
60///
61/// * `registry`: `FunctionRegistry` to register the UDFs
62///
63/// # Errors
64///
65/// Returns an error if the UDFs cannot be registered or if the rewriter cannot be registered.
66pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> {
67    let functions: Vec<Arc<ScalarUDF>> = vec![
68        json_get::json_get_udf(),
69        json_get_bool::json_get_bool_udf(),
70        json_get_float::json_get_float_udf(),
71        json_get_int::json_get_int_udf(),
72        json_get_json::json_get_json_udf(),
73        json_get_array::json_get_array_udf(),
74        json_as_text::json_as_text_udf(),
75        json_get_str::json_get_str_udf(),
76        json_contains::json_contains_udf(),
77        json_length::json_length_udf(),
78        json_object_keys::json_object_keys_udf(),
79        json_from_scalar::json_from_scalar_udf(),
80    ];
81    functions.into_iter().try_for_each(|udf| {
82        let existing_udf = registry.register_udf(udf)?;
83        if let Some(existing_udf) = existing_udf {
84            debug!("Overwrite existing UDF: {}", existing_udf.name());
85        }
86        Ok(()) as Result<()>
87    })?;
88    registry.register_function_rewrite(Arc::new(rewrite::JsonFunctionRewriter))?;
89    registry.register_expr_planner(Arc::new(rewrite::JsonExprPlanner))?;
90
91    Ok(())
92}
datafusion_functions_json/lib.rs

datafusion_functions_json/
lib.rs
