diff --git a/proto/expr.proto b/proto/expr.proto index 998ed63a4b08..9abb1d74f495 100644 --- a/proto/expr.proto +++ b/proto/expr.proto @@ -280,6 +280,7 @@ message ExprNode { JSONB_PATH_QUERY_FIRST = 623; JSONB_POPULATE_RECORD = 629; JSONB_TO_RECORD = 630; + JSONB_SET = 631; // Non-pure functions below (> 1000) // ------------------------ diff --git a/src/expr/impl/src/scalar/jsonb_set.rs b/src/expr/impl/src/scalar/jsonb_set.rs new file mode 100644 index 000000000000..e3efefb05d41 --- /dev/null +++ b/src/expr/impl/src/scalar/jsonb_set.rs @@ -0,0 +1,186 @@ +// Copyright 2024 RisingWave Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use jsonbb::ValueRef; +use risingwave_common::types::{JsonbRef, JsonbVal, ListRef}; +use risingwave_expr::{function, ExprError, Result}; + +/// Returns `target` with the item designated by `path` replaced by `new_value`, or with `new_value` +/// added if `create_if_missing` is true (which is the default) and the item designated by path does +/// not exist. All earlier steps in the path must exist, or the `target` is returned unchanged. As +/// with the path oriented operators, negative integers that appear in the path count from the end +/// of JSON arrays. +/// +/// If the last path step is an array index that is out of range, and `create_if_missing` is true, +/// the new value is added at the beginning of the array if the index is negative, or at the end of +/// the array if it is positive. +/// +/// # Examples +/// +/// ```slt +/// query T +/// SELECT jsonb_set('[{"f1":1,"f2":null},2,null,3]', '{0,f1}', '[2,3,4]', false); +/// ---- +/// [{"f1": [2, 3, 4], "f2": null}, 2, null, 3] +/// +/// query T +/// SELECT jsonb_set('[{"f1":1,"f2":null},2]', '{0,f3}', '[2,3,4]'); +/// ---- +/// [{"f1": 1, "f2": null, "f3": [2, 3, 4]}, 2] +/// ``` +#[function("jsonb_set(jsonb, varchar[], jsonb, boolean) -> jsonb")] +fn jsonb_set4( + target: JsonbRef<'_>, + path: ListRef<'_>, + new_value: JsonbRef<'_>, + create_if_missing: bool, +) -> Result { + if target.is_scalar() { + return Err(ExprError::InvalidParam { + name: "jsonb", + reason: "cannot set path in scalar".into(), + }); + } + let target: ValueRef<'_> = target.into(); + let new_value: ValueRef<'_> = new_value.into(); + let mut builder = jsonbb::Builder::>::with_capacity(target.capacity()); + jsonbb_set_path(target, path, 0, new_value, create_if_missing, &mut builder)?; + Ok(JsonbVal::from(builder.finish())) +} + +#[function("jsonb_set(jsonb, varchar[], jsonb) -> jsonb")] +fn jsonb_set3( + target: JsonbRef<'_>, + path: ListRef<'_>, + new_value: JsonbRef<'_>, +) -> Result { + jsonb_set4(target, path, new_value, true) +} + +/// Recursively set `path[i..]` in `target` to `new_value` and write the result to `builder`. +/// +/// Panics if `i` is out of bounds. +fn jsonbb_set_path( + target: ValueRef<'_>, + path: ListRef<'_>, + i: usize, + new_value: ValueRef<'_>, + create_if_missing: bool, + builder: &mut jsonbb::Builder, +) -> Result<()> { + let last_step = i == path.len() - 1; + match target { + ValueRef::Object(obj) => { + let key = path + .get(i) + .unwrap() + .ok_or_else(|| ExprError::InvalidParam { + name: "path", + reason: format!("path element at position {} is null", i + 1).into(), + })? + .into_utf8(); + builder.begin_object(); + for (k, v) in obj.iter() { + builder.add_string(k); + if k != key { + builder.add_value(v); + } else if last_step { + builder.add_value(new_value); + } else { + // recursively set path[i+1..] in v + jsonbb_set_path(v, path, i + 1, new_value, create_if_missing, builder)?; + } + } + if create_if_missing && last_step && !obj.contains_key(key) { + builder.add_string(key); + builder.add_value(new_value); + } + builder.end_object(); + Ok(()) + } + ValueRef::Array(array) => { + let key = path + .get(i) + .unwrap() + .ok_or_else(|| ExprError::InvalidParam { + name: "path", + reason: format!("path element at position {} is null", i + 1).into(), + })? + .into_utf8(); + let idx = key.parse::().map_err(|_| ExprError::InvalidParam { + name: "path", + reason: format!( + "path element at position {} is not an integer: \"{}\"", + i + 1, + key + ) + .into(), + })?; + let Some(idx) = normalize_array_index(array.len(), idx) else { + // out of bounds index + if create_if_missing { + builder.begin_array(); + // the new value is added at the beginning of the array if the index is negative + if idx < 0 { + builder.add_value(new_value); + } + for v in array.iter() { + builder.add_value(v); + } + // or at the end of the array if it is positive. + if idx >= 0 { + builder.add_value(new_value); + } + builder.end_array(); + } else { + builder.add_value(target); + } + return Ok(()); + }; + builder.begin_array(); + for (j, v) in array.iter().enumerate() { + if j != idx { + builder.add_value(v); + continue; + } + if last_step { + builder.add_value(new_value); + } else { + // recursively set path[i+1..] in v + jsonbb_set_path(v, path, i + 1, new_value, create_if_missing, builder)?; + } + } + builder.end_array(); + Ok(()) + } + _ => { + builder.add_value(target); + Ok(()) + } + } +} + +/// Normalizes an array index to `0..len`. +/// Negative indices count from the end. i.e. `-len..0 => 0..len`. +/// Returns `None` if index is out of bounds. +fn normalize_array_index(len: usize, index: i32) -> Option { + if index < -(len as i32) || index >= (len as i32) { + return None; + } + if index >= 0 { + Some(index as usize) + } else { + Some((len as i32 + index) as usize) + } +} diff --git a/src/expr/impl/src/scalar/mod.rs b/src/expr/impl/src/scalar/mod.rs index edbaaf4de01a..c4e7990de133 100644 --- a/src/expr/impl/src/scalar/mod.rs +++ b/src/expr/impl/src/scalar/mod.rs @@ -58,6 +58,7 @@ mod jsonb_info; mod jsonb_object; mod jsonb_path; mod jsonb_record; +mod jsonb_set; mod length; mod lower; mod make_time; diff --git a/src/frontend/src/binder/expr/function.rs b/src/frontend/src/binder/expr/function.rs index e1d8602c2487..74162c35e982 100644 --- a/src/frontend/src/binder/expr/function.rs +++ b/src/frontend/src/binder/expr/function.rs @@ -1160,6 +1160,7 @@ impl Binder { ("jsonb_path_exists", raw_call(ExprType::JsonbPathExists)), ("jsonb_path_query_array", raw_call(ExprType::JsonbPathQueryArray)), ("jsonb_path_query_first", raw_call(ExprType::JsonbPathQueryFirst)), + ("jsonb_set", raw_call(ExprType::JsonbSet)), // Functions that return a constant value ("pi", pi()), // greatest and least diff --git a/src/frontend/src/expr/pure.rs b/src/frontend/src/expr/pure.rs index 0199c86caeb3..d03b7507fdfb 100644 --- a/src/frontend/src/expr/pure.rs +++ b/src/frontend/src/expr/pure.rs @@ -209,6 +209,7 @@ impl ExprVisitor for ImpureAnalyzer { | Type::JsonbPathMatch | Type::JsonbPathQueryArray | Type::JsonbPathQueryFirst + | Type::JsonbSet | Type::IsJson | Type::ToJsonb | Type::Sind diff --git a/src/frontend/src/optimizer/plan_expr_visitor/strong.rs b/src/frontend/src/optimizer/plan_expr_visitor/strong.rs index ea55085f5a07..fefdd1e4547f 100644 --- a/src/frontend/src/optimizer/plan_expr_visitor/strong.rs +++ b/src/frontend/src/optimizer/plan_expr_visitor/strong.rs @@ -289,6 +289,7 @@ impl Strong { | ExprType::JsonbPathQueryFirst | ExprType::JsonbPopulateRecord | ExprType::JsonbToRecord + | ExprType::JsonbSet | ExprType::Vnode | ExprType::Proctime | ExprType::PgSleep diff --git a/src/tests/regress/Cargo.toml b/src/tests/regress/Cargo.toml index 8c5478d13d52..65248877adf6 100644 --- a/src/tests/regress/Cargo.toml +++ b/src/tests/regress/Cargo.toml @@ -26,6 +26,7 @@ tokio = { version = "0.2", package = "madsim-tokio", features = [ "time", "signal", "process", + "io-util", ] } tracing = "0.1" tracing-subscriber = "0.3.17" diff --git a/src/tests/regress/data/sql/jsonb.sql b/src/tests/regress/data/sql/jsonb.sql index b8d1e24654af..71ce107ef270 100644 --- a/src/tests/regress/data/sql/jsonb.sql +++ b/src/tests/regress/data/sql/jsonb.sql @@ -1114,18 +1114,18 @@ select '{"a":1 , "b":2, "c":3}'::jsonb - '{b}'::text[]; select '{"a":1 , "b":2, "c":3}'::jsonb - '{c,b}'::text[]; select '{"a":1 , "b":2, "c":3}'::jsonb - '{}'::text[]; ---@ select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{n}', '[1,2,3]'); ---@ select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{b,-1}', '[1,2,3]'); ---@ select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{d,1,0}', '[1,2,3]'); ---@ select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{d,NULL,0}', '[1,2,3]'); ---@ ---@ select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{n}', '{"1": 2}'); ---@ select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{b,-1}', '{"1": 2}'); ---@ select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{d,1,0}', '{"1": 2}'); ---@ select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{d,NULL,0}', '{"1": 2}'); ---@ ---@ select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{b,-1}', '"test"'); ---@ select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{b,-1}', '{"f": "test"}'); +select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{n}', '[1,2,3]'); +select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{b,-1}', '[1,2,3]'); +select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{d,1,0}', '[1,2,3]'); +select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{d,NULL,0}', '[1,2,3]'); + +select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{n}', '{"1": 2}'); +select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{b,-1}', '{"1": 2}'); +select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{d,1,0}', '{"1": 2}'); +select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{d,NULL,0}', '{"1": 2}'); + +select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{b,-1}', '"test"'); +select jsonb_set('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb, '{b,-1}', '{"f": "test"}'); select jsonb_delete_path('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}', '{n}'); select jsonb_delete_path('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}', '{b,-1}'); @@ -1148,33 +1148,33 @@ select '[]'::jsonb - 1; select '"a"'::jsonb #- '{a}'; -- error select '{}'::jsonb #- '{a}'; select '[]'::jsonb #- '{a}'; ---@ select jsonb_set('"a"','{a}','"b"'); --error ---@ select jsonb_set('{}','{a}','"b"', false); ---@ select jsonb_set('[]','{1}','"b"', false); ---@ select jsonb_set('[{"f1":1,"f2":null},2,null,3]', '{0}','[2,3,4]', false); ---@ ---@ -- jsonb_set adding instead of replacing ---@ ---@ -- prepend to array ---@ select jsonb_set('{"a":1,"b":[0,1,2],"c":{"d":4}}','{b,-33}','{"foo":123}'); ---@ -- append to array ---@ select jsonb_set('{"a":1,"b":[0,1,2],"c":{"d":4}}','{b,33}','{"foo":123}'); ---@ -- check nesting levels addition ---@ select jsonb_set('{"a":1,"b":[4,5,[0,1,2],6,7],"c":{"d":4}}','{b,2,33}','{"foo":123}'); ---@ -- add new key ---@ select jsonb_set('{"a":1,"b":[0,1,2],"c":{"d":4}}','{c,e}','{"foo":123}'); ---@ -- adding doesn't do anything if elements before last aren't present ---@ select jsonb_set('{"a":1,"b":[0,1,2],"c":{"d":4}}','{x,-33}','{"foo":123}'); ---@ select jsonb_set('{"a":1,"b":[0,1,2],"c":{"d":4}}','{x,y}','{"foo":123}'); ---@ -- add to empty object ---@ select jsonb_set('{}','{x}','{"foo":123}'); ---@ --add to empty array ---@ select jsonb_set('[]','{0}','{"foo":123}'); ---@ select jsonb_set('[]','{99}','{"foo":123}'); ---@ select jsonb_set('[]','{-99}','{"foo":123}'); ---@ select jsonb_set('{"a": [1, 2, 3]}', '{a, non_integer}', '"new_value"'); ---@ select jsonb_set('{"a": {"b": [1, 2, 3]}}', '{a, b, non_integer}', '"new_value"'); ---@ select jsonb_set('{"a": {"b": [1, 2, 3]}}', '{a, b, NULL}', '"new_value"'); +select jsonb_set('"a"','{a}','"b"'); --error +select jsonb_set('{}','{a}','"b"', false); +select jsonb_set('[]','{1}','"b"', false); +select jsonb_set('[{"f1":1,"f2":null},2,null,3]', '{0}','[2,3,4]', false); + +-- jsonb_set adding instead of replacing + +-- prepend to array +select jsonb_set('{"a":1,"b":[0,1,2],"c":{"d":4}}','{b,-33}','{"foo":123}'); +-- append to array +select jsonb_set('{"a":1,"b":[0,1,2],"c":{"d":4}}','{b,33}','{"foo":123}'); +-- check nesting levels addition +select jsonb_set('{"a":1,"b":[4,5,[0,1,2],6,7],"c":{"d":4}}','{b,2,33}','{"foo":123}'); +-- add new key +select jsonb_set('{"a":1,"b":[0,1,2],"c":{"d":4}}','{c,e}','{"foo":123}'); +-- adding doesn't do anything if elements before last aren't present +select jsonb_set('{"a":1,"b":[0,1,2],"c":{"d":4}}','{x,-33}','{"foo":123}'); +select jsonb_set('{"a":1,"b":[0,1,2],"c":{"d":4}}','{x,y}','{"foo":123}'); +-- add to empty object +select jsonb_set('{}','{x}','{"foo":123}'); +--add to empty array +select jsonb_set('[]','{0}','{"foo":123}'); +select jsonb_set('[]','{99}','{"foo":123}'); +select jsonb_set('[]','{-99}','{"foo":123}'); +select jsonb_set('{"a": [1, 2, 3]}', '{a, non_integer}', '"new_value"'); +select jsonb_set('{"a": {"b": [1, 2, 3]}}', '{a, b, non_integer}', '"new_value"'); +select jsonb_set('{"a": {"b": [1, 2, 3]}}', '{a, b, NULL}', '"new_value"'); -- jsonb_set_lax