table/
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
15use std::collections::{HashMap, HashSet};
16use std::sync::Arc;
17
18use common_recordbatch::SendableRecordBatchStream;
19use datafusion::execution::FunctionRegistry;
20use datafusion::logical_expr::Cast;
21use datafusion::logical_expr::expr::ScalarFunction;
22use datafusion::physical_plan::ExecutionPlan;
23use datafusion::prelude::SessionContext;
24use datafusion_expr::expr::Expr;
25use datatypes::data_type::DataType;
26use datatypes::prelude::ConcreteDataType;
27use datatypes::schema::constraint::{CURRENT_TIMESTAMP, CURRENT_TIMESTAMP_FN, NOW_FN};
28use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, SchemaRef};
29use lazy_static::lazy_static;
30use snafu::ResultExt;
31use store_api::data_source::DataSourceRef;
32use store_api::storage::ScanRequest;
33
34use crate::error::{Result, TablesRecordBatchSnafu};
35use crate::metadata::{FilterPushDownType, TableInfoRef, TableType};
36
37pub mod adapter;
38mod metrics;
39pub mod numbers;
40pub mod scan;
41
42lazy_static! {
43    /// The [`Expr`] to call UDF function `now()`.
44    static ref NOW_EXPR: Expr = {
45        let ctx = SessionContext::new();
46
47        let now_udf = ctx.udf("now").expect("now UDF not found");
48
49        Expr::ScalarFunction(ScalarFunction {
50            func: now_udf,
51            args: vec![],
52        })
53    };
54}
55
56/// Defines partition rules for a table.
57/// TODO(discord9): add the entire partition exprs rules here later
58pub struct PartitionRules {
59    /// The physical partition columns that are not in the logical table.
60    /// only used in kvbackend manager to store the physical partition columns that are not in the logical table.
61    /// This is used to avoid the partition columns in the physical table that are not in the logical table
62    /// to prevent certain optimizations, if table is not a logical table, this should be empty
63    pub extra_phy_cols_not_in_logical_table: Vec<String>,
64}
65
66pub type TableRef = Arc<Table>;
67
68/// Table handle.
69pub struct Table {
70    table_info: TableInfoRef,
71    filter_pushdown: FilterPushDownType,
72    data_source: DataSourceRef,
73    /// Columns default [`Expr`]
74    column_defaults: HashMap<String, Expr>,
75    partition_rules: Option<PartitionRules>,
76}
77
78impl Table {
79    pub fn new(
80        table_info: TableInfoRef,
81        filter_pushdown: FilterPushDownType,
82        data_source: DataSourceRef,
83    ) -> Self {
84        Self {
85            column_defaults: collect_column_defaults(table_info.meta.schema.column_schemas()),
86            table_info,
87            filter_pushdown,
88            data_source,
89            partition_rules: None,
90        }
91    }
92
93    pub fn new_partitioned(
94        table_info: TableInfoRef,
95        filter_pushdown: FilterPushDownType,
96        data_source: DataSourceRef,
97        partition_rules: Option<PartitionRules>,
98    ) -> Self {
99        Self {
100            column_defaults: collect_column_defaults(table_info.meta.schema.column_schemas()),
101            table_info,
102            filter_pushdown,
103            data_source,
104            partition_rules,
105        }
106    }
107
108    /// Get column default [`Expr`], if available.
109    pub fn get_column_default(&self, column: &str) -> Option<&Expr> {
110        self.column_defaults.get(column)
111    }
112
113    pub fn data_source(&self) -> DataSourceRef {
114        self.data_source.clone()
115    }
116
117    /// Get a reference to the schema for this table.
118    pub fn schema(&self) -> SchemaRef {
119        self.table_info.meta.schema.clone()
120    }
121
122    pub fn schema_ref(&self) -> &SchemaRef {
123        &self.table_info.meta.schema
124    }
125
126    /// Get a reference to the table info.
127    pub fn table_info(&self) -> TableInfoRef {
128        self.table_info.clone()
129    }
130
131    /// Get the type of this table for metadata/catalog purposes.
132    pub fn table_type(&self) -> TableType {
133        self.table_info.table_type
134    }
135
136    pub fn partition_rules(&self) -> Option<&PartitionRules> {
137        self.partition_rules.as_ref()
138    }
139
140    pub async fn scan_to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
141        self.data_source
142            .get_stream(request)
143            .context(TablesRecordBatchSnafu)
144    }
145
146    pub fn scan_to_plan(&self, request: ScanRequest) -> Result<Option<Arc<dyn ExecutionPlan>>> {
147        self.data_source
148            .get_physical_plan(request)
149            .context(TablesRecordBatchSnafu)
150    }
151
152    /// Tests whether the table provider can make use of any or all filter expressions
153    /// to optimise data retrieval.
154    pub fn supports_filters_pushdown(&self, filters: &[&Expr]) -> Result<Vec<FilterPushDownType>> {
155        Ok(vec![self.filter_pushdown; filters.len()])
156    }
157
158    /// Get primary key columns in the definition order.
159    pub fn primary_key_columns(&self) -> impl Iterator<Item = ColumnSchema> + '_ {
160        self.table_info
161            .meta
162            .primary_key_indices
163            .iter()
164            .map(|i| self.table_info.meta.schema.column_schemas()[*i].clone())
165    }
166
167    /// Get field columns in the definition order.
168    pub fn field_columns(&self) -> impl Iterator<Item = ColumnSchema> + '_ {
169        // `value_indices` in TableMeta is not reliable. Do a filter here.
170        let primary_keys = self
171            .table_info
172            .meta
173            .primary_key_indices
174            .iter()
175            .copied()
176            .collect::<HashSet<_>>();
177
178        self.table_info
179            .meta
180            .schema
181            .column_schemas()
182            .iter()
183            .enumerate()
184            .filter(move |(i, c)| !primary_keys.contains(i) && !c.is_time_index())
185            .map(|(_, c)| c.clone())
186    }
187}
188
189/// Collects column default [`Expr`] from column schemas.
190fn collect_column_defaults(column_schemas: &[ColumnSchema]) -> HashMap<String, Expr> {
191    column_schemas
192        .iter()
193        .filter_map(|column_schema| {
194            default_constraint_to_expr(
195                column_schema.default_constraint()?,
196                &column_schema.data_type,
197            )
198            .map(|expr| (column_schema.name.clone(), expr))
199        })
200        .collect()
201}
202
203/// Try to cast the [`ColumnDefaultConstraint`] to [`Expr`] by the target data type.
204fn default_constraint_to_expr(
205    default_constraint: &ColumnDefaultConstraint,
206    target_type: &ConcreteDataType,
207) -> Option<Expr> {
208    match default_constraint {
209        ColumnDefaultConstraint::Value(v) => v
210            .try_to_scalar_value(target_type)
211            .ok()
212            .map(|x| Expr::Literal(x, None)),
213
214        ColumnDefaultConstraint::Function(name)
215            if matches!(
216                name.as_str(),
217                CURRENT_TIMESTAMP | CURRENT_TIMESTAMP_FN | NOW_FN
218            ) =>
219        {
220            Some(Expr::Cast(Cast {
221                expr: Box::new(NOW_EXPR.clone()),
222                data_type: target_type.as_arrow_type(),
223            }))
224        }
225
226        ColumnDefaultConstraint::Function(_) => None,
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use datafusion_common::ScalarValue;
233    use datatypes::prelude::ConcreteDataType;
234    use datatypes::schema::ColumnDefaultConstraint;
235
236    use super::*;
237
238    #[test]
239    fn test_collect_columns_defaults() {
240        let column_schemas = [
241            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), false),
242            ColumnSchema::new("col2", ConcreteDataType::string_datatype(), true)
243                .with_default_constraint(Some(ColumnDefaultConstraint::Value("test".into())))
244                .unwrap(),
245            ColumnSchema::new(
246                "ts",
247                ConcreteDataType::timestamp_millisecond_datatype(),
248                false,
249            )
250            .with_time_index(true)
251            .with_default_constraint(Some(ColumnDefaultConstraint::Function(
252                "current_timestamp".to_string(),
253            )))
254            .unwrap(),
255        ];
256        let column_defaults = collect_column_defaults(&column_schemas[..]);
257
258        assert!(!column_defaults.contains_key("col1"));
259        assert!(matches!(column_defaults.get("col2").unwrap(),
260                         Expr::Literal(ScalarValue::Utf8(Some(s)), _) if s == "test"));
261        assert!(matches!(
262            column_defaults.get("ts").unwrap(),
263            Expr::Cast(Cast {
264                expr,
265                data_type
266            }) if **expr == *NOW_EXPR && *data_type == ConcreteDataType::timestamp_millisecond_datatype().as_arrow_type()
267        ));
268    }
269}
table/table.rs

table/
table.rs
