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