common_meta/ddl/create_logical_tables/
metadata.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::HashSet;
16use std::sync::Arc;
17
18use datatypes::prelude::ConcreteDataType;
19use datatypes::schema::{ColumnSchema, Schema};
20use snafu::OptionExt;
21use store_api::metric_engine_consts::is_metric_engine_internal_column;
22
23use crate::ddl::create_logical_tables::CreateLogicalTablesProcedure;
24use crate::error::Result;
25use crate::key::table_route::TableRouteValue;
26
27impl CreateLogicalTablesProcedure {
28    pub(crate) async fn fill_physical_table_info(&mut self) -> Result<()> {
29        let physical_region_numbers = self
30            .context
31            .table_metadata_manager
32            .table_route_manager()
33            .get_physical_table_route(self.data.physical_table_id)
34            .await
35            .map(|(_, route)| TableRouteValue::Physical(route).region_numbers())?;
36
37        self.data.physical_region_numbers = physical_region_numbers;
38
39        // Extract partition column names from the physical table
40        let physical_table_info = self
41            .context
42            .table_metadata_manager
43            .table_info_manager()
44            .get(self.data.physical_table_id)
45            .await?
46            .with_context(|| crate::error::TableInfoNotFoundSnafu {
47                table: format!("physical table {}", self.data.physical_table_id),
48            })?;
49
50        let physical_partition_columns: Vec<String> = physical_table_info
51            .table_info
52            .meta
53            .partition_key_indices
54            .iter()
55            .map(|&idx| {
56                physical_table_info.table_info.meta.schema.column_schemas()[idx]
57                    .name
58                    .clone()
59            })
60            .collect();
61
62        self.data.physical_partition_columns = physical_partition_columns;
63
64        Ok(())
65    }
66
67    pub(crate) fn merge_partition_columns_into_logical_tables(&mut self) -> Result<()> {
68        let partition_columns = &self.data.physical_partition_columns;
69
70        // Skip if no partition columns to add
71        if partition_columns.is_empty() {
72            return Ok(());
73        }
74
75        for task in &mut self.data.tasks {
76            // Get existing column names in the logical table
77            let column_schemas = task.table_info.meta.schema.column_schemas();
78            let existing_column_names: HashSet<_> =
79                column_schemas.iter().map(|c| &c.name).collect();
80
81            let mut new_columns = Vec::new();
82            let mut new_primary_key_indices = task.table_info.meta.primary_key_indices.clone();
83
84            // Add missing partition columns
85            for partition_column in partition_columns {
86                // The engine owns its identity columns: they are derived from the logical
87                // table at write time, so a logical schema must never declare them. Older
88                // procedures may have persisted an unfiltered partition column list, hence
89                // filtering here rather than only while collecting it.
90                if is_metric_engine_internal_column(partition_column) {
91                    continue;
92                }
93
94                if !existing_column_names.contains(partition_column) {
95                    let new_column_index = column_schemas.len() + new_columns.len();
96
97                    // Create new column schema for the partition column
98                    let column_schema = ColumnSchema::new(
99                        partition_column.clone(),
100                        ConcreteDataType::string_datatype(),
101                        true,
102                    );
103                    new_columns.push(column_schema);
104
105                    // Add to primary key indices (partition columns are part of primary key)
106                    new_primary_key_indices.push(new_column_index);
107                }
108            }
109
110            // If we added new columns, update the table info
111            if !new_columns.is_empty() {
112                let mut updated_columns = column_schemas.to_vec();
113                updated_columns.extend(new_columns);
114
115                // Create new schema with updated columns
116                let new_schema = Arc::new(Schema::new(updated_columns));
117
118                // Update the table info
119                task.table_info.meta.schema = new_schema;
120                task.table_info.meta.primary_key_indices = new_primary_key_indices;
121            }
122        }
123
124        Ok(())
125    }
126
127    pub(crate) async fn allocate_table_ids(&mut self) -> Result<()> {
128        for (task, table_id) in self
129            .data
130            .tasks
131            .iter_mut()
132            .zip(self.data.table_ids_already_exists.iter())
133        {
134            let table_id = if let Some(table_id) = table_id {
135                *table_id
136            } else {
137                self.context
138                    .table_metadata_allocator
139                    .allocate_table_id(&task.create_table.table_id)
140                    .await?
141            };
142            task.set_table_id(table_id);
143
144            // sort columns in task
145            task.sort_columns();
146        }
147
148        Ok(())
149    }
150}
common_meta/ddl/create_logical_tables/metadata.rs

common_meta/ddl/create_logical_tables/
metadata.rs
