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

mz_expr/scalar/func/impls/
array.rs
