diff --git a/e2e_test/batch/types/map.slt.part b/e2e_test/batch/types/map.slt.part index b4b4be7e5cba7..fe98fa3633000 100644 --- a/e2e_test/batch/types/map.slt.part +++ b/e2e_test/batch/types/map.slt.part @@ -122,6 +122,63 @@ select to_jsonb(m1), to_jsonb(m2), to_jsonb(m3), to_jsonb(l), to_jsonb(s) from t {"a": 1.0, "b": 2.0, "c": 3.0} null null null null {"a": 1.0, "b": 2.0, "c": 3.0} {"1": true, "2": false, "3": true} {"a": {"a1": "a2"}, "b": {"b1": "b2"}} [{"a": 1, "b": 2, "c": 3}, {"d": 4, "e": 5, "f": 6}] {"m": {"a": {"x": 1}, "b": {"x": 2}, "c": {"x": 3}}} +query ? +select jsonb_populate_map( + null::map(varchar, int), + '{"a": 1, "b": 2}'::jsonb +); +---- +{a:1,b:2} + + +query ? +select jsonb_populate_map( + MAP {'a': 1, 'b': 2}, + '{"b": 3, "c": 4}'::jsonb +); +---- +{a:1,b:3,c:4} + + +# implicit cast (int -> varchar) +query ? +select jsonb_populate_map( + MAP {'a': 'a', 'b': 'b'}, + '{"b": 3, "c": 4}'::jsonb +); +---- +{a:a,b:3,c:4} + + +query error +select jsonb_populate_map( + MAP {'a': 1, 'b': 2}, + '{"b": "3", "c": 4}'::jsonb +); +---- +db error: ERROR: Failed to run the query + +Caused by these errors (recent errors listed first): + 1: Expr error + 2: error while evaluating expression `jsonb_populate_map('{a:1,b:2}', '{"b": "3", "c": 4}')` + 3: Parse error: cannot cast jsonb string to type number + + +query error +select jsonb_populate_map( + null::map(int, int), + '{"a": 1, "b": 2}'::jsonb +); +---- +db error: ERROR: Failed to run the query + +Caused by these errors (recent errors listed first): + 1: Expr error + 2: error while evaluating expression `jsonb_populate_map(NULL, '{"a": 1, "b": 2}')` + 3: Parse error: cannot convert jsonb to a map with non-string keys + + + statement ok drop table t; diff --git a/proto/expr.proto b/proto/expr.proto index e5b5fb73ba8ff..53bba96cc587b 100644 --- a/proto/expr.proto +++ b/proto/expr.proto @@ -282,6 +282,7 @@ message ExprNode { JSONB_POPULATE_RECORD = 629; JSONB_TO_RECORD = 630; JSONB_SET = 631; + JSONB_POPULATE_MAP = 632; // Map functions MAP_FROM_ENTRIES = 700; diff --git a/src/common/src/types/jsonb.rs b/src/common/src/types/jsonb.rs index fa80069080ff4..6363864fd73e2 100644 --- a/src/common/src/types/jsonb.rs +++ b/src/common/src/types/jsonb.rs @@ -20,7 +20,9 @@ use jsonbb::{Value, ValueRef}; use postgres_types::{accepts, to_sql_checked, FromSql, IsNull, ToSql, Type}; use risingwave_common_estimate_size::EstimateSize; -use super::{Datum, IntoOrdered, ListValue, ScalarImpl, StructRef, ToOwnedDatum, F64}; +use super::{ + Datum, IntoOrdered, ListValue, MapType, MapValue, ScalarImpl, StructRef, ToOwnedDatum, F64, +}; use crate::types::{DataType, Scalar, ScalarRef, StructType, StructValue}; use crate::util::iter_util::ZipEqDebug; @@ -464,6 +466,28 @@ impl<'a> JsonbRef<'a> { Ok(StructValue::new(fields)) } + pub fn to_map(self, ty: &MapType) -> Result { + let object = self + .0 + .as_object() + .ok_or_else(|| format!("cannot convert to map from a jsonb {}", self.type_name()))?; + if !matches!(ty.key(), DataType::Varchar) { + return Err("cannot convert jsonb to a map with non-string keys".to_string()); + } + + let mut keys: Vec = Vec::with_capacity(object.len()); + let mut values: Vec = Vec::with_capacity(object.len()); + for (k, v) in object.iter() { + let v = Self(v).to_datum(ty.value())?; + keys.push(Some(ScalarImpl::Utf8(k.to_owned().into()))); + values.push(v); + } + MapValue::try_from_kv( + ListValue::from_datum_iter(ty.key(), keys), + ListValue::from_datum_iter(ty.value(), values), + ) + } + /// Expands the top-level JSON object to a row having the struct type of the `base` argument. pub fn populate_struct( self, diff --git a/src/expr/impl/src/scalar/jsonb_record.rs b/src/expr/impl/src/scalar/jsonb_record.rs index b85feb9190d2a..a6def7cb25643 100644 --- a/src/expr/impl/src/scalar/jsonb_record.rs +++ b/src/expr/impl/src/scalar/jsonb_record.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use risingwave_common::types::{JsonbRef, StructRef, StructValue}; +use risingwave_common::types::{JsonbRef, MapRef, MapValue, Scalar, StructRef, StructValue}; use risingwave_expr::expr::Context; use risingwave_expr::{function, ExprError, Result}; @@ -60,6 +60,22 @@ fn jsonb_populate_record( jsonb.populate_struct(output_type, base).map_err(parse_err) } +#[function("jsonb_populate_map(anymap, jsonb) -> anymap")] +pub fn jsonb_populate_map( + base: Option>, + v: JsonbRef<'_>, + ctx: &Context, +) -> Result { + let output_type = ctx.return_type.as_map(); + let jsonb_map = v + .to_map(output_type) + .map_err(|e| ExprError::Parse(e.into()))?; + match base { + Some(base) => Ok(MapValue::concat(base, jsonb_map.as_scalar_ref())), + None => Ok(jsonb_map), + } +} + /// Expands the top-level JSON array of objects to a set of rows having the composite type of the /// base argument. Each element of the JSON array is processed as described above for /// `jsonb_populate_record`. diff --git a/src/frontend/src/binder/expr/function/builtin_scalar.rs b/src/frontend/src/binder/expr/function/builtin_scalar.rs index 73eb722b26011..d46681c51ab3e 100644 --- a/src/frontend/src/binder/expr/function/builtin_scalar.rs +++ b/src/frontend/src/binder/expr/function/builtin_scalar.rs @@ -399,6 +399,7 @@ impl Binder { ("jsonb_path_query_array", raw_call(ExprType::JsonbPathQueryArray)), ("jsonb_path_query_first", raw_call(ExprType::JsonbPathQueryFirst)), ("jsonb_set", raw_call(ExprType::JsonbSet)), + ("jsonb_populate_map", raw_call(ExprType::JsonbPopulateMap)), // map ("map_from_entries", raw_call(ExprType::MapFromEntries)), ("map_access",raw_call(ExprType::MapAccess)), diff --git a/src/frontend/src/expr/pure.rs b/src/frontend/src/expr/pure.rs index 3e6c83d8330fb..d47cc3851f641 100644 --- a/src/frontend/src/expr/pure.rs +++ b/src/frontend/src/expr/pure.rs @@ -211,6 +211,7 @@ impl ExprVisitor for ImpureAnalyzer { | Type::JsonbPathQueryArray | Type::JsonbPathQueryFirst | Type::JsonbSet + | Type::JsonbPopulateMap | 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 2c14fc730877d..673a5f41746bb 100644 --- a/src/frontend/src/optimizer/plan_expr_visitor/strong.rs +++ b/src/frontend/src/optimizer/plan_expr_visitor/strong.rs @@ -291,6 +291,7 @@ impl Strong { | ExprType::JsonbPopulateRecord | ExprType::JsonbToRecord | ExprType::JsonbSet + | ExprType::JsonbPopulateMap | ExprType::MapFromEntries | ExprType::MapAccess | ExprType::MapKeys