mz_expr/scalar/func/impls/
array.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::array::{Array, ArrayDimension};
15use mz_repr::{Datum, DatumList, Row, RowArena, RowPacker, SqlColumnType, SqlScalarType};
16use serde::{Deserialize, Serialize};
17
18use crate::scalar::func::{LazyUnaryFunc, stringify_datum};
19use crate::{EvalError, MirScalarExpr};
20
21#[sqlfunc(
22    sqlname = "arraytolist",
23    preserves_uniqueness = true,
24    introduces_nulls = false,
25    output_type_expr = SqlScalarType::List {
26        element_type: Box::new(input_type.scalar_type.unwrap_array_element_type().clone()),
27        custom_id: None,
28    }.nullable(true)
29)]
30fn cast_array_to_list_one_dim<'a>(a: Array<'a>) -> Result<DatumList<'a>, EvalError> {
31    let ndims = a.dims().ndims();
32    if ndims > 1 {
33        return Err(EvalError::Unsupported {
34            feature: format!(
35                "casting multi-dimensional array to list; got array with {} dimensions",
36                ndims
37            )
38            .into(),
39            discussion_no: None,
40        });
41    }
42    Ok(a.elements())
43}
44
45#[derive(
46    Ord,
47    PartialOrd,
48    Clone,
49    Debug,
50    Eq,
51    PartialEq,
52    Serialize,
53    Deserialize,
54    Hash,
55    MzReflect
56)]
57pub struct CastArrayToString {
58    pub ty: SqlScalarType,
59}
60
61impl LazyUnaryFunc for CastArrayToString {
62    fn eval<'a>(
63        &'a self,
64        datums: &[Datum<'a>],
65        temp_storage: &'a RowArena,
66        a: &'a MirScalarExpr,
67    ) -> Result<Datum<'a>, EvalError> {
68        let a = a.eval(datums, temp_storage)?;
69        if a.is_null() {
70            return Ok(Datum::Null);
71        }
72        let mut buf = String::new();
73        stringify_datum(&mut buf, a, &self.ty)?;
74        Ok(Datum::String(temp_storage.push_string(buf)))
75    }
76
77    fn output_type(&self, input_type: SqlColumnType) -> SqlColumnType {
78        SqlScalarType::String.nullable(input_type.nullable)
79    }
80
81    fn propagates_nulls(&self) -> bool {
82        true
83    }
84
85    fn introduces_nulls(&self) -> bool {
86        false
87    }
88
89    fn preserves_uniqueness(&self) -> bool {
90        true
91    }
92
93    fn inverse(&self) -> Option<crate::UnaryFunc> {
94        // TODO? If we moved typeconv into `expr` we could determine the right
95        // inverse of this.
96        None
97    }
98
99    fn is_monotone(&self) -> bool {
100        false
101    }
102}
103
104impl fmt::Display for CastArrayToString {
105    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
106        f.write_str("arraytostr")
107    }
108}
109
110#[derive(
111    Ord,
112    PartialOrd,
113    Clone,
114    Debug,
115    Eq,
116    PartialEq,
117    Serialize,
118    Deserialize,
119    Hash,
120    MzReflect
121)]
122pub struct CastArrayToJsonb {
123    pub cast_element: Box<MirScalarExpr>,
124}
125
126impl LazyUnaryFunc for CastArrayToJsonb {
127    fn eval<'a>(
128        &'a self,
129        datums: &[Datum<'a>],
130        temp_storage: &'a RowArena,
131        a: &'a MirScalarExpr,
132    ) -> Result<Datum<'a>, EvalError> {
133        fn pack<'a>(
134            temp_storage: &RowArena,
135            elems: &mut impl Iterator<Item = Datum<'a>>,
136            dims: &[ArrayDimension],
137            cast_element: &MirScalarExpr,
138            packer: &mut RowPacker,
139        ) -> Result<(), EvalError> {
140            packer.push_list_with(|packer| match dims {
141                [] => Ok(()),
142                [dim] => {
143                    for _ in 0..dim.length {
144                        let elem = elems.next().unwrap();
145                        let elem = match cast_element.eval(&[elem], temp_storage)? {
146                            Datum::Null => Datum::JsonNull,
147                            d => d,
148                        };
149                        packer.push(elem);
150                    }
151                    Ok(())
152                }
153                [dim, rest @ ..] => {
154                    for _ in 0..dim.length {
155                        pack(temp_storage, elems, rest, cast_element, packer)?;
156                    }
157                    Ok(())
158                }
159            })
160        }
161
162        let a = a.eval(datums, temp_storage)?;
163        if a.is_null() {
164            return Ok(Datum::Null);
165        }
166        let a = a.unwrap_array();
167        let elements = a.elements();
168        let dims = a.dims().into_iter().collect::<Vec<_>>();
169        let mut row = Row::default();
170        pack(
171            temp_storage,
172            &mut elements.into_iter(),
173            &dims,
174            &self.cast_element,
175            &mut row.packer(),
176        )?;
177        Ok(temp_storage.push_unary_row(row))
178    }
179
180    fn output_type(&self, input_type: SqlColumnType) -> SqlColumnType {
181        SqlScalarType::Jsonb.nullable(input_type.nullable)
182    }
183
184    fn propagates_nulls(&self) -> bool {
185        true
186    }
187
188    fn introduces_nulls(&self) -> bool {
189        false
190    }
191
192    fn preserves_uniqueness(&self) -> bool {
193        true
194    }
195
196    fn inverse(&self) -> Option<crate::UnaryFunc> {
197        // TODO? If we moved typeconv into `expr` we could determine the right
198        // inverse of this.
199        None
200    }
201
202    fn is_monotone(&self) -> bool {
203        false
204    }
205}
206
207impl fmt::Display for CastArrayToJsonb {
208    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
209        f.write_str("arraytojsonb")
210    }
211}
212
213/// Casts an array of one type to an array of another type. Does so by casting
214/// each element of the first array to the desired inner type and collecting
215/// the results into a new array.
216#[derive(
217    Ord,
218    PartialOrd,
219    Clone,
220    Debug,
221    Eq,
222    PartialEq,
223    Serialize,
224    Deserialize,
225    Hash,
226    MzReflect
227)]
228pub struct CastArrayToArray {
229    pub return_ty: SqlScalarType,
230    pub cast_expr: Box<MirScalarExpr>,
231}
232
233impl LazyUnaryFunc for CastArrayToArray {
234    fn eval<'a>(
235        &'a self,
236        datums: &[Datum<'a>],
237        temp_storage: &'a RowArena,
238        a: &'a MirScalarExpr,
239    ) -> Result<Datum<'a>, EvalError> {
240        let a = a.eval(datums, temp_storage)?;
241        if a.is_null() {
242            return Ok(Datum::Null);
243        }
244
245        let arr = a.unwrap_array();
246        let dims = arr.dims().into_iter().collect::<Vec<ArrayDimension>>();
247
248        let casted_datums = arr
249            .elements()
250            .iter()
251            .map(|datum| self.cast_expr.eval(&[datum], temp_storage))
252            .collect::<Result<Vec<Datum<'a>>, EvalError>>()?;
253
254        Ok(temp_storage.try_make_datum(|packer| packer.try_push_array(&dims, casted_datums))?)
255    }
256
257    fn output_type(&self, _input_type: SqlColumnType) -> SqlColumnType {
258        self.return_ty.clone().nullable(true)
259    }
260
261    fn propagates_nulls(&self) -> bool {
262        true
263    }
264
265    fn introduces_nulls(&self) -> bool {
266        false
267    }
268
269    fn preserves_uniqueness(&self) -> bool {
270        false
271    }
272
273    fn inverse(&self) -> Option<crate::UnaryFunc> {
274        None
275    }
276
277    fn is_monotone(&self) -> bool {
278        false
279    }
280}
281
282impl fmt::Display for CastArrayToArray {
283    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
284        f.write_str("arraytoarray")
285    }
286}
mz_expr/scalar/func/impls/array.rs

mz_expr/scalar/func/impls/
array.rs
