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

mz_expr/scalar/func/impls/
array.rs
