datafusion_functions_json/
json_get_bool.rs

1use datafusion::arrow::array::BooleanArray;
2use datafusion::arrow::datatypes::DataType;
3use datafusion::common::Result as DataFusionResult;
4use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
5use jiter::Peek;
6
7use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
8use crate::common_macros::make_udf_function;
9
10make_udf_function!(
11    JsonGetBool,
12    json_get_bool,
13    json_data path,
14    r#"Get an boolean value from a JSON string by its "path""#
15);
16
17#[derive(Debug, PartialEq, Eq, Hash)]
18pub(super) struct JsonGetBool {
19    signature: Signature,
20    aliases: [String; 1],
21}
22
23impl Default for JsonGetBool {
24    fn default() -> Self {
25        Self {
26            signature: Signature::variadic_any(Volatility::Immutable),
27            aliases: ["json_get_bool".to_string()],
28        }
29    }
30}
31
32impl ScalarUDFImpl for JsonGetBool {
33    fn name(&self) -> &str {
34        self.aliases[0].as_str()
35    }
36
37    fn signature(&self) -> &Signature {
38        &self.signature
39    }
40
41    fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
42        return_type_check(arg_types, self.name(), DataType::Boolean).map(|_| DataType::Boolean)
43    }
44
45    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
46        invoke::<BooleanArray>(&args.args, jiter_json_get_bool)
47    }
48
49    fn aliases(&self) -> &[String] {
50        &self.aliases
51    }
52
53    fn placement(
54        &self,
55        args: &[datafusion::logical_expr::ExpressionPlacement],
56    ) -> datafusion::logical_expr::ExpressionPlacement {
57        // If the first argument is a column and the remaining arguments are literals (a path)
58        // then we can push this UDF down to the leaf nodes.
59        if args.len() >= 2
60            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
61            && args[1..]
62                .iter()
63                .all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
64        {
65            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
66        } else {
67            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
68        }
69    }
70}
71
72fn jiter_json_get_bool(json_data: Option<&str>, path: &[JsonPath]) -> Result<bool, GetError> {
73    if let Some((mut jiter, peek)) = jiter_json_find(json_data, path) {
74        match peek {
75            Peek::True | Peek::False => Ok(jiter.known_bool(peek)?),
76            Peek::String => {
77                let s = jiter.known_str()?;
78                s.parse::<bool>().map_err(|_| GetError)
79            }
80            _ => get_err!(),
81        }
82    } else {
83        get_err!()
84    }
85}
datafusion_functions_json/json_get_bool.rs

datafusion_functions_json/
json_get_bool.rs
