diff --git a/proto/expr.proto b/proto/expr.proto index 0c7290705f82..cceaecd64011 100644 --- a/proto/expr.proto +++ b/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) // ------------------------ diff --git a/src/common/src/types/jsonb.rs b/src/common/src/types/jsonb.rs index 2d0834701454..7f4c00203706 100644 --- a/src/common/src/types/jsonb.rs +++ b/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 diff --git a/src/expr/impl/src/scalar/jsonb_concat.rs b/src/expr/impl/src/scalar/jsonb_concat.rs new file mode 100644 index 000000000000..6277db8d5b98 --- /dev/null +++ b/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])), + } +} diff --git a/src/expr/impl/src/scalar/mod.rs b/src/expr/impl/src/scalar/mod.rs index 0d60bf04bed8..dd88a374ba96 100644 --- a/src/expr/impl/src/scalar/mod.rs +++ b/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; diff --git a/src/frontend/planner_test/tests/testdata/input/concat_op_dispatch.yaml b/src/frontend/planner_test/tests/testdata/input/concat_op_dispatch.yaml index 208eb6441b4f..e241638c8d15 100644 --- a/src/frontend/planner_test/tests/testdata/input/concat_op_dispatch.yaml +++ b/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 diff --git a/src/frontend/planner_test/tests/testdata/output/concat_op_dispatch.yaml b/src/frontend/planner_test/tests/testdata/output/concat_op_dispatch.yaml index 577c413051c2..41c3475ca1d3 100644 --- a/src/frontend/planner_test/tests/testdata/output/concat_op_dispatch.yaml +++ b/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; diff --git a/src/frontend/src/binder/expr/binary_op.rs b/src/frontend/src/binder/expr/binary_op.rs index 00b4a9f64a9c..f7c8a86144fc 100644 --- a/src/frontend/src/binder/expr/binary_op.rs +++ b/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!( diff --git a/src/frontend/src/expr/pure.rs b/src/frontend/src/expr/pure.rs index 91e92ee65120..4bd7a997ba41 100644 --- a/src/frontend/src/expr/pure.rs +++ b/src/frontend/src/expr/pure.rs @@ -169,6 +169,7 @@ impl ExprVisitor 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