Skip to main content

metric_engine/engine/alter/
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::{AlterKind, RegionAlterRequest};
22use store_api::storage::RegionId;
23
24use crate::error::{AddingFieldColumnSnafu, Result};
25
26/// Extract new columns from the create requests.
27///
28/// # Panics
29///
30/// This function will panic if the alter kind is not `AddColumns`.
31pub fn extract_new_columns<'a>(
32    requests: &'a [(RegionId, RegionAlterRequest)],
33    physical_columns: &HashMap<String, ColumnMetadata>,
34    new_column_indices: &mut HashMap<&'a str, usize>,
35    new_columns: &mut Vec<ColumnMetadata>,
36) -> Result<()> {
37    for (region_id, request) in requests {
38        let AlterKind::AddColumns { columns } = &request.kind else {
39            unreachable!()
40        };
41        for col in columns {
42            let column_name = col.column_metadata.column_schema.name.as_str();
43            if let Some(existing) = physical_columns.get(column_name).or_else(|| {
44                new_column_indices
45                    .get(column_name)
46                    .map(|&index| &new_columns[index])
47            }) {
48                crate::engine::validate_shared_column(*region_id, existing, &col.column_metadata)?;
49            } else {
50                ensure!(
51                    col.column_metadata.semantic_type != SemanticType::Field
52                        || is_native_histogram_value_type(
53                            &col.column_metadata.column_schema.data_type
54                        ),
55                    AddingFieldColumnSnafu {
56                        name: column_name.to_string(),
57                    }
58                );
59                new_column_indices.insert(column_name, new_columns.len());
60                // TODO(weny): avoid clone
61                new_columns.push(col.column_metadata.clone());
62            }
63        }
64    }
65
66    Ok(())
67}
68
69#[cfg(test)]
70mod tests {
71    use std::collections::HashMap;
72
73    use api::v1::SemanticType;
74    use datatypes::prelude::ConcreteDataType;
75    use datatypes::schema::ColumnSchema;
76    use store_api::metadata::ColumnMetadata;
77    use store_api::region_request::{AddColumn, AlterKind, RegionAlterRequest};
78    use store_api::storage::RegionId;
79
80    use super::*;
81    use crate::error::Error;
82
83    #[test]
84    fn test_alter_rejects_conflicting_shared_definitions() {
85        let column = ColumnMetadata {
86            column_schema: ColumnSchema::new("shared", ConcreteDataType::string_datatype(), true),
87            semantic_type: SemanticType::Tag,
88            column_id: 0,
89        };
90        let mut incompatible = column.clone();
91        incompatible.column_schema =
92            ColumnSchema::new("shared", ConcreteDataType::int64_datatype(), true);
93        for existing in [false, true] {
94            let physical = if existing {
95                HashMap::from([("shared".into(), column.clone())])
96            } else {
97                HashMap::new()
98            };
99            let requests = [(
100                RegionId::new(1, 1),
101                RegionAlterRequest {
102                    kind: AlterKind::AddColumns {
103                        columns: [column.clone(), incompatible.clone()]
104                            .into_iter()
105                            .map(|column_metadata| AddColumn {
106                                column_metadata,
107                                location: None,
108                            })
109                            .collect(),
110                    },
111                },
112            )];
113            let error =
114                extract_new_columns(&requests, &physical, &mut HashMap::new(), &mut Vec::new())
115                    .unwrap_err();
116            assert!(matches!(error, Error::InvalidRequest { .. }));
117        }
118    }
119
120    #[test]
121    fn test_extract_new_columns_preserves_first_seen_order() {
122        let requests = [["b", "a"], ["a", "b"]].map(|names| {
123            (
124                RegionId::new(1, 1),
125                RegionAlterRequest {
126                    kind: AlterKind::AddColumns {
127                        columns: names
128                            .into_iter()
129                            .map(|name| AddColumn {
130                                column_metadata: ColumnMetadata {
131                                    column_schema: ColumnSchema::new(
132                                        name,
133                                        ConcreteDataType::string_datatype(),
134                                        true,
135                                    ),
136                                    semantic_type: SemanticType::Tag,
137                                    column_id: 0,
138                                },
139                                location: None,
140                            })
141                            .collect(),
142                    },
143                },
144            )
145        });
146        let mut indices = HashMap::new();
147        let mut columns = Vec::new();
148        extract_new_columns(&requests, &HashMap::new(), &mut indices, &mut columns).unwrap();
149        assert_eq!(indices, HashMap::from([("b", 0), ("a", 1)]));
150        assert_eq!(
151            columns
152                .iter()
153                .map(|c| c.column_schema.name.as_str())
154                .collect::<Vec<_>>(),
155            ["b", "a"]
156        );
157    }
158
159    #[test]
160    fn test_extract_new_columns_with_field_type() {
161        let requests = vec![(
162            RegionId::new(1, 1),
163            RegionAlterRequest {
164                kind: AlterKind::AddColumns {
165                    columns: vec![AddColumn {
166                        column_metadata: ColumnMetadata {
167                            column_schema: ColumnSchema::new(
168                                "new_column".to_string(),
169                                ConcreteDataType::string_datatype(),
170                                false,
171                            ),
172                            semantic_type: SemanticType::Field,
173                            column_id: 0,
174                        },
175                        location: None,
176                    }],
177                },
178            },
179        )];
180
181        let physical_columns = HashMap::new();
182        let mut new_column_indices = HashMap::new();
183        let mut new_columns = Vec::new();
184
185        let err = extract_new_columns(
186            &requests,
187            &physical_columns,
188            &mut new_column_indices,
189            &mut new_columns,
190        )
191        .unwrap_err();
192
193        assert!(matches!(err, Error::AddingFieldColumn { .. }));
194    }
195}