diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index 3cd96b7525..3be41e0786 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -16,7 +16,7 @@ use std::collections::{BTreeSet, HashMap, HashSet}; use std::sync::{Arc, RwLock}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use api::v1::CreateTableExpr; +use api::v1::{CreateTableExpr, TableName}; use catalog::CatalogManagerRef; use common_error::ext::BoxedError; use common_query::logical_plan::breakup_insert_plan; @@ -25,8 +25,9 @@ use common_telemetry::{debug, info}; use common_time::Timestamp; use datafusion::datasource::DefaultTableSource; use datafusion::sql::unparser::expr_to_sql; -use datafusion_common::DFSchemaRef; use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::utils::quote_identifier; +use datafusion_common::{DFSchemaRef, TableReference}; use datafusion_expr::{DmlStatement, LogicalPlan, WriteOp, col, lit}; use datatypes::schema::Schema; use query::QueryEngineRef; @@ -51,7 +52,7 @@ use crate::batching_mode::state::{ use crate::batching_mode::table_creator::{QueryType, create_table_with_expr}; use crate::batching_mode::time_window::TimeWindowExpr; use crate::batching_mode::utils::{ - AddFilterRewriter, ColumnMatcherRewriter, gen_plan_with_matching_schema, + AddFilterRewriter, ColumnMatcherRewriter, df_plan_to_sql, gen_plan_with_matching_schema, get_table_info_df_schema, sql_to_df_plan, }; use crate::df_optimizer::apply_df_optimizer; @@ -105,6 +106,32 @@ fn is_merge_mode_last_non_null(options: &HashMap) -> bool { .unwrap_or(false) } +fn encode_insert_plan_request( + insert_to: TableName, + insert_input_plan: &LogicalPlan, +) -> Result { + let message = DFLogicalSubstraitConvertor {} + .encode(insert_input_plan, DefaultSerializer) + .context(SubstraitEncodeLogicalPlanSnafu)?; + Ok(api::v1::QueryRequest { + query: Some(api::v1::query_request::Query::InsertIntoPlan( + api::v1::InsertIntoPlan { + table_name: Some(insert_to), + logical_plan: message.to_vec(), + }, + )), + }) +} + +fn format_insert_target_columns(plan: &LogicalPlan) -> String { + plan.schema() + .fields() + .iter() + .map(|field| quote_identifier(field.name()).to_string()) + .collect::>() + .join(", ") +} + #[derive(Clone)] pub struct BatchingTask { pub config: Arc, @@ -549,19 +576,52 @@ impl BatchingTask { .with_label_values(&[flow_id.to_string().as_str()]) .start_timer(); - let req = if let Some((insert_to, insert_plan)) = + let req = if let Some((insert_to, insert_input_plan)) = breakup_insert_plan(&plan, catalog, schema) { - let message = DFLogicalSubstraitConvertor {} - .encode(&insert_plan, DefaultSerializer) - .context(SubstraitEncodeLogicalPlanSnafu)?; - api::v1::QueryRequest { - query: Some(api::v1::query_request::Query::InsertIntoPlan( - api::v1::InsertIntoPlan { - table_name: Some(insert_to), - logical_plan: message.to_vec(), - }, - )), + if query_mode == CheckpointMode::FullSnapshot + && matches!(self.config.query_type, QueryType::Sql) + && self.config.flow_eval_interval.is_some() + && self.config.time_window_expr.is_none() + { + // Evaluation-interval SQL flows without a time-window + // expression execute as full-query snapshots. Send these + // as SQL text instead of Substrait to avoid logical-plan + // round-trip issues around complex joins/unions/CTEs and + // duplicate field aliases. Keep ordinary SQL full snapshots + // on the existing InsertIntoPlan path because SQL unparsing + // is not valid for every planned aggregate shape yet. + // If the local SQL unparser does not support this plan, + // keep the previous InsertIntoPlan transport as a fallback. + match df_plan_to_sql(&insert_input_plan) { + Ok(select_sql) => { + let target_columns = format_insert_target_columns(&insert_input_plan); + let sql = format!( + "INSERT INTO {} ({}) {}", + TableReference::full( + insert_to.catalog_name.as_str(), + insert_to.schema_name.as_str(), + insert_to.table_name.as_str(), + ) + .to_quoted_string(), + target_columns, + select_sql + ); + api::v1::QueryRequest { + query: Some(api::v1::query_request::Query::Sql(sql)), + } + } + Err(err) => { + warn!( + "Failed to unparse full-snapshot SQL flow {} plan; \ + falling back to InsertIntoPlan: {:?}", + flow_id, err + ); + encode_insert_plan_request(insert_to, &insert_input_plan)? + } + } + } else { + encode_insert_plan_request(insert_to, &insert_input_plan)? } } else { let message = DFLogicalSubstraitConvertor {} diff --git a/src/flow/src/batching_mode/time_window.rs b/src/flow/src/batching_mode/time_window.rs index 410de4a81c..1c8ebea2a8 100644 --- a/src/flow/src/batching_mode/time_window.rs +++ b/src/flow/src/batching_mode/time_window.rs @@ -33,7 +33,9 @@ use datafusion::error::Result as DfResult; use datafusion::execution::SessionState; use datafusion::logical_expr::Expr; use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}; -use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter}; +use datafusion_common::tree_node::{ + Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, TreeNodeVisitor, +}; use datafusion_common::{DFSchema, TableReference}; use datafusion_expr::{ColumnarValue, LogicalPlan}; use datafusion_physical_expr::PhysicalExprRef; @@ -59,6 +61,84 @@ use crate::expr::error::DataTypeSnafu; /// Represents a test timestamp in seconds since the Unix epoch. const DEFAULT_TEST_TIMESTAMP: Timestamp = Timestamp::new_second(17_0000_0000); +#[derive(Default)] +struct TimeWindowPlanShape { + table_scan_count: usize, + aggregate_count: usize, + has_unsupported_pruning_node: bool, +} + +impl TimeWindowPlanShape { + fn should_skip_time_window_expr(&self) -> bool { + self.has_unsupported_pruning_node || self.table_scan_count != 1 || self.aggregate_count > 1 + } + + fn should_stop_inspection(&self) -> bool { + // This intentionally differs from the final skip predicate above: + // zero table scans make a fully inspected plan ineligible, but they are + // not an early-stop condition because a later subtree may still contain + // the first table scan. + self.has_unsupported_pruning_node || self.table_scan_count > 1 || self.aggregate_count > 1 + } +} + +impl TreeNodeVisitor<'_> for TimeWindowPlanShape { + type Node = LogicalPlan; + + fn f_down(&mut self, node: &Self::Node) -> DfResult { + match node { + LogicalPlan::TableScan(_) => { + self.table_scan_count += 1; + } + LogicalPlan::Aggregate(_) => { + self.aggregate_count += 1; + } + // The pinned DataFusion fork has no separate + // `LogicalPlan::CrossJoin` variant. SQL `CROSS JOIN` is represented + // as `LogicalPlan::Join(_)` here, so rejecting all joins also + // rejects cross joins. + LogicalPlan::Join(_) + | LogicalPlan::Window(_) + | LogicalPlan::Union(_) + | LogicalPlan::Distinct(_) + | LogicalPlan::Limit(_) + | LogicalPlan::Sort(_) + | LogicalPlan::Extension(_) + | LogicalPlan::Dml(_) + | LogicalPlan::Ddl(_) + | LogicalPlan::Unnest(_) + | LogicalPlan::RecursiveQuery(_) => { + self.has_unsupported_pruning_node = true; + } + _ => {} + } + + // These disqualifying conditions are monotonic. Once any of them is + // met, later traversal cannot make the plan eligible for TWE pruning + // again. Counts may be partial after `Stop`, but they are only used by + // the final skip predicate. + if self.should_stop_inspection() { + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + } +} + +fn inspect_time_window_plan_shape(plan: &LogicalPlan) -> DfResult { + let mut shape = TimeWindowPlanShape::default(); + plan.visit_with_subqueries(&mut shape)?; + Ok(shape) +} + +fn should_skip_time_window_expr(plan: &LogicalPlan) -> bool { + let Ok(shape) = inspect_time_window_plan_shape(plan) else { + return true; + }; + + shape.should_skip_time_window_expr() +} + /// Time window expr like `date_bin(INTERVAL '1' MINUTE, ts)`, this type help with /// evaluating the expr using given timestamp /// @@ -354,6 +434,21 @@ pub async fn find_time_window_expr( ) -> Result<(String, Option, TimeUnit, DFSchema), Error> { // TODO(discord9): find the expr that do time window + // Dirty-window pruning is only safe for simple single-source aggregate plans. + // For joins, window functions, set operations, distinct/sort/limit, extension + // nodes, or multi-scan plans, conservatively give up and let batching mode run + // the unfiltered full query instead. + if should_skip_time_window_expr(plan) { + // The column/schema/unit fields are placeholders when the TWE is None; + // callers must only use them when a TWE is present. + return Ok(( + String::new(), + None, + TimeUnit::Millisecond, + DFSchema::empty(), + )); + } + let mut table_name = None; // first find the table source in the logical plan @@ -948,4 +1043,154 @@ mod test { assert_eq!(expected_unparsed, new_sql); } } + + #[tokio::test] + async fn test_complex_plans_skip_time_window_expr() { + let query_engine = create_test_query_engine(); + let ctx = QueryContext::arc(); + + let testcases = [ + // A join may duplicate or drop rows across sources, so a time window + // found on one side is not safe to use as a full-query dirty-window + // pruning boundary. + r#" +SELECT + l.number, + date_bin('5 minutes', l.ts) AS time_window +FROM numbers_with_ts l +JOIN numbers_with_ts r ON l.number = r.number +GROUP BY l.number, time_window +"#, + // Window functions can depend on rows outside the dirty window, + // even if their input contains a group-by time window. + r#" +SELECT number, time_window +FROM ( + SELECT + number, + time_window, + row_number() OVER (PARTITION BY number ORDER BY time_window DESC) AS rn + FROM ( + SELECT number, date_bin('5 minutes', ts) AS time_window + FROM numbers_with_ts + GROUP BY number, time_window + ) +) +WHERE rn = 1 +"#, + // Set operations combine multiple query scopes/sources. + r#" +SELECT date_bin('5 minutes', ts) AS time_window +FROM numbers_with_ts +GROUP BY time_window +UNION ALL +SELECT date_bin('5 minutes', ts) AS time_window +FROM numbers_with_ts +GROUP BY time_window +"#, + // Nested aggregates are unsafe: pruning source rows by the inner + // time window is not equivalent for the outer/global aggregate. + r#" +SELECT max(cnt) +FROM ( + SELECT date_bin('5 minutes', ts) AS time_window, count(number) AS cnt + FROM numbers_with_ts + GROUP BY time_window +) +"#, + // Expression subqueries add another query scope/source and should + // not be treated as a simple single-source TWE plan. + r#" +SELECT date_bin('5 minutes', ts) AS time_window +FROM numbers_with_ts +WHERE number IN (SELECT number FROM numbers_with_ts) +GROUP BY time_window +"#, + // Sorting an otherwise valid TWE query is conservatively treated + // as full-query to avoid adding dirty-window predicates across + // post-aggregate plan nodes. + r#" +SELECT date_bin('5 minutes', ts) AS time_window +FROM numbers_with_ts +GROUP BY time_window +ORDER BY time_window +"#, + // DISTINCT is a post-query de-duplication boundary; keep it on the + // full-query path even if its input has a time window group key. + r#" +SELECT DISTINCT time_window +FROM ( + SELECT date_bin('5 minutes', ts) AS time_window + FROM numbers_with_ts + GROUP BY time_window +) +"#, + // LIMIT can change which rows are visible after pruning. + r#" +SELECT date_bin('5 minutes', ts) AS time_window +FROM numbers_with_ts +GROUP BY time_window +LIMIT 10 +"#, + // Cross joins may appear either as a dedicated node or as multiple + // table scans; either way they must not use source dirty-window + // pruning from one side only. + r#" +SELECT date_bin('5 minutes', l.ts) AS time_window +FROM numbers_with_ts l, numbers_with_ts r +GROUP BY time_window +"#, + // A cross join with a constant relation still has only one table + // scan, so it must be rejected by the join node instead of relying + // on the multi-scan guard. + r#" +SELECT date_bin('5 minutes', l.ts) AS time_window +FROM numbers_with_ts l CROSS JOIN (VALUES (1)) AS v(x) +GROUP BY time_window +"#, + ]; + + for sql in testcases { + let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), sql, true) + .await + .unwrap(); + let (_, lower, upper) = find_plan_time_window_bound( + &plan, + Timestamp::new(23, TimeUnit::Millisecond), + ctx.clone(), + query_engine.clone(), + ) + .await + .unwrap(); + + assert_eq!(None, lower, "query should not have TWE: {sql}"); + assert_eq!(None, upper, "query should not have TWE: {sql}"); + } + } + + #[tokio::test] + async fn test_simple_single_source_aggregate_keeps_time_window_expr() { + let query_engine = create_test_query_engine(); + let ctx = QueryContext::arc(); + + let sql = r#" +SELECT max(number) AS max_number, date_bin('5 minutes', ts) AS time_window +FROM numbers_with_ts +GROUP BY time_window +"#; + let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), sql, true) + .await + .unwrap(); + let (_, lower, upper) = find_plan_time_window_bound( + &plan, + Timestamp::new(23, TimeUnit::Millisecond), + ctx.clone(), + query_engine.clone(), + ) + .await + .unwrap(); + + assert_eq!(Some(Timestamp::new(0, TimeUnit::Millisecond)), lower); + assert_eq!(Some(Timestamp::new(300000, TimeUnit::Millisecond)), upper); + } } diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index b2614bb8d5..c4dfd8407d 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -410,7 +410,7 @@ impl<'a> ParserContext<'a> { .fail(); }; - if !utils::is_simple_tql_cte_query(query) { + if utils::has_tql_cte(query) && !utils::is_simple_tql_cte_query(query) { return InvalidFlowQuerySnafu { reason: "WITH is only supported for the simplest TQL CTE in CREATE FLOW" .to_string(), @@ -1962,7 +1962,7 @@ SELECT * FROM tql; } #[test] - fn test_parse_create_flow_with_sql_cte_is_unsupported() { + fn test_parse_create_flow_with_sql_cte_is_supported() { let sql = r#" CREATE FLOW f SINK TO s @@ -1970,11 +1970,17 @@ AS WITH cte AS (SELECT 1) SELECT * FROM cte; "#; - let err = + let stmts = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) - .unwrap_err(); - let msg = err.to_string(); - assert!(msg.to_uppercase().contains("WITH"), "err: {msg}"); + .unwrap(); + assert_eq!(1, stmts.len()); + let Statement::CreateFlow(create_flow) = &stmts[0] else { + panic!("unexpected stmt: {:?}", stmts[0]); + }; + assert_eq!( + "WITH cte AS (SELECT 1) SELECT * FROM cte", + create_flow.query.to_string() + ); } #[test] diff --git a/src/sql/src/parsers/utils.rs b/src/sql/src/parsers/utils.rs index 239f3155ca..e25a1d9522 100644 --- a/src/sql/src/parsers/utils.rs +++ b/src/sql/src/parsers/utils.rs @@ -98,6 +98,14 @@ pub(crate) fn is_simple_tql_cte_query(query: &Query) -> bool { reference == cte_name } +pub(crate) fn has_tql_cte(query: &Query) -> bool { + query.hybrid_cte.as_ref().is_some_and(|with| { + with.cte_tables + .iter() + .any(|cte| matches!(cte.content, CteContent::Tql(_))) + }) +} + fn has_only_hybrid_tql_cte(query: &Query) -> bool { query .inner diff --git a/src/sql/src/parsers/with_tql_parser.rs b/src/sql/src/parsers/with_tql_parser.rs index 609caaf3b0..382b2a9137 100644 --- a/src/sql/src/parsers/with_tql_parser.rs +++ b/src/sql/src/parsers/with_tql_parser.rs @@ -146,15 +146,14 @@ impl ParserContext<'_> { // Parse the main query let main_query = self.parser.parse_query().context(error::SyntaxSnafu)?; - // Convert the hybrid CTEs to a standard query with hybrid metadata - let hybrid_cte = HybridCteWith { - recursive, - cte_tables: tql_cte_tables, - }; - // Create a Query statement with hybrid CTE metadata let mut query = Query::try_from(*main_query)?; - query.hybrid_cte = Some(hybrid_cte); + if !tql_cte_tables.is_empty() { + query.hybrid_cte = Some(HybridCteWith { + recursive, + cte_tables: tql_cte_tables, + }); + } query.inner.with = Some(With { recursive, cte_tables: sql_cte_tables, diff --git a/src/sql/src/statements/query.rs b/src/sql/src/statements/query.rs index 85344bfb6d..6ea07fb253 100644 --- a/src/sql/src/statements/query.rs +++ b/src/sql/src/statements/query.rs @@ -108,5 +108,11 @@ mod test { .to_string(), "WITH tql_cte AS (TQL EVAL (0, 100, '5s') up) SELECT * FROM tql_cte" ); + assert_eq!( + create_query("WITH sql_cte AS (SELECT 1) SELECT * FROM sql_cte") + .unwrap() + .to_string(), + "WITH sql_cte AS (SELECT 1) SELECT * FROM sql_cte" + ); } }