From e5549ada08b6464e35d409d96560338d65ebbfef Mon Sep 17 00:00:00 2001 From: baishen Date: Thu, 24 Oct 2024 12:32:38 +0800 Subject: [PATCH] feat: support `object_delete` and `object_pick` function --- src/functions.rs | 64 +++++++++++++++++++++- tests/it/functions.rs | 123 ++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 182 insertions(+), 5 deletions(-) diff --git a/src/functions.rs b/src/functions.rs index f429621..e86c11f 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -2840,7 +2840,7 @@ fn array_overlap_jsonb(value1: &[u8], value2: &[u8]) -> Result { Ok(false) } -/// Insert a new value into a JSONB array value by the specified position. +/// Insert a new value into a JSONB object value by the new key and new value. pub fn object_insert( value: &[u8], new_key: &str, @@ -2924,6 +2924,68 @@ fn object_insert_jsonb( Ok(()) } +/// Delete keys and values from a JSONB object value by keys. +pub fn object_delete(value: &[u8], keys: &BTreeSet<&str>, buf: &mut Vec) -> Result<(), Error> { + if !is_jsonb(value) { + let value = parse_value(value)?; + let mut val_buf = Vec::new(); + value.write_to_vec(&mut val_buf); + return object_delete_jsonb(&val_buf, keys, buf); + } + object_delete_jsonb(value, keys, buf) +} + +fn object_delete_jsonb( + value: &[u8], + keys: &BTreeSet<&str>, + buf: &mut Vec, +) -> Result<(), Error> { + let header = read_u32(value, 0)?; + if header & CONTAINER_HEADER_TYPE_MASK != OBJECT_CONTAINER_TAG { + return Err(Error::InvalidObject); + } + + let mut builder = ObjectBuilder::new(); + for (key, jentry, item) in iterate_object_entries(value, header) { + if keys.contains(key) { + continue; + } + builder.push_raw(key, jentry, item); + } + builder.build_into(buf); + + Ok(()) +} + +/// Pick keys and values from a JSONB object value by keys. +pub fn object_pick(value: &[u8], keys: &BTreeSet<&str>, buf: &mut Vec) -> Result<(), Error> { + if !is_jsonb(value) { + let value = parse_value(value)?; + let mut val_buf = Vec::new(); + value.write_to_vec(&mut val_buf); + return object_pick_jsonb(&val_buf, keys, buf); + } + object_pick_jsonb(value, keys, buf) +} + +fn object_pick_jsonb(value: &[u8], keys: &BTreeSet<&str>, buf: &mut Vec) -> Result<(), Error> { + let header = read_u32(value, 0)?; + if header & CONTAINER_HEADER_TYPE_MASK != OBJECT_CONTAINER_TAG { + return Err(Error::InvalidObject); + } + + let mut builder = ObjectBuilder::new(); + for (key, jentry, item) in iterate_object_entries(value, header) { + if !keys.contains(key) { + continue; + } + builder.push_raw(key, jentry, item); + } + builder.build_into(buf); + + Ok(()) +} + /// Deletes all object fields that have null values from the given JSON value, recursively. /// Null values that are not object fields are untouched. pub fn strip_nulls(value: &[u8], buf: &mut Vec) -> Result<(), Error> { diff --git a/tests/it/functions.rs b/tests/it/functions.rs index bf089d9..d2119ac 100644 --- a/tests/it/functions.rs +++ b/tests/it/functions.rs @@ -15,16 +15,17 @@ use std::borrow::Cow; use std::cmp::Ordering; use std::collections::BTreeMap; +use std::collections::BTreeSet; use jsonb::{ array_distinct, array_except, array_insert, array_intersection, array_length, array_overlap, array_values, as_bool, as_null, as_number, as_str, build_array, build_object, compare, concat, contains, convert_to_comparable, delete_by_index, delete_by_keypath, delete_by_name, exists_all_keys, exists_any_keys, from_slice, get_by_index, get_by_keypath, get_by_name, - get_by_path, get_by_path_array, is_array, is_object, keypath::parse_key_paths, object_each, - object_insert, object_keys, parse_value, path_exists, path_match, strip_nulls, to_bool, to_f64, - to_i64, to_pretty_string, to_serde_json, to_serde_json_object, to_str, to_string, to_u64, - traverse_check_string, type_of, Error, Number, Object, Value, + get_by_path, get_by_path_array, is_array, is_object, keypath::parse_key_paths, object_delete, + object_each, object_insert, object_keys, object_pick, parse_value, path_exists, path_match, + strip_nulls, to_bool, to_f64, to_i64, to_pretty_string, to_serde_json, to_serde_json_object, + to_str, to_string, to_u64, traverse_check_string, type_of, Error, Number, Object, Value, }; use jsonb::jsonpath::parse_json_path; @@ -1818,6 +1819,120 @@ fn test_object_insert() { } } +#[test] +fn test_object_pick() { + let sources = vec![ + ( + r#"{"b":11,"d":22,"m":[1,2]}"#, + vec!["a", "b", "c"], + Some(r#"{"b":11}"#), + ), + ( + r#"{"b":11,"d":22,"m":[1,2]}"#, + vec!["a", "x", "y"], + Some(r#"{}"#), + ), + ( + r#"{"k1":"v1","k2":{"x":"y"}}"#, + vec!["k1"], + Some(r#"{"k1":"v1"}"#), + ), + (r#"1"#, vec!["a", "b"], None), + ]; + for (val, keys, result) in sources { + let keys = BTreeSet::from_iter(keys); + { + let val = val.as_bytes(); + let mut buf = Vec::new(); + let ret = object_pick(val, &keys, &mut buf); + match result { + Some(result) => { + assert!(ret.is_ok()); + let actual = from_slice(&buf).unwrap(); + let expected = parse_value(result.as_bytes()).unwrap(); + assert_eq!(actual, expected); + } + None => { + assert!(ret.is_err()); + } + } + } + { + let val = parse_value(val.as_bytes()).unwrap().to_vec(); + let mut buf = Vec::new(); + let ret = object_pick(&val, &keys, &mut buf); + match result { + Some(result) => { + assert!(ret.is_ok()); + let actual = from_slice(&buf).unwrap(); + let expected = parse_value(result.as_bytes()).unwrap(); + assert_eq!(actual, expected); + } + None => { + assert!(ret.is_err()); + } + } + } + } +} + +#[test] +fn test_object_delete() { + let sources = vec![ + ( + r#"{"b":11,"d":22,"m":[1,2]}"#, + vec!["a", "b", "c"], + Some(r#"{"d":22,"m":[1,2]}"#), + ), + ( + r#"{"b":11,"d":22,"m":[1,2]}"#, + vec!["a", "x", "y"], + Some(r#"{"b":11,"d":22,"m":[1,2]}"#), + ), + ( + r#"{"k1":"v1","k2":{"x":"y"}}"#, + vec!["k1"], + Some(r#"{"k2":{"x":"y"}}"#), + ), + (r#"1"#, vec!["a", "b"], None), + ]; + for (val, keys, result) in sources { + let keys = BTreeSet::from_iter(keys); + { + let val = val.as_bytes(); + let mut buf = Vec::new(); + let ret = object_delete(val, &keys, &mut buf); + match result { + Some(result) => { + assert!(ret.is_ok()); + let actual = from_slice(&buf).unwrap(); + let expected = parse_value(result.as_bytes()).unwrap(); + assert_eq!(actual, expected); + } + None => { + assert!(ret.is_err()); + } + } + } + { + let val = parse_value(val.as_bytes()).unwrap().to_vec(); + let mut buf = Vec::new(); + let ret = object_delete(&val, &keys, &mut buf); + match result { + Some(result) => { + assert!(ret.is_ok()); + let actual = from_slice(&buf).unwrap(); + let expected = parse_value(result.as_bytes()).unwrap(); + assert_eq!(actual, expected); + } + None => { + assert!(ret.is_err()); + } + } + } + } +} + #[test] fn test_to_serde_json() { let sources = vec![