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

datafusion_functions_json/
lib.rs
