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

mz_expr/scalar/func/impls/
list.rs
