mz_expr/scalar/func/impls/
map.rs

Help
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 itertools::Itertools;
13use mz_lowertest::MzReflect;
14use mz_repr::{ColumnType, Datum, RowArena, ScalarType};
15use proptest_derive::Arbitrary;
16use serde::{Deserialize, Serialize};
17
18use crate::scalar::func::{LazyUnaryFunc, stringify_datum};
19use crate::{EvalError, MirScalarExpr};
20
21#[derive(
22    Arbitrary, Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect,
23)]
24pub struct CastMapToString {
25    pub ty: ScalarType,
26}
27
28impl LazyUnaryFunc for CastMapToString {
29    fn eval<'a>(
30        &'a self,
31        datums: &[Datum<'a>],
32        temp_storage: &'a RowArena,
33        a: &'a MirScalarExpr,
34    ) -> Result<Datum<'a>, EvalError> {
35        let a = a.eval(datums, temp_storage)?;
36        if a.is_null() {
37            return Ok(Datum::Null);
38        }
39        let mut buf = String::new();
40        stringify_datum(&mut buf, a, &self.ty)?;
41        Ok(Datum::String(temp_storage.push_string(buf)))
42    }
43
44    fn output_type(&self, input_type: ColumnType) -> ColumnType {
45        ScalarType::String.nullable(input_type.nullable)
46    }
47
48    fn propagates_nulls(&self) -> bool {
49        true
50    }
51
52    fn introduces_nulls(&self) -> bool {
53        false
54    }
55
56    fn preserves_uniqueness(&self) -> bool {
57        true
58    }
59
60    fn inverse(&self) -> Option<crate::UnaryFunc> {
61        // TODO? If we moved typeconv into expr, we could evaluate this
62        None
63    }
64
65    fn is_monotone(&self) -> bool {
66        false
67    }
68}
69
70impl fmt::Display for CastMapToString {
71    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
72        f.write_str("maptostr")
73    }
74}
75
76#[derive(
77    Arbitrary, Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect,
78)]
79pub struct MapLength;
80
81impl LazyUnaryFunc for MapLength {
82    fn eval<'a>(
83        &'a self,
84        datums: &[Datum<'a>],
85        temp_storage: &'a RowArena,
86        a: &'a MirScalarExpr,
87    ) -> Result<Datum<'a>, EvalError> {
88        let a = a.eval(datums, temp_storage)?;
89        if a.is_null() {
90            return Ok(Datum::Null);
91        }
92        let count = a.unwrap_map().iter().count();
93        match count.try_into() {
94            Ok(c) => Ok(Datum::Int32(c)),
95            Err(_) => Err(EvalError::Int32OutOfRange(count.to_string().into())),
96        }
97    }
98
99    fn output_type(&self, input_type: ColumnType) -> ColumnType {
100        ScalarType::Int32.nullable(input_type.nullable)
101    }
102
103    fn propagates_nulls(&self) -> bool {
104        true
105    }
106
107    fn introduces_nulls(&self) -> bool {
108        false
109    }
110
111    fn preserves_uniqueness(&self) -> bool {
112        false
113    }
114
115    fn inverse(&self) -> Option<crate::UnaryFunc> {
116        None
117    }
118
119    fn is_monotone(&self) -> bool {
120        false
121    }
122}
123
124impl fmt::Display for MapLength {
125    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
126        f.write_str("map_length")
127    }
128}
129
130#[derive(
131    Arbitrary, Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, MzReflect,
132)]
133pub struct MapBuildFromRecordList {
134    pub value_type: ScalarType,
135}
136
137impl LazyUnaryFunc for MapBuildFromRecordList {
138    fn eval<'a>(
139        &'a self,
140        datums: &[Datum<'a>],
141        temp_storage: &'a RowArena,
142        a: &'a MirScalarExpr,
143    ) -> Result<Datum<'a>, EvalError> {
144        let a = a.eval(datums, temp_storage)?;
145        if a.is_null() {
146            return Ok(Datum::Null);
147        }
148        let list = a.unwrap_list();
149        let mut map = std::collections::BTreeMap::new();
150
151        for i in list.iter() {
152            if i.is_null() {
153                continue;
154            }
155
156            for (k, v) in i.unwrap_list().iter().tuples() {
157                if k.is_null() {
158                    continue;
159                }
160                map.insert(k.unwrap_str(), v);
161            }
162        }
163
164        let map = temp_storage.make_datum(|packer| packer.push_dict(map));
165        Ok(map)
166    }
167
168    fn output_type(&self, _input_type: ColumnType) -> ColumnType {
169        ScalarType::Map {
170            value_type: Box::new(self.value_type.clone()),
171            custom_id: None,
172        }
173        .nullable(true)
174    }
175
176    fn propagates_nulls(&self) -> bool {
177        true
178    }
179
180    fn introduces_nulls(&self) -> bool {
181        true
182    }
183
184    fn preserves_uniqueness(&self) -> bool {
185        false
186    }
187
188    fn inverse(&self) -> Option<crate::UnaryFunc> {
189        None
190    }
191
192    fn is_monotone(&self) -> bool {
193        false
194    }
195}
196
197impl fmt::Display for MapBuildFromRecordList {
198    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
199        f.write_str("map_build")
200    }
201}
mz_expr/scalar/func/impls/map.rs

mz_expr/scalar/func/impls/
map.rs
