Skip to content

Commit

Permalink
feat(expr): support jsonb_populate_map (#18378)
Browse files Browse the repository at this point in the history
Signed-off-by: xxchan <[email protected]>
  • Loading branch information
xxchan authored Sep 4, 2024
1 parent 0a4ccde commit c4b1dd4
Show file tree
Hide file tree
Showing 7 changed files with 103 additions and 2 deletions.
57 changes: 57 additions & 0 deletions e2e_test/batch/types/map.slt.part
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
1 change: 1 addition & 0 deletions proto/expr.proto
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
26 changes: 25 additions & 1 deletion src/common/src/types/jsonb.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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,
Expand Down
18 changes: 17 additions & 1 deletion src/expr/impl/src/scalar/jsonb_record.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand Down Expand Up @@ -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`.
Expand Down
1 change: 1 addition & 0 deletions src/frontend/src/binder/expr/function/builtin_scalar.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)),
Expand Down
1 change: 1 addition & 0 deletions src/frontend/src/expr/pure.rs
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,7 @@ impl ExprVisitor for ImpureAnalyzer {
| Type::JsonbPathQueryArray
| Type::JsonbPathQueryFirst
| Type::JsonbSet
| Type::JsonbPopulateMap
| Type::IsJson
| Type::ToJsonb
| Type::Sind
Expand Down
1 change: 1 addition & 0 deletions src/frontend/src/optimizer/plan_expr_visitor/strong.rs
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,7 @@ impl Strong {
| ExprType::JsonbPopulateRecord
| ExprType::JsonbToRecord
| ExprType::JsonbSet
| ExprType::JsonbPopulateMap
| ExprType::MapFromEntries
| ExprType::MapAccess
| ExprType::MapKeys
Expand Down

0 comments on commit c4b1dd4

Please sign in to comment.