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

metric_engine/engine/create/
extract_new_columns.rs
