risingwave_expr_impl/scalar/
to_jsonb.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 std::fmt::Debug;
16
17use jsonbb::Builder;
18use risingwave_common::types::{
19    DataType, Date, Decimal, F32, F64, Int256Ref, Interval, JsonbRef, JsonbVal, ListRef, MapRef,
20    ScalarRefImpl, Serial, StructRef, Time, Timestamp, Timestamptz, ToText,
21};
22use risingwave_common::util::iter_util::ZipEqDebug;
23use risingwave_expr::expr::Context;
24use risingwave_expr::{ExprError, Result, function};
25
26#[function("to_jsonb(*) -> jsonb")]
27fn to_jsonb(input: Option<impl ToJsonb>, ctx: &Context) -> Result<JsonbVal> {
28    let mut builder = Builder::default();
29    input.add_to(&ctx.arg_types[0], &mut builder)?;
30    Ok(builder.finish().into())
31}
32
33/// Values that can be converted to JSONB.
34///
35/// This trait is implemented for all scalar reference types.
36pub trait ToJsonb {
37    fn add_to(self, data_type: &DataType, builder: &mut Builder) -> Result<()>;
38}
39
40impl<T: ToJsonb> ToJsonb for Option<T> {
41    fn add_to(self, data_type: &DataType, builder: &mut Builder) -> Result<()> {
42        match self {
43            Some(inner) => inner.add_to(data_type, builder),
44            None => {
45                builder.add_null();
46                Ok(())
47            }
48        }
49    }
50}
51
52impl ToJsonb for ScalarRefImpl<'_> {
53    fn add_to(self, ty: &DataType, builder: &mut Builder) -> Result<()> {
54        use ScalarRefImpl::*;
55        match self {
56            Int16(v) => v.add_to(ty, builder),
57            Int32(v) => v.add_to(ty, builder),
58            Int64(v) => v.add_to(ty, builder),
59            Int256(v) => v.add_to(ty, builder),
60            Float32(v) => v.add_to(ty, builder),
61            Float64(v) => v.add_to(ty, builder),
62            Utf8(v) => v.add_to(ty, builder),
63            Bool(v) => v.add_to(ty, builder),
64            Decimal(v) => v.add_to(ty, builder),
65            Interval(v) => v.add_to(ty, builder),
66            Date(v) => v.add_to(ty, builder),
67            Time(v) => v.add_to(ty, builder),
68            Timestamp(v) => v.add_to(ty, builder),
69            Jsonb(v) => v.add_to(ty, builder),
70            Serial(v) => v.add_to(ty, builder),
71            Bytea(v) => v.add_to(ty, builder),
72            Timestamptz(v) => v.add_to(ty, builder),
73            Struct(v) => v.add_to(ty, builder),
74            List(v) => v.add_to(ty, builder),
75            Map(v) => v.add_to(ty, builder),
76        }
77    }
78}
79
80impl ToJsonb for bool {
81    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
82        builder.add_bool(self);
83        Ok(())
84    }
85}
86
87impl ToJsonb for i16 {
88    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
89        builder.add_i64(self as _);
90        Ok(())
91    }
92}
93
94impl ToJsonb for i32 {
95    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
96        builder.add_i64(self as _);
97        Ok(())
98    }
99}
100
101impl ToJsonb for i64 {
102    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
103        builder.add_i64(self as _);
104        Ok(())
105    }
106}
107
108impl ToJsonb for F32 {
109    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
110        if self.0 == f32::INFINITY {
111            builder.add_string("Infinity");
112        } else if self.0 == f32::NEG_INFINITY {
113            builder.add_string("-Infinity");
114        } else if self.0.is_nan() {
115            builder.add_string("NaN");
116        } else {
117            builder.add_f64(self.0 as f64);
118        }
119        Ok(())
120    }
121}
122
123impl ToJsonb for F64 {
124    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
125        if self.0 == f64::INFINITY {
126            builder.add_string("Infinity");
127        } else if self.0 == f64::NEG_INFINITY {
128            builder.add_string("-Infinity");
129        } else if self.0.is_nan() {
130            builder.add_string("NaN");
131        } else {
132            builder.add_f64(self.0);
133        }
134        Ok(())
135    }
136}
137
138impl ToJsonb for Decimal {
139    fn add_to(self, t: &DataType, builder: &mut Builder) -> Result<()> {
140        let res: F64 = self
141            .try_into()
142            .map_err(|_| ExprError::CastOutOfRange("IEEE 754 double"))?;
143        res.add_to(t, builder)?;
144        Ok(())
145    }
146}
147
148impl ToJsonb for Int256Ref<'_> {
149    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
150        builder.display(ToTextDisplay(self));
151        Ok(())
152    }
153}
154
155impl ToJsonb for &str {
156    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
157        builder.add_string(self);
158        Ok(())
159    }
160}
161
162impl ToJsonb for &[u8] {
163    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
164        builder.display(ToTextDisplay(self));
165        Ok(())
166    }
167}
168
169impl ToJsonb for Date {
170    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
171        builder.display(ToTextDisplay(self));
172        Ok(())
173    }
174}
175
176impl ToJsonb for Time {
177    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
178        builder.display(ToTextDisplay(self));
179        Ok(())
180    }
181}
182
183impl ToJsonb for Interval {
184    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
185        builder.display(ToTextDisplay(self));
186        Ok(())
187    }
188}
189
190impl ToJsonb for Timestamp {
191    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
192        builder.display(format_args!("{}T{}", self.0.date(), self.0.time()));
193        Ok(())
194    }
195}
196
197impl ToJsonb for Timestamptz {
198    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
199        let instant_local = self.to_datetime_utc();
200        builder.display(instant_local.to_rfc3339().as_str());
201        Ok(())
202    }
203}
204
205impl ToJsonb for Serial {
206    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
207        builder.display(ToTextDisplay(self));
208        Ok(())
209    }
210}
211
212impl ToJsonb for JsonbRef<'_> {
213    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
214        builder.add_value(self.into());
215        Ok(())
216    }
217}
218
219impl ToJsonb for ListRef<'_> {
220    fn add_to(self, data_type: &DataType, builder: &mut Builder) -> Result<()> {
221        let elem_type = data_type.as_list();
222        builder.begin_array();
223        for value in self.iter() {
224            value.add_to(elem_type, builder)?;
225        }
226        builder.end_array();
227        Ok(())
228    }
229}
230
231impl ToJsonb for MapRef<'_> {
232    fn add_to(self, data_type: &DataType, builder: &mut Builder) -> Result<()> {
233        let value_type = data_type.as_map().value();
234        builder.begin_object();
235        for (k, v) in self.iter() {
236            // XXX: is to_text here reasonable?
237            builder.add_string(&k.to_text());
238            v.add_to(value_type, builder)?;
239        }
240        builder.end_object();
241        Ok(())
242    }
243}
244
245impl ToJsonb for StructRef<'_> {
246    fn add_to(self, data_type: &DataType, builder: &mut Builder) -> Result<()> {
247        builder.begin_object();
248        for (value, (field_name, field_type)) in self
249            .iter_fields_ref()
250            .zip_eq_debug(data_type.as_struct().iter())
251        {
252            builder.add_string(field_name);
253            value.add_to(field_type, builder)?;
254        }
255        builder.end_object();
256        Ok(())
257    }
258}
259
260/// A wrapper type to implement `Display` for `ToText`.
261pub struct ToTextDisplay<T>(pub T);
262
263impl<T: ToText> std::fmt::Display for ToTextDisplay<T> {
264    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
265        self.0.write(f)
266    }
267}