feat(expr): support jsonb_populate_map (#18378) (#18399)

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

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;
 

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;

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<MapValue, String> {
+        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<Datum> = Vec::with_capacity(object.len());
+        let mut values: Vec<Datum> = 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,

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<MapRef<'_>>,
+    v: JsonbRef<'_>,
+    ctx: &Context,
+) -> Result<MapValue> {
+    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`.

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

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

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