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;
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_indices: &mut HashMap<&'a str, usize>,
31    new_columns: &mut Vec<ColumnMetadata>,
32) -> Result<()> {
33    for (region_id, request) in requests {
34        for col in &request.column_metadatas {
35            if let Some(existing) = physical_columns.get(&col.column_schema.name).or_else(|| {
36                new_column_indices
37                    .get(col.column_schema.name.as_str())
38                    .map(|&index| &new_columns[index])
39            }) {
40                crate::engine::validate_shared_column(*region_id, existing, col)?;
41            } else {
42                ensure!(
43                    col.semantic_type != SemanticType::Field
44                        || is_native_histogram_value_type(&col.column_schema.data_type),
45                    AddingFieldColumnSnafu {
46                        name: col.column_schema.name.clone(),
47                    }
48                );
49                new_column_indices.insert(&col.column_schema.name, new_columns.len());
50                // TODO(weny): avoid clone
51                new_columns.push(col.clone());
52            }
53        }
54    }
55
56    Ok(())
57}
58
59#[cfg(test)]
60mod tests {
61    use std::assert_matches;
62    use std::collections::HashMap;
63
64    use api::v1::SemanticType;
65    use datatypes::prelude::ConcreteDataType;
66    use datatypes::schema::ColumnSchema;
67    use store_api::metadata::ColumnMetadata;
68    use store_api::region_request::{PathType, RegionCreateRequest};
69    use store_api::storage::RegionId;
70
71    use super::*;
72    use crate::error::Error;
73
74    #[test]
75    fn test_conflicting_shared_definitions_rejected() {
76        let column = ColumnMetadata {
77            column_schema: ColumnSchema::new("shared", ConcreteDataType::string_datatype(), true),
78            semantic_type: SemanticType::Tag,
79            column_id: 0,
80        };
81        let mut incompatible = column.clone();
82        incompatible.semantic_type = SemanticType::Field;
83        let region = RegionId::new(1, 1);
84        assert!(crate::engine::validate_shared_column(region, &column, &incompatible).is_err());
85        incompatible.semantic_type = SemanticType::Tag;
86        incompatible.column_schema =
87            ColumnSchema::new("shared", ConcreteDataType::int64_datatype(), true);
88        assert!(crate::engine::validate_shared_column(region, &column, &incompatible).is_err());
89        assert!(crate::engine::validate_shared_column(region, &column, &column).is_ok());
90        for existing in [false, true] {
91            let physical = if existing {
92                HashMap::from([("shared".into(), column.clone())])
93            } else {
94                HashMap::new()
95            };
96            let requests = [(
97                region,
98                RegionCreateRequest {
99                    column_metadatas: vec![column.clone(), incompatible.clone()],
100                    engine: "metric".into(),
101                    primary_key: vec![0],
102                    options: HashMap::new(),
103                    table_dir: "test".into(),
104                    path_type: PathType::Bare,
105                    partition_expr_json: None,
106                    requirements: Default::default(),
107                },
108            )];
109            let error =
110                extract_new_columns(&requests, &physical, &mut HashMap::new(), &mut Vec::new())
111                    .unwrap_err();
112            assert!(matches!(error, Error::InvalidRequest { .. }));
113        }
114    }
115
116    #[test]
117    fn test_extract_new_columns() {
118        let mut requests = vec![
119            (
120                RegionId::new(1, 1),
121                RegionCreateRequest {
122                    column_metadatas: vec![
123                        ColumnMetadata {
124                            column_schema: ColumnSchema::new(
125                                "existing_column".to_string(),
126                                ConcreteDataType::string_datatype(),
127                                false,
128                            ),
129                            semantic_type: SemanticType::Tag,
130                            column_id: 0,
131                        },
132                        ColumnMetadata {
133                            column_schema: ColumnSchema::new(
134                                "new_column".to_string(),
135                                ConcreteDataType::string_datatype(),
136                                false,
137                            ),
138                            semantic_type: SemanticType::Tag,
139                            column_id: 0,
140                        },
141                    ],
142                    engine: "test".to_string(),
143                    primary_key: vec![],
144                    options: HashMap::new(),
145                    table_dir: "test".to_string(),
146                    path_type: PathType::Bare,
147                    partition_expr_json: Some("".to_string()),
148                    requirements: Default::default(),
149                },
150            ),
151            (
152                RegionId::new(1, 2),
153                RegionCreateRequest {
154                    column_metadatas: vec![ColumnMetadata {
155                        // Duplicate column name
156                        column_schema: ColumnSchema::new(
157                            "new_column".to_string(),
158                            ConcreteDataType::string_datatype(),
159                            false,
160                        ),
161                        semantic_type: SemanticType::Tag,
162                        column_id: 0,
163                    }],
164                    engine: "test".to_string(),
165                    primary_key: vec![],
166                    options: HashMap::new(),
167                    table_dir: "test".to_string(),
168                    path_type: PathType::Bare,
169                    partition_expr_json: Some("".to_string()),
170                    requirements: Default::default(),
171                },
172            ),
173        ];
174
175        let mut another = requests[0].1.column_metadatas[1].clone();
176        another.column_schema.name = "another_column".into();
177        requests[0].1.column_metadatas.push(another.clone());
178        requests[1].1.column_metadatas.push(another);
179
180        let mut physical_columns = HashMap::new();
181        physical_columns.insert(
182            "existing_column".to_string(),
183            ColumnMetadata {
184                column_schema: ColumnSchema::new(
185                    "existing_column".to_string(),
186                    ConcreteDataType::string_datatype(),
187                    false,
188                ),
189                semantic_type: SemanticType::Tag,
190                column_id: 0,
191            },
192        );
193        let mut new_column_indices = HashMap::new();
194        let mut new_columns = Vec::new();
195
196        let result = extract_new_columns(
197            &requests,
198            &physical_columns,
199            &mut new_column_indices,
200            &mut new_columns,
201        );
202
203        assert!(result.is_ok());
204        assert_eq!(new_column_indices["new_column"], 0);
205        assert_eq!(new_column_indices["another_column"], 1);
206        assert_eq!(new_columns.len(), 2);
207        assert_eq!(new_columns[0].column_schema.name, "new_column");
208        assert_eq!(new_columns[1].column_schema.name, "another_column");
209    }
210
211    #[test]
212    fn test_extract_new_columns_with_field_type() {
213        let requests = vec![(
214            RegionId::new(1, 1),
215            RegionCreateRequest {
216                column_metadatas: vec![ColumnMetadata {
217                    column_schema: ColumnSchema::new(
218                        "new_column".to_string(),
219                        ConcreteDataType::string_datatype(),
220                        false,
221                    ),
222                    semantic_type: SemanticType::Field,
223                    column_id: 0,
224                }],
225                engine: "test".to_string(),
226                primary_key: vec![],
227                options: HashMap::new(),
228                table_dir: "test".to_string(),
229                path_type: PathType::Bare,
230                partition_expr_json: Some("".to_string()),
231                requirements: Default::default(),
232            },
233        )];
234
235        let physical_columns = HashMap::new();
236        let mut new_column_indices = HashMap::new();
237        let mut new_columns = Vec::new();
238
239        let err = extract_new_columns(
240            &requests,
241            &physical_columns,
242            &mut new_column_indices,
243            &mut new_columns,
244        )
245        .unwrap_err();
246
247        assert_matches!(err, Error::AddingFieldColumn { .. });
248    }
249}