Skip to content

Commit c512972

Browse files
feat(expr): support jsonb_populate_map (#18378) (#18399)
Signed-off-by: xxchan <[email protected]> Co-authored-by: xxchan <[email protected]>
1 parent 6eaf8bf commit c512972

File tree

7 files changed

+103
-2
lines changed

7 files changed

+103
-2
lines changed

e2e_test/batch/types/map.slt.part

+57
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,63 @@ select to_jsonb(m1), to_jsonb(m2), to_jsonb(m3), to_jsonb(l), to_jsonb(s) from t
122122
{"a": 1.0, "b": 2.0, "c": 3.0} null null null null
123123
{"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}}}
124124

125+
query ?
126+
select jsonb_populate_map(
127+
null::map(varchar, int),
128+
'{"a": 1, "b": 2}'::jsonb
129+
);
130+
----
131+
{a:1,b:2}
132+
133+
134+
query ?
135+
select jsonb_populate_map(
136+
MAP {'a': 1, 'b': 2},
137+
'{"b": 3, "c": 4}'::jsonb
138+
);
139+
----
140+
{a:1,b:3,c:4}
141+
142+
143+
# implicit cast (int -> varchar)
144+
query ?
145+
select jsonb_populate_map(
146+
MAP {'a': 'a', 'b': 'b'},
147+
'{"b": 3, "c": 4}'::jsonb
148+
);
149+
----
150+
{a:a,b:3,c:4}
151+
152+
153+
query error
154+
select jsonb_populate_map(
155+
MAP {'a': 1, 'b': 2},
156+
'{"b": "3", "c": 4}'::jsonb
157+
);
158+
----
159+
db error: ERROR: Failed to run the query
160+
161+
Caused by these errors (recent errors listed first):
162+
1: Expr error
163+
2: error while evaluating expression `jsonb_populate_map('{a:1,b:2}', '{"b": "3", "c": 4}')`
164+
3: Parse error: cannot cast jsonb string to type number
165+
166+
167+
query error
168+
select jsonb_populate_map(
169+
null::map(int, int),
170+
'{"a": 1, "b": 2}'::jsonb
171+
);
172+
----
173+
db error: ERROR: Failed to run the query
174+
175+
Caused by these errors (recent errors listed first):
176+
1: Expr error
177+
2: error while evaluating expression `jsonb_populate_map(NULL, '{"a": 1, "b": 2}')`
178+
3: Parse error: cannot convert jsonb to a map with non-string keys
179+
180+
181+
125182
statement ok
126183
drop table t;
127184

proto/expr.proto

+1
Original file line numberDiff line numberDiff line change
@@ -282,6 +282,7 @@ message ExprNode {
282282
JSONB_POPULATE_RECORD = 629;
283283
JSONB_TO_RECORD = 630;
284284
JSONB_SET = 631;
285+
JSONB_POPULATE_MAP = 632;
285286

286287
// Map functions
287288
MAP_FROM_ENTRIES = 700;

src/common/src/types/jsonb.rs

+25-1
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,9 @@ use jsonbb::{Value, ValueRef};
2020
use postgres_types::{accepts, to_sql_checked, FromSql, IsNull, ToSql, Type};
2121
use risingwave_common_estimate_size::EstimateSize;
2222

23-
use super::{Datum, IntoOrdered, ListValue, ScalarImpl, StructRef, ToOwnedDatum, F64};
23+
use super::{
24+
Datum, IntoOrdered, ListValue, MapType, MapValue, ScalarImpl, StructRef, ToOwnedDatum, F64,
25+
};
2426
use crate::types::{DataType, Scalar, ScalarRef, StructType, StructValue};
2527
use crate::util::iter_util::ZipEqDebug;
2628

@@ -464,6 +466,28 @@ impl<'a> JsonbRef<'a> {
464466
Ok(StructValue::new(fields))
465467
}
466468

469+
pub fn to_map(self, ty: &MapType) -> Result<MapValue, String> {
470+
let object = self
471+
.0
472+
.as_object()
473+
.ok_or_else(|| format!("cannot convert to map from a jsonb {}", self.type_name()))?;
474+
if !matches!(ty.key(), DataType::Varchar) {
475+
return Err("cannot convert jsonb to a map with non-string keys".to_string());
476+
}
477+
478+
let mut keys: Vec<Datum> = Vec::with_capacity(object.len());
479+
let mut values: Vec<Datum> = Vec::with_capacity(object.len());
480+
for (k, v) in object.iter() {
481+
let v = Self(v).to_datum(ty.value())?;
482+
keys.push(Some(ScalarImpl::Utf8(k.to_owned().into())));
483+
values.push(v);
484+
}
485+
MapValue::try_from_kv(
486+
ListValue::from_datum_iter(ty.key(), keys),
487+
ListValue::from_datum_iter(ty.value(), values),
488+
)
489+
}
490+
467491
/// Expands the top-level JSON object to a row having the struct type of the `base` argument.
468492
pub fn populate_struct(
469493
self,

src/expr/impl/src/scalar/jsonb_record.rs

+17-1
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15-
use risingwave_common::types::{JsonbRef, StructRef, StructValue};
15+
use risingwave_common::types::{JsonbRef, MapRef, MapValue, Scalar, StructRef, StructValue};
1616
use risingwave_expr::expr::Context;
1717
use risingwave_expr::{function, ExprError, Result};
1818

@@ -60,6 +60,22 @@ fn jsonb_populate_record(
6060
jsonb.populate_struct(output_type, base).map_err(parse_err)
6161
}
6262

63+
#[function("jsonb_populate_map(anymap, jsonb) -> anymap")]
64+
pub fn jsonb_populate_map(
65+
base: Option<MapRef<'_>>,
66+
v: JsonbRef<'_>,
67+
ctx: &Context,
68+
) -> Result<MapValue> {
69+
let output_type = ctx.return_type.as_map();
70+
let jsonb_map = v
71+
.to_map(output_type)
72+
.map_err(|e| ExprError::Parse(e.into()))?;
73+
match base {
74+
Some(base) => Ok(MapValue::concat(base, jsonb_map.as_scalar_ref())),
75+
None => Ok(jsonb_map),
76+
}
77+
}
78+
6379
/// Expands the top-level JSON array of objects to a set of rows having the composite type of the
6480
/// base argument. Each element of the JSON array is processed as described above for
6581
/// `jsonb_populate_record`.

src/frontend/src/binder/expr/function/builtin_scalar.rs

+1
Original file line numberDiff line numberDiff line change
@@ -399,6 +399,7 @@ impl Binder {
399399
("jsonb_path_query_array", raw_call(ExprType::JsonbPathQueryArray)),
400400
("jsonb_path_query_first", raw_call(ExprType::JsonbPathQueryFirst)),
401401
("jsonb_set", raw_call(ExprType::JsonbSet)),
402+
("jsonb_populate_map", raw_call(ExprType::JsonbPopulateMap)),
402403
// map
403404
("map_from_entries", raw_call(ExprType::MapFromEntries)),
404405
("map_access",raw_call(ExprType::MapAccess)),

src/frontend/src/expr/pure.rs

+1
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,7 @@ impl ExprVisitor for ImpureAnalyzer {
211211
| Type::JsonbPathQueryArray
212212
| Type::JsonbPathQueryFirst
213213
| Type::JsonbSet
214+
| Type::JsonbPopulateMap
214215
| Type::IsJson
215216
| Type::ToJsonb
216217
| Type::Sind

src/frontend/src/optimizer/plan_expr_visitor/strong.rs

+1
Original file line numberDiff line numberDiff line change
@@ -291,6 +291,7 @@ impl Strong {
291291
| ExprType::JsonbPopulateRecord
292292
| ExprType::JsonbToRecord
293293
| ExprType::JsonbSet
294+
| ExprType::JsonbPopulateMap
294295
| ExprType::MapFromEntries
295296
| ExprType::MapAccess
296297
| ExprType::MapKeys

0 commit comments

Comments
 (0)