risingwave_common/types/
map_type.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::Formatter;
16
17use anyhow::Context;
18
19use super::*;
20
21/// Refer to [`super::super::array::MapArray`] for the invariants of a map value.
22#[derive(Debug, Clone, PartialEq, Eq, Hash)]
23pub struct MapType(Box<(DataType, DataType)>);
24
25impl From<MapType> for DataType {
26    fn from(value: MapType) -> Self {
27        DataType::Map(value)
28    }
29}
30
31impl MapType {
32    /// # Panics
33    /// Panics if the key type is not valid for a map.
34    pub fn from_kv(key: DataType, value: DataType) -> Self {
35        Self::check_key_type_valid(&key).unwrap();
36        Self(Box::new((key, value)))
37    }
38
39    pub fn try_from_kv(key: DataType, value: DataType) -> Result<Self, String> {
40        Self::check_key_type_valid(&key)?;
41        Ok(Self(Box::new((key, value))))
42    }
43
44    pub fn try_from_entries(list_entries_type: DataType) -> Result<Self, String> {
45        match list_entries_type {
46            DataType::Struct(s) => {
47                let Some((k, v)) = s.iter().collect_tuple() else {
48                    return Err(format!(
49                        "the underlying struct for map must have exactly two fields, got: {s:?}"
50                    ));
51                };
52                // the field names are not strictly enforced
53                // Currently this panics for SELECT * FROM t
54                // if cfg!(debug_assertions) {
55                //     itertools::assert_equal(struct_type.names(), ["key", "value"]);
56                // }
57                Self::try_from_kv(k.1.clone(), v.1.clone())
58            }
59            _ => Err(format!(
60                "invalid map entries type, expected struct, got: {list_entries_type}"
61            )),
62        }
63    }
64
65    /// # Panics
66    /// Panics if the key type is not valid for a map, or the
67    /// entries type is not a valid struct type.
68    pub fn from_entries(list_entries_type: DataType) -> Self {
69        Self::try_from_entries(list_entries_type).unwrap()
70    }
71
72    /// # Panics
73    /// Panics if the key type is not valid for a map.
74    pub fn struct_type_for_map(key_type: DataType, value_type: DataType) -> StructType {
75        MapType::check_key_type_valid(&key_type).unwrap();
76        StructType::new(vec![("key", key_type), ("value", value_type)])
77    }
78
79    pub fn key(&self) -> &DataType {
80        &self.0.0
81    }
82
83    pub fn value(&self) -> &DataType {
84        &self.0.1
85    }
86
87    pub fn into_kv(self) -> (DataType, DataType) {
88        *self.0
89    }
90
91    pub fn into_struct(self) -> DataType {
92        let (key, value) = *self.0;
93        DataType::Struct(Self::struct_type_for_map(key, value))
94    }
95
96    pub fn into_list(self) -> DataType {
97        DataType::list(self.into_struct())
98    }
99
100    /// Sames as [`Self::into_list`] but returns [`ListType`].
101    pub fn into_list_type(self) -> ListType {
102        ListType::new(self.into_struct())
103    }
104
105    /// String and integral types are allowed.
106    ///
107    /// This is similar to [Protobuf](https://protobuf.dev/programming-guides/proto3/#maps)'s
108    /// decision.
109    ///
110    /// Note that this isn't definitive.
111    /// Just be conservative at the beginning, but not too restrictive (like only allowing strings).
112    pub fn check_key_type_valid(data_type: &DataType) -> Result<(), String> {
113        let ok = match data_type {
114            DataType::Int16 | DataType::Int32 | DataType::Int64 => true,
115            DataType::Varchar => true,
116            DataType::Boolean
117            | DataType::Float32
118            | DataType::Float64
119            | DataType::Decimal
120            | DataType::Date
121            | DataType::Time
122            | DataType::Timestamp
123            | DataType::Timestamptz
124            | DataType::Interval
125            | DataType::Struct(_)
126            | DataType::List(_)
127            | DataType::Bytea
128            | DataType::Jsonb
129            | DataType::Serial
130            | DataType::Int256
131            | DataType::Vector(_)
132            | DataType::Map(_) => false,
133        };
134        if !ok {
135            Err(format!("invalid map key type: {data_type}"))
136        } else {
137            Ok(())
138        }
139    }
140}
141
142impl FromStr for MapType {
143    type Err = anyhow::Error;
144
145    fn from_str(s: &str) -> Result<Self, Self::Err> {
146        if !(s.starts_with("map(") && s.ends_with(')')) {
147            return Err(anyhow::anyhow!("expect map(...,...)"));
148        };
149        if let Some((key, value)) = s[4..s.len() - 1].split(',').collect_tuple() {
150            let key = key.parse().context("failed to parse map key type")?;
151            let value = value.parse().context("failed to parse map value type")?;
152            MapType::try_from_kv(key, value).map_err(|e| anyhow::anyhow!(e))
153        } else {
154            Err(anyhow::anyhow!("expect map(...,...)"))
155        }
156    }
157}
158
159impl std::fmt::Display for MapType {
160    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
161        write!(f, "map({},{})", self.key(), self.value())
162    }
163}