mz_expr/scalar/func/impls/
array.rs

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

mz_expr/scalar/func/impls/
array.rs
