rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-07-08 06:50:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(flow): support complex SQL full-query flows (#8242)

Browse Source 
* fix(flow): skip TWE detection for complex plans

Signed-off-by: discord9 <discord9@163.com>

* chore: helper

Signed-off-by: discord9 <discord9@163.com>

* fix(flow): allow SQL CTE flow parsing

Signed-off-by: discord9 <discord9@163.com>

* fix(flow): execute unscoped SQL snapshots as SQL

Signed-off-by: discord9 <discord9@163.com>

* fix: column list

Signed-off-by: discord9 <discord9@163.com>

* refactor: per review

Signed-off-by: discord9 <discord9@163.com>

---------

Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-06-08 21:07:43 +08:00
committed by GitHub
parent e403133eb2
commit 1c7b4c11fc
6 changed files with 352 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
src/flow/src/batching_mode/task.rs
View File

@@ -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<String, String>) -> bool {
.unwrap_or(false)
}
fn encode_insert_plan_request(
insert_to: TableName,
insert_input_plan: &LogicalPlan,
) -> Result<api::v1::QueryRequest, Error> {
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::<Vec<_>>()
.join(", ")
}
#[derive(Clone)]
pub struct BatchingTask {
pub config: Arc<TaskConfig>,
@@ -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 {}

247
src/flow/src/batching_mode/time_window.rs
View File

@@ -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<TreeNodeRecursion> {
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<TimeWindowPlanShape> {
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<datafusion_expr::Expr>, 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);
}
}

18
src/sql/src/parsers/create_parser.rs
View File

@@ -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]

8
src/sql/src/parsers/utils.rs
View File

@@ -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

13
src/sql/src/parsers/with_tql_parser.rs
View File

@@ -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,

6
src/sql/src/statements/query.rs
View File

@@ -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"
);
}
}