feat: jsonb || operator (#12502)

Co-authored-by: xxchan <[email protected]>

proto/expr.proto

@@ -214,6 +214,7 @@ message ExprNode {
     JSONB_TYPEOF = 602;
     JSONB_ARRAY_LENGTH = 603;
     IS_JSON = 604;
+    JSONB_CAT = 605;
 
     // Non-pure functions below (> 1000)
     // ------------------------

src/common/src/types/jsonb.rs

@@ -345,6 +345,10 @@ impl<'a> JsonbRef<'a> {
             _ => Err(format!("cannot deconstruct a jsonb {}", self.type_name())),
         }
     }
+
+    pub fn value(&self) -> &'a Value {
+        self.0
+    }
 }
 
 /// A custom implementation for [`serde_json::ser::Formatter`] to match PostgreSQL, which adds extra

src/expr/impl/src/scalar/jsonb_concat.rs

@@ -0,0 +1,101 @@
+// 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 risingwave_common::types::{JsonbRef, JsonbVal};
+use risingwave_expr::function;
+use serde_json::{json, Value};
+
+/// Concatenates the two jsonbs.
+///
+/// Examples:
+///
+/// ```slt
+/// # concat
+/// query T
+/// SELECT '[1,2]'::jsonb || '[3,4]'::jsonb;
+/// ----
+/// [1, 2, 3, 4]
+///
+/// query T
+/// SELECT '{"a": 1}'::jsonb || '{"b": 2}'::jsonb;
+/// ----
+/// {"a": 1, "b": 2}
+///
+/// query T
+/// SELECT '[1,2]'::jsonb || '{"a": 1}'::jsonb;
+/// ----
+/// [1, 2, {"a": 1}]
+///
+/// query T
+/// SELECT '1'::jsonb || '2'::jsonb;
+/// ----
+/// [1, 2]
+///
+/// query T
+/// SELECT '[1,2]'::jsonb || 'null'::jsonb;
+/// ----
+/// [1, 2, null]
+///
+/// query T
+/// SELECT 'null'::jsonb || '[1,2]'::jsonb;
+/// ----
+/// [null, 1, 2]
+///
+/// query T
+/// SELECT 'null'::jsonb || '1'::jsonb;
+/// ----
+/// [null, 1]
+/// ```
+#[function("jsonb_cat(jsonb, jsonb) -> jsonb")]
+pub fn jsonb_cat(left: JsonbRef<'_>, right: JsonbRef<'_>) -> JsonbVal {
+    let left_val = left.value().clone();
+    let right_val = right.value().clone();
+    match (left_val, right_val) {
+        // left and right are object based.
+        // This would have left:{'a':1}, right:{'b':2} -> {'a':1,'b':2}
+        (Value::Object(mut left_map), Value::Object(right_map)) => {
+            left_map.extend(right_map);
+            JsonbVal::from(Value::Object(left_map))
+        }
+
+        // left and right are array-based.
+        // This would merge both arrays into one array.
+        // This would have left:[1,2], right:[3,4] -> [1,2,3,4]
+        (Value::Array(mut left_arr), Value::Array(right_arr)) => {
+            left_arr.extend(right_arr);
+            JsonbVal::from(Value::Array(left_arr))
+        }
+
+        // One operand is an array, and the other is a single element.
+        // This would insert the non-array value as another element into the array
+        // Eg left:[1,2] right: {'a':1} -> [1,2,{'a':1}]
+        (Value::Array(mut left_arr), single_val) => {
+            left_arr.push(single_val);
+            JsonbVal::from(Value::Array(left_arr))
+        }
+
+        // One operand is an array, and the other is a single element.
+        // This would insert the non-array value as another element into the array
+        // Eg left:{'a':1} right:[1,2] -> [{'a':1},1,2]
+        (single_val, Value::Array(mut right_arr)) => {
+            right_arr.insert(0, single_val);
+            JsonbVal::from(Value::Array(right_arr))
+        }
+
+        // Both are non-array inputs.
+        // Both elements would be placed together in an array
+        // Eg left:1 right: 2 -> [1,2]
+        (left, right) => JsonbVal::from(json!([left, right])),
+    }
+}

src/expr/impl/src/scalar/mod.rs

@@ -43,6 +43,7 @@ mod format;
 mod format_type;
 mod int256;
 mod jsonb_access;
+mod jsonb_concat;
 mod jsonb_info;
 mod length;
 mod lower;

src/frontend/planner_test/tests/testdata/input/concat_op_dispatch.yaml

@@ -37,12 +37,12 @@
     select '1'::jsonb || '2'::jsonb;
   name: jsonb || jsonb -> jsonb
   expected_outputs:
-  - binder_error
+  - batch_plan
 - sql: |
     select '1'::jsonb || '2';
   name: jsonb || unknown (as jsonb) -> jsonb
   expected_outputs:
-  - binder_error
+  - batch_plan
 - sql: |
     with t(s) as (select '2') select '1'::jsonb || s from t;
   name: jsonb || text -> text

src/frontend/planner_test/tests/testdata/output/concat_op_dispatch.yaml

@@ -34,19 +34,11 @@
 - name: jsonb || jsonb -> jsonb
   sql: |
     select '1'::jsonb || '2'::jsonb;
-  binder_error: |-
-    Bind error: failed to bind expression: CAST('1' AS jsonb) || CAST('2' AS jsonb)
-
-    Caused by:
-      Bind error: operator not implemented yet: jsonb || jsonb
+  batch_plan: 'BatchValues { rows: [[''[1, 2]'':Jsonb]] }'
 - name: jsonb || unknown (as jsonb) -> jsonb
   sql: |
     select '1'::jsonb || '2';
-  binder_error: |-
-    Bind error: failed to bind expression: CAST('1' AS jsonb) || '2'
-
-    Caused by:
-      Bind error: operator not implemented yet: jsonb || jsonb
+  batch_plan: 'BatchValues { rows: [[''[1, 2]'':Jsonb]] }'
 - name: jsonb || text -> text
   sql: |
     with t(s) as (select '2') select '1'::jsonb || s from t;

src/frontend/src/binder/expr/binary_op.rs

@@ -108,11 +108,12 @@ impl Binder {
                         ExprType::ConcatOp
                     }
 
-                    // jsonb, bytea (and varbit, tsvector, tsquery)
-                    (Some(t @ DataType::Jsonb), Some(DataType::Jsonb))
-                    | (Some(t @ DataType::Jsonb), None)
-                    | (None, Some(t @ DataType::Jsonb))
-                    | (Some(t @ DataType::Bytea), Some(DataType::Bytea))
+                    (Some(DataType::Jsonb), Some(DataType::Jsonb))
+                    | (Some(DataType::Jsonb), None)
+                    | (None, Some(DataType::Jsonb)) => ExprType::JsonbCat,
+
+                    // bytea (and varbit, tsvector, tsquery)
+                    (Some(t @ DataType::Bytea), Some(DataType::Bytea))
                     | (Some(t @ DataType::Bytea), None)
                     | (None, Some(t @ DataType::Bytea)) => {
                         return Err(ErrorCode::BindError(format!(

src/frontend/src/expr/pure.rs

@@ -169,6 +169,7 @@ impl ExprVisitor<bool> for ImpureAnalyzer {
             | expr_node::Type::ArrayReplace
             | expr_node::Type::ArrayPosition
             | expr_node::Type::HexToInt256
+            | expr_node::Type::JsonbCat
             | expr_node::Type::JsonbAccessInner
             | expr_node::Type::JsonbAccessStr
             | expr_node::Type::JsonbTypeof