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

datafusion_functions_json/
lib.rs
