feat(expr): add jsonb_object function (risingwavelabs#13036)

Signed-off-by: Runji Wang <[email protected]>

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)

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;

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<JsonbVal> {
+    if array.len() % 2 == 1 {
+        return Err(ExprError::InvalidParam {
+            name: "array",
+            reason: "array must have even number of elements".into(),
+        });
+    }
+    let mut builder = Builder::<Vec<u8>>::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<JsonbVal> {
+    let mut builder = Builder::<Vec<u8>>::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<JsonbVal> {
+    if keys.len() != values.len() {
+        return Err(ExprError::InvalidParam {
+            name: "values",
+            reason: "mismatched array dimensions".into(),
+        });
+    }
+    let mut builder = Builder::<Vec<u8>>::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())
+}

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;

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()),

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