From 162392863732525b02c4fcdb2541e24ca4081d06 Mon Sep 17 00:00:00 2001 From: Andrey Koshchiy Date: Sat, 30 Sep 2023 18:12:43 +0400 Subject: [PATCH] feat: add object_each --- src/functions.rs | 66 +++++++++++++++++++++++++++++++++++++ tests/it/functions.rs | 75 +++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 138 insertions(+), 3 deletions(-) diff --git a/src/functions.rs b/src/functions.rs index 64e7cd5..254f0e1 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -390,6 +390,72 @@ pub fn object_keys(value: &[u8]) -> Option> { } } +/// Convert the values of a `JSONB` object to vector of key-value pairs. +pub fn object_each(value: &[u8]) -> Option, Vec)>> { + 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, Vec)> = 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> = 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>> { if !is_jsonb(value) { diff --git a/tests/it/functions.rs b/tests/it/functions.rs index 297cf7b..cd3c5d9 100644 --- a/tests/it/functions.rs +++ b/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), + } + } + } +}