feat: add get_by_keypath

Author: akoshchiy

Cargo.toml

@@ -30,6 +30,7 @@ byteorder = "1.5.0"
 fast-float = "0.2.0"
 nom = "7.1.3"
 ordered-float = { version = "4.1.1", default-features = false }
+
 rand = { version = "0.8.5", features = ["small_rng"] }
 serde_json = { version = "1.0.107", default-features = false, features = [
   "preserve_order",
@@ -40,6 +41,7 @@ goldenfile = "1.5.2"
 serde_json = "1.0.107"
 json-deserializer = "0.4.4"
 simd-json = {version = "0.11.1", features = ["allow-non-simd"]}
+
 mockalloc = "0.1.2"
 criterion = "0.5.1"
 

src/functions.rs

@@ -16,6 +16,7 @@ use core::convert::TryInto;
 use std::borrow::Cow;
 use std::cmp::Ordering;
 use std::collections::VecDeque;
+use std::str::from_utf8;
 
 use crate::constants::*;
 use crate::error::*;
@@ -232,22 +233,13 @@ pub fn get_by_index(value: &[u8], index: usize) -> Option<Vec<u8>> {
     let header = read_u32(value, 0).unwrap();
     match header & CONTAINER_HEADER_TYPE_MASK {
         ARRAY_CONTAINER_TAG => {
-            let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
-            if index >= length {
-                return None;
-            }
-            let mut jentry_offset = 4;
-            let mut val_offset = 4 * length + 4;
-            for i in 0..length {
+            let offsets = get_offsets_by_index(value, 0, header, index);
+
+            offsets.map(|(jentry_offset, val_offset)| {
                 let encoded = read_u32(value, jentry_offset).unwrap();
                 let jentry = JEntry::decode_jentry(encoded);
                 let val_length = jentry.length as usize;
-                if i < index {
-                    jentry_offset += 4;
-                    val_offset += val_length;
-                    continue;
-                }
-                let val = match jentry.type_code {
+                match jentry.type_code {
                     CONTAINER_TAG => value[val_offset..val_offset + val_length].to_vec(),
                     _ => {
                         let mut buf = Vec::with_capacity(8 + val_length);
@@ -258,10 +250,8 @@ pub fn get_by_index(value: &[u8], index: usize) -> Option<Vec<u8>> {
                         }
                         buf
                     }
-                };
-                return Some(val);
-            }
-            None
+                }
+            })
         }
         _ => None,
     }
@@ -289,40 +279,8 @@ pub fn get_by_name(value: &[u8], name: &str, ignore_case: bool) -> Option<Vec<u8
     let header = read_u32(value, 0).unwrap();
     match header & CONTAINER_HEADER_TYPE_MASK {
         OBJECT_CONTAINER_TAG => {
-            let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
-            let mut jentry_offset = 4;
-            let mut val_offset = 8 * length + 4;
-
-            let mut key_jentries: VecDeque<JEntry> = VecDeque::with_capacity(length);
-            for _ in 0..length {
-                let encoded = read_u32(value, jentry_offset).unwrap();
-                let key_jentry = JEntry::decode_jentry(encoded);
-
-                jentry_offset += 4;
-                val_offset += key_jentry.length as usize;
-                key_jentries.push_back(key_jentry);
-            }
+            let offsets = get_offsets_by_name(value, 0, header, name, ignore_case);
 
-            let mut offsets = None;
-            let mut key_offset = 8 * length + 4;
-            while let Some(key_jentry) = key_jentries.pop_front() {
-                let prev_key_offset = key_offset;
-                key_offset += key_jentry.length as usize;
-                let key =
-                    unsafe { std::str::from_utf8_unchecked(&value[prev_key_offset..key_offset]) };
-                // first match the value with the same name, if not found,
-                // then match the value with the ignoring case name.
-                if name.eq(key) {
-                    offsets = Some((jentry_offset, val_offset));
-                    break;
-                } else if ignore_case && name.eq_ignore_ascii_case(key) && offsets.is_none() {
-                    offsets = Some((jentry_offset, val_offset));
-                }
-                let val_encoded = read_u32(value, jentry_offset).unwrap();
-                let val_jentry = JEntry::decode_jentry(val_encoded);
-                jentry_offset += 4;
-                val_offset += val_jentry.length as usize;
-            }
             if let Some((jentry_offset, val_offset)) = offsets {
                 let encoded = read_u32(value, jentry_offset).unwrap();
                 let jentry = JEntry::decode_jentry(encoded);
@@ -348,6 +306,172 @@ pub fn get_by_name(value: &[u8], name: &str, ignore_case: bool) -> Option<Vec<u8
     }
 }
 
+/// Extracts JSON sub-object at the specified path,
+/// where path elements can be either field keys or array indexes encoded in utf-8 string.
+pub fn get_by_keypath<'a, I: Iterator<Item = &'a [u8]>>(
+    value: &[u8],
+    keypath: I,
+) -> Option<Vec<u8>> {
+    if !is_jsonb(value) {
+        return match parse_value(value) {
+            Ok(val) => {
+                let mut current_val = &val;
+                for key in keypath {
+                    match from_utf8(key) {
+                        Ok(k) => {
+                            let res = match current_val {
+                                Value::Array(arr) => match k.parse::<usize>() {
+                                    Ok(idx) => arr.get(idx),
+                                    Err(_) => None,
+                                },
+                                Value::Object(obj) => obj.get(k),
+                                _ => None,
+                            };
+                            match res {
+                                Some(v) => current_val = v,
+                                None => return None,
+                            };
+                        }
+                        Err(_) => return None,
+                    }
+                }
+                Some(current_val.to_vec())
+            }
+            Err(_) => None,
+        };
+    }
+
+    let mut curr_val_offset = 0;
+    let mut curr_jentry_encoded = 0;
+    let mut curr_jentry: Option<JEntry> = None;
+
+    for key in keypath {
+        match from_utf8(key) {
+            Ok(k) => {
+                if let Some(ref jentry) = curr_jentry {
+                    if jentry.type_code != CONTAINER_TAG {
+                        return None;
+                    }
+                };
+                let header = read_u32(value, curr_val_offset).unwrap();
+                match header & CONTAINER_HEADER_TYPE_MASK {
+                    OBJECT_CONTAINER_TAG => {
+                        match get_offsets_by_name(value, curr_val_offset, header, k, false) {
+                            Some((jentry_offset, value_offset)) => {
+                                curr_jentry_encoded = read_u32(value, jentry_offset).unwrap();
+                                curr_jentry = Some(JEntry::decode_jentry(curr_jentry_encoded));
+                                curr_val_offset = value_offset;
+                            }
+                            None => return None,
+                        };
+                    }
+                    ARRAY_CONTAINER_TAG => match k.parse::<usize>() {
+                        Ok(idx) => {
+                            match get_offsets_by_index(value, curr_val_offset, header, idx) {
+                                Some((jentry_offset, value_offset)) => {
+                                    curr_jentry_encoded = read_u32(value, jentry_offset).unwrap();
+                                    curr_jentry = Some(JEntry::decode_jentry(curr_jentry_encoded));
+                                    curr_val_offset = value_offset;
+                                }
+                                None => return None,
+                            }
+                        }
+                        Err(_) => return None,
+                    },
+                    _ => return None,
+                }
+            }
+            Err(_) => return None,
+        }
+    }
+    curr_jentry.map(|jentry| {
+        let val_length = jentry.length as usize;
+        match jentry.type_code {
+            CONTAINER_TAG => value[curr_val_offset..curr_val_offset + val_length].to_vec(),
+            _ => {
+                let mut buf: Vec<u8> = Vec::with_capacity(val_length + 8);
+                let scalar_header = SCALAR_CONTAINER_TAG;
+                buf.extend_from_slice(&scalar_header.to_be_bytes());
+                buf.extend_from_slice(&curr_jentry_encoded.to_be_bytes());
+                if val_length > 0 {
+                    buf.extend_from_slice(&value[curr_val_offset..curr_val_offset + val_length]);
+                }
+                buf
+            }
+        }
+    })
+}
+
+fn get_offsets_by_name(
+    value: &[u8],
+    offset: usize,
+    header: u32,
+    name: &str,
+    ignore_case: bool,
+) -> Option<(usize, usize)> {
+    let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
+    let mut jentry_offset = offset + 4;
+    let mut val_offset = offset + 8 * length + 4;
+
+    let mut key_jentries: VecDeque<JEntry> = VecDeque::with_capacity(length);
+    for _ in 0..length {
+        let encoded = read_u32(value, jentry_offset).unwrap();
+        let key_jentry = JEntry::decode_jentry(encoded);
+
+        jentry_offset += 4;
+        val_offset += key_jentry.length as usize;
+        key_jentries.push_back(key_jentry);
+    }
+
+    let mut offsets = None;
+    let mut key_offset = offset + 8 * length + 4;
+    while let Some(key_jentry) = key_jentries.pop_front() {
+        let prev_key_offset = key_offset;
+        key_offset += key_jentry.length as usize;
+        let key = unsafe { std::str::from_utf8_unchecked(&value[prev_key_offset..key_offset]) };
+        // first match the value with the same name, if not found,
+        // then match the value with the ignoring case name.
+        if name.eq(key) {
+            offsets = Some((jentry_offset, val_offset));
+            break;
+        } else if ignore_case && name.eq_ignore_ascii_case(key) && offsets.is_none() {
+            offsets = Some((jentry_offset, val_offset));
+        }
+        let val_encoded = read_u32(value, jentry_offset).unwrap();
+        let val_jentry = JEntry::decode_jentry(val_encoded);
+        jentry_offset += 4;
+        val_offset += val_jentry.length as usize;
+    }
+    offsets
+}
+
+fn get_offsets_by_index(
+    value: &[u8],
+    offset: usize,
+    header: u32,
+    index: usize,
+) -> Option<(usize, usize)> {
+    let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
+    if index >= length {
+        return None;
+    }
+    let mut jentry_offset = offset + 4;
+    let mut val_offset = offset + 4 * length + 4;
+
+    for i in 0..length {
+        let encoded = read_u32(value, jentry_offset).unwrap();
+        let jentry = JEntry::decode_jentry(encoded);
+        let val_length = jentry.length as usize;
+        if i < index {
+            jentry_offset += 4;
+            val_offset += val_length;
+            continue;
+        }
+        return Some((jentry_offset, val_offset));
+    }
+    None
+}
+
 /// Get the keys of a `JSONB` object.
 pub fn object_keys(value: &[u8]) -> Option<Vec<u8>> {
     if !is_jsonb(value) {

tests/it/functions.rs

@@ -18,13 +18,14 @@ use std::collections::BTreeMap;
 
 use jsonb::{
     array_length, array_values, as_bool, as_null, as_number, as_str, build_array, build_object,
-    compare, convert_to_comparable, from_slice, get_by_index, get_by_name, get_by_path, is_array,
-    is_object, object_each, object_keys, parse_value, path_exists, strip_nulls, to_bool, to_f64,
-    to_i64, to_pretty_string, to_str, to_string, to_u64, traverse_check_string, type_of, Number,
-    Object, Value,
+    compare, convert_to_comparable, from_slice, get_by_index, get_by_keypath, get_by_name,
+    get_by_path, is_array, is_object, object_each, object_keys, parse_value, path_exists,
+    strip_nulls, to_bool, to_f64, to_i64, to_pretty_string, to_str, to_string, to_u64,
+    traverse_check_string, type_of, Number, Object, Value,
 };
 
 use jsonb::jsonpath::parse_json_path;
+use nom::AsBytes;
 
 #[test]
 fn test_build_array() {
@@ -979,13 +980,6 @@ fn test_type_of() {
 
 #[test]
 fn test_object_each() {
-    fn init_object<'a>(entries: Vec<(&str, Value<'a>)>) -> Value<'a> {
-        let mut map = BTreeMap::new();
-        for (key, val) in entries {
-            map.insert(key.to_string(), val);
-        }
-        Value::Object(map)
-    }
     let sources = vec![
         ("true", None),
         (r#"[1,2,3]"#, None),
@@ -1044,3 +1038,58 @@ fn test_object_each() {
         }
     }
 }
+
+#[test]
+fn test_get_by_keypath() {
+    let sources = vec![
+        ("null", vec!["a", "b"], None),
+        ("true", vec!["a", "b"], None),
+        (r#""sdasd""#, vec!["1"], None),
+        (
+            "[10,20,30]",
+            vec!["1"],
+            Some(Value::Number(Number::UInt64(20))),
+        ),
+        (
+            r#"[10,20,["a","b","c"]]"#,
+            vec!["2", "0"],
+            Some(Value::String(Cow::from("a"))),
+        ),
+        (r#"[10,20,["a","b","c"]]"#, vec!["2", "a"], None),
+        (
+            r#"[10,20,[{"k1":[1,2,3],"k2":{"w":1,"z":2}},"b","c"]]"#,
+            vec!["2", "0", "k2"],
+            Some(init_object(vec![
+                ("w", Value::Number(Number::UInt64(1))),
+                ("z", Value::Number(Number::UInt64(2))),
+            ])),
+        ),
+    ];
+    for (json_str, path_str, expected) in sources {
+        let path = path_str.into_iter().map(|p| p.as_bytes());
+        {
+            let json = parse_value(json_str.as_bytes()).unwrap().to_vec();
+            let result = get_by_keypath(&json, path.clone());
+            match expected.clone() {
+                Some(e) => assert_eq!(e, from_slice(&result.unwrap()).unwrap()),
+                None => assert_eq!(result, None),
+            }
+        }
+        {
+            let json = json_str.as_bytes();
+            let result = get_by_keypath(json, path);
+            match expected {
+                Some(e) => assert_eq!(e, from_slice(&result.unwrap()).unwrap()),
+                None => assert_eq!(result, None),
+            }
+        }
+    }
+}
+
+fn init_object<'a>(entries: Vec<(&str, Value<'a>)>) -> Value<'a> {
+    let mut map = BTreeMap::new();
+    for (key, val) in entries {
+        map.insert(key.to_string(), val);
+    }
+    Value::Object(map)
+}