feat: implement json_get_array UDF

Author: vigimite

src/json_get_array.rs

@@ -0,0 +1,130 @@
+use std::any::Any;
+use std::sync::Arc;
+
+use arrow::array::{GenericListArray, ListBuilder, StringBuilder};
+use arrow_schema::{DataType, Field};
+use datafusion_common::arrow::array::ArrayRef;
+use datafusion_common::{Result as DataFusionResult, ScalarValue};
+use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
+use jiter::Peek;
+
+use crate::common::{check_args, get_err, invoke, jiter_json_find, GetError, JsonPath};
+use crate::common_macros::make_udf_function;
+
+struct StrArrayColumn {
+    rows: GenericListArray<i32>,
+}
+
+impl FromIterator<Option<Vec<String>>> for StrArrayColumn {
+    fn from_iter<T: IntoIterator<Item = Option<Vec<String>>>>(iter: T) -> Self {
+        let string_builder = StringBuilder::new();
+        let mut list_builder = ListBuilder::new(string_builder);
+
+        for row in iter {
+            if let Some(row) = row {
+                for elem in row {
+                    list_builder.values().append_value(elem);
+                }
+
+                list_builder.append(true);
+            } else {
+                list_builder.append(false);
+            }
+        }
+
+        Self {
+            rows: list_builder.finish(),
+        }
+    }
+}
+
+make_udf_function!(
+    JsonGetArray,
+    json_get_array,
+    json_data path,
+    r#"Get an arrow array value from a JSON string by its "path""#
+);
+
+#[derive(Debug)]
+pub(super) struct JsonGetArray {
+    signature: Signature,
+    aliases: [String; 1],
+}
+
+impl Default for JsonGetArray {
+    fn default() -> Self {
+        Self {
+            signature: Signature::variadic_any(Volatility::Immutable),
+            aliases: ["json_get_array".to_string()],
+        }
+    }
+}
+
+impl ScalarUDFImpl for JsonGetArray {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        self.aliases[0].as_str()
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
+        check_args(arg_types, self.name()).map(|()| DataType::List(Field::new("item", DataType::Utf8, true).into()))
+    }
+
+    fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
+        invoke::<StrArrayColumn, Vec<String>>(
+            args,
+            jiter_json_get_array,
+            |c| Ok(Arc::new(c.rows) as ArrayRef),
+            |i| {
+                let string_builder = StringBuilder::new();
+                let mut list_builder = ListBuilder::new(string_builder);
+
+                if let Some(row) = i {
+                    for elem in row {
+                        list_builder.values().append_value(elem);
+                    }
+                }
+
+                ScalarValue::List(list_builder.finish().into())
+            },
+        )
+    }
+
+    fn aliases(&self) -> &[String] {
+        &self.aliases
+    }
+}
+
+fn jiter_json_get_array(json_data: Option<&str>, path: &[JsonPath]) -> Result<Vec<String>, GetError> {
+    if let Some((mut jiter, peek)) = jiter_json_find(json_data, path) {
+        match peek {
+            Peek::Array => {
+                let mut peek_opt = jiter.known_array()?;
+                let mut array_values = Vec::new();
+
+                while let Some(peek) = peek_opt {
+                    let start = jiter.current_index();
+                    jiter.known_skip(peek)?;
+                    let object_slice = jiter.slice_to_current(start);
+                    let object_string = std::str::from_utf8(object_slice)?;
+
+                    array_values.push(object_string.to_owned());
+
+                    peek_opt = jiter.array_step()?;
+                }
+
+                Ok(array_values)
+            }
+            _ => get_err!(),
+        }
+    } else {
+        get_err!()
+    }
+}

src/lib.rs

@@ -10,6 +10,7 @@ mod common_union;
 mod json_as_text;
 mod json_contains;
 mod json_get;
+mod json_get_array;
 mod json_get_bool;
 mod json_get_float;
 mod json_get_int;
@@ -22,6 +23,7 @@ pub mod functions {
     pub use crate::json_as_text::json_as_text;
     pub use crate::json_contains::json_contains;
     pub use crate::json_get::json_get;
+    pub use crate::json_get_array::json_get_array;
     pub use crate::json_get_bool::json_get_bool;
     pub use crate::json_get_float::json_get_float;
     pub use crate::json_get_int::json_get_int;
@@ -34,6 +36,7 @@ pub mod udfs {
     pub use crate::json_as_text::json_as_text_udf;
     pub use crate::json_contains::json_contains_udf;
     pub use crate::json_get::json_get_udf;
+    pub use crate::json_get_array::json_get_array_udf;
     pub use crate::json_get_bool::json_get_bool_udf;
     pub use crate::json_get_float::json_get_float_udf;
     pub use crate::json_get_int::json_get_int_udf;
@@ -54,6 +57,7 @@ pub mod udfs {
 pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> {
     let functions: Vec<Arc<ScalarUDF>> = vec![
         json_get::json_get_udf(),
+        json_get_array::json_get_array_udf(),
         json_get_bool::json_get_bool_udf(),
         json_get_float::json_get_float_udf(),
         json_get_int::json_get_int_udf(),

tests/main.rs

@@ -476,6 +476,31 @@ async fn test_json_length_vec() {
     assert_batches_eq!(expected, &batches);
 }
 
+#[tokio::test]
+async fn test_json_get_arrow_array() {
+    let sql = r#"select name, json_get_array(json_data, 'foo') from test"#;
+    let batches = run_query(sql).await.unwrap();
+
+    let expected = [
+        "+------------------+--------------------------------------------+",
+        "| name             | json_get_array(test.json_data,Utf8(\"foo\")) |",
+        "+------------------+--------------------------------------------+",
+        "| object_foo       |                                            |",
+        "| object_foo_array | [1]                                        |",
+        "| object_foo_obj   |                                            |",
+        "| object_foo_null  |                                            |",
+        "| object_bar       |                                            |",
+        "| list_foo         |                                            |",
+        "| invalid_json     |                                            |",
+        "+------------------+--------------------------------------------+",
+    ];
+
+    assert_batches_eq!(expected, &batches);
+
+    let batches = run_query_large(sql).await.unwrap();
+    assert_batches_eq!(expected, &batches);
+}
+
 #[tokio::test]
 async fn test_no_args() {
     let err = run_query(r#"select json_len()"#).await.unwrap_err();
@@ -1131,6 +1156,7 @@ async fn test_long_arrow_cast() {
     assert_batches_eq!(expected, &batches);
 }
 
+#[tokio::test]
 async fn test_arrow_cast_numeric() {
     let sql = r#"select ('{"foo": 420}'->'foo')::numeric = 420"#;
     let batches = run_query(sql).await.unwrap();