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 = true)]
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        true
96    }
97
98    fn inverse(&self) -> Option<crate::UnaryFunc> {
99        // TODO? If we moved typeconv into `expr` we could determine the right
100        // inverse of this.
101        None
102    }
103
104    fn is_monotone(&self) -> bool {
105        false
106    }
107
108    fn is_eliminable_cast(&self) -> bool {
109        false
110    }
111}
112
113impl<E> CastListToJsonb<E> {
114    /// Rebuilds this function with the element cast expression converted to
115    /// `E2`.
116    pub fn try_map_expr<'a, E2: TryFrom<&'a E>>(
117        &'a self,
118    ) -> Result<CastListToJsonb<E2>, E2::Error> {
119        Ok(CastListToJsonb {
120            cast_element: Box::new(E2::try_from(&*self.cast_element)?),
121        })
122    }
123
124    /// Rebuilds this function with the element cast expression converted to
125    /// `E2`.
126    pub fn map_expr<'a, E2: From<&'a E>>(&'a self) -> CastListToJsonb<E2> {
127        CastListToJsonb {
128            cast_element: Box::new(E2::from(&*self.cast_element)),
129        }
130    }
131}
132
133impl<E> fmt::Display for CastListToJsonb<E> {
134    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
135        f.write_str("listtojsonb")
136    }
137}
138
139/// Casts between two list types by casting each element of `a` ("list1") using
140/// `cast_expr` and collecting the results into a new list ("list2").
141#[derive(
142    Ord,
143    PartialOrd,
144    Clone,
145    Debug,
146    Eq,
147    PartialEq,
148    Serialize,
149    Deserialize,
150    Hash
151)]
152pub struct CastList1ToList2<E = MirScalarExpr> {
153    /// List2's type
154    pub return_ty: SqlScalarType,
155    /// The expression to cast List1's elements to List2's elements' type
156    pub cast_expr: Box<E>,
157}
158
159impl<E: Eval> LazyUnaryFunc for CastList1ToList2<E> {
160    fn eval<'a>(
161        &'a self,
162        datums: &[Datum<'a>],
163        temp_storage: &'a RowArena,
164        a: &'a impl Eval,
165    ) -> Result<Datum<'a>, EvalError> {
166        let a = a.eval(datums, temp_storage)?;
167        if a.is_null() {
168            return Ok(Datum::Null);
169        }
170        let mut cast_datums = Vec::new();
171        for el in a.unwrap_list().iter() {
172            // `cast_expr` is evaluated as an expression that casts the
173            // first column in `datums` (i.e. `datums[0]`) from the list elements'
174            // current type to a target type.
175            cast_datums.push(self.cast_expr.eval(&[el], temp_storage)?);
176        }
177
178        Ok(temp_storage.make_datum(|packer| packer.push_list(cast_datums)))
179    }
180
181    fn output_sql_type(&self, input_type: SqlColumnType) -> SqlColumnType {
182        self.return_ty
183            .without_modifiers()
184            .nullable(input_type.nullable)
185    }
186
187    fn propagates_nulls(&self) -> bool {
188        true
189    }
190
191    fn introduces_nulls(&self) -> bool {
192        false
193    }
194
195    fn preserves_uniqueness(&self) -> bool {
196        false
197    }
198
199    fn inverse(&self) -> Option<crate::UnaryFunc> {
200        // TODO: this could be figured out--might be easier after enum dispatch?
201        None
202    }
203
204    fn is_monotone(&self) -> bool {
205        false
206    }
207
208    fn is_eliminable_cast(&self) -> bool {
209        false
210    }
211}
212
213impl<E> CastList1ToList2<E> {
214    /// Rebuilds this function with the element cast expression converted to
215    /// `E2`.
216    pub fn try_map_expr<'a, E2: TryFrom<&'a E>>(
217        &'a self,
218    ) -> Result<CastList1ToList2<E2>, E2::Error> {
219        Ok(CastList1ToList2 {
220            return_ty: self.return_ty.clone(),
221            cast_expr: Box::new(E2::try_from(&*self.cast_expr)?),
222        })
223    }
224
225    /// Rebuilds this function with the element cast expression converted to
226    /// `E2`.
227    pub fn map_expr<'a, E2: From<&'a E>>(&'a self) -> CastList1ToList2<E2> {
228        CastList1ToList2 {
229            return_ty: self.return_ty.clone(),
230            cast_expr: Box::new(E2::from(&*self.cast_expr)),
231        }
232    }
233}
234
235impl<E> fmt::Display for CastList1ToList2<E> {
236    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
237        f.write_str("list1tolist2")
238    }
239}
240
241#[sqlfunc(sqlname = "list_length")]
242fn list_length<'a>(a: DatumList<'a>) -> Result<i32, EvalError> {
243    let count = a.iter().count();
244    count
245        .try_into()
246        .map_err(|_| EvalError::Int32OutOfRange(count.to_string().into()))
247}
248
249/// The `list_length_max` implementation.
250#[derive(
251    Ord,
252    PartialOrd,
253    Clone,
254    Debug,
255    Eq,
256    PartialEq,
257    Serialize,
258    Deserialize,
259    Hash
260)]
261pub struct ListLengthMax {
262    /// Maximal allowed layer to query.
263    pub max_layer: usize,
264}
265#[sqlfunc(ListLengthMax)]
266// TODO(benesch): remove potentially dangerous usage of `as`.
267#[allow(clippy::as_conversions)]
268fn list_length_max<'a>(&self, a: DatumList<'a>, b: i64) -> Result<Option<i32>, EvalError> {
269    fn max_len_on_layer(i: DatumList<'_>, on_layer: i64) -> Option<usize> {
270        let i = i.iter();
271        if on_layer > 1 {
272            let mut max_len = None;
273            for d in i {
274                if let Datum::List(i) = d {
275                    max_len = std::cmp::max(max_len_on_layer(i, on_layer - 1), max_len);
276                }
277            }
278            max_len
279        } else {
280            Some(i.count())
281        }
282    }
283    if b as usize > self.max_layer || b < 1 {
284        Err(EvalError::InvalidLayer {
285            max_layer: self.max_layer,
286            val: b,
287        })
288    } else {
289        match max_len_on_layer(a, b) {
290            Some(l) => match l.try_into() {
291                Ok(c) => Ok(Some(c)),
292                Err(_) => Err(EvalError::Int32OutOfRange(l.to_string().into())),
293            },
294            None => Ok(None),
295        }
296    }
297}
mz_expr/scalar/func/impls/list.rs

mz_expr/scalar/func/impls/
list.rs
