risingwave_expr_impl/scalar/
jsonb_build.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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/// Builds a possibly-heterogeneously-typed JSON array out of a variadic argument list.
26/// Each argument is converted as per `to_jsonb`.
27///
28/// # Examples
29///
30/// ```slt
31/// query T
32/// select jsonb_build_array(1, 2, 'foo', 4, 5);
33/// ----
34/// [1, 2, "foo", 4, 5]
35///
36/// query T
37/// select jsonb_build_array(variadic array[1, 2, 4, 5]);
38/// ----
39/// [1, 2, 4, 5]
40/// ```
41#[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_element_type();
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/// Builds a JSON object out of a variadic argument list.
60/// By convention, the argument list consists of alternating keys and values.
61/// Key arguments are coerced to text; value arguments are converted as per `to_jsonb`.
62///
63/// # Examples
64///
65/// ```slt
66/// query T
67/// select jsonb_build_object('foo', 1, 2, 'bar');
68/// ----
69/// {"2": "bar", "foo": 1}
70///
71/// query T
72/// select jsonb_build_object(variadic array['foo', '1', '2', 'bar']);
73/// ----
74/// {"2": "bar", "foo": "1"}
75/// ```
76#[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_element_type(),
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            // special treatment for bool, `false` & `true` rather than `f` & `t`.
107            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}
risingwave_expr_impl/scalar/jsonb_build.rs

risingwave_expr_impl/scalar/
jsonb_build.rs
