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