diff --git a/proto/expr.proto b/proto/expr.proto index 2fcc905c4b2ab..7998f2fe8128a 100644 --- a/proto/expr.proto +++ b/proto/expr.proto @@ -220,6 +220,7 @@ message ExprNode { JSONB_ARRAY_LENGTH = 603; IS_JSON = 604; JSONB_CAT = 605; + JSONB_OBJECT = 606; JSONB_PRETTY = 607; // Non-pure functions below (> 1000) diff --git a/src/expr/impl/src/lib.rs b/src/expr/impl/src/lib.rs index 6ea82d30ac5f1..51b9a20a75c46 100644 --- a/src/expr/impl/src/lib.rs +++ b/src/expr/impl/src/lib.rs @@ -31,6 +31,7 @@ #![feature(coroutines)] #![feature(test)] #![feature(arc_unwrap_or_clone)] +#![feature(iter_array_chunks)] mod aggregate; mod scalar; diff --git a/src/expr/impl/src/scalar/jsonb_object.rs b/src/expr/impl/src/scalar/jsonb_object.rs new file mode 100644 index 0000000000000..3eb99cbaae615 --- /dev/null +++ b/src/expr/impl/src/scalar/jsonb_object.rs @@ -0,0 +1,178 @@ +// Copyright 2023 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::Builder; +use risingwave_common::types::{JsonbVal, ListRef}; +use risingwave_common::util::iter_util::ZipEqFast; +use risingwave_expr::{function, ExprError, Result}; + +/// Builds a JSON object out of a text array. +/// +/// The array must have either exactly one dimension with an even number of members, +/// in which case they are taken as alternating key/value pairs, or two dimensions +/// such that each inner array has exactly two elements, which are taken as a key/value pair. +/// All values are converted to JSON strings. +/// +/// # Examples +/// +/// ```slt +/// query T +/// select jsonb_object('{a, 1, b, def, c, 3.5}' :: text[]); +/// ---- +/// {"a": "1", "b": "def", "c": "3.5"} +/// +/// query error array must have even number of elements +/// select jsonb_object('{a, 1, b, "def", c}' :: text[]); +/// +/// query error null value not allowed for object key +/// select jsonb_object(array[null, 'b']); +/// +/// query T +/// select jsonb_object(array['a', null]); +/// ---- +/// {"a": null} +/// ``` +#[function("jsonb_object(varchar[]) -> jsonb")] +fn jsonb_object_1d(array: ListRef<'_>) -> Result { + if array.len() % 2 == 1 { + return Err(ExprError::InvalidParam { + name: "array", + reason: "array must have even number of elements".into(), + }); + } + let mut builder = Builder::>::new(); + builder.begin_object(); + for [key, value] in array.iter().array_chunks() { + match key { + Some(s) => builder.add_string(s.into_utf8()), + None => { + return Err(ExprError::InvalidParam { + name: "array", + reason: "null value not allowed for object key".into(), + }) + } + } + match value { + Some(s) => builder.add_string(s.into_utf8()), + None => builder.add_null(), + } + } + builder.end_object(); + Ok(builder.finish().into()) +} + +/// Builds a JSON object out of a text array. +/// +/// The array must have either exactly one dimension with an even number of members, +/// in which case they are taken as alternating key/value pairs, or two dimensions +/// such that each inner array has exactly two elements, which are taken as a key/value pair. +/// All values are converted to JSON strings. +/// +/// # Examples +/// +/// ```slt +/// query T +/// select jsonb_object('{{a, 1}, {b, def}, {c, 3.5}}' :: text[][]); +/// ---- +/// {"a": "1", "b": "def", "c": "3.5"} +/// +/// # FIXME: `null` should be parsed as a null value instead of a "null" string. +/// # query error null value not allowed for object key +/// # select jsonb_object('{{a, 1}, {null, "def"}, {c, 3.5}}' :: text[][]); +/// +/// query error array must have two columns +/// select jsonb_object('{{a, 1, 2}, {b, "def"}, {c, 3.5}}' :: text[][]); +/// ``` +#[function("jsonb_object(varchar[][]) -> jsonb")] +fn jsonb_object_2d(array: ListRef<'_>) -> Result { + let mut builder = Builder::>::new(); + builder.begin_object(); + for kv in array.iter() { + let Some(kv) = kv else { + return Err(ExprError::InvalidParam { + name: "array", + reason: "Unexpected array element.".into(), + }); + }; + let kv = kv.into_list(); + if kv.len() != 2 { + return Err(ExprError::InvalidParam { + name: "array", + reason: "array must have two columns".into(), + }); + } + match kv.get(0).unwrap() { + Some(s) => builder.add_string(s.into_utf8()), + None => { + return Err(ExprError::InvalidParam { + name: "array", + reason: "null value not allowed for object key".into(), + }) + } + } + match kv.get(1).unwrap() { + Some(s) => builder.add_string(s.into_utf8()), + None => builder.add_null(), + } + } + builder.end_object(); + Ok(builder.finish().into()) +} + +/// This form of `jsonb_object` takes keys and values pairwise from separate text arrays. +/// Otherwise it is identical to the one-argument form. +/// +/// # Examples +/// +/// ```slt +/// query T +/// select jsonb_object('{a,b}', '{1,2}'); +/// ---- +/// {"a": "1", "b": "2"} +/// +/// query error mismatched array dimensions +/// select jsonb_object('{a,b}', '{1,2,3}'); +/// +/// # FIXME: `null` should be parsed as a null value instead of a "null" string. +/// # query error null value not allowed for object key +/// # select jsonb_object('{a,null}', '{1,2}'); +/// ``` +#[function("jsonb_object(varchar[], varchar[]) -> jsonb")] +fn jsonb_object_kv(keys: ListRef<'_>, values: ListRef<'_>) -> Result { + if keys.len() != values.len() { + return Err(ExprError::InvalidParam { + name: "values", + reason: "mismatched array dimensions".into(), + }); + } + let mut builder = Builder::>::new(); + builder.begin_object(); + for (key, value) in keys.iter().zip_eq_fast(values.iter()) { + match key { + Some(s) => builder.add_string(s.into_utf8()), + None => { + return Err(ExprError::InvalidParam { + name: "keys", + reason: "null value not allowed for object key".into(), + }) + } + } + match value { + Some(s) => builder.add_string(s.into_utf8()), + None => builder.add_null(), + } + } + builder.end_object(); + Ok(builder.finish().into()) +} diff --git a/src/expr/impl/src/scalar/mod.rs b/src/expr/impl/src/scalar/mod.rs index dd88a374ba966..d9d10e4548aee 100644 --- a/src/expr/impl/src/scalar/mod.rs +++ b/src/expr/impl/src/scalar/mod.rs @@ -45,6 +45,7 @@ mod int256; mod jsonb_access; mod jsonb_concat; mod jsonb_info; +mod jsonb_object; mod length; mod lower; mod md5; diff --git a/src/frontend/src/binder/expr/function.rs b/src/frontend/src/binder/expr/function.rs index cc2519fb4371f..50a621c0879d5 100644 --- a/src/frontend/src/binder/expr/function.rs +++ b/src/frontend/src/binder/expr/function.rs @@ -877,6 +877,7 @@ impl Binder { ("jsonb_array_element_text", raw_call(ExprType::JsonbAccessStr)), ("jsonb_typeof", raw_call(ExprType::JsonbTypeof)), ("jsonb_array_length", raw_call(ExprType::JsonbArrayLength)), + ("jsonb_object", raw_call(ExprType::JsonbObject)), ("jsonb_pretty", raw_call(ExprType::JsonbPretty)), // Functions that return a constant value ("pi", pi()), diff --git a/src/frontend/src/expr/pure.rs b/src/frontend/src/expr/pure.rs index 42813ca07b354..a3c7abf1ef482 100644 --- a/src/frontend/src/expr/pure.rs +++ b/src/frontend/src/expr/pure.rs @@ -177,6 +177,7 @@ impl ExprVisitor for ImpureAnalyzer { | expr_node::Type::JsonbAccessStr | expr_node::Type::JsonbTypeof | expr_node::Type::JsonbArrayLength + | expr_node::Type::JsonbObject | expr_node::Type::JsonbPretty | expr_node::Type::IsJson | expr_node::Type::Sind