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

mz_expr/scalar/func/impls/
array.rs
