fix(flow): bind scheduled now in dist plan (#8389)

Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-07-01 10:56:30 +00:00
committed by GitHub
parent 3aa82567c4
commit 5f6b88b782
9 changed files with 565 additions and 4 deletions
+12
View File
@@ -67,6 +67,7 @@ use crate::metrics::{
OnDone, QUERY_STAGE_ELAPSED, maybe_attach_region_watermark_metrics,
should_collect_region_watermark_from_query_ctx,
};
use crate::options::ScheduledTimeExtension;
use crate::physical_wrapper::PhysicalPlanWrapperRef;
use crate::planner::{DfLogicalPlanner, LogicalPlanner};
use crate::query_engine::{DescribeResult, QueryEngineContext, QueryEngineState};
@@ -559,6 +560,17 @@ impl QueryEngine for DatafusionQueryEngine {
state: self.engine_state().function_state(),
});
// Carry scheduled Flow time through ConfigOptions.extensions so that
// the distributed plan analyzer can read it during expression
// simplification (preventing wall-clock constant-folding of `now()`).
state
.config_mut()
.options_mut()
.extensions
.insert(ScheduledTimeExtension {
scheduled_time: crate::options::scheduled_time_from_ctx(&query_ctx),
});
let config_options = state.config_options().clone();
let _ = state
.execution_props_mut()
+11
View File
@@ -52,6 +52,7 @@ use crate::dist_plan::commutativity::{
use crate::dist_plan::merge_scan::MergeScanLogicalPlan;
use crate::dist_plan::merge_sort::MergeSortLogicalPlan;
use crate::metrics::PUSH_DOWN_FALLBACK_ERRORS_TOTAL;
use crate::options::ScheduledTimeExtension;
use crate::plan::ExtractExpr;
use crate::query_engine::DefaultSerializer;
@@ -123,9 +124,18 @@ impl AnalyzerRule for DistPlannerAnalyzer {
let opt = config.extensions.get::<DistPlannerOptions>();
let allow_fallback = opt.map(|o| o.allow_query_fallback).unwrap_or(false);
// When the query is running under a scheduled Flow context, carry the
// logical "now" so that `SimplifyExpressions` does not constant-fold
// `now()` into wall-clock literals on the remote sub-plans.
let scheduled_time = config
.extensions
.get::<ScheduledTimeExtension>()
.and_then(|ext| ext.scheduled_time);
let optimizer_context = PatchOptimizerContext {
inner: datafusion_optimizer::OptimizerContext::new(),
config: config.clone(),
scheduled_time,
};
let plan = plan
@@ -142,6 +152,7 @@ impl AnalyzerRule for DistPlannerAnalyzer {
let optimizer_context = PatchOptimizerContext {
inner: datafusion_optimizer::OptimizerContext::new(),
config: config.clone(),
scheduled_time,
};
let plan = match pre_merge_scan_optimizer().optimize(plan, &optimizer_context, |_, _| {}) {
Ok(plan) => plan,
+200
View File
@@ -2454,3 +2454,203 @@ fn test_join_cross_table_predicate_not_pushed_to_single_side() {
"Cross-table predicate should not become a single-side TableScan filter:\n{plan_str}"
);
}
/// When `ScheduledTimeExtension` is injected into `ConfigOptions`, the
/// `SimplifyExpressions` pass (driven by `PatchOptimizerContext`) uses the
/// scheduled time instead of wall-clock time. The remote sub-plan must contain
/// the scheduled literal — not a variable wall-clock value.
#[test]
fn scheduled_now_yields_stable_literal_in_remote_plan() {
init_default_ut_logging();
let scheduled_time =
chrono::DateTime::<chrono::Utc>::from_timestamp_millis(1700000000000).unwrap();
let scheduled_ns = 1700000000000000000i64;
let test_table = TestTable::table_with_name(0, "t".to_string());
let table_source = Arc::new(DefaultTableSource::new(Arc::new(
DfTableProviderAdapter::new(test_table),
)));
// Build a plan with `now()` in filter and both `now()` and its
// `current_timestamp()` alias in projection.
let plan = LogicalPlanBuilder::scan_with_filters("t", table_source.clone(), None, vec![])
.unwrap()
.filter(binary_expr(now(), Operator::LtEq, col("ts")))
.unwrap()
.project(vec![now(), now().alias("current_timestamp()")])
.unwrap()
.build()
.unwrap();
let mut config = ConfigOptions::default();
config.extensions.insert(ScheduledTimeExtension {
scheduled_time: Some(scheduled_time),
});
let result = DistPlannerAnalyzer {}
.analyze(plan.clone(), &config)
.unwrap();
let result_str = result.to_string();
common_telemetry::info!("Analyzed plan with scheduled time: {}", result_str);
// The top-level should still say `Projection: now()` (schema recovery).
assert!(
result_str.contains("Projection: now()"),
"Expected top-level Projection: now(), got:\n{result_str}"
);
assert!(
result_str.contains("current_timestamp()"),
"Expected top-level current_timestamp() alias, got:\n{result_str}"
);
// The remote sub-plan must contain the scheduled-time literal, i.e.
// `TimestampNanosecond(1700000000000000000, None)`.
let expected_literal = format!("TimestampNanosecond({scheduled_ns}, None)");
assert!(
result_str.contains(&expected_literal),
"Expected remote plan literal '{expected_literal}', got:\n{result_str}"
);
// The remote sub-plan must contain the simplified literal (TimestampNanosecond).
let remote_section = if let Some(idx) = result_str.find("remote_input=[") {
&result_str[idx..]
} else {
""
};
assert!(
remote_section.contains("TimestampNanosecond("),
"Remote plan should contain TimestampNanosecond literal:\n{result_str}"
);
// Absent the extension (default config), the same plan simplifies to a
// wall-clock literal — the remote_input contains a variable timestamp.
let default_config = ConfigOptions::default();
let wall_result = DistPlannerAnalyzer {}
.analyze(plan, &default_config)
.unwrap();
let wall_str = wall_result.to_string();
// The wall-clock result must NOT contain the scheduled literal.
assert!(
!wall_str.contains(&expected_literal),
"Wall-clock result should not contain scheduled literal {expected_literal}:\n{wall_str}"
);
// It should still simplify to a literal (TimestampNanosecond present in remote).
let wall_remote = if let Some(idx) = wall_str.find("remote_input=[") {
&wall_str[idx..]
} else {
""
};
assert!(
wall_remote.contains("TimestampNanosecond("),
"Wall-clock remote plan should contain TimestampNanosecond:\n{wall_str}"
);
}
/// `current_timestamp()` is an alias of `now()` in DataFusion, but keep a
/// dedicated SQL-level regression: when it appears in a side-local filter, it
/// must be folded to the scheduled literal before the filter interacts with the
/// pre-MergeScan pushdown / remote planning path.
#[test]
fn scheduled_current_timestamp_filter_folds_before_remote_pushdown() {
init_default_ut_logging();
let scheduled_time =
chrono::DateTime::<chrono::Utc>::from_timestamp_millis(1700000000000).unwrap();
let left_table =
TestTable::table_with_filter_pushdown(0, "t1".to_string(), FilterPushDownType::Inexact);
let right_table =
TestTable::table_with_filter_pushdown(1, "t2".to_string(), FilterPushDownType::Inexact);
let left_provider = Arc::new(DfTableProviderAdapter::new(left_table));
let right_provider = Arc::new(DfTableProviderAdapter::new(right_table));
let ctx = SessionContext::new();
ctx.register_table(TableReference::bare("t1"), left_provider)
.unwrap();
ctx.register_table(TableReference::bare("t2"), right_provider)
.unwrap();
let plan = futures::executor::block_on(async {
ctx.sql(
"SELECT t1.number \
FROM t1 JOIN t2 ON t1.number = t2.number \
WHERE t1.ts < date_trunc('second', current_timestamp())",
)
.await
.unwrap()
})
.into_unoptimized_plan();
assert!(
plan.to_string().contains("current_timestamp") || plan.to_string().contains("now()"),
"Unoptimized plan should contain the time function spelling before analysis:\n{plan}"
);
let mut config = ConfigOptions::default();
config.extensions.insert(ScheduledTimeExtension {
scheduled_time: Some(scheduled_time),
});
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
assert_remote_table_scan_filters_are_safe(&result);
let result_str = result.to_string();
common_telemetry::info!("Analyzed current_timestamp filter plan: {}", result_str);
let remote_section = result_str
.find("remote_input=[")
.map(|idx| &result_str[idx..])
.unwrap_or("");
assert!(
remote_section.contains("TableScan: t1") && remote_section.contains("partial_filters="),
"Expected left-side filter to be pushed into the remote TableScan:\n{result_str}"
);
assert!(
remote_section.contains("1700000000000000000") || remote_section.contains("1700000000000"),
"Expected scheduled timestamp literal in remote filter:\n{result_str}"
);
assert!(
!remote_section.contains("current_timestamp") && !remote_section.contains("now()"),
"Remote filter should be folded before pushdown, got:\n{result_str}"
);
}
/// When `ScheduledTimeExtension.scheduled_time` is `None`, the analyzer
/// must fall back to wall-clock behavior (same as no extension at all).
#[test]
fn scheduled_none_falls_back_to_wall_clock() {
init_default_ut_logging();
let test_table = TestTable::table_with_name(0, "t".to_string());
let table_source = Arc::new(DefaultTableSource::new(Arc::new(
DfTableProviderAdapter::new(test_table),
)));
let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
.unwrap()
.project(vec![now()])
.unwrap()
.build()
.unwrap();
let mut config = ConfigOptions::default();
config.extensions.insert(ScheduledTimeExtension {
scheduled_time: None,
});
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
let result_str = result.to_string();
// Must still simplify now() — remotely a literal, top is now().
assert!(
result_str.contains("Projection: now()"),
"Top-level projection: {result_str}"
);
let remote_section = if let Some(idx) = result_str.find("remote_input=[") {
&result_str[idx..]
} else {
""
};
// The literal must be TimestampNanosecond (since fallback uses wall clock).
assert!(
remote_section.contains("TimestampNanosecond("),
"Remote should contain TimestampNanosecond:\n{result_str}"
);
}
+6 -1
View File
@@ -43,11 +43,16 @@ use crate::plan::ExtractExpr as _;
pub(crate) struct PatchOptimizerContext {
pub(crate) inner: datafusion_optimizer::OptimizerContext,
pub(crate) config: Arc<ConfigOptions>,
/// Override for `query_execution_start_time()` — used during scheduled Flow
/// evaluation so that `SimplifyExpressions` does not constant-fold `now()`
/// into wall-clock literals. When `None`, falls back to the inner context.
pub(crate) scheduled_time: Option<DateTime<Utc>>,
}
impl OptimizerConfig for PatchOptimizerContext {
fn query_execution_start_time(&self) -> Option<DateTime<Utc>> {
self.inner.query_execution_start_time()
self.scheduled_time
.or_else(|| self.inner.query_execution_start_time())
}
fn alias_generator(&self) -> &Arc<AliasGenerator> {
+40 -3
View File
@@ -16,6 +16,7 @@ use std::collections::HashMap;
use chrono::{DateTime, Utc};
use common_base::memory_limit::MemoryLimit;
use datafusion::config::{ConfigEntry, ConfigExtension, ExtensionOptions};
use serde::{Deserialize, Serialize};
use session::context::QueryContextRef;
use store_api::storage::RegionId;
@@ -326,15 +327,51 @@ pub fn parse_scheduled_time_datetime(
}
}
/// Convenience helper: extract scheduled time from a [`QueryContextRef`] as a [`DateTime<Utc>`].
/// Best-effort helper: extract scheduled time from a [`QueryContextRef`] as a [`DateTime<Utc>`].
///
/// Errors are silently swallowed, returning `None`.
/// Use [`parse_scheduled_time_datetime`] when a `Result` is needed.
/// Errors are silently swallowed, returning `None`; use [`parse_scheduled_time_datetime`]
/// at call sites that need to reject malformed scheduled time.
pub fn scheduled_time_from_ctx(query_ctx: &QueryContextRef) -> Option<DateTime<Utc>> {
let extensions = query_ctx.extensions();
parse_scheduled_time_datetime(&extensions).ok().flatten()
}
/// Carries the scheduled logical "now" through [`ConfigOptions::extensions`] so
/// that the distributed plan analyzer can apply it during expression
/// simplification (preventing wall-clock constant-folding of `now()`).
#[derive(Debug, Clone)]
pub(crate) struct ScheduledTimeExtension {
pub(crate) scheduled_time: Option<DateTime<Utc>>,
}
impl ConfigExtension for ScheduledTimeExtension {
const PREFIX: &'static str = "flow_scheduled_time";
}
impl ExtensionOptions for ScheduledTimeExtension {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
self
}
fn cloned(&self) -> Box<dyn ExtensionOptions> {
Box::new(self.clone())
}
fn set(&mut self, key: &str, value: &str) -> datafusion::error::Result<()> {
Err(datafusion_common::DataFusionError::NotImplemented(format!(
"ScheduledTimeExtension does not support set key: {key} with value: {value}"
)))
}
fn entries(&self) -> Vec<ConfigEntry> {
vec![]
}
}
fn invalid_query_context_extension(reason: String) -> Error {
InvalidQueryContextExtensionSnafu { reason }.build()
}
@@ -0,0 +1,91 @@
CREATE TABLE scheduled_late_overlap_input (
ts TIMESTAMP(3) TIME INDEX,
series STRING,
reading DOUBLE,
PRIMARY KEY(series)
);
Affected Rows: 0
CREATE FLOW scheduled_late_overlap_flow
SINK TO scheduled_late_overlap_sink
EVAL INTERVAL '1s'
AS
WITH
target_offsets(delta) AS (
VALUES
(INTERVAL '1 second'),
(INTERVAL '2 seconds'),
(INTERVAL '3 seconds'),
(INTERVAL '4 seconds'),
(INTERVAL '5 seconds'),
(INTERVAL '6 seconds'),
(INTERVAL '7 seconds'),
(INTERVAL '8 seconds'),
(INTERVAL '9 seconds'),
(INTERVAL '10 seconds')
),
target_seconds AS (
SELECT date_trunc('second', now()) - delta AS target_ts
FROM target_offsets
),
bucketed AS (
SELECT
series,
date_bin(INTERVAL '1 second', ts) AS bucket_ts,
last_value(reading ORDER BY ts) AS reading
FROM scheduled_late_overlap_input
WHERE ts >= date_trunc('second', now()) - INTERVAL '20 seconds'
AND ts < date_trunc('second', now())
GROUP BY series, date_bin(INTERVAL '1 second', ts)
)
SELECT
target_seconds.target_ts AS ts,
bucketed.series,
bucketed.reading,
now() AS create_time
FROM target_seconds
JOIN bucketed
ON bucketed.bucket_ts = target_seconds.target_ts;
Affected Rows: 0
INSERT INTO scheduled_late_overlap_input VALUES
(date_trunc('second', now()), 'seed', 1.0);
Affected Rows: 1
-- SQLNESS SLEEP 3s
INSERT INTO scheduled_late_overlap_input
SELECT min(ts), 'late_a', 10.0 FROM scheduled_late_overlap_input
UNION ALL
SELECT min(ts), 'late_b', 11.0 FROM scheduled_late_overlap_input;
Affected Rows: 2
-- SQLNESS SLEEP 4s
SELECT
count(DISTINCT series) >= 3 AS replayed_late_series,
bool_and(ts < create_time) AS all_ts_before_create_time,
bool_and(ts >= create_time - INTERVAL '10 seconds') AS all_ts_within_overlap_window
FROM scheduled_late_overlap_sink
WHERE ts = (SELECT min(ts) FROM scheduled_late_overlap_input);
+----------------------+---------------------------+------------------------------+
| replayed_late_series | all_ts_before_create_time | all_ts_within_overlap_window |
+----------------------+---------------------------+------------------------------+
| true | true | true |
+----------------------+---------------------------+------------------------------+
DROP FLOW scheduled_late_overlap_flow;
Affected Rows: 0
DROP TABLE scheduled_late_overlap_sink;
Affected Rows: 0
DROP TABLE scheduled_late_overlap_input;
Affected Rows: 0
@@ -0,0 +1,68 @@
CREATE TABLE scheduled_late_overlap_input (
ts TIMESTAMP(3) TIME INDEX,
series STRING,
reading DOUBLE,
PRIMARY KEY(series)
);
CREATE FLOW scheduled_late_overlap_flow
SINK TO scheduled_late_overlap_sink
EVAL INTERVAL '1s'
AS
WITH
target_offsets(delta) AS (
VALUES
(INTERVAL '1 second'),
(INTERVAL '2 seconds'),
(INTERVAL '3 seconds'),
(INTERVAL '4 seconds'),
(INTERVAL '5 seconds'),
(INTERVAL '6 seconds'),
(INTERVAL '7 seconds'),
(INTERVAL '8 seconds'),
(INTERVAL '9 seconds'),
(INTERVAL '10 seconds')
),
target_seconds AS (
SELECT date_trunc('second', now()) - delta AS target_ts
FROM target_offsets
),
bucketed AS (
SELECT
series,
date_bin(INTERVAL '1 second', ts) AS bucket_ts,
last_value(reading ORDER BY ts) AS reading
FROM scheduled_late_overlap_input
WHERE ts >= date_trunc('second', now()) - INTERVAL '20 seconds'
AND ts < date_trunc('second', now())
GROUP BY series, date_bin(INTERVAL '1 second', ts)
)
SELECT
target_seconds.target_ts AS ts,
bucketed.series,
bucketed.reading,
now() AS create_time
FROM target_seconds
JOIN bucketed
ON bucketed.bucket_ts = target_seconds.target_ts;
INSERT INTO scheduled_late_overlap_input VALUES
(date_trunc('second', now()), 'seed', 1.0);
-- SQLNESS SLEEP 3s
INSERT INTO scheduled_late_overlap_input
SELECT min(ts), 'late_a', 10.0 FROM scheduled_late_overlap_input
UNION ALL
SELECT min(ts), 'late_b', 11.0 FROM scheduled_late_overlap_input;
-- SQLNESS SLEEP 4s
SELECT
count(DISTINCT series) >= 3 AS replayed_late_series,
bool_and(ts < create_time) AS all_ts_before_create_time,
bool_and(ts >= create_time - INTERVAL '10 seconds') AS all_ts_within_overlap_window
FROM scheduled_late_overlap_sink
WHERE ts = (SELECT min(ts) FROM scheduled_late_overlap_input);
DROP FLOW scheduled_late_overlap_flow;
DROP TABLE scheduled_late_overlap_sink;
DROP TABLE scheduled_late_overlap_input;
@@ -0,0 +1,82 @@
CREATE TABLE now_boundary_input (
ts TIMESTAMP(3) TIME INDEX,
v DOUBLE,
PRIMARY KEY(v)
);
Affected Rows: 0
CREATE FLOW now_boundary_flow
SINK TO now_boundary_sink
EVAL INTERVAL '1s'
AS
SELECT
date_bin(INTERVAL '1 second', ts) AS window_start,
count(v) AS value_count,
now() AS create_time,
current_timestamp() AS cur_ts
FROM now_boundary_input
WHERE ts >= date_trunc('second', now()) - INTERVAL '1 second'
AND ts < date_trunc('second', current_timestamp())
GROUP BY date_bin(INTERVAL '1 second', ts);
Affected Rows: 0
INSERT INTO now_boundary_input VALUES
(date_trunc('second', now()) - INTERVAL '1 second', 0.0),
(date_trunc('second', now()), 1.0),
(date_trunc('second', now()) + INTERVAL '1 second', 2.0),
(date_trunc('second', now()) + INTERVAL '2 seconds', 3.0),
(date_trunc('second', now()) + INTERVAL '3 seconds', 4.0),
(date_trunc('second', now()) + INTERVAL '4 seconds', 5.0),
(date_trunc('second', now()) + INTERVAL '5 seconds', 6.0),
(date_trunc('second', now()) + INTERVAL '6 seconds', 7.0),
(date_trunc('second', now()) + INTERVAL '7 seconds', 8.0);
Affected Rows: 9
-- SQLNESS SLEEP 4s
SELECT
count(DISTINCT window_start) >= 1 AS has_selected_window,
min(value_count) AS min_value_count,
max(value_count) AS max_value_count,
bool_and(create_time = date_trunc('second', create_time)) AS all_create_time_at_second_boundary,
bool_and(cur_ts = create_time) AS cur_ts_equals_create_time,
bool_and(window_start = create_time - INTERVAL '1 second') AS all_windows_match_scheduled_previous_second
FROM now_boundary_sink
WHERE value_count > 0;
+---------------------+-----------------+-----------------+------------------------------------+---------------------------+---------------------------------------------+
| has_selected_window | min_value_count | max_value_count | all_create_time_at_second_boundary | cur_ts_equals_create_time | all_windows_match_scheduled_previous_second |
+---------------------+-----------------+-----------------+------------------------------------+---------------------------+---------------------------------------------+
| true | 1 | 1 | true | true | true |
+---------------------+-----------------+-----------------+------------------------------------+---------------------------+---------------------------------------------+
-- The WHERE clause should select exactly the previous scheduled second.
-- It intentionally uses both now() and current_timestamp() in filter bounds.
-- If the distributed analyzer folds now()/current_timestamp() with wall-clock time, the remote
-- filter shifts forward and this invariant fails: window_start becomes equal
-- to create_time instead of create_time - 1s.
SELECT
bool_and(value_count = 1) AS all_value_count_equals_one
FROM now_boundary_sink
WHERE value_count > 0;
+----------------------------+
| all_value_count_equals_one |
+----------------------------+
| true |
+----------------------------+
DROP FLOW now_boundary_flow;
Affected Rows: 0
DROP TABLE now_boundary_sink;
Affected Rows: 0
DROP TABLE now_boundary_input;
Affected Rows: 0
@@ -0,0 +1,55 @@
CREATE TABLE now_boundary_input (
ts TIMESTAMP(3) TIME INDEX,
v DOUBLE,
PRIMARY KEY(v)
);
CREATE FLOW now_boundary_flow
SINK TO now_boundary_sink
EVAL INTERVAL '1s'
AS
SELECT
date_bin(INTERVAL '1 second', ts) AS window_start,
count(v) AS value_count,
now() AS create_time,
current_timestamp() AS cur_ts
FROM now_boundary_input
WHERE ts >= date_trunc('second', now()) - INTERVAL '1 second'
AND ts < date_trunc('second', current_timestamp())
GROUP BY date_bin(INTERVAL '1 second', ts);
INSERT INTO now_boundary_input VALUES
(date_trunc('second', now()) - INTERVAL '1 second', 0.0),
(date_trunc('second', now()), 1.0),
(date_trunc('second', now()) + INTERVAL '1 second', 2.0),
(date_trunc('second', now()) + INTERVAL '2 seconds', 3.0),
(date_trunc('second', now()) + INTERVAL '3 seconds', 4.0),
(date_trunc('second', now()) + INTERVAL '4 seconds', 5.0),
(date_trunc('second', now()) + INTERVAL '5 seconds', 6.0),
(date_trunc('second', now()) + INTERVAL '6 seconds', 7.0),
(date_trunc('second', now()) + INTERVAL '7 seconds', 8.0);
-- SQLNESS SLEEP 4s
SELECT
count(DISTINCT window_start) >= 1 AS has_selected_window,
min(value_count) AS min_value_count,
max(value_count) AS max_value_count,
bool_and(create_time = date_trunc('second', create_time)) AS all_create_time_at_second_boundary,
bool_and(cur_ts = create_time) AS cur_ts_equals_create_time,
bool_and(window_start = create_time - INTERVAL '1 second') AS all_windows_match_scheduled_previous_second
FROM now_boundary_sink
WHERE value_count > 0;
-- The WHERE clause should select exactly the previous scheduled second.
-- It intentionally uses both now() and current_timestamp() in filter bounds.
-- If the distributed analyzer folds now()/current_timestamp() with wall-clock time, the remote
-- filter shifts forward and this invariant fails: window_start becomes equal
-- to create_time instead of create_time - 1s.
SELECT
bool_and(value_count = 1) AS all_value_count_equals_one
FROM now_boundary_sink
WHERE value_count > 0;
DROP FLOW now_boundary_flow;
DROP TABLE now_boundary_sink;
DROP TABLE now_boundary_input;