mz_expr/scalar/func/impls/
jsonb.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::fmt;
11
12use mz_lowertest::MzReflect;
13use mz_repr::adt::jsonb::{Jsonb, JsonbRef};
14use mz_repr::adt::numeric::{self, Numeric, NumericMaxScale};
15use mz_repr::{ColumnType, Datum, Row, RowPacker, ScalarType, strconv};
16use proptest_derive::Arbitrary;
17use serde::{Deserialize, Serialize};
18
19use crate::EvalError;
20use crate::scalar::func::EagerUnaryFunc;
21use crate::scalar::func::impls::numeric::*;
22
23sqlfunc!(
24    #[sqlname = "jsonb_to_text"]
25    #[preserves_uniqueness = false]
26    #[inverse = to_unary!(super::CastStringToJsonb)]
27    fn cast_jsonb_to_string<'a>(a: JsonbRef<'a>) -> String {
28        let mut buf = String::new();
29        strconv::format_jsonb(&mut buf, a);
30        buf
31    }
32);
33
34sqlfunc!(
35    #[sqlname = "jsonb_to_smallint"]
36    #[is_monotone = true]
37    fn cast_jsonb_to_int16<'a>(a: JsonbRef<'a>) -> Result<i16, EvalError> {
38        match a.into_datum() {
39            Datum::Numeric(a) => cast_numeric_to_int16(a.into_inner()),
40            datum => Err(EvalError::InvalidJsonbCast {
41                from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
42                to: "smallint".into(),
43            }),
44        }
45    }
46);
47
48sqlfunc!(
49    #[sqlname = "jsonb_to_integer"]
50    #[is_monotone = true]
51    fn cast_jsonb_to_int32<'a>(a: JsonbRef<'a>) -> Result<i32, EvalError> {
52        match a.into_datum() {
53            Datum::Numeric(a) => cast_numeric_to_int32(a.into_inner()),
54            datum => Err(EvalError::InvalidJsonbCast {
55                from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
56                to: "integer".into(),
57            }),
58        }
59    }
60);
61
62sqlfunc!(
63    #[sqlname = "jsonb_to_bigint"]
64    #[is_monotone = true]
65    fn cast_jsonb_to_int64<'a>(a: JsonbRef<'a>) -> Result<i64, EvalError> {
66        match a.into_datum() {
67            Datum::Numeric(a) => cast_numeric_to_int64(a.into_inner()),
68            datum => Err(EvalError::InvalidJsonbCast {
69                from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
70                to: "bigint".into(),
71            }),
72        }
73    }
74);
75
76sqlfunc!(
77    #[sqlname = "jsonb_to_real"]
78    #[is_monotone = true]
79    fn cast_jsonb_to_float32<'a>(a: JsonbRef<'a>) -> Result<f32, EvalError> {
80        match a.into_datum() {
81            Datum::Numeric(a) => cast_numeric_to_float32(a.into_inner()),
82            datum => Err(EvalError::InvalidJsonbCast {
83                from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
84                to: "real".into(),
85            }),
86        }
87    }
88);
89
90sqlfunc!(
91    #[sqlname = "jsonb_to_double"]
92    #[is_monotone = true]
93    fn cast_jsonb_to_float64<'a>(a: JsonbRef<'a>) -> Result<f64, EvalError> {
94        match a.into_datum() {
95            Datum::Numeric(a) => cast_numeric_to_float64(a.into_inner()),
96            datum => Err(EvalError::InvalidJsonbCast {
97                from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
98                to: "double precision".into(),
99            }),
100        }
101    }
102);
103
104#[derive(
105    Arbitrary, Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect,
106)]
107pub struct CastJsonbToNumeric(pub Option<NumericMaxScale>);
108
109impl<'a> EagerUnaryFunc<'a> for CastJsonbToNumeric {
110    type Input = JsonbRef<'a>;
111    type Output = Result<Numeric, EvalError>;
112
113    fn call(&self, a: JsonbRef<'a>) -> Result<Numeric, EvalError> {
114        match a.into_datum() {
115            Datum::Numeric(mut num) => match self.0 {
116                None => Ok(num.into_inner()),
117                Some(scale) => {
118                    if numeric::rescale(&mut num.0, scale.into_u8()).is_err() {
119                        return Err(EvalError::NumericFieldOverflow);
120                    };
121                    Ok(num.into_inner())
122                }
123            },
124            datum => Err(EvalError::InvalidJsonbCast {
125                from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
126                to: "numeric".into(),
127            }),
128        }
129    }
130
131    fn output_type(&self, input: ColumnType) -> ColumnType {
132        ScalarType::Numeric { max_scale: self.0 }.nullable(input.nullable)
133    }
134
135    fn is_monotone(&self) -> bool {
136        true
137    }
138}
139
140impl fmt::Display for CastJsonbToNumeric {
141    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
142        f.write_str("jsonb_to_numeric")
143    }
144}
145
146sqlfunc!(
147    #[sqlname = "jsonb_to_boolean"]
148    #[is_monotone = true]
149    fn cast_jsonb_to_bool<'a>(a: JsonbRef<'a>) -> Result<bool, EvalError> {
150        match a.into_datum() {
151            Datum::True => Ok(true),
152            Datum::False => Ok(false),
153            datum => Err(EvalError::InvalidJsonbCast {
154                from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
155                to: "boolean".into(),
156            }),
157        }
158    }
159);
160
161sqlfunc!(
162    #[sqlname = "jsonbable_to_jsonb"]
163    fn cast_jsonbable_to_jsonb<'a>(a: JsonbRef<'a>) -> JsonbRef<'a> {
164        match a.into_datum() {
165            Datum::Numeric(n) => {
166                let n = n.into_inner();
167                let datum = if n.is_finite() {
168                    Datum::from(n)
169                } else if n.is_nan() {
170                    Datum::String("NaN")
171                } else if n.is_negative() {
172                    Datum::String("-Infinity")
173                } else {
174                    Datum::String("Infinity")
175                };
176                JsonbRef::from_datum(datum)
177            }
178            datum => JsonbRef::from_datum(datum),
179        }
180    }
181);
182
183sqlfunc!(
184    fn jsonb_array_length<'a>(a: JsonbRef<'a>) -> Result<Option<i32>, EvalError> {
185        match a.into_datum() {
186            Datum::List(list) => {
187                let count = list.iter().count();
188                match i32::try_from(count) {
189                    Ok(len) => Ok(Some(len)),
190                    Err(_) => Err(EvalError::Int32OutOfRange(count.to_string().into())),
191                }
192            }
193            _ => Ok(None),
194        }
195    }
196);
197
198sqlfunc!(
199    fn jsonb_typeof<'a>(a: JsonbRef<'a>) -> &'a str {
200        match a.into_datum() {
201            Datum::Map(_) => "object",
202            Datum::List(_) => "array",
203            Datum::String(_) => "string",
204            Datum::Numeric(_) => "number",
205            Datum::True | Datum::False => "boolean",
206            Datum::JsonNull => "null",
207            d => panic!("Not jsonb: {:?}", d),
208        }
209    }
210);
211
212sqlfunc!(
213    fn jsonb_strip_nulls<'a>(a: JsonbRef<'a>) -> Jsonb {
214        fn strip_nulls(a: Datum, row: &mut RowPacker) {
215            match a {
216                Datum::Map(dict) => row.push_dict_with(|row| {
217                    for (k, v) in dict.iter() {
218                        match v {
219                            Datum::JsonNull => (),
220                            _ => {
221                                row.push(Datum::String(k));
222                                strip_nulls(v, row);
223                            }
224                        }
225                    }
226                }),
227                Datum::List(list) => row.push_list_with(|row| {
228                    for elem in list.iter() {
229                        strip_nulls(elem, row);
230                    }
231                }),
232                _ => row.push(a),
233            }
234        }
235        let mut row = Row::default();
236        strip_nulls(a.into_datum(), &mut row.packer());
237        Jsonb::from_row(row)
238    }
239);
240
241sqlfunc!(
242    fn jsonb_pretty<'a>(a: JsonbRef<'a>) -> String {
243        let mut buf = String::new();
244        strconv::format_jsonb_pretty(&mut buf, a);
245        buf
246    }
247);