flow/batching_mode/
table_creator.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 api::v1::CreateTableExpr;
16use common_recordbatch::map_dictionary_to_values_data_type;
17use datafusion_common::tree_node::TreeNode;
18use datafusion_expr::LogicalPlan;
19use datatypes::prelude::ConcreteDataType;
20use datatypes::schema::ColumnSchema;
21use operator::expr_helper::column_schemas_to_defs;
22use snafu::ResultExt;
23
24use crate::Error;
25use crate::adapter::{AUTO_CREATED_PLACEHOLDER_TS_COL, AUTO_CREATED_UPDATE_AT_TS_COL};
26use crate::batching_mode::utils::FindGroupByFinalName;
27use crate::error::{ConvertColumnSchemaSnafu, DatafusionSnafu};
28
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub enum QueryType {
31    /// query is a tql query
32    Tql,
33    /// query is a sql query
34    Sql,
35}
36
37// auto created table have a auto added column `update_at`, and optional have a `AUTO_CREATED_PLACEHOLDER_TS_COL` column for time index placeholder if no timestamp column is specified
38// TODO(discord9): for now no default value is set for auto added column for compatibility reason with streaming mode, but this might change in favor of simpler code?
39pub(super) fn create_table_with_expr(
40    plan: &LogicalPlan,
41    sink_table_name: &[String; 3],
42    query_type: &QueryType,
43) -> Result<CreateTableExpr, Error> {
44    let table_def = match query_type {
45        &QueryType::Sql => {
46            if let Some(def) = build_pk_from_aggr(plan)? {
47                def
48            } else {
49                build_by_sql_schema(plan)?
50            }
51        }
52        QueryType::Tql => {
53            // first try build from aggr, then from tql schema because tql query might not have aggr node
54            if let Some(table_def) = build_pk_from_aggr(plan)? {
55                table_def
56            } else {
57                build_by_tql_schema(plan)?
58            }
59        }
60    };
61    let first_time_stamp = table_def.ts_col;
62    let primary_keys = table_def.pks;
63
64    let mut column_schemas = Vec::new();
65    for field in plan.schema().fields() {
66        let name = field.name();
67        let ty = map_dictionary_to_values_data_type(&ConcreteDataType::from_arrow_type(
68            field.data_type(),
69        ));
70        let col_schema = if first_time_stamp == Some(name.clone()) {
71            ColumnSchema::new(name, ty, false).with_time_index(true)
72        } else {
73            ColumnSchema::new(name, ty, true)
74        };
75
76        match query_type {
77            QueryType::Sql => {
78                column_schemas.push(col_schema);
79            }
80            QueryType::Tql => {
81                // if is val column, need to rename as val DOUBLE NULL
82                // if is tag column, need to cast type as STRING NULL
83                let is_tag_column = primary_keys.contains(name);
84                let is_val_column = !is_tag_column && first_time_stamp.as_ref() != Some(name);
85                if is_val_column {
86                    let col_schema =
87                        ColumnSchema::new(name, ConcreteDataType::float64_datatype(), true);
88                    column_schemas.push(col_schema);
89                } else if is_tag_column {
90                    let col_schema =
91                        ColumnSchema::new(name, ConcreteDataType::string_datatype(), true);
92                    column_schemas.push(col_schema);
93                } else {
94                    // time index column
95                    column_schemas.push(col_schema);
96                }
97            }
98        }
99    }
100
101    if query_type == &QueryType::Sql {
102        let update_at_schema = ColumnSchema::new(
103            AUTO_CREATED_UPDATE_AT_TS_COL,
104            ConcreteDataType::timestamp_millisecond_datatype(),
105            true,
106        );
107        column_schemas.push(update_at_schema);
108    }
109
110    let time_index = if let Some(time_index) = first_time_stamp {
111        time_index
112    } else {
113        column_schemas.push(
114            ColumnSchema::new(
115                AUTO_CREATED_PLACEHOLDER_TS_COL,
116                ConcreteDataType::timestamp_millisecond_datatype(),
117                false,
118            )
119            .with_time_index(true),
120        );
121        AUTO_CREATED_PLACEHOLDER_TS_COL.to_string()
122    };
123
124    let column_defs =
125        column_schemas_to_defs(column_schemas, &primary_keys).context(ConvertColumnSchemaSnafu)?;
126    Ok(CreateTableExpr {
127        catalog_name: sink_table_name[0].clone(),
128        schema_name: sink_table_name[1].clone(),
129        table_name: sink_table_name[2].clone(),
130        desc: "Auto created table by flow engine".to_string(),
131        column_defs,
132        time_index,
133        primary_keys,
134        create_if_not_exists: true,
135        table_options: Default::default(),
136        table_id: None,
137        engine: "mito".to_string(),
138    })
139}
140
141/// simply build by schema, return first timestamp column and no primary key
142fn build_by_sql_schema(plan: &LogicalPlan) -> Result<TableDef, Error> {
143    let first_time_stamp = plan.schema().fields().iter().find_map(|f| {
144        if ConcreteDataType::from_arrow_type(f.data_type()).is_timestamp() {
145            Some(f.name().clone())
146        } else {
147            None
148        }
149    });
150    Ok(TableDef {
151        ts_col: first_time_stamp,
152        pks: vec![],
153    })
154}
155
156/// Return first timestamp column found in output schema and all string columns
157fn build_by_tql_schema(plan: &LogicalPlan) -> Result<TableDef, Error> {
158    let first_time_stamp = plan.schema().fields().iter().find_map(|f| {
159        if ConcreteDataType::from_arrow_type(f.data_type()).is_timestamp() {
160            Some(f.name().clone())
161        } else {
162            None
163        }
164    });
165    let string_columns = plan
166        .schema()
167        .fields()
168        .iter()
169        .filter_map(|f| {
170            if map_dictionary_to_values_data_type(&ConcreteDataType::from_arrow_type(f.data_type()))
171                .is_string()
172            {
173                Some(f.name().clone())
174            } else {
175                None
176            }
177        })
178        .collect::<Vec<_>>();
179
180    Ok(TableDef {
181        ts_col: first_time_stamp,
182        pks: string_columns,
183    })
184}
185
186struct TableDef {
187    ts_col: Option<String>,
188    pks: Vec<String>,
189}
190
191/// Return first timestamp column which is in group by clause and other columns which are also in group by clause
192///
193/// # Returns
194///
195/// * `Option<String>` - first timestamp column which is in group by clause
196/// * `Vec<String>` - other columns which are also in group by clause
197///
198/// if no aggregation found, return None
199fn build_pk_from_aggr(plan: &LogicalPlan) -> Result<Option<TableDef>, Error> {
200    let fields = plan.schema().fields();
201    let mut pk_names = FindGroupByFinalName::default();
202
203    plan.visit(&mut pk_names)
204        .with_context(|_| DatafusionSnafu {
205            context: format!("Can't find aggr expr in plan {plan:?}"),
206        })?;
207
208    // if no group by clause, return empty with first timestamp column found in output schema
209    let Some(pk_final_names) = pk_names.get_group_expr_names() else {
210        return Ok(None);
211    };
212    if pk_final_names.is_empty() {
213        let first_ts_col = fields
214            .iter()
215            .find(|f| ConcreteDataType::from_arrow_type(f.data_type()).is_timestamp())
216            .map(|f| f.name().clone());
217        return Ok(Some(TableDef {
218            ts_col: first_ts_col,
219            pks: vec![],
220        }));
221    }
222
223    let all_pk_cols: Vec<_> = fields
224        .iter()
225        .filter(|f| pk_final_names.contains(f.name()))
226        .map(|f| f.name().clone())
227        .collect();
228    // Auto-created tables use the first timestamp column in the group-by keys
229    // as the time index. It is possible that timestamp columns appear only as
230    // aggregate outputs (for example `max(ts)`) and are not group-by keys; in
231    // that case `first_time_stamp` stays `None` and the caller falls back to a
232    // placeholder time index column.
233    let first_time_stamp = fields
234        .iter()
235        .find(|f| {
236            all_pk_cols.contains(&f.name().clone())
237                && ConcreteDataType::from_arrow_type(f.data_type()).is_timestamp()
238        })
239        .map(|f| f.name().clone());
240
241    let all_pk_cols: Vec<_> = all_pk_cols
242        .into_iter()
243        .filter(|col| first_time_stamp.as_ref() != Some(col))
244        .collect();
245
246    Ok(Some(TableDef {
247        ts_col: first_time_stamp,
248        pks: all_pk_cols,
249    }))
250}
251
252#[cfg(test)]
253mod test {
254    use std::sync::Arc;
255
256    use api::v1::column_def::try_as_column_schema;
257    use datafusion::arrow::datatypes::{
258        DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit,
259    };
260    use datafusion_common::DFSchema;
261    use datafusion_expr::logical_plan::EmptyRelation;
262    use datatypes::prelude::ConcreteDataType;
263    use datatypes::schema::ColumnSchema;
264    use pretty_assertions::assert_eq;
265    use session::context::QueryContext;
266
267    use super::*;
268    use crate::adapter::{AUTO_CREATED_PLACEHOLDER_TS_COL, AUTO_CREATED_UPDATE_AT_TS_COL};
269    use crate::batching_mode::utils::sql_to_df_plan;
270    use crate::test_utils::create_test_query_engine;
271
272    #[test]
273    fn test_tql_dictionary_string_is_label() {
274        let arrow_schema = Arc::new(ArrowSchema::new(vec![
275            Field::new_dictionary("host", ArrowDataType::UInt32, ArrowDataType::Utf8, true),
276            Field::new("value", ArrowDataType::Float64, true),
277            Field::new(
278                "ts",
279                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
280                false,
281            ),
282        ]));
283        let plan = LogicalPlan::EmptyRelation(EmptyRelation {
284            produce_one_row: false,
285            schema: Arc::new(DFSchema::try_from(arrow_schema).unwrap()),
286        });
287
288        let expr = create_table_with_expr(
289            &plan,
290            &[
291                "greptime".to_string(),
292                "public".to_string(),
293                "sink".to_string(),
294            ],
295            &QueryType::Tql,
296        )
297        .unwrap();
298        let columns = expr
299            .column_defs
300            .iter()
301            .map(|column| try_as_column_schema(column).unwrap())
302            .collect::<Vec<_>>();
303
304        assert_eq!(vec!["host".to_string()], expr.primary_keys);
305        assert_eq!("ts", expr.time_index);
306        assert_eq!(ConcreteDataType::string_datatype(), columns[0].data_type);
307        assert_eq!(ConcreteDataType::float64_datatype(), columns[1].data_type);
308        assert!(columns[2].is_time_index());
309    }
310
311    #[tokio::test]
312    async fn test_gen_create_table_sql() {
313        let query_engine = create_test_query_engine();
314        let ctx = QueryContext::arc();
315        struct TestCase {
316            sql: String,
317            sink_table_name: String,
318            column_schemas: Vec<ColumnSchema>,
319            primary_keys: Vec<String>,
320            time_index: String,
321        }
322
323        let update_at_schema = ColumnSchema::new(
324            AUTO_CREATED_UPDATE_AT_TS_COL,
325            ConcreteDataType::timestamp_millisecond_datatype(),
326            true,
327        );
328
329        let ts_placeholder_schema = ColumnSchema::new(
330            AUTO_CREATED_PLACEHOLDER_TS_COL,
331            ConcreteDataType::timestamp_millisecond_datatype(),
332            false,
333        )
334        .with_time_index(true);
335
336        let testcases = vec![
337            TestCase {
338                sql: "SELECT number, ts FROM numbers_with_ts".to_string(),
339                sink_table_name: "new_table".to_string(),
340                column_schemas: vec![
341                    ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true),
342                    ColumnSchema::new(
343                        "ts",
344                        ConcreteDataType::timestamp_millisecond_datatype(),
345                        false,
346                    )
347                    .with_time_index(true),
348                    update_at_schema.clone(),
349                ],
350                primary_keys: vec![],
351                time_index: "ts".to_string(),
352            },
353            TestCase {
354                sql: "SELECT number, max(ts) FROM numbers_with_ts GROUP BY number".to_string(),
355                sink_table_name: "new_table".to_string(),
356                column_schemas: vec![
357                    ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true),
358                    ColumnSchema::new(
359                        "max(numbers_with_ts.ts)",
360                        ConcreteDataType::timestamp_millisecond_datatype(),
361                        true,
362                    ),
363                    update_at_schema.clone(),
364                    ts_placeholder_schema.clone(),
365                ],
366                primary_keys: vec!["number".to_string()],
367                time_index: AUTO_CREATED_PLACEHOLDER_TS_COL.to_string(),
368            },
369            TestCase {
370                sql: "SELECT max(number), ts FROM numbers_with_ts GROUP BY ts".to_string(),
371                sink_table_name: "new_table".to_string(),
372                column_schemas: vec![
373                    ColumnSchema::new(
374                        "max(numbers_with_ts.number)",
375                        ConcreteDataType::uint32_datatype(),
376                        true,
377                    ),
378                    ColumnSchema::new(
379                        "ts",
380                        ConcreteDataType::timestamp_millisecond_datatype(),
381                        false,
382                    )
383                    .with_time_index(true),
384                    update_at_schema.clone(),
385                ],
386                primary_keys: vec![],
387                time_index: "ts".to_string(),
388            },
389            TestCase {
390                sql: "SELECT number, ts FROM numbers_with_ts GROUP BY ts, number".to_string(),
391                sink_table_name: "new_table".to_string(),
392                column_schemas: vec![
393                    ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true),
394                    ColumnSchema::new(
395                        "ts",
396                        ConcreteDataType::timestamp_millisecond_datatype(),
397                        false,
398                    )
399                    .with_time_index(true),
400                    update_at_schema.clone(),
401                ],
402                primary_keys: vec!["number".to_string()],
403                time_index: "ts".to_string(),
404            },
405        ];
406
407        for tc in testcases {
408            let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), &tc.sql, true)
409                .await
410                .unwrap();
411            let expr = create_table_with_expr(
412                &plan,
413                &[
414                    "greptime".to_string(),
415                    "public".to_string(),
416                    tc.sink_table_name.clone(),
417                ],
418                &QueryType::Sql,
419            )
420            .unwrap();
421            // TODO(discord9): assert expr
422            let column_schemas = expr
423                .column_defs
424                .iter()
425                .map(|c| try_as_column_schema(c).unwrap())
426                .collect::<Vec<_>>();
427            assert_eq!(tc.column_schemas, column_schemas, "{:?}", tc.sql);
428            assert_eq!(tc.primary_keys, expr.primary_keys, "{:?}", tc.sql);
429            assert_eq!(tc.time_index, expr.time_index, "{:?}", tc.sql);
430        }
431    }
432}
flow/batching_mode/table_creator.rs

flow/batching_mode/
table_creator.rs
