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

mz_expr/scalar/func/impls/
list.rs
