catalog/system_schema/information_schema/
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::sync::{Arc, Weak};
16
17use arrow_schema::SchemaRef as ArrowSchemaRef;
18use common_catalog::consts::{
19    INFORMATION_SCHEMA_COLUMNS_TABLE_ID, SEMANTIC_TYPE_FIELD, SEMANTIC_TYPE_PRIMARY_KEY,
20    SEMANTIC_TYPE_TIME_INDEX,
21};
22use common_error::ext::BoxedError;
23use common_recordbatch::adapter::RecordBatchStreamAdapter;
24use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
25use datafusion::execution::TaskContext;
26use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
27use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
28use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream;
29use datatypes::prelude::{ConcreteDataType, DataType, MutableVector};
30use datatypes::scalars::ScalarVectorBuilder;
31use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
32use datatypes::value::Value;
33use datatypes::vectors::{
34    Int64Vector, Int64VectorBuilder, StringVector, StringVectorBuilder, VectorRef,
35};
36use futures::TryStreamExt;
37use snafu::{OptionExt, ResultExt};
38use sql::statements;
39use store_api::storage::{ScanRequest, TableId};
40
41use crate::CatalogManager;
42use crate::error::{
43    CreateRecordBatchSnafu, InternalSnafu, Result, UpgradeWeakCatalogManagerRefSnafu,
44};
45use crate::information_schema::Predicates;
46use crate::system_schema::SystemTable;
47use crate::system_schema::information_schema::{
48    COLUMNS, InformationSchemaTableScope, InformationTable,
49};
50
51#[derive(Debug)]
52pub(super) struct InformationSchemaColumns {
53    schema: SchemaRef,
54    catalog_name: String,
55    catalog_manager: Weak<dyn CatalogManager>,
56}
57
58pub const TABLE_CATALOG: &str = "table_catalog";
59pub const TABLE_SCHEMA: &str = "table_schema";
60pub const TABLE_NAME: &str = "table_name";
61pub const COLUMN_NAME: &str = "column_name";
62pub const REGION_ID: &str = "region_id";
63pub const PEER_ID: &str = "peer_id";
64const ORDINAL_POSITION: &str = "ordinal_position";
65const CHARACTER_MAXIMUM_LENGTH: &str = "character_maximum_length";
66const CHARACTER_OCTET_LENGTH: &str = "character_octet_length";
67const NUMERIC_PRECISION: &str = "numeric_precision";
68const NUMERIC_SCALE: &str = "numeric_scale";
69const DATETIME_PRECISION: &str = "datetime_precision";
70const CHARACTER_SET_NAME: &str = "character_set_name";
71pub const COLLATION_NAME: &str = "collation_name";
72pub const COLUMN_KEY: &str = "column_key";
73pub const EXTRA: &str = "extra";
74pub const PRIVILEGES: &str = "privileges";
75const GENERATION_EXPRESSION: &str = "generation_expression";
76// Extension field to keep greptime data type name
77pub const GREPTIME_DATA_TYPE: &str = "greptime_data_type";
78pub const DATA_TYPE: &str = "data_type";
79pub const SEMANTIC_TYPE: &str = "semantic_type";
80pub const COLUMN_DEFAULT: &str = "column_default";
81pub const IS_NULLABLE: &str = "is_nullable";
82const COLUMN_TYPE: &str = "column_type";
83pub const COLUMN_COMMENT: &str = "column_comment";
84const SRS_ID: &str = "srs_id";
85const INIT_CAPACITY: usize = 42;
86
87// The maximum length of string type
88const MAX_STRING_LENGTH: i64 = 2147483647;
89const UTF8_CHARSET_NAME: &str = "utf8";
90const UTF8_COLLATE_NAME: &str = "utf8_bin";
91const PRI_COLUMN_KEY: &str = "PRI";
92const TIME_INDEX_COLUMN_KEY: &str = "TIME INDEX";
93const DEFAULT_PRIVILEGES: &str = "select,insert";
94const EMPTY_STR: &str = "";
95const YES: &str = "YES";
96const NO: &str = "NO";
97
98impl InformationSchemaColumns {
99    pub(super) fn new(catalog_name: String, catalog_manager: Weak<dyn CatalogManager>) -> Self {
100        Self {
101            schema: Self::schema(),
102            catalog_name,
103            catalog_manager,
104        }
105    }
106
107    fn schema() -> SchemaRef {
108        Arc::new(Schema::new(vec![
109            ColumnSchema::new(TABLE_CATALOG, ConcreteDataType::string_datatype(), false),
110            ColumnSchema::new(TABLE_SCHEMA, ConcreteDataType::string_datatype(), false),
111            ColumnSchema::new(TABLE_NAME, ConcreteDataType::string_datatype(), false),
112            ColumnSchema::new(COLUMN_NAME, ConcreteDataType::string_datatype(), false),
113            ColumnSchema::new(ORDINAL_POSITION, ConcreteDataType::int64_datatype(), false),
114            ColumnSchema::new(
115                CHARACTER_MAXIMUM_LENGTH,
116                ConcreteDataType::int64_datatype(),
117                true,
118            ),
119            ColumnSchema::new(
120                CHARACTER_OCTET_LENGTH,
121                ConcreteDataType::int64_datatype(),
122                true,
123            ),
124            ColumnSchema::new(NUMERIC_PRECISION, ConcreteDataType::int64_datatype(), true),
125            ColumnSchema::new(NUMERIC_SCALE, ConcreteDataType::int64_datatype(), true),
126            ColumnSchema::new(DATETIME_PRECISION, ConcreteDataType::int64_datatype(), true),
127            ColumnSchema::new(
128                CHARACTER_SET_NAME,
129                ConcreteDataType::string_datatype(),
130                true,
131            ),
132            ColumnSchema::new(COLLATION_NAME, ConcreteDataType::string_datatype(), true),
133            ColumnSchema::new(COLUMN_KEY, ConcreteDataType::string_datatype(), false),
134            ColumnSchema::new(EXTRA, ConcreteDataType::string_datatype(), false),
135            ColumnSchema::new(PRIVILEGES, ConcreteDataType::string_datatype(), false),
136            ColumnSchema::new(
137                GENERATION_EXPRESSION,
138                ConcreteDataType::string_datatype(),
139                false,
140            ),
141            ColumnSchema::new(
142                GREPTIME_DATA_TYPE,
143                ConcreteDataType::string_datatype(),
144                false,
145            ),
146            ColumnSchema::new(DATA_TYPE, ConcreteDataType::string_datatype(), false),
147            ColumnSchema::new(SEMANTIC_TYPE, ConcreteDataType::string_datatype(), false),
148            ColumnSchema::new(COLUMN_DEFAULT, ConcreteDataType::string_datatype(), true),
149            ColumnSchema::new(IS_NULLABLE, ConcreteDataType::string_datatype(), false),
150            ColumnSchema::new(COLUMN_TYPE, ConcreteDataType::string_datatype(), false),
151            ColumnSchema::new(COLUMN_COMMENT, ConcreteDataType::string_datatype(), true),
152            ColumnSchema::new(SRS_ID, ConcreteDataType::int64_datatype(), true),
153        ]))
154    }
155
156    fn builder(&self) -> InformationSchemaColumnsBuilder {
157        InformationSchemaColumnsBuilder::new(
158            self.schema.clone(),
159            self.catalog_name.clone(),
160            self.catalog_manager.clone(),
161        )
162    }
163}
164
165impl InformationTable for InformationSchemaColumns {
166    fn scope(&self) -> InformationSchemaTableScope {
167        InformationSchemaTableScope::Catalog
168    }
169}
170
171impl SystemTable for InformationSchemaColumns {
172    fn table_id(&self) -> TableId {
173        INFORMATION_SCHEMA_COLUMNS_TABLE_ID
174    }
175
176    fn table_name(&self) -> &'static str {
177        COLUMNS
178    }
179
180    fn schema(&self) -> SchemaRef {
181        self.schema.clone()
182    }
183
184    fn to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
185        let schema = self.schema.arrow_schema().clone();
186        let mut builder = self.builder();
187        let stream = Box::pin(DfRecordBatchStreamAdapter::new(
188            schema,
189            futures::stream::once(async move {
190                builder
191                    .make_columns(Some(request))
192                    .await
193                    .map(|x| x.into_df_record_batch())
194                    .map_err(Into::into)
195            }),
196        ));
197        Ok(Box::pin(
198            RecordBatchStreamAdapter::try_new(stream)
199                .map_err(BoxedError::new)
200                .context(InternalSnafu)?,
201        ))
202    }
203}
204
205struct InformationSchemaColumnsBuilder {
206    schema: SchemaRef,
207    catalog_name: String,
208    catalog_manager: Weak<dyn CatalogManager>,
209
210    catalog_names: StringVectorBuilder,
211    schema_names: StringVectorBuilder,
212    table_names: StringVectorBuilder,
213    column_names: StringVectorBuilder,
214    ordinal_positions: Int64VectorBuilder,
215    character_maximum_lengths: Int64VectorBuilder,
216    character_octet_lengths: Int64VectorBuilder,
217    numeric_precisions: Int64VectorBuilder,
218    numeric_scales: Int64VectorBuilder,
219    datetime_precisions: Int64VectorBuilder,
220    character_set_names: StringVectorBuilder,
221    collation_names: StringVectorBuilder,
222    column_keys: StringVectorBuilder,
223    greptime_data_types: StringVectorBuilder,
224    data_types: StringVectorBuilder,
225    semantic_types: StringVectorBuilder,
226    column_defaults: StringVectorBuilder,
227    is_nullables: StringVectorBuilder,
228    column_types: StringVectorBuilder,
229    column_comments: StringVectorBuilder,
230}
231
232impl InformationSchemaColumnsBuilder {
233    fn new(
234        schema: SchemaRef,
235        catalog_name: String,
236        catalog_manager: Weak<dyn CatalogManager>,
237    ) -> Self {
238        Self {
239            schema,
240            catalog_name,
241            catalog_manager,
242            catalog_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
243            schema_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
244            table_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
245            column_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
246            ordinal_positions: Int64VectorBuilder::with_capacity(INIT_CAPACITY),
247            character_maximum_lengths: Int64VectorBuilder::with_capacity(INIT_CAPACITY),
248            character_octet_lengths: Int64VectorBuilder::with_capacity(INIT_CAPACITY),
249            numeric_precisions: Int64VectorBuilder::with_capacity(INIT_CAPACITY),
250            numeric_scales: Int64VectorBuilder::with_capacity(INIT_CAPACITY),
251            datetime_precisions: Int64VectorBuilder::with_capacity(INIT_CAPACITY),
252            character_set_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
253            collation_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
254            column_keys: StringVectorBuilder::with_capacity(INIT_CAPACITY),
255            greptime_data_types: StringVectorBuilder::with_capacity(INIT_CAPACITY),
256            data_types: StringVectorBuilder::with_capacity(INIT_CAPACITY),
257            semantic_types: StringVectorBuilder::with_capacity(INIT_CAPACITY),
258            column_defaults: StringVectorBuilder::with_capacity(INIT_CAPACITY),
259            is_nullables: StringVectorBuilder::with_capacity(INIT_CAPACITY),
260            column_types: StringVectorBuilder::with_capacity(INIT_CAPACITY),
261            column_comments: StringVectorBuilder::with_capacity(INIT_CAPACITY),
262        }
263    }
264
265    /// Construct the `information_schema.columns` virtual table
266    async fn make_columns(&mut self, request: Option<ScanRequest>) -> Result<RecordBatch> {
267        let catalog_name = self.catalog_name.clone();
268        let catalog_manager = self
269            .catalog_manager
270            .upgrade()
271            .context(UpgradeWeakCatalogManagerRefSnafu)?;
272        let predicates = Predicates::from_scan_request(&request);
273
274        for schema_name in catalog_manager.schema_names(&catalog_name, None).await? {
275            let mut stream = catalog_manager.tables(&catalog_name, &schema_name, None);
276
277            while let Some(table) = stream.try_next().await? {
278                let keys = &table.table_info().meta.primary_key_indices;
279                let schema = table.schema();
280
281                for (idx, column) in schema.column_schemas().iter().enumerate() {
282                    let semantic_type = if column.is_time_index() {
283                        SEMANTIC_TYPE_TIME_INDEX
284                    } else if keys.contains(&idx) {
285                        SEMANTIC_TYPE_PRIMARY_KEY
286                    } else {
287                        SEMANTIC_TYPE_FIELD
288                    };
289
290                    self.add_column(
291                        &predicates,
292                        idx,
293                        &catalog_name,
294                        &schema_name,
295                        &table.table_info().name,
296                        semantic_type,
297                        column,
298                    );
299                }
300            }
301        }
302
303        self.finish()
304    }
305
306    #[allow(clippy::too_many_arguments)]
307    fn add_column(
308        &mut self,
309        predicates: &Predicates,
310        index: usize,
311        catalog_name: &str,
312        schema_name: &str,
313        table_name: &str,
314        semantic_type: &str,
315        column_schema: &ColumnSchema,
316    ) {
317        // Use sql data type name
318        let data_type = statements::concrete_data_type_to_sql_data_type(&column_schema.data_type)
319            .map(|dt| dt.to_string().to_lowercase())
320            .unwrap_or_else(|_| column_schema.data_type.name());
321
322        let column_key = match semantic_type {
323            SEMANTIC_TYPE_PRIMARY_KEY => PRI_COLUMN_KEY,
324            SEMANTIC_TYPE_TIME_INDEX => TIME_INDEX_COLUMN_KEY,
325            _ => EMPTY_STR,
326        };
327
328        let row = [
329            (TABLE_CATALOG, &Value::from(catalog_name)),
330            (TABLE_SCHEMA, &Value::from(schema_name)),
331            (TABLE_NAME, &Value::from(table_name)),
332            (COLUMN_NAME, &Value::from(column_schema.name.as_str())),
333            (DATA_TYPE, &Value::from(data_type.as_str())),
334            (SEMANTIC_TYPE, &Value::from(semantic_type)),
335            (ORDINAL_POSITION, &Value::from((index + 1) as i64)),
336            (COLUMN_KEY, &Value::from(column_key)),
337        ];
338
339        if !predicates.eval(&row) {
340            return;
341        }
342
343        self.catalog_names.push(Some(catalog_name));
344        self.schema_names.push(Some(schema_name));
345        self.table_names.push(Some(table_name));
346        self.column_names.push(Some(&column_schema.name));
347        // Starts from 1
348        self.ordinal_positions.push(Some((index + 1) as i64));
349
350        if column_schema.data_type.is_string() {
351            self.character_maximum_lengths.push(Some(MAX_STRING_LENGTH));
352            self.character_octet_lengths.push(Some(MAX_STRING_LENGTH));
353            self.numeric_precisions.push(None);
354            self.numeric_scales.push(None);
355            self.datetime_precisions.push(None);
356            self.character_set_names.push(Some(UTF8_CHARSET_NAME));
357            self.collation_names.push(Some(UTF8_COLLATE_NAME));
358        } else if column_schema.data_type.is_numeric() || column_schema.data_type.is_decimal() {
359            self.character_maximum_lengths.push(None);
360            self.character_octet_lengths.push(None);
361
362            self.numeric_precisions.push(
363                column_schema
364                    .data_type
365                    .numeric_precision()
366                    .map(|x| x as i64),
367            );
368            self.numeric_scales
369                .push(column_schema.data_type.numeric_scale().map(|x| x as i64));
370
371            self.datetime_precisions.push(None);
372            self.character_set_names.push(None);
373            self.collation_names.push(None);
374        } else {
375            self.character_maximum_lengths.push(None);
376            self.character_octet_lengths.push(None);
377            self.numeric_precisions.push(None);
378            self.numeric_scales.push(None);
379
380            match &column_schema.data_type {
381                ConcreteDataType::Timestamp(ts_type) => {
382                    self.datetime_precisions
383                        .push(Some(ts_type.precision() as i64));
384                }
385                ConcreteDataType::Time(time_type) => {
386                    self.datetime_precisions
387                        .push(Some(time_type.precision() as i64));
388                }
389                _ => self.datetime_precisions.push(None),
390            }
391
392            self.character_set_names.push(None);
393            self.collation_names.push(None);
394        }
395
396        self.column_keys.push(Some(column_key));
397        self.greptime_data_types
398            .push(Some(&column_schema.data_type.name()));
399        self.data_types.push(Some(&data_type));
400        self.semantic_types.push(Some(semantic_type));
401        self.column_defaults.push(
402            column_schema
403                .default_constraint()
404                .map(|s| format!("{}", s))
405                .as_deref(),
406        );
407        if column_schema.is_nullable() {
408            self.is_nullables.push(Some(YES));
409        } else {
410            self.is_nullables.push(Some(NO));
411        }
412        self.column_types.push(Some(&data_type));
413        let column_comment = column_schema.column_comment().map(|x| x.as_ref());
414        self.column_comments.push(column_comment);
415    }
416
417    fn finish(&mut self) -> Result<RecordBatch> {
418        let rows_num = self.collation_names.len();
419
420        let privileges = Arc::new(StringVector::from(vec![DEFAULT_PRIVILEGES; rows_num]));
421        let empty_string = Arc::new(StringVector::from(vec![EMPTY_STR; rows_num]));
422        let srs_ids = Arc::new(Int64Vector::from(vec![None; rows_num]));
423
424        let columns: Vec<VectorRef> = vec![
425            Arc::new(self.catalog_names.finish()),
426            Arc::new(self.schema_names.finish()),
427            Arc::new(self.table_names.finish()),
428            Arc::new(self.column_names.finish()),
429            Arc::new(self.ordinal_positions.finish()),
430            Arc::new(self.character_maximum_lengths.finish()),
431            Arc::new(self.character_octet_lengths.finish()),
432            Arc::new(self.numeric_precisions.finish()),
433            Arc::new(self.numeric_scales.finish()),
434            Arc::new(self.datetime_precisions.finish()),
435            Arc::new(self.character_set_names.finish()),
436            Arc::new(self.collation_names.finish()),
437            Arc::new(self.column_keys.finish()),
438            empty_string.clone(),
439            privileges,
440            empty_string,
441            Arc::new(self.greptime_data_types.finish()),
442            Arc::new(self.data_types.finish()),
443            Arc::new(self.semantic_types.finish()),
444            Arc::new(self.column_defaults.finish()),
445            Arc::new(self.is_nullables.finish()),
446            Arc::new(self.column_types.finish()),
447            Arc::new(self.column_comments.finish()),
448            srs_ids,
449        ];
450
451        RecordBatch::new(self.schema.clone(), columns).context(CreateRecordBatchSnafu)
452    }
453}
454
455impl DfPartitionStream for InformationSchemaColumns {
456    fn schema(&self) -> &ArrowSchemaRef {
457        self.schema.arrow_schema()
458    }
459
460    fn execute(&self, _: Arc<TaskContext>) -> DfSendableRecordBatchStream {
461        let schema = self.schema.arrow_schema().clone();
462        let mut builder = self.builder();
463        Box::pin(DfRecordBatchStreamAdapter::new(
464            schema,
465            futures::stream::once(async move {
466                builder
467                    .make_columns(None)
468                    .await
469                    .map(|x| x.into_df_record_batch())
470                    .map_err(Into::into)
471            }),
472        ))
473    }
474}
catalog/system_schema/information_schema/columns.rs

catalog/system_schema/information_schema/
columns.rs
