Skip to main content

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 mz_expr_derive::sqlfunc;
11use mz_repr::adt::jsonb::{Jsonb, JsonbRef};
12use mz_repr::adt::numeric::{self, Numeric, NumericMaxScale};
13use mz_repr::{Datum, Row, RowPacker, SqlScalarType, strconv};
14use serde::{Deserialize, Serialize};
15
16use crate::EvalError;
17use crate::scalar::func::impls::numeric::*;
18
19#[sqlfunc(
20    sqlname = "jsonb_to_text",
21    preserves_uniqueness = false,
22    inverse = super::CastStringToJsonb
23)]
24pub fn cast_jsonb_to_string<'a>(a: JsonbRef<'a>) -> String {
25    let mut buf = String::new();
26    strconv::format_jsonb(&mut buf, a);
27    buf
28}
29
30#[sqlfunc(sqlname = "jsonb_to_smallint", is_monotone = true)]
31fn cast_jsonb_to_int16<'a>(a: JsonbRef<'a>) -> Result<i16, EvalError> {
32    match a.into_datum() {
33        Datum::Numeric(a) => cast_numeric_to_int16(a.into_inner()),
34        datum => Err(EvalError::InvalidJsonbCast {
35            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
36            to: "smallint".into(),
37        }),
38    }
39}
40
41#[sqlfunc(sqlname = "jsonb_to_integer", is_monotone = true)]
42fn cast_jsonb_to_int32<'a>(a: JsonbRef<'a>) -> Result<i32, EvalError> {
43    match a.into_datum() {
44        Datum::Numeric(a) => cast_numeric_to_int32(a.into_inner()),
45        datum => Err(EvalError::InvalidJsonbCast {
46            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
47            to: "integer".into(),
48        }),
49    }
50}
51
52#[sqlfunc(sqlname = "jsonb_to_bigint", is_monotone = true)]
53fn cast_jsonb_to_int64<'a>(a: JsonbRef<'a>) -> Result<i64, EvalError> {
54    match a.into_datum() {
55        Datum::Numeric(a) => cast_numeric_to_int64(a.into_inner()),
56        datum => Err(EvalError::InvalidJsonbCast {
57            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
58            to: "bigint".into(),
59        }),
60    }
61}
62
63#[sqlfunc(sqlname = "jsonb_to_real", is_monotone = true)]
64fn cast_jsonb_to_float32<'a>(a: JsonbRef<'a>) -> Result<f32, EvalError> {
65    match a.into_datum() {
66        Datum::Numeric(a) => cast_numeric_to_float32(a.into_inner()),
67        datum => Err(EvalError::InvalidJsonbCast {
68            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
69            to: "real".into(),
70        }),
71    }
72}
73
74#[sqlfunc(sqlname = "jsonb_to_double", is_monotone = true)]
75fn cast_jsonb_to_float64<'a>(a: JsonbRef<'a>) -> Result<f64, EvalError> {
76    match a.into_datum() {
77        Datum::Numeric(a) => cast_numeric_to_float64(a.into_inner()),
78        datum => Err(EvalError::InvalidJsonbCast {
79            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
80            to: "double precision".into(),
81        }),
82    }
83}
84
85#[derive(
86    Ord,
87    PartialOrd,
88    Clone,
89    Debug,
90    Eq,
91    PartialEq,
92    Serialize,
93    Deserialize,
94    Hash
95)]
96pub struct CastJsonbToNumeric(pub Option<NumericMaxScale>);
97
98#[sqlfunc(
99    CastJsonbToNumeric,
100    sqlname = "jsonb_to_numeric",
101    is_monotone = true,
102    output_type_expr = SqlScalarType::Numeric { max_scale: self.0 }
103        .nullable(input_type.nullable)
104)]
105fn cast_jsonb_to_numeric<'a>(&self, a: JsonbRef<'a>) -> Result<Numeric, EvalError> {
106    match a.into_datum() {
107        Datum::Numeric(mut num) => match self.0 {
108            None => Ok(num.into_inner()),
109            Some(scale) => {
110                if numeric::rescale(&mut num.0, scale.into_u8()).is_err() {
111                    return Err(EvalError::NumericFieldOverflow);
112                };
113                Ok(num.into_inner())
114            }
115        },
116        datum => Err(EvalError::InvalidJsonbCast {
117            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
118            to: "numeric".into(),
119        }),
120    }
121}
122
123#[sqlfunc(sqlname = "jsonb_to_boolean", is_monotone = true)]
124fn cast_jsonb_to_bool<'a>(a: JsonbRef<'a>) -> Result<bool, EvalError> {
125    match a.into_datum() {
126        Datum::True => Ok(true),
127        Datum::False => Ok(false),
128        datum => Err(EvalError::InvalidJsonbCast {
129            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
130            to: "boolean".into(),
131        }),
132    }
133}
134
135#[sqlfunc(sqlname = "jsonbable_to_jsonb")]
136fn cast_jsonbable_to_jsonb<'a>(a: JsonbRef<'a>) -> JsonbRef<'a> {
137    match a.into_datum() {
138        Datum::Numeric(n) => {
139            let n = n.into_inner();
140            let datum = if n.is_finite() {
141                Datum::from(n)
142            } else if n.is_nan() {
143                Datum::String("NaN")
144            } else if n.is_negative() {
145                Datum::String("-Infinity")
146            } else {
147                Datum::String("Infinity")
148            };
149            JsonbRef::from_datum(datum)
150        }
151        datum => JsonbRef::from_datum(datum),
152    }
153}
154
155#[sqlfunc]
156fn jsonb_array_length<'a>(a: JsonbRef<'a>) -> Result<Option<i32>, EvalError> {
157    match a.into_datum() {
158        Datum::List(list) => {
159            let count = list.iter().count();
160            match i32::try_from(count) {
161                Ok(len) => Ok(Some(len)),
162                Err(_) => Err(EvalError::Int32OutOfRange(count.to_string().into())),
163            }
164        }
165        _ => Ok(None),
166    }
167}
168
169#[sqlfunc]
170fn jsonb_typeof<'a>(a: JsonbRef<'a>) -> &'a str {
171    match a.into_datum() {
172        Datum::Map(_) => "object",
173        Datum::List(_) => "array",
174        Datum::String(_) => "string",
175        Datum::Numeric(_) => "number",
176        Datum::True | Datum::False => "boolean",
177        Datum::JsonNull => "null",
178        d => panic!("Not jsonb: {:?}", d),
179    }
180}
181
182#[sqlfunc]
183fn jsonb_strip_nulls<'a>(a: JsonbRef<'a>) -> Jsonb {
184    fn strip_nulls(a: Datum, row: &mut RowPacker) {
185        match a {
186            Datum::Map(dict) => row.push_dict_with(|row| {
187                for (k, v) in dict.iter() {
188                    match v {
189                        Datum::JsonNull => (),
190                        _ => {
191                            row.push(Datum::String(k));
192                            strip_nulls(v, row);
193                        }
194                    }
195                }
196            }),
197            Datum::List(list) => row.push_list_with(|row| {
198                for elem in list.iter() {
199                    strip_nulls(elem, row);
200                }
201            }),
202            _ => row.push(a),
203        }
204    }
205    let mut row = Row::default();
206    strip_nulls(a.into_datum(), &mut row.packer());
207    Jsonb::from_row(row)
208}
209
210// NOTE: no budget pre-check, see the exception on `crate::func::check_build_fits_budget`.
211#[sqlfunc]
212fn jsonb_pretty<'a>(a: JsonbRef<'a>) -> String {
213    let mut buf = String::new();
214    strconv::format_jsonb_pretty(&mut buf, a);
215    buf
216}