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