table/table/
numbers.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::pin::Pin;
16use std::sync::Arc;
17
18use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
19use common_error::ext::BoxedError;
20use common_recordbatch::adapter::RecordBatchMetrics;
21use common_recordbatch::error::Result as RecordBatchResult;
22use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream};
23use datafusion::arrow::record_batch::RecordBatch as DfRecordBatch;
24use datatypes::arrow::array::UInt32Array;
25use datatypes::data_type::ConcreteDataType;
26use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef};
27use futures::Stream;
28use futures::task::{Context, Poll};
29use store_api::data_source::DataSource;
30use store_api::storage::ScanRequest;
31
32use crate::metadata::{
33    FilterPushDownType, TableId, TableInfoBuilder, TableInfoRef, TableMeta, TableType,
34};
35use crate::{Table, TableRef};
36
37const NUMBER_COLUMN: &str = "number";
38
39pub const NUMBERS_TABLE_NAME: &str = "numbers";
40
41/// numbers table for test
42#[derive(Debug, Clone)]
43pub struct NumbersTable;
44
45impl NumbersTable {
46    pub fn table(table_id: TableId) -> TableRef {
47        Self::table_with_name(table_id, NUMBERS_TABLE_NAME.to_string())
48    }
49
50    pub fn table_with_name(table_id: TableId, name: String) -> TableRef {
51        let data_source = Arc::new(NumbersDataSource::new(Self::schema()));
52        let table = Table::new(
53            Self::table_info(table_id, name, "test_engine".to_string()),
54            FilterPushDownType::Unsupported,
55            data_source,
56        );
57        Arc::new(table)
58    }
59
60    pub fn schema() -> SchemaRef {
61        let column_schemas = vec![ColumnSchema::new(
62            NUMBER_COLUMN,
63            ConcreteDataType::uint32_datatype(),
64            false,
65        )];
66        let schema = SchemaBuilder::try_from_columns(column_schemas)
67            .unwrap()
68            .build()
69            .unwrap();
70        Arc::new(schema)
71    }
72
73    pub fn table_info(table_id: TableId, name: String, engine: String) -> TableInfoRef {
74        let table_meta = TableMeta {
75            schema: Self::schema(),
76            primary_key_indices: vec![0],
77            value_indices: vec![],
78            engine,
79            next_column_id: 1,
80            options: Default::default(),
81            created_on: Default::default(),
82            updated_on: Default::default(),
83            partition_key_indices: vec![],
84            column_ids: vec![],
85        };
86
87        let table_info = TableInfoBuilder::default()
88            .table_id(table_id)
89            .name(name)
90            .catalog_name(DEFAULT_CATALOG_NAME)
91            .schema_name(DEFAULT_SCHEMA_NAME)
92            .table_version(0)
93            .table_type(TableType::Temporary)
94            .meta(table_meta)
95            .build()
96            .unwrap();
97        Arc::new(table_info)
98    }
99}
100
101struct NumbersDataSource {
102    schema: SchemaRef,
103}
104
105impl NumbersDataSource {
106    pub fn new(schema: SchemaRef) -> Self {
107        Self { schema }
108    }
109}
110
111impl DataSource for NumbersDataSource {
112    fn get_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream, BoxedError> {
113        let projection = request.projection_input.map(|input| input.projection);
114        let projected_schema = match &projection {
115            Some(projection) => Arc::new(self.schema.try_project(projection).unwrap()),
116            None => self.schema.clone(),
117        };
118        Ok(Box::pin(NumbersStream {
119            limit: request.limit.unwrap_or(100) as u32,
120            schema: self.schema.clone(),
121            already_run: false,
122            projection,
123            projected_schema,
124        }))
125    }
126}
127
128// Limited numbers stream
129struct NumbersStream {
130    limit: u32,
131    schema: SchemaRef,
132    already_run: bool,
133    projection: Option<Vec<usize>>,
134    projected_schema: SchemaRef,
135}
136
137impl RecordBatchStream for NumbersStream {
138    fn schema(&self) -> SchemaRef {
139        self.projected_schema.clone()
140    }
141
142    fn output_ordering(&self) -> Option<&[OrderOption]> {
143        None
144    }
145
146    fn metrics(&self) -> Option<RecordBatchMetrics> {
147        None
148    }
149}
150
151impl Stream for NumbersStream {
152    type Item = RecordBatchResult<RecordBatch>;
153
154    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
155        if self.already_run {
156            return Poll::Ready(None);
157        }
158        self.already_run = true;
159        let numbers: Vec<u32> = (0..self.limit).collect();
160        let mut batch = DfRecordBatch::try_new(
161            self.schema.arrow_schema().clone(),
162            vec![Arc::new(UInt32Array::from(numbers))],
163        )
164        .unwrap();
165        if let Some(projection) = &self.projection {
166            batch = batch.project(projection).unwrap();
167        }
168
169        Poll::Ready(Some(Ok(RecordBatch::from_df_record_batch(
170            self.projected_schema.clone(),
171            batch,
172        ))))
173    }
174}
table/table/numbers.rs

table/table/
numbers.rs
