Add json_extract_scalar

Author: Fantayeneh Asres Gizaw

src/common.rs

@@ -1,6 +1,9 @@
 use std::str::Utf8Error;
 use std::sync::Arc;
 
+use crate::common_union::{
+    is_json_union, json_from_union_scalar, nested_json_array, nested_json_array_ref, TYPE_ID_NULL,
+};
 use datafusion::arrow::array::{
     downcast_array, AnyDictionaryArray, Array, ArrayAccessor, ArrayRef, AsArray, DictionaryArray, LargeStringArray,
     PrimitiveArray, PrimitiveBuilder, RunArray, StringArray, StringViewArray,
@@ -11,10 +14,8 @@ use datafusion::arrow::datatypes::{ArrowNativeType, DataType, Int64Type, UInt64T
 use datafusion::common::{exec_err, plan_err, Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::ColumnarValue;
 use jiter::{Jiter, JiterError, Peek};
-
-use crate::common_union::{
-    is_json_union, json_from_union_scalar, nested_json_array, nested_json_array_ref, TYPE_ID_NULL,
-};
+use jsonpath_rust::parser::model::{Segment, Selector};
+use jsonpath_rust::parser::parse_json_path;
 
 /// General implementation of `ScalarUDFImpl::return_type`.
 ///
@@ -140,6 +141,22 @@ impl<'s> JsonPathArgs<'s> {
     }
 }
 
+pub(crate) 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<_>>()
+}
+
 pub trait InvokeResult {
     type Item;
     type Builder;
@@ -585,3 +602,21 @@ fn mask_dictionary_keys(keys: &PrimitiveArray<Int64Type>, type_ids: &[i8]) -> Pr
     }
     PrimitiveArray::new(keys.values().clone(), Some(null_mask.into()))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+
+    // 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_extract.rs

@@ -1,13 +1,11 @@
-use std::any::Any;
+use crate::common::{invoke, parse_jsonpath, return_type_check};
+use crate::common_macros::make_udf_function;
+use crate::json_get_json::jiter_json_get_json;
 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;
+use std::any::Any;
 
 make_udf_function!(
     JsonExtract,
@@ -65,55 +63,20 @@ impl ScalarUDFImpl for JsonExtract {
 
         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()),
+            _ => {
+                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)
-        })
+        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_extract_scalar.rs

@@ -0,0 +1,83 @@
+use std::any::Any;
+
+use datafusion::arrow::datatypes::DataType;
+use datafusion::common::{exec_err, Result as DataFusionResult};
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
+use datafusion::scalar::ScalarValue;
+
+use crate::common::parse_jsonpath;
+use crate::common::{invoke, return_type_check};
+use crate::common_macros::make_udf_function;
+use crate::common_union::JsonUnion;
+use crate::json_get::jiter_json_get_union;
+
+make_udf_function!(
+    JsonExtractScalar,
+    json_extract_scalar,
+    json_data path,
+    r#"Get a value from a JSON string by its "path""#
+);
+
+#[derive(Debug)]
+pub(super) struct JsonExtractScalar {
+    signature: Signature,
+    aliases: [String; 1],
+}
+
+impl Default for JsonExtractScalar {
+    fn default() -> Self {
+        Self {
+            signature: Signature::variadic_any(Volatility::Immutable),
+            aliases: ["json_extract_scalar".to_string()],
+        }
+    }
+}
+
+impl ScalarUDFImpl for JsonExtractScalar {
+    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(), JsonUnion::data_type())
+    }
+
+    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::<JsonUnion>(&[json_arg.clone()], |json, _| jiter_json_get_union(json, &path))
+    }
+
+    fn aliases(&self) -> &[String] {
+        &self.aliases
+    }
+}

src/lib.rs

@@ -10,6 +10,8 @@ mod common_macros;
 mod common_union;
 mod json_as_text;
 mod json_contains;
+mod json_extract;
+mod json_extract_scalar;
 mod json_get;
 mod json_get_bool;
 mod json_get_float;
@@ -19,15 +21,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_extract_scalar::json_extract_scalar;
+    pub use crate::json_get::json_get;
     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;
@@ -63,6 +65,7 @@ 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_extract_scalar::json_extract_scalar_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_scalar_test.rs

@@ -0,0 +1,43 @@
+use crate::utils::{display_val, run_query};
+use rstest::{fixture, rstest};
+
+mod utils;
+
+#[fixture]
+fn json_data() -> String {
+    let json = r#"
+        {
+          "store": {
+            "book name": "My Favorite Books",
+            "book": [
+              {"title": "1984", "author": "George Orwell"},
+              {"title": "Pride and Prejudice", "author": "Jane Austen"}
+            ]
+          }
+        }
+    "#;
+    json.to_string()
+}
+
+#[rstest]
+#[case("$.store.book[0].author", "{str=George Orwell}")]
+#[tokio::test]
+async fn test_json_extract_scalar(json_data: String, #[case] path: &str, #[case] expected: &str) {
+    let result = json_extract_scalar(&json_data, path).await;
+    assert_eq!(result, expected.to_string());
+}
+
+#[rstest]
+#[case("[1, 2, 3]", "$[2]", "{int=3}")]
+#[case("[1, 2, 3]", "$[3]", "{null=}")]
+#[tokio::test]
+async fn test_json_extract_scalar_simple(#[case] json: String, #[case] path: &str, #[case] expected: &str) {
+    let result = json_extract_scalar(&json, path).await;
+    assert_eq!(result, expected.to_string());
+}
+
+async fn json_extract_scalar(json: &str, path: &str) -> String {
+    let sql = format!("select json_extract_scalar('{}', '{}')", json, path);
+    let batches = run_query(sql.as_str()).await.unwrap();
+    display_val(batches).await.1
+}

tests/json_extract_test.rs

@@ -1,5 +1,5 @@
-use rstest::{fixture, rstest};
 use crate::utils::{display_val, run_query};
+use rstest::{fixture, rstest};
 
 mod utils;
 
@@ -10,13 +10,12 @@ fn json_data() -> String {
 }
 
 #[rstest]
-#[case("$.a.ab", "[{\"ac\": \"Dune\", \"ca\": \"Frank Herbert\"},{\"ad\": \"Foundation\", \"da\": \"Isaac Asimov\"}]")]
+#[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,
-) {
+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());
 }
@@ -29,10 +28,8 @@ async fn test_invalid_json_path(json_data: String) {
     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
 }
-