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                    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("existing_column".to_string(), 0);
127        let mut new_column_names = HashSet::new();
128        let mut new_columns = Vec::new();
129
130        let result = extract_new_columns(
131            &requests,
132            &physical_columns,
133            &mut new_column_names,
134            &mut new_columns,
135        );
136
137        assert!(result.is_ok());
138        assert!(new_column_names.contains("new_column"));
139        assert_eq!(new_columns.len(), 1);
140        assert_eq!(new_columns[0].column_schema.name, "new_column");
141    }
142
143    #[test]
144    fn test_extract_new_columns_with_field_type() {
145        let requests = vec![(
146            RegionId::new(1, 1),
147            RegionCreateRequest {
148                column_metadatas: vec![ColumnMetadata {
149                    column_schema: ColumnSchema::new(
150                        "new_column".to_string(),
151                        ConcreteDataType::string_datatype(),
152                        false,
153                    ),
154                    semantic_type: SemanticType::Field,
155                    column_id: 0,
156                }],
157                engine: "test".to_string(),
158                primary_key: vec![],
159                options: HashMap::new(),
160                table_dir: "test".to_string(),
161                path_type: PathType::Bare,
162                partition_expr_json: Some("".to_string()),
163            },
164        )];
165
166        let physical_columns = HashMap::new();
167        let mut new_column_names = HashSet::new();
168        let mut new_columns = Vec::new();
169
170        let err = extract_new_columns(
171            &requests,
172            &physical_columns,
173            &mut new_column_names,
174            &mut new_columns,
175        )
176        .unwrap_err();
177
178        assert_matches!(err, Error::AddingFieldColumn { .. });
179    }
180}
metric_engine/engine/create/extract_new_columns.rs

metric_engine/engine/create/
extract_new_columns.rs
