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(Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect)]
90pub struct CastJsonbToNumeric(pub Option<NumericMaxScale>);
91
92impl<'a> EagerUnaryFunc<'a> for CastJsonbToNumeric {
93    type Input = JsonbRef<'a>;
94    type Output = Result<Numeric, EvalError>;
95
96    fn call(&self, a: JsonbRef<'a>) -> Result<Numeric, EvalError> {
97        match a.into_datum() {
98            Datum::Numeric(mut num) => match self.0 {
99                None => Ok(num.into_inner()),
100                Some(scale) => {
101                    if numeric::rescale(&mut num.0, scale.into_u8()).is_err() {
102                        return Err(EvalError::NumericFieldOverflow);
103                    };
104                    Ok(num.into_inner())
105                }
106            },
107            datum => Err(EvalError::InvalidJsonbCast {
108                from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
109                to: "numeric".into(),
110            }),
111        }
112    }
113
114    fn output_type(&self, input: SqlColumnType) -> SqlColumnType {
115        SqlScalarType::Numeric { max_scale: self.0 }.nullable(input.nullable)
116    }
117
118    fn is_monotone(&self) -> bool {
119        true
120    }
121}
122
123impl fmt::Display for CastJsonbToNumeric {
124    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
125        f.write_str("jsonb_to_numeric")
126    }
127}
128
129#[sqlfunc(sqlname = "jsonb_to_boolean", is_monotone = true)]
130fn cast_jsonb_to_bool<'a>(a: JsonbRef<'a>) -> Result<bool, EvalError> {
131    match a.into_datum() {
132        Datum::True => Ok(true),
133        Datum::False => Ok(false),
134        datum => Err(EvalError::InvalidJsonbCast {
135            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
136            to: "boolean".into(),
137        }),
138    }
139}
140
141#[sqlfunc(sqlname = "jsonbable_to_jsonb")]
142fn cast_jsonbable_to_jsonb<'a>(a: JsonbRef<'a>) -> JsonbRef<'a> {
143    match a.into_datum() {
144        Datum::Numeric(n) => {
145            let n = n.into_inner();
146            let datum = if n.is_finite() {
147                Datum::from(n)
148            } else if n.is_nan() {
149                Datum::String("NaN")
150            } else if n.is_negative() {
151                Datum::String("-Infinity")
152            } else {
153                Datum::String("Infinity")
154            };
155            JsonbRef::from_datum(datum)
156        }
157        datum => JsonbRef::from_datum(datum),
158    }
159}
160
161#[sqlfunc]
162fn jsonb_array_length<'a>(a: JsonbRef<'a>) -> Result<Option<i32>, EvalError> {
163    match a.into_datum() {
164        Datum::List(list) => {
165            let count = list.iter().count();
166            match i32::try_from(count) {
167                Ok(len) => Ok(Some(len)),
168                Err(_) => Err(EvalError::Int32OutOfRange(count.to_string().into())),
169            }
170        }
171        _ => Ok(None),
172    }
173}
174
175#[sqlfunc]
176fn jsonb_typeof<'a>(a: JsonbRef<'a>) -> &'a str {
177    match a.into_datum() {
178        Datum::Map(_) => "object",
179        Datum::List(_) => "array",
180        Datum::String(_) => "string",
181        Datum::Numeric(_) => "number",
182        Datum::True | Datum::False => "boolean",
183        Datum::JsonNull => "null",
184        d => panic!("Not jsonb: {:?}", d),
185    }
186}
187
188#[sqlfunc]
189fn jsonb_strip_nulls<'a>(a: JsonbRef<'a>) -> Jsonb {
190    fn strip_nulls(a: Datum, row: &mut RowPacker) {
191        match a {
192            Datum::Map(dict) => row.push_dict_with(|row| {
193                for (k, v) in dict.iter() {
194                    match v {
195                        Datum::JsonNull => (),
196                        _ => {
197                            row.push(Datum::String(k));
198                            strip_nulls(v, row);
199                        }
200                    }
201                }
202            }),
203            Datum::List(list) => row.push_list_with(|row| {
204                for elem in list.iter() {
205                    strip_nulls(elem, row);
206                }
207            }),
208            _ => row.push(a),
209        }
210    }
211    let mut row = Row::default();
212    strip_nulls(a.into_datum(), &mut row.packer());
213    Jsonb::from_row(row)
214}
215
216#[sqlfunc]
217fn jsonb_pretty<'a>(a: JsonbRef<'a>) -> String {
218    let mut buf = String::new();
219    strconv::format_jsonb_pretty(&mut buf, a);
220    buf
221}