catalog/
system_schema.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
15pub mod information_schema;
16mod memory_table;
17pub mod numbers_table_provider;
18pub mod pg_catalog;
19pub mod predicate;
20mod utils;
21
22use std::collections::HashMap;
23use std::sync::Arc;
24
25use common_error::ext::BoxedError;
26use common_recordbatch::{RecordBatchStreamWrapper, SendableRecordBatchStream};
27use common_telemetry::tracing::Span;
28use datafusion::physical_plan::ExecutionPlan;
29use datatypes::schema::SchemaRef;
30use futures_util::StreamExt;
31use snafu::ResultExt;
32use store_api::data_source::DataSource;
33use store_api::storage::ScanRequest;
34use table::error::{SchemaConversionSnafu, TablesRecordBatchSnafu};
35use table::metadata::{
36    FilterPushDownType, TableId, TableInfoBuilder, TableInfoRef, TableMetaBuilder, TableType,
37};
38use table::{Table, TableRef};
39
40use crate::error::Result;
41
42pub trait SystemSchemaProvider {
43    /// Returns a map of [TableRef] in information schema.
44    fn tables(&self) -> &HashMap<String, TableRef>;
45
46    /// Returns the [TableRef] by table name.
47    fn table(&self, name: &str) -> Option<TableRef> {
48        self.tables().get(name).cloned()
49    }
50
51    /// Returns table names in the order of table id.
52    fn table_names(&self) -> Vec<String> {
53        let mut tables = self.tables().values().clone().collect::<Vec<_>>();
54
55        tables.sort_by(|t1, t2| {
56            t1.table_info()
57                .table_id()
58                .partial_cmp(&t2.table_info().table_id())
59                .unwrap()
60        });
61        tables
62            .into_iter()
63            .map(|t| t.table_info().name.clone())
64            .collect()
65    }
66}
67
68trait SystemSchemaProviderInner {
69    fn catalog_name(&self) -> &str;
70    fn schema_name() -> &'static str;
71    fn build_table(&self, name: &str) -> Option<TableRef> {
72        self.system_table(name).map(|table| {
73            let table_info = Self::table_info(self.catalog_name().to_string(), &table);
74            let filter_pushdown = FilterPushDownType::Inexact;
75            let data_source = Arc::new(SystemTableDataSource::new(table));
76            let table = Table::new(table_info, filter_pushdown, data_source);
77            Arc::new(table)
78        })
79    }
80    fn system_table(&self, name: &str) -> Option<SystemTableRef>;
81
82    fn table_info(catalog_name: String, table: &SystemTableRef) -> TableInfoRef {
83        let table_meta = TableMetaBuilder::empty()
84            .schema(table.schema())
85            .primary_key_indices(vec![])
86            .next_column_id(0)
87            .build()
88            .unwrap();
89        let table_info = TableInfoBuilder::default()
90            .table_id(table.table_id())
91            .name(table.table_name().to_string())
92            .catalog_name(catalog_name)
93            .schema_name(Self::schema_name().to_string())
94            .meta(table_meta)
95            .table_type(table.table_type())
96            .build()
97            .unwrap();
98        Arc::new(table_info)
99    }
100}
101
102pub trait SystemTable {
103    fn table_id(&self) -> TableId;
104
105    fn table_name(&self) -> &'static str;
106
107    fn schema(&self) -> SchemaRef;
108
109    fn to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream>;
110
111    fn scan_plan(&self, _request: ScanRequest) -> Result<Option<Arc<dyn ExecutionPlan>>> {
112        Ok(None)
113    }
114
115    fn table_type(&self) -> TableType {
116        TableType::Temporary
117    }
118}
119
120pub type SystemTableRef = Arc<dyn SystemTable + Send + Sync>;
121
122struct SystemTableDataSource {
123    table: SystemTableRef,
124}
125
126impl SystemTableDataSource {
127    fn new(table: SystemTableRef) -> Self {
128        Self { table }
129    }
130
131    fn try_project(&self, projection: &[usize]) -> std::result::Result<SchemaRef, BoxedError> {
132        let schema = self
133            .table
134            .schema()
135            .try_project(projection)
136            .context(SchemaConversionSnafu)
137            .map_err(BoxedError::new)?;
138        Ok(Arc::new(schema))
139    }
140}
141
142impl DataSource for SystemTableDataSource {
143    fn get_stream(
144        &self,
145        request: ScanRequest,
146    ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
147        let projection = request
148            .projection_input
149            .as_ref()
150            .map(|input| input.projection.clone());
151
152        let projected_schema = match projection.as_ref() {
153            Some(projection) => self.try_project(projection)?,
154            None => self.table.schema(),
155        };
156
157        let stream = self
158            .table
159            .to_stream(request)
160            .map_err(BoxedError::new)
161            .context(TablesRecordBatchSnafu)
162            .map_err(BoxedError::new)?
163            .map(move |batch| match (&projection, batch) {
164                // Some tables (e.g., inspect tables) already honor projection in their inner stream;
165                // others ignore it and return full rows. We will only apply projection here if the
166                // inner batch width doesn't match the projection size.
167                (Some(p), Ok(b)) if b.num_columns() != p.len() => b.try_project(p),
168                (_, res) => res,
169            });
170
171        let stream = RecordBatchStreamWrapper {
172            schema: projected_schema,
173            stream: Box::pin(stream),
174            output_ordering: None,
175            metrics: Default::default(),
176            span: Span::current(),
177        };
178
179        Ok(Box::pin(stream))
180    }
181
182    fn get_physical_plan(
183        &self,
184        request: ScanRequest,
185    ) -> std::result::Result<Option<Arc<dyn ExecutionPlan>>, BoxedError> {
186        self.table
187            .scan_plan(request)
188            .map_err(BoxedError::new)
189            .context(TablesRecordBatchSnafu)
190            .map_err(BoxedError::new)
191    }
192}
catalog/system_schema.rs

catalog/
system_schema.rs
