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, VectorRef,
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            Vector(v) => v.add_to(ty, builder),
77        }
78    }
79}
80
81impl ToJsonb for bool {
82    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
83        builder.add_bool(self);
84        Ok(())
85    }
86}
87
88impl ToJsonb for i16 {
89    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
90        builder.add_i64(self as _);
91        Ok(())
92    }
93}
94
95impl ToJsonb for i32 {
96    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
97        builder.add_i64(self as _);
98        Ok(())
99    }
100}
101
102impl ToJsonb for i64 {
103    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
104        builder.add_i64(self as _);
105        Ok(())
106    }
107}
108
109impl ToJsonb for F32 {
110    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
111        if self.0 == f32::INFINITY {
112            builder.add_string("Infinity");
113        } else if self.0 == f32::NEG_INFINITY {
114            builder.add_string("-Infinity");
115        } else if self.0.is_nan() {
116            builder.add_string("NaN");
117        } else {
118            builder.add_f64(self.0 as f64);
119        }
120        Ok(())
121    }
122}
123
124impl ToJsonb for F64 {
125    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
126        if self.0 == f64::INFINITY {
127            builder.add_string("Infinity");
128        } else if self.0 == f64::NEG_INFINITY {
129            builder.add_string("-Infinity");
130        } else if self.0.is_nan() {
131            builder.add_string("NaN");
132        } else {
133            builder.add_f64(self.0);
134        }
135        Ok(())
136    }
137}
138
139impl ToJsonb for Decimal {
140    fn add_to(self, t: &DataType, builder: &mut Builder) -> Result<()> {
141        let res: F64 = self
142            .try_into()
143            .map_err(|_| ExprError::CastOutOfRange("IEEE 754 double"))?;
144        res.add_to(t, builder)?;
145        Ok(())
146    }
147}
148
149impl ToJsonb for Int256Ref<'_> {
150    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
151        builder.display(ToTextDisplay(self));
152        Ok(())
153    }
154}
155
156impl ToJsonb for &str {
157    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
158        builder.add_string(self);
159        Ok(())
160    }
161}
162
163impl ToJsonb for &[u8] {
164    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
165        builder.display(ToTextDisplay(self));
166        Ok(())
167    }
168}
169
170impl ToJsonb for Date {
171    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
172        builder.display(ToTextDisplay(self));
173        Ok(())
174    }
175}
176
177impl ToJsonb for Time {
178    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
179        builder.display(ToTextDisplay(self));
180        Ok(())
181    }
182}
183
184impl ToJsonb for Interval {
185    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
186        builder.display(ToTextDisplay(self));
187        Ok(())
188    }
189}
190
191impl ToJsonb for Timestamp {
192    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
193        builder.display(format_args!("{}T{}", self.0.date(), self.0.time()));
194        Ok(())
195    }
196}
197
198impl ToJsonb for Timestamptz {
199    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
200        let instant_local = self.to_datetime_utc();
201        builder.display(instant_local.to_rfc3339().as_str());
202        Ok(())
203    }
204}
205
206impl ToJsonb for Serial {
207    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
208        builder.display(ToTextDisplay(self));
209        Ok(())
210    }
211}
212
213impl ToJsonb for VectorRef<'_> {
214    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
215        builder.display(ToTextDisplay(self));
216        Ok(())
217    }
218}
219
220impl ToJsonb for JsonbRef<'_> {
221    fn add_to(self, _: &DataType, builder: &mut Builder) -> Result<()> {
222        builder.add_value(self.into());
223        Ok(())
224    }
225}
226
227impl ToJsonb for ListRef<'_> {
228    fn add_to(self, data_type: &DataType, builder: &mut Builder) -> Result<()> {
229        let elem_type = data_type.as_list_element_type();
230        builder.begin_array();
231        for value in self.iter() {
232            value.add_to(elem_type, builder)?;
233        }
234        builder.end_array();
235        Ok(())
236    }
237}
238
239impl ToJsonb for MapRef<'_> {
240    fn add_to(self, data_type: &DataType, builder: &mut Builder) -> Result<()> {
241        let value_type = data_type.as_map().value();
242        builder.begin_object();
243        for (k, v) in self.iter() {
244            // XXX: is to_text here reasonable?
245            builder.add_string(&k.to_text());
246            v.add_to(value_type, builder)?;
247        }
248        builder.end_object();
249        Ok(())
250    }
251}
252
253impl ToJsonb for StructRef<'_> {
254    fn add_to(self, data_type: &DataType, builder: &mut Builder) -> Result<()> {
255        builder.begin_object();
256        for (value, (field_name, field_type)) in self
257            .iter_fields_ref()
258            .zip_eq_debug(data_type.as_struct().iter())
259        {
260            builder.add_string(field_name);
261            value.add_to(field_type, builder)?;
262        }
263        builder.end_object();
264        Ok(())
265    }
266}
267
268/// A wrapper type to implement `Display` for `ToText`.
269pub struct ToTextDisplay<T>(pub T);
270
271impl<T: ToText> std::fmt::Display for ToTextDisplay<T> {
272    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
273        self.0.write(f)
274    }
275}