Skip to main content

metric_engine/engine/create/
extract_new_columns.rs

1// Copyright 2023 Greptime Team
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::collections::{HashMap, HashSet};
16
17use api::v1::SemanticType;
18use snafu::ensure;
19use store_api::metadata::ColumnMetadata;
20use store_api::region_request::RegionCreateRequest;
21use store_api::storage::RegionId;
22
23use crate::error::{AddingFieldColumnSnafu, Result};
24
25/// Extract new columns from the create requests.
26pub fn extract_new_columns<'a>(
27    requests: &'a [(RegionId, RegionCreateRequest)],
28    physical_columns: &HashMap<String, ColumnMetadata>,
29    new_column_names: &mut HashSet<&'a str>,
30    new_columns: &mut Vec<ColumnMetadata>,
31) -> Result<()> {
32    for (_, request) in requests {
33        for col in &request.column_metadatas {
34            if !physical_columns.contains_key(&col.column_schema.name)
35                && !new_column_names.contains(&col.column_schema.name.as_str())
36            {
37                ensure!(
38                    col.semantic_type != SemanticType::Field,
39                    AddingFieldColumnSnafu {
40                        name: col.column_schema.name.clone(),
41                    }
42                );
43                new_column_names.insert(&col.column_schema.name);
44                // TODO(weny): avoid clone
45                new_columns.push(col.clone());
46            }
47        }
48    }
49
50    Ok(())
51}
52
53#[cfg(test)]
54mod tests {
55    use std::assert_matches;
56    use std::collections::{HashMap, HashSet};
57
58    use api::v1::SemanticType;
59    use datatypes::prelude::ConcreteDataType;
60    use datatypes::schema::ColumnSchema;
61    use store_api::metadata::ColumnMetadata;
62    use store_api::region_request::{PathType, RegionCreateRequest};
63    use store_api::storage::RegionId;
64
65    use super::*;
66    use crate::error::Error;
67
68    #[test]
69    fn test_extract_new_columns() {
70        let requests = vec![
71            (
72                RegionId::new(1, 1),
73                RegionCreateRequest {
74                    column_metadatas: vec![
75                        ColumnMetadata {
76                            column_schema: ColumnSchema::new(
77                                "existing_column".to_string(),
78                                ConcreteDataType::string_datatype(),
79                                false,
80                            ),
81                            semantic_type: SemanticType::Tag,
82                            column_id: 0,
83                        },
84                        ColumnMetadata {
85                            column_schema: ColumnSchema::new(
86                                "new_column".to_string(),
87                                ConcreteDataType::string_datatype(),
88                                false,
89                            ),
90                            semantic_type: SemanticType::Tag,
91                            column_id: 0,
92                        },
93                    ],
94                    engine: "test".to_string(),
95                    primary_key: vec![],
96                    options: HashMap::new(),
97                    table_dir: "test".to_string(),
98                    path_type: PathType::Bare,
99                    partition_expr_json: Some("".to_string()),
100                    requirements: Default::default(),
101                },
102            ),
103            (
104                RegionId::new(1, 2),
105                RegionCreateRequest {
106                    column_metadatas: vec![ColumnMetadata {
107                        // Duplicate column name
108                        column_schema: ColumnSchema::new(
109                            "new_column".to_string(),
110                            ConcreteDataType::string_datatype(),
111                            false,
112                        ),
113                        semantic_type: SemanticType::Tag,
114                        column_id: 0,
115                    }],
116                    engine: "test".to_string(),
117                    primary_key: vec![],
118                    options: HashMap::new(),
119                    table_dir: "test".to_string(),
120                    path_type: PathType::Bare,
121                    partition_expr_json: Some("".to_string()),
122                    requirements: Default::default(),
123                },
124            ),
125        ];
126
127        let mut physical_columns = HashMap::new();
128        physical_columns.insert(
129            "existing_column".to_string(),
130            ColumnMetadata {
131                column_schema: ColumnSchema::new(
132                    "existing_column".to_string(),
133                    ConcreteDataType::string_datatype(),
134                    false,
135                ),
136                semantic_type: SemanticType::Tag,
137                column_id: 0,
138            },
139        );
140        let mut new_column_names = HashSet::new();
141        let mut new_columns = Vec::new();
142
143        let result = extract_new_columns(
144            &requests,
145            &physical_columns,
146            &mut new_column_names,
147            &mut new_columns,
148        );
149
150        assert!(result.is_ok());
151        assert!(new_column_names.contains("new_column"));
152        assert_eq!(new_columns.len(), 1);
153        assert_eq!(new_columns[0].column_schema.name, "new_column");
154    }
155
156    #[test]
157    fn test_extract_new_columns_with_field_type() {
158        let requests = vec![(
159            RegionId::new(1, 1),
160            RegionCreateRequest {
161                column_metadatas: vec![ColumnMetadata {
162                    column_schema: ColumnSchema::new(
163                        "new_column".to_string(),
164                        ConcreteDataType::string_datatype(),
165                        false,
166                    ),
167                    semantic_type: SemanticType::Field,
168                    column_id: 0,
169                }],
170                engine: "test".to_string(),
171                primary_key: vec![],
172                options: HashMap::new(),
173                table_dir: "test".to_string(),
174                path_type: PathType::Bare,
175                partition_expr_json: Some("".to_string()),
176                requirements: Default::default(),
177            },
178        )];
179
180        let physical_columns = HashMap::new();
181        let mut new_column_names = HashSet::new();
182        let mut new_columns = Vec::new();
183
184        let err = extract_new_columns(
185            &requests,
186            &physical_columns,
187            &mut new_column_names,
188            &mut new_columns,
189        )
190        .unwrap_err();
191
192        assert_matches!(err, Error::AddingFieldColumn { .. });
193    }
194}