datafusion_functions_json/
json_get_str.rs

1use std::any::Any;
2
3use datafusion::arrow::array::StringArray;
4use datafusion::arrow::datatypes::DataType;
5use datafusion::common::Result as DataFusionResult;
6use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
7use jiter::Peek;
8
9use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
10use crate::common_macros::make_udf_function;
11
12make_udf_function!(
13    JsonGetStr,
14    json_get_str,
15    json_data path,
16    r#"Get a string value from a JSON string by its "path""#
17);
18
19#[derive(Debug, PartialEq, Eq, Hash)]
20pub(super) struct JsonGetStr {
21    signature: Signature,
22    aliases: [String; 1],
23}
24
25impl Default for JsonGetStr {
26    fn default() -> Self {
27        Self {
28            signature: Signature::variadic_any(Volatility::Immutable),
29            aliases: ["json_get_str".to_string()],
30        }
31    }
32}
33
34impl ScalarUDFImpl for JsonGetStr {
35    fn as_any(&self) -> &dyn Any {
36        self
37    }
38
39    fn name(&self) -> &str {
40        self.aliases[0].as_str()
41    }
42
43    fn signature(&self) -> &Signature {
44        &self.signature
45    }
46
47    fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
48        return_type_check(arg_types, self.name(), DataType::Utf8)
49    }
50
51    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
52        invoke::<StringArray>(&args.args, jiter_json_get_str)
53    }
54
55    fn aliases(&self) -> &[String] {
56        &self.aliases
57    }
58}
59
60fn jiter_json_get_str(json_data: Option<&str>, path: &[JsonPath]) -> Result<String, GetError> {
61    if let Some((mut jiter, peek)) = jiter_json_find(json_data, path) {
62        match peek {
63            Peek::String => Ok(jiter.known_str()?.to_owned()),
64            _ => get_err!(),
65        }
66    } else {
67        get_err!()
68    }
69}