mz_expr/scalar/func/impls/
list.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::{Datum, DatumList, ExcludeNull, Row, RowArena, SqlColumnType, SqlScalarType};
14use serde::{Deserialize, Serialize};
15
16use crate::scalar::func::{LazyUnaryFunc, stringify_datum};
17use crate::{Eval, EvalError, MirScalarExpr};
18
19#[derive(
20    Ord,
21    PartialOrd,
22    Clone,
23    Debug,
24    Eq,
25    PartialEq,
26    Serialize,
27    Deserialize,
28    Hash
29)]
30pub struct CastListToString {
31    pub ty: SqlScalarType,
32}
33
34// TODO? if typeconv was in expr, we could determine the inverse of this cast
35#[sqlfunc(CastListToString, sqlname = "listtostr", preserves_uniqueness = super::text_preserves_uniqueness(&self.ty))]
36fn cast_list_to_string<'a>(&self, a: ExcludeNull<Datum<'a>>) -> Result<String, EvalError> {
37    let mut buf = String::new();
38    stringify_datum(&mut buf, *a, &self.ty)?;
39    Ok(buf)
40}
41
42#[derive(
43    Ord,
44    PartialOrd,
45    Clone,
46    Debug,
47    Eq,
48    PartialEq,
49    Serialize,
50    Deserialize,
51    Hash
52)]
53pub struct CastListToJsonb<E = MirScalarExpr> {
54    pub cast_element: Box<E>,
55}
56
57impl<E: Eval> LazyUnaryFunc for CastListToJsonb<E> {
58    fn eval<'a>(
59        &'a self,
60        datums: &[Datum<'a>],
61        temp_storage: &'a RowArena,
62        a: &'a impl Eval,
63    ) -> Result<Datum<'a>, EvalError> {
64        let a = a.eval(datums, temp_storage)?;
65        if a.is_null() {
66            return Ok(Datum::Null);
67        }
68        let mut row = Row::default();
69        row.packer().push_list_with(|packer| {
70            for elem in a.unwrap_list().iter() {
71                let elem = match self.cast_element.eval(&[elem], temp_storage)? {
72                    Datum::Null => Datum::JsonNull,
73                    d => d,
74                };
75                packer.push(elem);
76            }
77            Ok::<_, EvalError>(())
78        })?;
79        Ok(temp_storage.push_unary_row(row))
80    }
81
82    fn output_sql_type(&self, input_type: SqlColumnType) -> SqlColumnType {
83        SqlScalarType::Jsonb.nullable(input_type.nullable)
84    }
85
86    fn propagates_nulls(&self) -> bool {
87        true
88    }
89
90    fn introduces_nulls(&self) -> bool {
91        false
92    }
93
94    fn preserves_uniqueness(&self) -> bool {
95        // NOTE: `to_jsonb` maps distinct elements to the same JSONB value, for
96        // example `time` and timestamp elements through their text form, which
97        // rounds nanoseconds away. This cast is therefore not
98        // uniqueness-preserving.
99        false
100    }
101
102    fn inverse(&self) -> Option<crate::UnaryFunc> {
103        // TODO? If we moved typeconv into `expr` we could determine the right
104        // inverse of this.
105        None
106    }
107
108    fn is_monotone(&self) -> bool {
109        false
110    }
111
112    fn is_eliminable_cast(&self) -> bool {
113        false
114    }
115}
116
117impl<E> CastListToJsonb<E> {
118    /// Rebuilds this function with the element cast expression converted to
119    /// `E2`.
120    pub fn try_map_expr<'a, E2: TryFrom<&'a E>>(
121        &'a self,
122    ) -> Result<CastListToJsonb<E2>, E2::Error> {
123        Ok(CastListToJsonb {
124            cast_element: Box::new(E2::try_from(&*self.cast_element)?),
125        })
126    }
127
128    /// Rebuilds this function with the element cast expression converted to
129    /// `E2`.
130    pub fn map_expr<'a, E2: From<&'a E>>(&'a self) -> CastListToJsonb<E2> {
131        CastListToJsonb {
132            cast_element: Box::new(E2::from(&*self.cast_element)),
133        }
134    }
135}
136
137impl<E> fmt::Display for CastListToJsonb<E> {
138    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
139        f.write_str("listtojsonb")
140    }
141}
142
143/// Casts between two list types by casting each element of `a` ("list1") using
144/// `cast_expr` and collecting the results into a new list ("list2").
145#[derive(
146    Ord,
147    PartialOrd,
148    Clone,
149    Debug,
150    Eq,
151    PartialEq,
152    Serialize,
153    Deserialize,
154    Hash
155)]
156pub struct CastList1ToList2<E = MirScalarExpr> {
157    /// List2's type
158    pub return_ty: SqlScalarType,
159    /// The expression to cast List1's elements to List2's elements' type
160    pub cast_expr: Box<E>,
161}
162
163impl<E: Eval> LazyUnaryFunc for CastList1ToList2<E> {
164    fn eval<'a>(
165        &'a self,
166        datums: &[Datum<'a>],
167        temp_storage: &'a RowArena,
168        a: &'a impl Eval,
169    ) -> Result<Datum<'a>, EvalError> {
170        let a = a.eval(datums, temp_storage)?;
171        if a.is_null() {
172            return Ok(Datum::Null);
173        }
174        let mut cast_datums = Vec::new();
175        for el in a.unwrap_list().iter() {
176            // `cast_expr` is evaluated as an expression that casts the
177            // first column in `datums` (i.e. `datums[0]`) from the list elements'
178            // current type to a target type.
179            cast_datums.push(self.cast_expr.eval(&[el], temp_storage)?);
180        }
181
182        Ok(temp_storage.make_datum(|packer| packer.push_list(cast_datums)))
183    }
184
185    fn output_sql_type(&self, input_type: SqlColumnType) -> SqlColumnType {
186        self.return_ty
187            .without_modifiers()
188            .nullable(input_type.nullable)
189    }
190
191    fn propagates_nulls(&self) -> bool {
192        true
193    }
194
195    fn introduces_nulls(&self) -> bool {
196        false
197    }
198
199    fn preserves_uniqueness(&self) -> bool {
200        false
201    }
202
203    fn inverse(&self) -> Option<crate::UnaryFunc> {
204        // TODO: this could be figured out--might be easier after enum dispatch?
205        None
206    }
207
208    fn is_monotone(&self) -> bool {
209        false
210    }
211
212    fn is_eliminable_cast(&self) -> bool {
213        false
214    }
215}
216
217impl<E> CastList1ToList2<E> {
218    /// Rebuilds this function with the element cast expression converted to
219    /// `E2`.
220    pub fn try_map_expr<'a, E2: TryFrom<&'a E>>(
221        &'a self,
222    ) -> Result<CastList1ToList2<E2>, E2::Error> {
223        Ok(CastList1ToList2 {
224            return_ty: self.return_ty.clone(),
225            cast_expr: Box::new(E2::try_from(&*self.cast_expr)?),
226        })
227    }
228
229    /// Rebuilds this function with the element cast expression converted to
230    /// `E2`.
231    pub fn map_expr<'a, E2: From<&'a E>>(&'a self) -> CastList1ToList2<E2> {
232        CastList1ToList2 {
233            return_ty: self.return_ty.clone(),
234            cast_expr: Box::new(E2::from(&*self.cast_expr)),
235        }
236    }
237}
238
239impl<E> fmt::Display for CastList1ToList2<E> {
240    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
241        f.write_str("list1tolist2")
242    }
243}
244
245#[sqlfunc(sqlname = "list_length")]
246fn list_length<'a>(a: DatumList<'a>) -> Result<i32, EvalError> {
247    let count = a.iter().count();
248    count
249        .try_into()
250        .map_err(|_| EvalError::Int32OutOfRange(count.to_string().into()))
251}
252
253/// The `list_length_max` implementation.
254#[derive(
255    Ord,
256    PartialOrd,
257    Clone,
258    Debug,
259    Eq,
260    PartialEq,
261    Serialize,
262    Deserialize,
263    Hash
264)]
265pub struct ListLengthMax {
266    /// Maximal allowed layer to query.
267    pub max_layer: usize,
268}
269#[sqlfunc(ListLengthMax)]
270// TODO(benesch): remove potentially dangerous usage of `as`.
271#[allow(clippy::as_conversions)]
272fn list_length_max<'a>(&self, a: DatumList<'a>, b: i64) -> Result<Option<i32>, EvalError> {
273    fn max_len_on_layer(i: DatumList<'_>, on_layer: i64) -> Option<usize> {
274        let i = i.iter();
275        if on_layer > 1 {
276            let mut max_len = None;
277            for d in i {
278                if let Datum::List(i) = d {
279                    max_len = std::cmp::max(max_len_on_layer(i, on_layer - 1), max_len);
280                }
281            }
282            max_len
283        } else {
284            Some(i.count())
285        }
286    }
287    if b as usize > self.max_layer || b < 1 {
288        Err(EvalError::InvalidLayer {
289            max_layer: self.max_layer,
290            val: b,
291        })
292    } else {
293        match max_len_on_layer(a, b) {
294            Some(l) => match l.try_into() {
295                Ok(c) => Ok(Some(c)),
296                Err(_) => Err(EvalError::Int32OutOfRange(l.to_string().into())),
297            },
298            None => Ok(None),
299        }
300    }
301}
mz_expr/scalar/func/impls/list.rs

mz_expr/scalar/func/impls/
list.rs
