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

catalog/system_schema/
memory_table.rs
