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