risingwave_expr_impl/scalar/
jsonb_build.rs1use itertools::Either;
16use jsonbb::Builder;
17use risingwave_common::row::Row;
18use risingwave_common::types::{JsonbVal, ScalarRefImpl};
19use risingwave_common::util::iter_util::ZipEqDebug;
20use risingwave_expr::expr::Context;
21use risingwave_expr::{ExprError, Result, function};
22
23use super::{ToJsonb, ToTextDisplay};
24
25#[function("jsonb_build_array(variadic anyarray) -> jsonb")]
42fn jsonb_build_array(args: impl Row, ctx: &Context) -> Result<JsonbVal> {
43    let mut builder = Builder::<Vec<u8>>::new();
44    builder.begin_array();
45    if ctx.variadic {
46        for (value, ty) in args.iter().zip_eq_debug(&ctx.arg_types) {
47            value.add_to(ty, &mut builder)?;
48        }
49    } else {
50        let ty = ctx.arg_types[0].as_list_elem();
51        for value in args.iter() {
52            value.add_to(ty, &mut builder)?;
53        }
54    }
55    builder.end_array();
56    Ok(builder.finish().into())
57}
58
59#[function("jsonb_build_object(variadic anyarray) -> jsonb")]
77fn jsonb_build_object(args: impl Row, ctx: &Context) -> Result<JsonbVal> {
78    if args.len() % 2 == 1 {
79        return Err(ExprError::InvalidParam {
80            name: "args",
81            reason: "argument list must have even number of elements".into(),
82        });
83    }
84    let mut builder = Builder::<Vec<u8>>::new();
85    builder.begin_object();
86    let arg_types = match ctx.variadic {
87        true => Either::Left(ctx.arg_types.iter()),
88        false => Either::Right(itertools::repeat_n(
89            ctx.arg_types[0].as_list_elem(),
90            args.len(),
91        )),
92    };
93    for (i, [(key, _), (value, value_type)]) in args
94        .iter()
95        .zip_eq_debug(arg_types)
96        .array_chunks()
97        .enumerate()
98    {
99        match key {
100            Some(ScalarRefImpl::List(_) | ScalarRefImpl::Struct(_) | ScalarRefImpl::Jsonb(_)) => {
101                return Err(ExprError::InvalidParam {
102                    name: "args",
103                    reason: "key value must be scalar, not array, composite, or json".into(),
104                });
105            }
106            Some(ScalarRefImpl::Bool(b)) => builder.display(b),
108            Some(s) => builder.display(ToTextDisplay(s)),
109            None => {
110                return Err(ExprError::InvalidParam {
111                    name: "args",
112                    reason: format!("argument {}: key must not be null", i * 2 + 1).into(),
113                });
114            }
115        }
116        value.add_to(value_type, &mut builder)?;
117    }
118    builder.end_object();
119    Ok(builder.finish().into())
120}