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, PartialOrd, Ord, 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_struct(self) -> DataType {
88        let (key, value) = *self.0;
89        DataType::Struct(Self::struct_type_for_map(key, value))
90    }
91
92    pub fn into_list(self) -> DataType {
93        DataType::List(Box::new(self.into_struct()))
94    }
95
96    /// String and integral types are allowed.
97    ///
98    /// This is similar to [Protobuf](https://protobuf.dev/programming-guides/proto3/#maps)'s
99    /// decision.
100    ///
101    /// Note that this isn't definitive.
102    /// Just be conservative at the beginning, but not too restrictive (like only allowing strings).
103    pub fn check_key_type_valid(data_type: &DataType) -> Result<(), String> {
104        let ok = match data_type {
105            DataType::Int16 | DataType::Int32 | DataType::Int64 => true,
106            DataType::Varchar => true,
107            DataType::Boolean
108            | DataType::Float32
109            | DataType::Float64
110            | DataType::Decimal
111            | DataType::Date
112            | DataType::Time
113            | DataType::Timestamp
114            | DataType::Timestamptz
115            | DataType::Interval
116            | DataType::Struct(_)
117            | DataType::List(_)
118            | DataType::Bytea
119            | DataType::Jsonb
120            | DataType::Serial
121            | DataType::Int256
122            | DataType::Map(_) => false,
123        };
124        if !ok {
125            Err(format!("invalid map key type: {data_type}"))
126        } else {
127            Ok(())
128        }
129    }
130}
131
132impl FromStr for MapType {
133    type Err = anyhow::Error;
134
135    fn from_str(s: &str) -> Result<Self, Self::Err> {
136        if !(s.starts_with("map(") && s.ends_with(')')) {
137            return Err(anyhow::anyhow!("expect map(...,...)"));
138        };
139        if let Some((key, value)) = s[4..s.len() - 1].split(',').collect_tuple() {
140            let key = key.parse().context("failed to parse map key type")?;
141            let value = value.parse().context("failed to parse map value type")?;
142            MapType::try_from_kv(key, value).map_err(|e| anyhow::anyhow!(e))
143        } else {
144            Err(anyhow::anyhow!("expect map(...,...)"))
145        }
146    }
147}
148
149impl std::fmt::Display for MapType {
150    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
151        write!(f, "map({},{})", self.key(), self.value())
152    }
153}