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 risingwave_common::row::Row;
17use risingwave_common::types::ScalarRefImpl;
18use risingwave_common::util::iter_util::ZipEqDebug;
19use risingwave_expr::expr::Context;
20use risingwave_expr::{ExprError, Result, function};
21
22use super::{ToJsonb, ToTextDisplay};
23
24/// Builds a possibly-heterogeneously-typed JSON array out of a variadic argument list.
25/// Each argument is converted as per `to_jsonb`.
26///
27/// # Examples
28///
29/// ```slt
30/// query T
31/// select jsonb_build_array(1, 2, 'foo', 4, 5);
32/// ----
33/// [1, 2, "foo", 4, 5]
34///
35/// query T
36/// select jsonb_build_array(variadic array[1, 2, 4, 5]);
37/// ----
38/// [1, 2, 4, 5]
39/// ```
40#[function("jsonb_build_array(variadic anyarray) -> jsonb")]
41fn jsonb_build_array(args: impl Row, ctx: &Context, writer: &mut jsonbb::Builder) -> Result<()> {
42    writer.begin_array();
43    if ctx.variadic {
44        for (value, ty) in args.iter().zip_eq_debug(&ctx.arg_types) {
45            value.add_to(ty, writer)?;
46        }
47    } else {
48        let ty = ctx.arg_types[0].as_list_elem();
49        for value in args.iter() {
50            value.add_to(ty, writer)?;
51        }
52    }
53    writer.end_array();
54    Ok(())
55}
56
57/// Builds a JSON object out of a variadic argument list.
58/// By convention, the argument list consists of alternating keys and values.
59/// Key arguments are coerced to text; value arguments are converted as per `to_jsonb`.
60///
61/// # Examples
62///
63/// ```slt
64/// query T
65/// select jsonb_build_object('foo', 1, 2, 'bar');
66/// ----
67/// {"2": "bar", "foo": 1}
68///
69/// query T
70/// select jsonb_build_object(variadic array['foo', '1', '2', 'bar']);
71/// ----
72/// {"2": "bar", "foo": "1"}
73/// ```
74#[function("jsonb_build_object(variadic anyarray) -> jsonb")]
75fn jsonb_build_object(args: impl Row, ctx: &Context, writer: &mut jsonbb::Builder) -> Result<()> {
76    if args.len() % 2 == 1 {
77        return Err(ExprError::InvalidParam {
78            name: "args",
79            reason: "argument list must have even number of elements".into(),
80        });
81    }
82    writer.begin_object();
83    let arg_types = match ctx.variadic {
84        true => Either::Left(ctx.arg_types.iter()),
85        false => Either::Right(itertools::repeat_n(
86            ctx.arg_types[0].as_list_elem(),
87            args.len(),
88        )),
89    };
90    for (i, [(key, _), (value, value_type)]) in args
91        .iter()
92        .zip_eq_debug(arg_types)
93        .array_chunks()
94        .enumerate()
95    {
96        match key {
97            Some(ScalarRefImpl::List(_) | ScalarRefImpl::Struct(_) | ScalarRefImpl::Jsonb(_)) => {
98                return Err(ExprError::InvalidParam {
99                    name: "args",
100                    reason: "key value must be scalar, not array, composite, or json".into(),
101                });
102            }
103            // special treatment for bool, `false` & `true` rather than `f` & `t`.
104            Some(ScalarRefImpl::Bool(b)) => writer.display(b),
105            Some(s) => writer.display(ToTextDisplay(s)),
106            None => {
107                return Err(ExprError::InvalidParam {
108                    name: "args",
109                    reason: format!("argument {}: key must not be null", i * 2 + 1).into(),
110                });
111            }
112        }
113        value.add_to(value_type, writer)?;
114    }
115    writer.end_object();
116    Ok(())
117}