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::{ColumnId, 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, ColumnId>,
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.to_string(),
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::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                },
100            ),
101            (
102                RegionId::new(1, 2),
103                RegionCreateRequest {
104                    column_metadatas: vec![ColumnMetadata {
105                        // Duplicate column name
106                        column_schema: ColumnSchema::new(
107                            "new_column".to_string(),
108                            ConcreteDataType::string_datatype(),
109                            false,
110                        ),
111                        semantic_type: SemanticType::Tag,
112                        column_id: 0,
113                    }],
114                    engine: "test".to_string(),
115                    primary_key: vec![],
116                    options: HashMap::new(),
117                    table_dir: "test".to_string(),
118                    path_type: PathType::Bare,
119                },
120            ),
121        ];
122
123        let mut physical_columns = HashMap::new();
124        physical_columns.insert("existing_column".to_string(), 0);
125        let mut new_column_names = HashSet::new();
126        let mut new_columns = Vec::new();
127
128        let result = extract_new_columns(
129            &requests,
130            &physical_columns,
131            &mut new_column_names,
132            &mut new_columns,
133        );
134
135        assert!(result.is_ok());
136        assert!(new_column_names.contains("new_column"));
137        assert_eq!(new_columns.len(), 1);
138        assert_eq!(new_columns[0].column_schema.name, "new_column");
139    }
140
141    #[test]
142    fn test_extract_new_columns_with_field_type() {
143        let requests = vec![(
144            RegionId::new(1, 1),
145            RegionCreateRequest {
146                column_metadatas: vec![ColumnMetadata {
147                    column_schema: ColumnSchema::new(
148                        "new_column".to_string(),
149                        ConcreteDataType::string_datatype(),
150                        false,
151                    ),
152                    semantic_type: SemanticType::Field,
153                    column_id: 0,
154                }],
155                engine: "test".to_string(),
156                primary_key: vec![],
157                options: HashMap::new(),
158                table_dir: "test".to_string(),
159                path_type: PathType::Bare,
160            },
161        )];
162
163        let physical_columns = HashMap::new();
164        let mut new_column_names = HashSet::new();
165        let mut new_columns = Vec::new();
166
167        let err = extract_new_columns(
168            &requests,
169            &physical_columns,
170            &mut new_column_names,
171            &mut new_columns,
172        )
173        .unwrap_err();
174
175        assert_matches!(err, Error::AddingFieldColumn { .. });
176    }
177}