catalog/system_schema/
memory_table.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
15mod table_columns;
16
17use std::sync::Arc;
18
19use arrow_schema::SchemaRef as ArrowSchemaRef;
20use common_error::ext::BoxedError;
21use common_recordbatch::adapter::RecordBatchStreamAdapter;
22use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
23use datafusion::execution::TaskContext;
24use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
25use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
26use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream;
27use datatypes::schema::SchemaRef;
28use datatypes::vectors::VectorRef;
29use snafu::ResultExt;
30use store_api::storage::{ScanRequest, TableId};
31
32use crate::error::{CreateRecordBatchSnafu, InternalSnafu, Result};
33use crate::system_schema::SystemTable;
34use crate::system_schema::information_schema::{InformationSchemaTableScope, InformationTable};
35
36/// A memory table with specified schema and columns.
37#[derive(Debug)]
38pub(crate) struct MemoryTable {
39    pub(crate) table_id: TableId,
40    pub(crate) table_name: &'static str,
41    pub(crate) schema: SchemaRef,
42    pub(crate) columns: Vec<VectorRef>,
43    scope: InformationSchemaTableScope,
44}
45
46impl MemoryTable {
47    /// Creates a memory table with table id, name, schema and columns.
48    pub fn new(
49        table_id: TableId,
50        table_name: &'static str,
51        schema: SchemaRef,
52        columns: Vec<VectorRef>,
53        scope: InformationSchemaTableScope,
54    ) -> Self {
55        Self {
56            table_id,
57            table_name,
58            schema,
59            columns,
60            scope,
61        }
62    }
63
64    pub fn builder(&self) -> MemoryTableBuilder {
65        MemoryTableBuilder::new(self.schema.clone(), self.columns.clone())
66    }
67}
68
69pub(crate) struct MemoryTableBuilder {
70    schema: SchemaRef,
71    columns: Vec<VectorRef>,
72}
73
74impl MemoryTableBuilder {
75    fn new(schema: SchemaRef, columns: Vec<VectorRef>) -> Self {
76        Self { schema, columns }
77    }
78
79    /// Construct the `information_schema.{table_name}` virtual table
80    pub async fn memory_records(&mut self) -> Result<RecordBatch> {
81        if self.columns.is_empty() {
82            Ok(RecordBatch::new_empty(self.schema.clone()))
83        } else {
84            RecordBatch::new(self.schema.clone(), std::mem::take(&mut self.columns))
85                .context(CreateRecordBatchSnafu)
86        }
87    }
88}
89
90impl DfPartitionStream for MemoryTable {
91    fn schema(&self) -> &ArrowSchemaRef {
92        self.schema.arrow_schema()
93    }
94
95    fn execute(&self, _: Arc<TaskContext>) -> DfSendableRecordBatchStream {
96        let schema = self.schema.arrow_schema().clone();
97        let mut builder = self.builder();
98        Box::pin(DfRecordBatchStreamAdapter::new(
99            schema,
100            futures::stream::once(async move {
101                builder
102                    .memory_records()
103                    .await
104                    .map(|x| x.into_df_record_batch())
105                    .map_err(Into::into)
106            }),
107        ))
108    }
109}
110
111impl InformationTable for MemoryTable {
112    fn scope(&self) -> InformationSchemaTableScope {
113        self.scope
114    }
115}
116
117impl SystemTable for MemoryTable {
118    fn table_id(&self) -> TableId {
119        self.table_id
120    }
121
122    fn table_name(&self) -> &'static str {
123        self.table_name
124    }
125
126    fn schema(&self) -> SchemaRef {
127        self.schema.clone()
128    }
129
130    fn to_stream(&self, _request: ScanRequest) -> Result<SendableRecordBatchStream> {
131        let schema = self.schema.arrow_schema().clone();
132        let mut builder = self.builder();
133        let stream = Box::pin(DfRecordBatchStreamAdapter::new(
134            schema,
135            futures::stream::once(async move {
136                builder
137                    .memory_records()
138                    .await
139                    .map(|x| x.into_df_record_batch())
140                    .map_err(Into::into)
141            }),
142        ));
143        Ok(Box::pin(
144            RecordBatchStreamAdapter::try_new(stream)
145                .map_err(BoxedError::new)
146                .context(InternalSnafu)?,
147        ))
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use std::sync::Arc;
154
155    use common_recordbatch::RecordBatches;
156    use datatypes::prelude::ConcreteDataType;
157    use datatypes::schema::{ColumnSchema, Schema};
158    use datatypes::vectors::StringVector;
159
160    use super::*;
161    use crate::system_schema::SystemTable;
162
163    #[tokio::test]
164    async fn test_memory_table() {
165        let schema = Arc::new(Schema::new(vec![
166            ColumnSchema::new("a", ConcreteDataType::string_datatype(), false),
167            ColumnSchema::new("b", ConcreteDataType::string_datatype(), false),
168        ]));
169
170        let table = MemoryTable::new(
171            42,
172            "test",
173            schema.clone(),
174            vec![
175                Arc::new(StringVector::from(vec!["a1", "a2"])),
176                Arc::new(StringVector::from(vec!["b1", "b2"])),
177            ],
178            InformationSchemaTableScope::Global,
179        );
180
181        assert_eq!(42, table.table_id());
182        assert_eq!("test", table.table_name);
183        assert_eq!(schema, SystemTable::schema(&table));
184
185        let stream = table.to_stream(ScanRequest::default()).unwrap();
186
187        let batches = RecordBatches::try_collect(stream).await.unwrap();
188
189        assert_eq!(
190            "\
191+----+----+
192| a  | b  |
193+----+----+
194| a1 | b1 |
195| a2 | b2 |
196+----+----+",
197            batches.pretty_print().unwrap()
198        );
199    }
200
201    #[tokio::test]
202    async fn test_empty_memory_table() {
203        let schema = Arc::new(Schema::new(vec![
204            ColumnSchema::new("a", ConcreteDataType::string_datatype(), false),
205            ColumnSchema::new("b", ConcreteDataType::string_datatype(), false),
206        ]));
207
208        let table = MemoryTable::new(
209            42,
210            "test",
211            schema.clone(),
212            vec![],
213            InformationSchemaTableScope::Global,
214        );
215
216        assert_eq!(42, table.table_id());
217        assert_eq!("test", table.table_name());
218        assert_eq!(schema, SystemTable::schema(&table));
219
220        let stream = table.to_stream(ScanRequest::default()).unwrap();
221
222        let batches = RecordBatches::try_collect(stream).await.unwrap();
223
224        assert_eq!(
225            "\
226+---+---+
227| a | b |
228+---+---+
229+---+---+",
230            batches.pretty_print().unwrap()
231        );
232    }
233}
catalog/system_schema/memory_table.rs

catalog/system_schema/
memory_table.rs
