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

datafusion_functions_json/
lib.rs
