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 projected_schema = match &request.projection {
114            Some(projection) => Arc::new(self.schema.try_project(projection).unwrap()),
115            None => self.schema.clone(),
116        };
117        Ok(Box::pin(NumbersStream {
118            limit: request.limit.unwrap_or(100) as u32,
119            schema: self.schema.clone(),
120            already_run: false,
121            projection: request.projection,
122            projected_schema,
123        }))
124    }
125}
126
127// Limited numbers stream
128struct NumbersStream {
129    limit: u32,
130    schema: SchemaRef,
131    already_run: bool,
132    projection: Option<Vec<usize>>,
133    projected_schema: SchemaRef,
134}
135
136impl RecordBatchStream for NumbersStream {
137    fn schema(&self) -> SchemaRef {
138        self.projected_schema.clone()
139    }
140
141    fn output_ordering(&self) -> Option<&[OrderOption]> {
142        None
143    }
144
145    fn metrics(&self) -> Option<RecordBatchMetrics> {
146        None
147    }
148}
149
150impl Stream for NumbersStream {
151    type Item = RecordBatchResult<RecordBatch>;
152
153    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
154        if self.already_run {
155            return Poll::Ready(None);
156        }
157        self.already_run = true;
158        let numbers: Vec<u32> = (0..self.limit).collect();
159        let mut batch = DfRecordBatch::try_new(
160            self.schema.arrow_schema().clone(),
161            vec![Arc::new(UInt32Array::from(numbers))],
162        )
163        .unwrap();
164        if let Some(projection) = &self.projection {
165            batch = batch.project(projection).unwrap();
166        }
167
168        Poll::Ready(Some(Ok(RecordBatch::from_df_record_batch(
169            self.projected_schema.clone(),
170            batch,
171        ))))
172    }
173}
table/table/numbers.rs

table/table/
numbers.rs
