common_meta/ddl/create_table/
template.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::column_def::try_as_column_def;
18use api::v1::meta::Partition;
19use api::v1::region::{CreateRequest, RegionColumnDef};
20use api::v1::{ColumnDataType, ColumnDef, CreateTableExpr, SemanticType};
21use common_catalog::consts::{MITO_ENGINE, MITO2_ENGINE};
22use common_telemetry::warn;
23use snafu::{OptionExt, ResultExt};
24use store_api::metric_engine_consts::{
25    DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
26    PRIMARY_KEY_ENCODING,
27};
28use store_api::region_request::RegionRequirements;
29use store_api::storage::consts::ReservedColumnId;
30use store_api::storage::{RegionId, RegionNumber};
31use table::metadata::{TableId, TableInfo};
32
33use crate::error::{self, Result, SerializeWalOptionsSnafu};
34use crate::reconciliation::utils::build_column_metadata_from_table_info;
35use crate::wal_provider::{RegionWalOptions, serialize_wal_options};
36
37/// Constructs a [CreateRequest] based on the provided [TableInfo].
38///
39/// Note: This function is primarily intended for creating logical tables.
40///
41/// Logical table templates keep the original column order and primary key indices from
42/// `TableInfo` (including internal columns when present), because these are used to
43/// reconstruct the logical schema on the engine side.
44pub fn build_template_from_raw_table_info(table_info: &TableInfo) -> Result<CreateRequest> {
45    let primary_key_indices = &table_info.meta.primary_key_indices;
46    let column_defs = table_info
47        .meta
48        .schema
49        .column_schemas()
50        .iter()
51        .enumerate()
52        .map(|(i, c)| {
53            let is_primary_key = primary_key_indices.contains(&i);
54            let column_def = try_as_column_def(c, is_primary_key)
55                .context(error::ConvertColumnDefSnafu { column: &c.name })?;
56            Ok(RegionColumnDef {
57                column_def: Some(column_def),
58                // The column id will be overridden by the metric engine.
59                // So we just use the index as the column id.
60                column_id: i as u32,
61            })
62        })
63        .collect::<Result<Vec<_>>>()?;
64
65    let options = HashMap::from(&table_info.meta.options);
66    let template = CreateRequest {
67        region_id: 0,
68        engine: table_info.meta.engine.clone(),
69        column_defs,
70        primary_key: table_info
71            .meta
72            .primary_key_indices
73            .iter()
74            .map(|i| *i as u32)
75            .collect(),
76        path: String::new(),
77        options,
78        partition: None,
79        requirements: None,
80    };
81
82    Ok(template)
83}
84
85/// Constructs a [CreateRequest] based on the provided [TableInfo] for physical table.
86///
87/// Note: This function is primarily intended for creating physical table.
88///
89/// Physical table templates mark primary
90/// keys by tag semantic type to match the physical storage layout.
91pub fn build_template_from_raw_table_info_for_physical_table(
92    table_info: &TableInfo,
93) -> Result<CreateRequest> {
94    let name_to_ids = table_info
95        .name_to_ids()
96        .context(error::MissingColumnIdsSnafu)?;
97    let column_metadatas = build_column_metadata_from_table_info(
98        table_info.meta.schema.column_schemas(),
99        &table_info.meta.primary_key_indices,
100        &name_to_ids,
101    )?;
102    let column_defs = column_metadatas
103        .iter()
104        .map(|c| {
105            let column_def =
106                try_as_column_def(&c.column_schema, c.semantic_type == SemanticType::Tag).context(
107                    error::ConvertColumnDefSnafu {
108                        column: &c.column_schema.name,
109                    },
110                )?;
111            let region_column_def = RegionColumnDef {
112                column_def: Some(column_def),
113                column_id: c.column_id,
114            };
115
116            Ok(region_column_def)
117        })
118        .collect::<Result<Vec<_>>>()?;
119    // Preserve the order of primary key indices as defined in the original table info.
120    let primary_key = table_info
121        .meta
122        .primary_key_indices
123        .iter()
124        .map(|idx| column_metadatas[*idx].column_id)
125        .collect();
126
127    let options = HashMap::from(&table_info.meta.options);
128    let template = CreateRequest {
129        region_id: 0,
130        engine: table_info.meta.engine.clone(),
131        column_defs,
132        primary_key,
133        path: String::new(),
134        options,
135        partition: None,
136        requirements: None,
137    };
138
139    Ok(template)
140}
141
142pub(crate) fn build_template(create_table_expr: &CreateTableExpr) -> Result<CreateRequest> {
143    // Explicit sparse identity columns use the same IDs as metric physical regions.
144    let reserved_prefix = matches!(
145        create_table_expr.engine.as_str(),
146        MITO_ENGINE | MITO2_ENGINE
147    ) && create_table_expr
148        .table_options
149        .get(PRIMARY_KEY_ENCODING)
150        .is_some_and(|value| value.eq_ignore_ascii_case("sparse"))
151        && create_table_expr.primary_keys.starts_with(&[
152            DATA_SCHEMA_TABLE_ID_COLUMN_NAME.into(),
153            DATA_SCHEMA_TSID_COLUMN_NAME.into(),
154        ])
155        && [
156            (DATA_SCHEMA_TABLE_ID_COLUMN_NAME, ColumnDataType::Uint32),
157            (DATA_SCHEMA_TSID_COLUMN_NAME, ColumnDataType::Uint64),
158        ]
159        .into_iter()
160        .all(|(name, data_type)| {
161            create_table_expr.column_defs.iter().any(|column| {
162                column.name == name && column.data_type == data_type as i32 && !column.is_nullable
163            })
164        });
165    let column_defs = create_table_expr
166        .column_defs
167        .iter()
168        .enumerate()
169        .map(|(i, c)| {
170            let semantic_type = if create_table_expr.time_index == c.name {
171                SemanticType::Timestamp
172            } else if create_table_expr.primary_keys.contains(&c.name) {
173                SemanticType::Tag
174            } else {
175                SemanticType::Field
176            };
177
178            RegionColumnDef {
179                column_def: Some(ColumnDef {
180                    name: c.name.clone(),
181                    data_type: c.data_type,
182                    is_nullable: c.is_nullable,
183                    default_constraint: c.default_constraint.clone(),
184                    semantic_type: semantic_type as i32,
185                    comment: String::new(),
186                    datatype_extension: c.datatype_extension.clone(),
187                    options: c.options.clone(),
188                }),
189                column_id: if reserved_prefix && c.name == DATA_SCHEMA_TABLE_ID_COLUMN_NAME {
190                    ReservedColumnId::table_id()
191                } else if reserved_prefix && c.name == DATA_SCHEMA_TSID_COLUMN_NAME {
192                    ReservedColumnId::tsid()
193                } else {
194                    i as u32
195                },
196            }
197        })
198        .collect::<Vec<_>>();
199
200    let primary_key = create_table_expr
201        .primary_keys
202        .iter()
203        .map(|key| {
204            column_defs
205                .iter()
206                .find_map(|c| {
207                    c.column_def.as_ref().and_then(|x| {
208                        if &x.name == key {
209                            Some(c.column_id)
210                        } else {
211                            None
212                        }
213                    })
214                })
215                .context(error::PrimaryKeyNotFoundSnafu { key })
216        })
217        .collect::<Result<_>>()?;
218
219    let template = CreateRequest {
220        region_id: 0,
221        engine: create_table_expr.engine.clone(),
222        column_defs,
223        primary_key,
224        path: String::new(),
225        options: create_table_expr.table_options.clone(),
226        partition: None,
227        requirements: None,
228    };
229
230    Ok(template)
231}
232
233/// Builder for [PbCreateRegionRequest].
234pub struct CreateRequestBuilder {
235    template: CreateRequest,
236    /// Optional. Only for metric engine.
237    physical_table_id: Option<TableId>,
238    requirements: RegionRequirements,
239}
240
241impl CreateRequestBuilder {
242    pub fn new(template: CreateRequest, physical_table_id: Option<TableId>) -> Self {
243        Self {
244            template,
245            physical_table_id,
246            requirements: RegionRequirements::empty(),
247        }
248    }
249
250    pub fn template(&self) -> &CreateRequest {
251        &self.template
252    }
253
254    pub fn with_requirements(mut self, requirements: RegionRequirements) -> Self {
255        self.requirements = requirements;
256        self
257    }
258
259    pub fn build_one(
260        &self,
261        region_id: RegionId,
262        storage_path: String,
263        region_wal_options: &RegionWalOptions,
264        partition_exprs: &HashMap<RegionNumber, String>,
265    ) -> Result<CreateRequest> {
266        let mut request = self.template.clone();
267
268        request.region_id = region_id.as_u64();
269        request.path = storage_path;
270        request.requirements = Some(self.requirements.into());
271        // Stores the encoded wal options into the request options.
272        serialize_wal_options(&mut request.options, region_id, region_wal_options)
273            .context(SerializeWalOptionsSnafu { region_id })?;
274        request.partition = Some(prepare_partition_expr(region_id, partition_exprs));
275
276        if let Some(physical_table_id) = self.physical_table_id {
277            // Logical table has the same region numbers with physical table, and they have a one-to-one mapping.
278            // For example, region 0 of logical table must resides with region 0 of physical table. So here we can
279            // simply concat the physical table id and the logical region number to get the physical region id.
280            let physical_region_id = RegionId::new(physical_table_id, region_id.region_number());
281
282            request.options.insert(
283                LOGICAL_TABLE_METADATA_KEY.to_string(),
284                physical_region_id.as_u64().to_string(),
285            );
286        }
287
288        Ok(request)
289    }
290}
291
292fn prepare_partition_expr(
293    region_id: RegionId,
294    partition_exprs: &HashMap<RegionNumber, String>,
295) -> Partition {
296    let expr = partition_exprs.get(&region_id.region_number()).cloned();
297    if expr.is_none() {
298        warn!("region {} has no partition expr", region_id);
299    }
300
301    Partition {
302        expression: expr.unwrap_or_default(),
303        ..Default::default()
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use std::collections::HashMap;
310
311    use store_api::storage::{RegionId, RegionNumber};
312
313    use super::*;
314    use crate::key::test_utils;
315
316    #[test]
317    fn test_create_sparse_reserved_column_ids() {
318        let mut expr = CreateTableExpr {
319            engine: MITO_ENGINE.into(),
320            table_options: HashMap::from([(PRIMARY_KEY_ENCODING.into(), "sparse".into())]),
321            primary_keys: vec![
322                DATA_SCHEMA_TABLE_ID_COLUMN_NAME.into(),
323                DATA_SCHEMA_TSID_COLUMN_NAME.into(),
324                "job".into(),
325            ],
326            time_index: "ts".into(),
327            column_defs: [
328                ("ts", ColumnDataType::TimestampMillisecond),
329                (DATA_SCHEMA_TABLE_ID_COLUMN_NAME, ColumnDataType::Uint32),
330                (DATA_SCHEMA_TSID_COLUMN_NAME, ColumnDataType::Uint64),
331                ("job", ColumnDataType::String),
332            ]
333            .into_iter()
334            .map(|(name, data_type)| ColumnDef {
335                name: name.into(),
336                data_type: data_type as i32,
337                is_nullable: false,
338                ..Default::default()
339            })
340            .collect(),
341            ..Default::default()
342        };
343        let template = build_template(&expr).unwrap();
344        assert_eq!(
345            template.primary_key,
346            vec![ReservedColumnId::table_id(), ReservedColumnId::tsid(), 3]
347        );
348        assert_eq!(
349            template.column_defs[1].column_id,
350            ReservedColumnId::table_id()
351        );
352        assert_eq!(template.column_defs[2].column_id, ReservedColumnId::tsid());
353        // The names alone must not change ordinary dense-table column allocation.
354        expr.table_options.clear();
355        assert_eq!(build_template(&expr).unwrap().primary_key, vec![1, 2, 3]);
356    }
357
358    #[test]
359    fn test_build_one_sets_partition_expr_per_region() {
360        // minimal template
361        let template = CreateRequest {
362            region_id: 0,
363            engine: "mito".to_string(),
364            column_defs: vec![],
365            primary_key: vec![],
366            path: String::new(),
367            options: Default::default(),
368            partition: None,
369            requirements: None,
370        };
371        let builder = CreateRequestBuilder::new(template, None);
372
373        let mut partition_exprs: HashMap<RegionNumber, String> = HashMap::new();
374        let expr_a =
375            r#"{"Expr":{"lhs":{"Column":"a"},"op":"Eq","rhs":{"Value":{"UInt32":1}}}}"#.to_string();
376        partition_exprs.insert(0, expr_a.clone());
377
378        let r0 = builder
379            .build_one(
380                RegionId::new(42, 0),
381                "/p".to_string(),
382                &Default::default(),
383                &partition_exprs,
384            )
385            .unwrap();
386        assert_eq!(r0.partition.as_ref().unwrap().expression, expr_a);
387        assert_eq!(
388            r0.requirements.map(RegionRequirements::from),
389            Some(RegionRequirements::empty())
390        );
391    }
392
393    #[test]
394    fn test_build_one_sets_explicit_requirements() {
395        let template = CreateRequest {
396            region_id: 0,
397            engine: "mito".to_string(),
398            column_defs: vec![],
399            primary_key: vec![],
400            path: String::new(),
401            options: Default::default(),
402            partition: None,
403            requirements: None,
404        };
405        let builder = CreateRequestBuilder::new(template, None)
406            .with_requirements(RegionRequirements::object_storage());
407
408        let request = builder
409            .build_one(
410                RegionId::new(42, 0),
411                "/p".to_string(),
412                &Default::default(),
413                &Default::default(),
414            )
415            .unwrap();
416
417        assert_eq!(
418            request.requirements.map(RegionRequirements::from),
419            Some(RegionRequirements::object_storage())
420        );
421    }
422
423    #[test]
424    fn test_build_template_for_physical_table_primary_key_matches_indices() {
425        let mut table_info = test_utils::new_test_table_info(42);
426        table_info.meta.primary_key_indices = vec![0, 2];
427        table_info.meta.column_ids = vec![10, 20, 30];
428
429        let template = build_template_from_raw_table_info_for_physical_table(&table_info).unwrap();
430        assert_eq!(template.primary_key, vec![10, 30]);
431    }
432}
common_meta/ddl/create_table/template.rs

common_meta/ddl/create_table/
template.rs
