Skip to content

Commit 4737ea0

Browse files
vigimitevigimite
committed
feat: implement json_get_array udf v2
1 parent b9dfacd commit 4737ea0

File tree

3 files changed

+202
-1
lines changed

3 files changed

+202
-1
lines changed

‎src/json_get_array.rs‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
use std::any::Any;
2+
use std::sync::Arc;
3+
4+
use datafusion::arrow::array::{ArrayRef, ListBuilder, StringBuilder};
5+
use datafusion::arrow::datatypes::DataType;
6+
use datafusion::common::{Result as DataFusionResult, ScalarValue};
7+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
8+
use jiter::Peek;
9+
10+
use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, InvokeResult, JsonPath};
11+
use crate::common_macros::make_udf_function;
12+
13+
make_udf_function!(
14+
JsonGetArray,
15+
json_get_array,
16+
json_data path,
17+
r#"Get an arrow array from a JSON string by its "path""#
18+
);
19+
20+
#[derive(Debug)]
21+
pub(super) struct JsonGetArray {
22+
signature: Signature,
23+
aliases: [String; 1],
24+
}
25+
26+
impl Default for JsonGetArray {
27+
fn default() -> Self {
28+
Self {
29+
signature: Signature::variadic_any(Volatility::Immutable),
30+
aliases: ["json_get_array".to_string()],
31+
}
32+
}
33+
}
34+
35+
impl ScalarUDFImpl for JsonGetArray {
36+
fn as_any(&self) -> &dyn Any {
37+
self
38+
}
39+
40+
fn name(&self) -> &str {
41+
self.aliases[0].as_str()
42+
}
43+
44+
fn signature(&self) -> &Signature {
45+
&self.signature
46+
}
47+
48+
fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
49+
return_type_check(
50+
arg_types,
51+
self.name(),
52+
DataType::List(Arc::new(datafusion::arrow::datatypes::Field::new(
53+
"item",
54+
DataType::Utf8,
55+
true,
56+
))),
57+
)
58+
}
59+
60+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
61+
invoke::<BuildArrayList>(&args.args, jiter_json_get_array)
62+
}
63+
64+
fn aliases(&self) -> &[String] {
65+
&self.aliases
66+
}
67+
}
68+
69+
#[derive(Debug)]
70+
struct BuildArrayList;
71+
72+
impl InvokeResult for BuildArrayList {
73+
type Item = Vec<String>;
74+
75+
type Builder = ListBuilder<StringBuilder>;
76+
77+
const ACCEPT_DICT_RETURN: bool = true;
78+
79+
fn builder(capacity: usize) -> Self::Builder {
80+
let values_builder = StringBuilder::new();
81+
ListBuilder::with_capacity(values_builder, capacity)
82+
}
83+
84+
fn append_value(builder: &mut Self::Builder, value: Option<Self::Item>) {
85+
builder.append_option(value.map(|v| v.into_iter().map(Some)));
86+
}
87+
88+
fn finish(mut builder: Self::Builder) -> DataFusionResult<ArrayRef> {
89+
Ok(Arc::new(builder.finish()))
90+
}
91+
92+
fn scalar(value: Option<Self::Item>) -> ScalarValue {
93+
let mut builder = ListBuilder::new(StringBuilder::new());
94+
95+
if let Some(array_items) = value {
96+
for item in array_items {
97+
builder.values().append_value(item);
98+
}
99+
100+
builder.append(true);
101+
} else {
102+
builder.append(false);
103+
}
104+
let array = builder.finish();
105+
ScalarValue::List(Arc::new(array))
106+
}
107+
}
108+
109+
fn jiter_json_get_array(opt_json: Option<&str>, path: &[JsonPath]) -> Result<Vec<String>, GetError> {
110+
if let Some((mut jiter, peek)) = jiter_json_find(opt_json, path) {
111+
match peek {
112+
Peek::Array => {
113+
let mut peek_opt = jiter.known_array()?;
114+
let mut array_items: Vec<String> = Vec::new();
115+
116+
while let Some(element_peek) = peek_opt {
117+
// Get the raw JSON slice for each array element
118+
let start = jiter.current_index();
119+
jiter.known_skip(element_peek)?;
120+
let slice = jiter.slice_to_current(start);
121+
let element_str = std::str::from_utf8(slice)?.to_string();
122+
123+
array_items.push(element_str);
124+
peek_opt = jiter.array_step()?;
125+
}
126+
127+
Ok(array_items)
128+
}
129+
_ => get_err!(),
130+
}
131+
} else {
132+
get_err!()
133+
}
134+
}

‎src/lib.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ mod common_union;
1111
mod json_as_text;
1212
mod json_contains;
1313
mod json_get;
14+
mod json_get_array;
1415
mod json_get_bool;
1516
mod json_get_float;
1617
mod json_get_int;
@@ -26,6 +27,7 @@ pub mod functions {
2627
pub use crate::json_as_text::json_as_text;
2728
pub use crate::json_contains::json_contains;
2829
pub use crate::json_get::json_get;
30+
pub use crate::json_get_array::json_get_array;
2931
pub use crate::json_get_bool::json_get_bool;
3032
pub use crate::json_get_float::json_get_float;
3133
pub use crate::json_get_int::json_get_int;
@@ -39,6 +41,7 @@ pub mod udfs {
3941
pub use crate::json_as_text::json_as_text_udf;
4042
pub use crate::json_contains::json_contains_udf;
4143
pub use crate::json_get::json_get_udf;
44+
pub use crate::json_get_array::json_get_array_udf;
4245
pub use crate::json_get_bool::json_get_bool_udf;
4346
pub use crate::json_get_float::json_get_float_udf;
4447
pub use crate::json_get_int::json_get_int_udf;
@@ -64,6 +67,7 @@ pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> {
6467
json_get_float::json_get_float_udf(),
6568
json_get_int::json_get_int_udf(),
6669
json_get_json::json_get_json_udf(),
70+
json_get_array::json_get_array_udf(),
6771
json_as_text::json_as_text_udf(),
6872
json_get_str::json_get_str_udf(),
6973
json_contains::json_contains_udf(),

‎tests/main.rs‎

Lines changed: 64 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,14 +83,77 @@ async fn test_json_get_union() {
8383
}
8484

8585
#[tokio::test]
86-
async fn test_json_get_array() {
86+
async fn test_json_get_array_elem() {
8787
let sql = "select json_get('[1, 2, 3]', 2)";
8888
let batches = run_query(sql).await.unwrap();
8989
let (value_type, value_repr) = display_val(batches).await;
9090
assert!(matches!(value_type, DataType::Union(_, _)));
9191
assert_eq!(value_repr, "{int=3}");
9292
}
9393

94+
#[tokio::test]
95+
async fn test_json_get_array_basic_numbers() {
96+
let sql = "select json_get_array('[1, 2, 3]')";
97+
let batches = run_query(sql).await.unwrap();
98+
let (value_type, value_repr) = display_val(batches).await;
99+
assert!(matches!(value_type, DataType::List(_)));
100+
assert_eq!(value_repr, "[1, 2, 3]");
101+
}
102+
103+
#[tokio::test]
104+
async fn test_json_get_array_mixed_types() {
105+
let sql = r#"select json_get_array('["hello", 42, true, null, 3.14]')"#;
106+
let batches = run_query(sql).await.unwrap();
107+
let (value_type, value_repr) = display_val(batches).await;
108+
assert!(matches!(value_type, DataType::List(_)));
109+
assert_eq!(value_repr, r#"["hello", 42, true, null, 3.14]"#);
110+
}
111+
112+
#[tokio::test]
113+
async fn test_json_get_array_nested_objects() {
114+
let sql = r#"select json_get_array('[{"name": "John"}, {"age": 30}]')"#;
115+
let batches = run_query(sql).await.unwrap();
116+
let (value_type, value_repr) = display_val(batches).await;
117+
assert!(matches!(value_type, DataType::List(_)));
118+
assert_eq!(value_repr, r#"[{"name": "John"}, {"age": 30}]"#);
119+
}
120+
121+
#[tokio::test]
122+
async fn test_json_get_array_nested_arrays() {
123+
let sql = r#"select json_get_array('[[1, 2], [3, 4]]')"#;
124+
let batches = run_query(sql).await.unwrap();
125+
let (value_type, value_repr) = display_val(batches).await;
126+
assert!(matches!(value_type, DataType::List(_)));
127+
assert_eq!(value_repr, "[[1, 2], [3, 4]]");
128+
}
129+
130+
#[tokio::test]
131+
async fn test_json_get_array_empty() {
132+
let sql = "select json_get_array('[]')";
133+
let batches = run_query(sql).await.unwrap();
134+
let (value_type, value_repr) = display_val(batches).await;
135+
assert!(matches!(value_type, DataType::List(_)));
136+
assert_eq!(value_repr, "[]");
137+
}
138+
139+
#[tokio::test]
140+
async fn test_json_get_array_invalid_json() {
141+
let sql = "select json_get_array('')";
142+
let batches = run_query(sql).await.unwrap();
143+
let (value_type, value_repr) = display_val(batches).await;
144+
assert!(matches!(value_type, DataType::List(_)));
145+
assert_eq!(value_repr, "");
146+
}
147+
148+
#[tokio::test]
149+
async fn test_json_get_array_with_path() {
150+
let sql = r#"select json_get_array('{"items": [1, 2, 3]}', 'items')"#;
151+
let batches = run_query(sql).await.unwrap();
152+
let (value_type, value_repr) = display_val(batches).await;
153+
assert!(matches!(value_type, DataType::List(_)));
154+
assert_eq!(value_repr, "[1, 2, 3]");
155+
}
156+
94157
#[tokio::test]
95158
async fn test_json_get_equals() {
96159
let e = run_query(r"select name, json_get(json_data, 'foo')='abc' from test")

0 commit comments

Comments
 (0)