datafusion_functions_json/
json_get.rs

1use std::any::Any;
2use std::sync::Arc;
3
4use datafusion::arrow::array::ArrayRef;
5use datafusion::arrow::array::UnionArray;
6use datafusion::arrow::datatypes::DataType;
7use datafusion::common::Result as DataFusionResult;
8use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
9use datafusion::scalar::ScalarValue;
10use jiter::{Jiter, NumberAny, NumberInt, Peek};
11
12use crate::common::InvokeResult;
13use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
14use crate::common_macros::make_udf_function;
15use crate::common_union::{JsonUnion, JsonUnionField};
16
17make_udf_function!(
18    JsonGet,
19    json_get,
20    json_data path,
21    r#"Get a value from a JSON string by its "path""#
22);
23
24// build_typed_get!(JsonGet, "json_get", Union, Float64Array, jiter_json_get_float);
25
26#[derive(Debug)]
27pub(super) struct JsonGet {
28    signature: Signature,
29    aliases: [String; 1],
30}
31
32impl Default for JsonGet {
33    fn default() -> Self {
34        Self {
35            signature: Signature::variadic_any(Volatility::Immutable),
36            aliases: ["json_get".to_string()],
37        }
38    }
39}
40
41impl ScalarUDFImpl for JsonGet {
42    fn as_any(&self) -> &dyn Any {
43        self
44    }
45
46    fn name(&self) -> &str {
47        self.aliases[0].as_str()
48    }
49
50    fn signature(&self) -> &Signature {
51        &self.signature
52    }
53
54    fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
55        return_type_check(arg_types, self.name(), JsonUnion::data_type())
56    }
57
58    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
59        invoke::<JsonUnion>(&args.args, jiter_json_get_union)
60    }
61
62    fn aliases(&self) -> &[String] {
63        &self.aliases
64    }
65}
66
67impl InvokeResult for JsonUnion {
68    type Item = JsonUnionField;
69
70    type Builder = JsonUnion;
71
72    const ACCEPT_DICT_RETURN: bool = true;
73
74    fn builder(capacity: usize) -> Self::Builder {
75        JsonUnion::new(capacity)
76    }
77
78    fn append_value(builder: &mut Self::Builder, value: Option<Self::Item>) {
79        if let Some(value) = value {
80            builder.push(value);
81        } else {
82            builder.push_none();
83        }
84    }
85
86    fn finish(builder: Self::Builder) -> DataFusionResult<ArrayRef> {
87        let array: UnionArray = builder.try_into()?;
88        Ok(Arc::new(array) as ArrayRef)
89    }
90
91    fn scalar(value: Option<Self::Item>) -> ScalarValue {
92        JsonUnionField::scalar_value(value)
93    }
94}
95
96fn jiter_json_get_union(opt_json: Option<&str>, path: &[JsonPath]) -> Result<JsonUnionField, GetError> {
97    if let Some((mut jiter, peek)) = jiter_json_find(opt_json, path) {
98        build_union(&mut jiter, peek)
99    } else {
100        get_err!()
101    }
102}
103
104fn build_union(jiter: &mut Jiter, peek: Peek) -> Result<JsonUnionField, GetError> {
105    match peek {
106        Peek::Null => {
107            jiter.known_null()?;
108            Ok(JsonUnionField::JsonNull)
109        }
110        Peek::True | Peek::False => {
111            let value = jiter.known_bool(peek)?;
112            Ok(JsonUnionField::Bool(value))
113        }
114        Peek::String => {
115            let value = jiter.known_str()?;
116            Ok(JsonUnionField::Str(value.to_owned()))
117        }
118        Peek::Array => {
119            let start = jiter.current_index();
120            jiter.known_skip(peek)?;
121            let array_slice = jiter.slice_to_current(start);
122            let array_string = std::str::from_utf8(array_slice)?;
123            Ok(JsonUnionField::Array(array_string.to_owned()))
124        }
125        Peek::Object => {
126            let start = jiter.current_index();
127            jiter.known_skip(peek)?;
128            let object_slice = jiter.slice_to_current(start);
129            let object_string = std::str::from_utf8(object_slice)?;
130            Ok(JsonUnionField::Object(object_string.to_owned()))
131        }
132        _ => match jiter.known_number(peek)? {
133            NumberAny::Int(NumberInt::Int(value)) => Ok(JsonUnionField::Int(value)),
134            NumberAny::Int(NumberInt::BigInt(_)) => todo!("BigInt not supported yet"),
135            NumberAny::Float(value) => Ok(JsonUnionField::Float(value)),
136        },
137    }
138}
datafusion_functions_json/json_get.rs

datafusion_functions_json/
json_get.rs
