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