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<E = MirScalarExpr> {
99    pub cast_element: Box<E>,
100}
101
102impl<E: Eval> LazyUnaryFunc for CastListToJsonb<E> {
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<E> CastListToJsonb<E> {
159    /// Rebuilds this function with the element cast expression converted to
160    /// `E2`.
161    pub fn try_map_expr<'a, E2: TryFrom<&'a E>>(
162        &'a self,
163    ) -> Result<CastListToJsonb<E2>, E2::Error> {
164        Ok(CastListToJsonb {
165            cast_element: Box::new(E2::try_from(&*self.cast_element)?),
166        })
167    }
168
169    /// Rebuilds this function with the element cast expression converted to
170    /// `E2`.
171    pub fn map_expr<'a, E2: From<&'a E>>(&'a self) -> CastListToJsonb<E2> {
172        CastListToJsonb {
173            cast_element: Box::new(E2::from(&*self.cast_element)),
174        }
175    }
176}
177
178impl<E> fmt::Display for CastListToJsonb<E> {
179    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
180        f.write_str("listtojsonb")
181    }
182}
183
184/// Casts between two list types by casting each element of `a` ("list1") using
185/// `cast_expr` and collecting the results into a new list ("list2").
186#[derive(
187    Ord,
188    PartialOrd,
189    Clone,
190    Debug,
191    Eq,
192    PartialEq,
193    Serialize,
194    Deserialize,
195    Hash
196)]
197pub struct CastList1ToList2<E = MirScalarExpr> {
198    /// List2's type
199    pub return_ty: SqlScalarType,
200    /// The expression to cast List1's elements to List2's elements' type
201    pub cast_expr: Box<E>,
202}
203
204impl<E: Eval> LazyUnaryFunc for CastList1ToList2<E> {
205    fn eval<'a>(
206        &'a self,
207        datums: &[Datum<'a>],
208        temp_storage: &'a RowArena,
209        a: &'a impl Eval,
210    ) -> Result<Datum<'a>, EvalError> {
211        let a = a.eval(datums, temp_storage)?;
212        if a.is_null() {
213            return Ok(Datum::Null);
214        }
215        let mut cast_datums = Vec::new();
216        for el in a.unwrap_list().iter() {
217            // `cast_expr` is evaluated as an expression that casts the
218            // first column in `datums` (i.e. `datums[0]`) from the list elements'
219            // current type to a target type.
220            cast_datums.push(self.cast_expr.eval(&[el], temp_storage)?);
221        }
222
223        Ok(temp_storage.make_datum(|packer| packer.push_list(cast_datums)))
224    }
225
226    fn output_sql_type(&self, input_type: SqlColumnType) -> SqlColumnType {
227        self.return_ty
228            .without_modifiers()
229            .nullable(input_type.nullable)
230    }
231
232    fn propagates_nulls(&self) -> bool {
233        true
234    }
235
236    fn introduces_nulls(&self) -> bool {
237        false
238    }
239
240    fn preserves_uniqueness(&self) -> bool {
241        false
242    }
243
244    fn inverse(&self) -> Option<crate::UnaryFunc> {
245        // TODO: this could be figured out--might be easier after enum dispatch?
246        None
247    }
248
249    fn is_monotone(&self) -> bool {
250        false
251    }
252
253    fn is_eliminable_cast(&self) -> bool {
254        false
255    }
256}
257
258impl<E> CastList1ToList2<E> {
259    /// Rebuilds this function with the element cast expression converted to
260    /// `E2`.
261    pub fn try_map_expr<'a, E2: TryFrom<&'a E>>(
262        &'a self,
263    ) -> Result<CastList1ToList2<E2>, E2::Error> {
264        Ok(CastList1ToList2 {
265            return_ty: self.return_ty.clone(),
266            cast_expr: Box::new(E2::try_from(&*self.cast_expr)?),
267        })
268    }
269
270    /// Rebuilds this function with the element cast expression converted to
271    /// `E2`.
272    pub fn map_expr<'a, E2: From<&'a E>>(&'a self) -> CastList1ToList2<E2> {
273        CastList1ToList2 {
274            return_ty: self.return_ty.clone(),
275            cast_expr: Box::new(E2::from(&*self.cast_expr)),
276        }
277    }
278}
279
280impl<E> fmt::Display for CastList1ToList2<E> {
281    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
282        f.write_str("list1tolist2")
283    }
284}
285
286#[sqlfunc(sqlname = "list_length")]
287fn list_length<'a>(a: DatumList<'a>) -> Result<i32, EvalError> {
288    let count = a.iter().count();
289    count
290        .try_into()
291        .map_err(|_| EvalError::Int32OutOfRange(count.to_string().into()))
292}
293
294/// The `list_length_max` implementation.
295///
296/// We're not deriving `sqlfunc` here because we need to pass in the `max_layer` parameter.
297#[derive(
298    Ord,
299    PartialOrd,
300    Clone,
301    Debug,
302    Eq,
303    PartialEq,
304    Serialize,
305    Deserialize,
306    Hash
307)]
308pub struct ListLengthMax {
309    /// Maximal allowed layer to query.
310    pub max_layer: usize,
311}
312impl EagerBinaryFunc for ListLengthMax {
313    type Input<'a> = (DatumList<'a>, i64);
314    type Output<'a> = Result<Option<i32>, EvalError>;
315    // TODO(benesch): remove potentially dangerous usage of `as`.
316    #[allow(clippy::as_conversions)]
317    fn call<'a>(&self, (a, b): Self::Input<'a>, _: &'a RowArena) -> Self::Output<'a> {
318        fn max_len_on_layer(i: DatumList<'_>, on_layer: i64) -> Option<usize> {
319            let i = i.iter();
320            if on_layer > 1 {
321                let mut max_len = None;
322                for d in i {
323                    if let Datum::List(i) = d {
324                        max_len = std::cmp::max(max_len_on_layer(i, on_layer - 1), max_len);
325                    }
326                }
327                max_len
328            } else {
329                Some(i.count())
330            }
331        }
332        if b as usize > self.max_layer || b < 1 {
333            Err(EvalError::InvalidLayer {
334                max_layer: self.max_layer,
335                val: b,
336            })
337        } else {
338            match max_len_on_layer(a, b) {
339                Some(l) => match l.try_into() {
340                    Ok(c) => Ok(Some(c)),
341                    Err(_) => Err(EvalError::Int32OutOfRange(l.to_string().into())),
342                },
343                None => Ok(None),
344            }
345        }
346    }
347    fn output_sql_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
348        let output = Self::Output::as_column_type();
349        let propagates_nulls = self.propagates_nulls();
350        let nullable = output.nullable;
351        let input_nullable = input_types.iter().any(|t| t.nullable);
352        output.nullable(nullable || (propagates_nulls && input_nullable))
353    }
354}
355impl fmt::Display for ListLengthMax {
356    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
357        f.write_str("list_length_max")
358    }
359}
mz_expr/scalar/func/impls/list.rs

mz_expr/scalar/func/impls/
list.rs
