Merge pull request #33 from akoshchiy/11270-json_object_each

feat: add object_each

src/functions.rs

@@ -390,6 +390,72 @@ pub fn object_keys(value: &[u8]) -> Option<Vec<u8>> {
     }
 }
 
+/// Convert the values of a `JSONB` object to vector of key-value pairs.
+pub fn object_each(value: &[u8]) -> Option<Vec<(Vec<u8>, Vec<u8>)>> {
+    if !is_jsonb(value) {
+        return match parse_value(value) {
+            Ok(val) => match val {
+                Value::Object(obj) => {
+                    let mut result = Vec::with_capacity(obj.len());
+                    for (k, v) in obj {
+                        result.push((k.as_bytes().to_vec(), v.to_vec()));
+                    }
+                    Some(result)
+                }
+                _ => None,
+            },
+            Err(_) => None,
+        };
+    }
+
+    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 items: Vec<(Vec<u8>, Vec<u8>)> = Vec::with_capacity(length);
+            let mut jentries: VecDeque<(JEntry, u32)> = VecDeque::with_capacity(length * 2);
+            let mut offset = 4;
+
+            for _ in 0..length * 2 {
+                let encoded = read_u32(value, offset).unwrap();
+                offset += 4;
+                jentries.push_back((JEntry::decode_jentry(encoded), encoded));
+            }
+
+            let mut keys: VecDeque<Vec<u8>> = VecDeque::with_capacity(length);
+            for _ in 0..length {
+                let (jentry, _) = jentries.pop_front().unwrap();
+                let key_len = jentry.length as usize;
+                keys.push_back(value[offset..offset + key_len].to_vec());
+                offset += key_len;
+            }
+
+            for _ in 0..length {
+                let (jentry, encoded) = jentries.pop_front().unwrap();
+                let key = keys.pop_front().unwrap();
+                let val_length = jentry.length as usize;
+                let val = match jentry.type_code {
+                    CONTAINER_TAG => value[offset..offset + val_length].to_vec(),
+                    _ => {
+                        let mut buf = Vec::with_capacity(val_length + 8);
+                        buf.extend_from_slice(&SCALAR_CONTAINER_TAG.to_be_bytes());
+                        buf.extend_from_slice(&encoded.to_be_bytes());
+                        if jentry.length > 0 {
+                            buf.extend_from_slice(&value[offset..offset + val_length]);
+                        }
+                        buf
+                    }
+                };
+                offset += val_length;
+                items.push((key, val));
+            }
+            Some(items)
+        }
+        _ => None,
+    }
+}
+
 /// Convert the values of a `JSONB` array to vector.
 pub fn array_values(value: &[u8]) -> Option<Vec<Vec<u8>>> {
     if !is_jsonb(value) {

tests/it/functions.rs

@@ -14,13 +14,14 @@
 
 use std::borrow::Cow;
 use std::cmp::Ordering;
+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_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,
+    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;
@@ -975,3 +976,71 @@ 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),
+        (
+            r#"{"a":1,"b":false}"#,
+            Some(vec![
+                ("a", Value::Number(Number::Int64(1))),
+                ("b", Value::Bool(false)),
+            ]),
+        ),
+        (
+            r#"{"a":[1,2,3],"b":{"k":1}}"#,
+            Some(vec![
+                (
+                    "a",
+                    Value::Array(vec![
+                        Value::Number(Number::Int64(1)),
+                        Value::Number(Number::Int64(2)),
+                        Value::Number(Number::Int64(3)),
+                    ]),
+                ),
+                (
+                    "b",
+                    init_object(vec![("k", Value::Number(Number::Int64(1)))]),
+                ),
+            ]),
+        ),
+    ];
+    for (src, expected) in sources {
+        {
+            let res = object_each(src.as_bytes());
+            match expected.clone() {
+                Some(expected) => {
+                    let arr = res.unwrap();
+                    for (v, e) in arr.iter().zip(expected.iter()) {
+                        assert_eq!(v.0, e.0.as_bytes().to_vec());
+                        assert_eq!(from_slice(&v.1).unwrap(), e.1);
+                    }
+                }
+                None => assert_eq!(res, None),
+            }
+        }
+        {
+            let jsonb = parse_value(src.as_bytes()).unwrap().to_vec();
+            let res = object_each(&jsonb);
+            match expected {
+                Some(expected) => {
+                    let arr = res.unwrap();
+                    for (v, e) in arr.iter().zip(expected.iter()) {
+                        assert_eq!(v.0, e.0.as_bytes().to_vec());
+                        assert_eq!(from_slice(&v.1).unwrap(), e.1);
+                    }
+                }
+                None => assert_eq!(res, None),
+            }
+        }
+    }
+}