Skip to main content

risingwave_expr_impl/scalar/
to_jsonb.rs

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