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

catalog/
system_schema.rs
