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_lowertest::MzReflect;
13use mz_repr::{ColumnType, Datum, Row, RowArena, ScalarType};
14use proptest_derive::Arbitrary;
15use serde::{Deserialize, Serialize};
16
17use crate::scalar::func::{LazyUnaryFunc, stringify_datum};
18use crate::{EvalError, MirScalarExpr};
19
20#[derive(
21    Arbitrary, Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect,
22)]
23pub struct CastListToString {
24    pub ty: ScalarType,
25}
26
27impl LazyUnaryFunc for CastListToString {
28    fn eval<'a>(
29        &'a self,
30        datums: &[Datum<'a>],
31        temp_storage: &'a RowArena,
32        a: &'a MirScalarExpr,
33    ) -> Result<Datum<'a>, EvalError> {
34        let a = a.eval(datums, temp_storage)?;
35        if a.is_null() {
36            return Ok(Datum::Null);
37        }
38        let mut buf = String::new();
39        stringify_datum(&mut buf, a, &self.ty)?;
40        Ok(Datum::String(temp_storage.push_string(buf)))
41    }
42
43    fn output_type(&self, input_type: ColumnType) -> ColumnType {
44        ScalarType::String.nullable(input_type.nullable)
45    }
46
47    fn propagates_nulls(&self) -> bool {
48        true
49    }
50
51    fn introduces_nulls(&self) -> bool {
52        false
53    }
54
55    fn preserves_uniqueness(&self) -> bool {
56        true
57    }
58
59    fn inverse(&self) -> Option<crate::UnaryFunc> {
60        // TODO? if typeconv was in expr, we could determine this
61        None
62    }
63
64    fn is_monotone(&self) -> bool {
65        false
66    }
67}
68
69impl fmt::Display for CastListToString {
70    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
71        f.write_str("listtostr")
72    }
73}
74
75#[derive(
76    Arbitrary, Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect,
77)]
78pub struct CastListToJsonb {
79    pub cast_element: Box<MirScalarExpr>,
80}
81
82impl LazyUnaryFunc for CastListToJsonb {
83    fn eval<'a>(
84        &'a self,
85        datums: &[Datum<'a>],
86        temp_storage: &'a RowArena,
87        a: &'a MirScalarExpr,
88    ) -> Result<Datum<'a>, EvalError> {
89        let a = a.eval(datums, temp_storage)?;
90        if a.is_null() {
91            return Ok(Datum::Null);
92        }
93        let mut row = Row::default();
94        row.packer().push_list_with(|packer| {
95            for elem in a.unwrap_list().iter() {
96                let elem = match self.cast_element.eval(&[elem], temp_storage)? {
97                    Datum::Null => Datum::JsonNull,
98                    d => d,
99                };
100                packer.push(elem);
101            }
102            Ok::<_, EvalError>(())
103        })?;
104        Ok(temp_storage.push_unary_row(row))
105    }
106
107    fn output_type(&self, input_type: ColumnType) -> ColumnType {
108        ScalarType::Jsonb.nullable(input_type.nullable)
109    }
110
111    fn propagates_nulls(&self) -> bool {
112        true
113    }
114
115    fn introduces_nulls(&self) -> bool {
116        false
117    }
118
119    fn preserves_uniqueness(&self) -> bool {
120        true
121    }
122
123    fn inverse(&self) -> Option<crate::UnaryFunc> {
124        // TODO? If we moved typeconv into `expr` we could determine the right
125        // inverse of this.
126        None
127    }
128
129    fn is_monotone(&self) -> bool {
130        false
131    }
132}
133
134impl fmt::Display for CastListToJsonb {
135    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
136        f.write_str("listtojsonb")
137    }
138}
139
140/// Casts between two list types by casting each element of `a` ("list1") using
141/// `cast_expr` and collecting the results into a new list ("list2").
142#[derive(Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect)]
143pub struct CastList1ToList2 {
144    /// List2's type
145    pub return_ty: ScalarType,
146    /// The expression to cast List1's elements to List2's elements' type
147    pub cast_expr: Box<MirScalarExpr>,
148}
149
150impl LazyUnaryFunc for CastList1ToList2 {
151    fn eval<'a>(
152        &'a self,
153        datums: &[Datum<'a>],
154        temp_storage: &'a RowArena,
155        a: &'a MirScalarExpr,
156    ) -> Result<Datum<'a>, EvalError> {
157        let a = a.eval(datums, temp_storage)?;
158        if a.is_null() {
159            return Ok(Datum::Null);
160        }
161        let mut cast_datums = Vec::new();
162        for el in a.unwrap_list().iter() {
163            // `cast_expr` is evaluated as an expression that casts the
164            // first column in `datums` (i.e. `datums[0]`) from the list elements'
165            // current type to a target type.
166            cast_datums.push(self.cast_expr.eval(&[el], temp_storage)?);
167        }
168
169        Ok(temp_storage.make_datum(|packer| packer.push_list(cast_datums)))
170    }
171
172    fn output_type(&self, input_type: ColumnType) -> ColumnType {
173        self.return_ty
174            .without_modifiers()
175            .nullable(input_type.nullable)
176    }
177
178    fn propagates_nulls(&self) -> bool {
179        true
180    }
181
182    fn introduces_nulls(&self) -> bool {
183        false
184    }
185
186    fn preserves_uniqueness(&self) -> bool {
187        false
188    }
189
190    fn inverse(&self) -> Option<crate::UnaryFunc> {
191        // TODO: this could be figured out--might be easier after enum dispatch?
192        None
193    }
194
195    fn is_monotone(&self) -> bool {
196        false
197    }
198}
199
200impl fmt::Display for CastList1ToList2 {
201    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
202        f.write_str("list1tolist2")
203    }
204}
205
206#[derive(
207    Arbitrary, Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect,
208)]
209pub struct ListLength;
210
211impl LazyUnaryFunc for ListLength {
212    fn eval<'a>(
213        &'a self,
214        datums: &[Datum<'a>],
215        temp_storage: &'a RowArena,
216        a: &'a MirScalarExpr,
217    ) -> Result<Datum<'a>, EvalError> {
218        let a = a.eval(datums, temp_storage)?;
219        if a.is_null() {
220            return Ok(Datum::Null);
221        }
222        let count = a.unwrap_list().iter().count();
223        match count.try_into() {
224            Ok(c) => Ok(Datum::Int32(c)),
225            Err(_) => Err(EvalError::Int32OutOfRange(count.to_string().into())),
226        }
227    }
228
229    fn output_type(&self, input_type: ColumnType) -> ColumnType {
230        ScalarType::Int32.nullable(input_type.nullable)
231    }
232
233    fn propagates_nulls(&self) -> bool {
234        true
235    }
236
237    fn introduces_nulls(&self) -> bool {
238        false
239    }
240
241    fn preserves_uniqueness(&self) -> bool {
242        false
243    }
244
245    fn inverse(&self) -> Option<crate::UnaryFunc> {
246        None
247    }
248
249    fn is_monotone(&self) -> bool {
250        false
251    }
252}
253
254impl fmt::Display for ListLength {
255    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
256        f.write_str("list_length")
257    }
258}
mz_expr/scalar/func/impls/list.rs

mz_expr/scalar/func/impls/
list.rs
