Merge pull request #1 from featuremesh/feature/json-extract

json extract

Author: fantayeneh

Cargo.toml

@@ -15,13 +15,15 @@ datafusion = { version = "46", default-features = false }
 jiter = "0.9"
 paste = "1"
 log = "0.4"
+jsonpath-rust = "1.0.0"
 
 [dev-dependencies]
 datafusion = { version = "46", default-features = false, features = ["nested_expressions"] }
 codspeed-criterion-compat = "2.6"
 criterion = "0.5.1"
 clap = "4"
 tokio = { version = "1.43", features = ["full"] }
+rstest = "0.25.0"
 
 [lints.clippy]
 dbg_macro = "deny"

src/common.rs

@@ -68,7 +68,7 @@ fn dict_key_type(d: &DataType) -> Option<DataType> {
     None
 }
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub enum JsonPath<'s> {
     Key(&'s str),
     Index(usize),

src/json_extract.rs

@@ -0,0 +1,119 @@
+use std::any::Any;
+use datafusion::arrow::array::StringArray;
+use datafusion::arrow::datatypes::{DataType, DataType::Utf8};
+use datafusion::common::{exec_err, Result as DataFusionResult, ScalarValue};
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
+use jsonpath_rust::parser::model::{Segment, Selector};
+use jsonpath_rust::parser::parse_json_path;
+use crate::common::{invoke, return_type_check, JsonPath};
+use crate::common_macros::make_udf_function;
+use crate::json_get_json::jiter_json_get_json;
+
+make_udf_function!(
+    JsonExtract,
+    json_extract,
+    json_data path,
+    r#"Get a value from a JSON string by its "path" in JSONPath format"#
+);
+
+#[derive(Debug)]
+pub(super) struct JsonExtract {
+    signature: Signature,
+    aliases: [String; 1],
+}
+
+impl Default for JsonExtract {
+    fn default() -> Self {
+        Self {
+            signature: Signature::exact(
+                vec![Utf8, Utf8], // JSON data and JSONPath as strings
+                Volatility::Immutable,
+            ),
+            aliases: ["json_extract".to_string()],
+        }
+    }
+}
+
+impl ScalarUDFImpl for JsonExtract {
+    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> {
+        return_type_check(arg_types, self.name(), Utf8)
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
+        if args.args.len() != 2 {
+            return exec_err!(
+                "'{}' expects exactly 2 arguments (JSON data, path), got {}",
+                self.name(),
+                args.args.len()
+            );
+        }
+
+        let json_arg = &args.args[0];
+        let path_arg = &args.args[1];
+
+        let path_str = match path_arg {
+            ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => s,
+            _ => return exec_err!("'{}' expects a valid JSONPath string (e.g., '$.key[0]') as second argument", self.name()),
+        };
+
+        let path = parse_jsonpath(path_str);
+
+        invoke::<StringArray>(&[json_arg.clone()], |json, _| {
+            jiter_json_get_json(json, &path)
+        })
+    }
+
+    fn aliases(&self) -> &[String] {
+        &self.aliases
+    }
+}
+
+fn parse_jsonpath(path: &str) -> Vec<JsonPath<'static>> {
+    let segments = parse_json_path(path)
+        .map(|it| it.segments)
+        .unwrap_or(Vec::new());
+
+    segments.into_iter().map(|segment| {
+        match segment {
+            Segment::Selector(s) => match s {
+                Selector::Name(name) => JsonPath::Key(Box::leak(name.into_boxed_str())),
+                Selector::Index(idx) => JsonPath::Index(idx as usize),
+                _ => JsonPath::None,
+            },
+            _ => JsonPath::None,
+        }
+    }).collect::<Vec<_>>()
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+    use super::*;
+
+    // Test cases for parse_jsonpath
+    #[rstest]
+    #[case("$.a.aa", vec![JsonPath::Key("a"), JsonPath::Key("aa")])]
+    #[case("$.a.ab[0].ac", vec![JsonPath::Key("a"), JsonPath::Key("ab"), JsonPath::Index(0), JsonPath::Key("ac")])]
+    #[case("$.a.ab[1].ad", vec![JsonPath::Key("a"), JsonPath::Key("ab"), JsonPath::Index(1), JsonPath::Key("ad")])]
+    #[case(r#"$.a["a b"].ad"#, vec![JsonPath::Key("a"), JsonPath::Key("\"a b\""), JsonPath::Key("ad")])]
+    #[tokio::test]
+    async fn test_parse_jsonpath(
+        #[case] path: &str,
+        #[case] expected: Vec<JsonPath<'static>>,
+    ) {
+        let result = parse_jsonpath(path);
+        assert_eq!(result, expected);
+    }
+}

src/json_get.rs

@@ -93,7 +93,7 @@ impl InvokeResult for JsonUnion {
     }
 }
 
-fn jiter_json_get_union(opt_json: Option<&str>, path: &[JsonPath]) -> Result<JsonUnionField, GetError> {
+pub(crate) fn jiter_json_get_union(opt_json: Option<&str>, path: &[JsonPath]) -> Result<JsonUnionField, GetError> {
     if let Some((mut jiter, peek)) = jiter_json_find(opt_json, path) {
         build_union(&mut jiter, peek)
     } else {

src/json_get_json.rs

@@ -56,7 +56,7 @@ impl ScalarUDFImpl for JsonGetJson {
     }
 }
 
-fn jiter_json_get_json(opt_json: Option<&str>, path: &[JsonPath]) -> Result<String, GetError> {
+pub(crate) fn jiter_json_get_json(opt_json: Option<&str>, path: &[JsonPath]) -> Result<String, GetError> {
     if let Some((mut jiter, peek)) = jiter_json_find(opt_json, path) {
         let start = jiter.current_index();
         jiter.known_skip(peek)?;

src/lib.rs

@@ -19,13 +19,15 @@ mod json_get_str;
 mod json_length;
 mod json_object_keys;
 mod rewrite;
+mod json_extract;
 
 pub use common_union::{JsonUnionEncoder, JsonUnionValue};
 
 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_extract::json_extract;
     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;
@@ -60,6 +62,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_extract::json_extract_udf(),
         json_get_bool::json_get_bool_udf(),
         json_get_float::json_get_float_udf(),
         json_get_int::json_get_int_udf(),

tests/json_extract_test.rs

@@ -0,0 +1,38 @@
+use rstest::{fixture, rstest};
+use crate::utils::{display_val, run_query};
+
+mod utils;
+
+#[fixture]
+fn json_data() -> String {
+    let json = r#"{"a": {"a a": "My Collection","ab": [{"ac": "Dune", "ca": "Frank Herbert"},{"ad": "Foundation", "da": "Isaac Asimov"}]}}"#;
+    json.to_string()
+}
+
+#[rstest]
+#[case("$.a.ab", "[{\"ac\": \"Dune\", \"ca\": \"Frank Herbert\"},{\"ad\": \"Foundation\", \"da\": \"Isaac Asimov\"}]")]
+#[tokio::test]
+async fn test_json_paths(
+    json_data: String,
+    #[case] path: &str,
+    #[case] expected: &str,
+) {
+    let result = json_extract(&json_data, path).await;
+    assert_eq!(result, expected.to_string());
+}
+
+#[rstest]
+#[tokio::test]
+#[ignore]
+async fn test_invalid_json_path(json_data: String) {
+    let result = json_extract(&json_data, "store.invalid.path").await;
+    assert_eq!(result, "".to_string());
+}
+
+
+async fn json_extract(json: &str, path: &str) -> String {
+    let sql = format!("select json_extract('{}', '{}')", json, path);
+    let batches = run_query(sql.as_str()).await.unwrap();
+    display_val(batches).await.1
+}
+

tests/main.rs

@@ -82,6 +82,28 @@ async fn test_json_get_union() {
     assert_batches_eq!(expected, &batches);
 }
 
+#[tokio::test]
+async fn test_json_extract_union() {
+    let batches = run_query("select name, json_extract(json_data, '$.foo') as foo from test")
+        .await
+        .unwrap();
+
+    let expected = [
+        "+------------------+-------------+",
+        "| name             | foo         |",
+        "+------------------+-------------+",
+        "| object_foo       | {str=abc}   |",
+        "| object_foo_array | {array=[1]} |",
+        "| object_foo_obj   | {object={}} |",
+        "| object_foo_null  | {null=}     |",
+        "| object_bar       | {null=}     |",
+        "| list_foo         | {null=}     |",
+        "| invalid_json     | {null=}     |",
+        "+------------------+-------------+",
+    ];
+    assert_batches_eq!(expected, &batches);
+}
+
 #[tokio::test]
 async fn test_json_get_array() {
     let sql = "select json_get('[1, 2, 3]', 2)";