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, PartialEq, Eq, Hash)]
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    fn placement(
67        &self,
68        args: &[datafusion::logical_expr::ExpressionPlacement],
69    ) -> datafusion::logical_expr::ExpressionPlacement {
70        // If the first argument is a column and the remaining arguments are literals (a path)
71        // then we can push this UDF down to the leaf nodes.
72        if args.len() >= 2
73            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
74            && args[1..]
75                .iter()
76                .all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
77        {
78            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
79        } else {
80            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
81        }
82    }
83}
84
85impl InvokeResult for JsonUnion {
86    type Item = JsonUnionField;
87
88    type Builder = JsonUnion;
89
90    const ACCEPT_DICT_RETURN: bool = true;
91
92    fn builder(capacity: usize) -> Self::Builder {
93        JsonUnion::new(capacity)
94    }
95
96    fn append_value(builder: &mut Self::Builder, value: Option<Self::Item>) {
97        if let Some(value) = value {
98            builder.push(value);
99        } else {
100            builder.push_none();
101        }
102    }
103
104    fn finish(builder: Self::Builder) -> DataFusionResult<ArrayRef> {
105        let array: UnionArray = builder.try_into()?;
106        Ok(Arc::new(array) as ArrayRef)
107    }
108
109    fn scalar(value: Option<Self::Item>) -> ScalarValue {
110        JsonUnionField::scalar_value(value)
111    }
112}
113
114fn jiter_json_get_union(opt_json: Option<&str>, path: &[JsonPath]) -> Result<JsonUnionField, GetError> {
115    if let Some((mut jiter, peek)) = jiter_json_find(opt_json, path) {
116        build_union(&mut jiter, peek)
117    } else {
118        get_err!()
119    }
120}
121
122fn build_union(jiter: &mut Jiter, peek: Peek) -> Result<JsonUnionField, GetError> {
123    match peek {
124        Peek::Null => {
125            jiter.known_null()?;
126            Ok(JsonUnionField::JsonNull)
127        }
128        Peek::True | Peek::False => {
129            let value = jiter.known_bool(peek)?;
130            Ok(JsonUnionField::Bool(value))
131        }
132        Peek::String => {
133            let value = jiter.known_str()?;
134            Ok(JsonUnionField::Str(value.to_owned()))
135        }
136        Peek::Array => {
137            let start = jiter.current_index();
138            jiter.known_skip(peek)?;
139            let array_slice = jiter.slice_to_current(start);
140            let array_string = std::str::from_utf8(array_slice)?;
141            Ok(JsonUnionField::Array(array_string.to_owned()))
142        }
143        Peek::Object => {
144            let start = jiter.current_index();
145            jiter.known_skip(peek)?;
146            let object_slice = jiter.slice_to_current(start);
147            let object_string = std::str::from_utf8(object_slice)?;
148            Ok(JsonUnionField::Object(object_string.to_owned()))
149        }
150        _ => match jiter.known_number(peek)? {
151            NumberAny::Int(NumberInt::Int(value)) => Ok(JsonUnionField::Int(value)),
152            NumberAny::Int(NumberInt::BigInt(_)) => todo!("BigInt not supported yet"),
153            NumberAny::Float(value) => Ok(JsonUnionField::Float(value)),
154        },
155    }
156}
datafusion_functions_json/json_get.rs

datafusion_functions_json/
json_get.rs
