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