datafusion_functions_json/
json_get.rs

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

datafusion_functions_json/
json_get.rs
