1use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
16use std::sync::{Arc, RwLock};
17use std::time::{Duration, SystemTime, UNIX_EPOCH};
18
19use api::v1::{CreateTableExpr, TableName};
20use catalog::CatalogManagerRef;
21use common_error::ext::BoxedError;
22use common_query::logical_plan::breakup_insert_plan;
23use common_telemetry::tracing::warn;
24use common_telemetry::{debug, info};
25use common_time::Timestamp;
26use datafusion::datasource::DefaultTableSource;
27use datafusion::sql::unparser::expr_to_sql;
28use datafusion_common::tree_node::{Transformed, TreeNode};
29use datafusion_common::utils::quote_identifier;
30use datafusion_common::{DFSchemaRef, TableReference};
31use datafusion_expr::{DmlStatement, LogicalPlan, WriteOp, col, lit};
32use datatypes::schema::Schema;
33use query::QueryEngineRef;
34use query::options::FLOW_INCREMENTAL_MODE;
35use query::query_engine::DefaultSerializer;
36use session::context::QueryContextRef;
37use snafu::{OptionExt, ResultExt};
38use sql::parsers::utils::is_tql;
39use store_api::mito_engine_options::MERGE_MODE_KEY;
40use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
41use table::table::adapter::DfTableProviderAdapter;
42use tokio::sync::oneshot::error::TryRecvError;
43use tokio::sync::{Mutex, oneshot};
44use tokio::time::Instant;
45
46use crate::batching_mode::BatchingModeOptions;
47use crate::batching_mode::checkpoint::checkpoint_mode_label;
48use crate::batching_mode::eval_schedule::{EvalSchedule, select_due_scheduled_times};
49use crate::batching_mode::frontend_client::{FrontendClient, PeerDesc};
50use crate::batching_mode::state::{
51 CheckpointMode, DirtyTimeWindows, FilterExprInfo, TaskState, to_df_literal,
52};
53use crate::batching_mode::table_creator::{QueryType, create_table_with_expr};
54use crate::batching_mode::time_window::TimeWindowExpr;
55use crate::batching_mode::utils::{
56 AddFilterRewriter, ColumnMatcherRewriter, df_plan_to_sql, gen_plan_with_matching_schema,
57 get_table_info_df_schema, sql_to_df_plan,
58};
59use crate::df_optimizer::apply_df_optimizer;
60use crate::error::{
61 DatafusionSnafu, ExternalSnafu, InvalidQuerySnafu, SubstraitEncodeLogicalPlanSnafu,
62 UnexpectedSnafu,
63};
64use crate::metrics::{
65 METRIC_FLOW_BATCHING_ENGINE_ERROR_CNT, METRIC_FLOW_BATCHING_ENGINE_QUERY_TIME,
66 METRIC_FLOW_BATCHING_ENGINE_SLOW_QUERY, METRIC_FLOW_BATCHING_ENGINE_START_QUERY_CNT,
67 METRIC_FLOW_ROWS,
68};
69use crate::{Error, FlowId};
70
71mod ckpt;
72mod inc;
73
74fn wall_clock_unix_secs() -> i64 {
76 SystemTime::now()
77 .duration_since(UNIX_EPOCH)
78 .unwrap_or_default()
79 .as_secs() as i64
80}
81
82#[derive(Clone)]
84pub struct TaskConfig {
85 pub flow_id: FlowId,
86 pub query: String,
87 pub output_schema: DFSchemaRef,
89 pub time_window_expr: Option<TimeWindowExpr>,
90 pub expire_after: Option<i64>,
92 pub sink_table_name: [String; 3],
93 pub source_table_names: HashSet<[String; 3]>,
94 pub catalog_manager: CatalogManagerRef,
95 pub query_type: QueryType,
96 pub batch_opts: Arc<BatchingModeOptions>,
97 pub flow_eval_interval: Option<Duration>,
98 pub eval_schedule: Option<EvalSchedule>,
100}
101
102fn determine_query_type(query: &str, query_ctx: &QueryContextRef) -> Result<QueryType, Error> {
103 let is_tql = is_tql(query_ctx.sql_dialect(), query)
104 .map_err(BoxedError::new)
105 .context(ExternalSnafu)?;
106 Ok(if is_tql {
107 QueryType::Tql
108 } else {
109 QueryType::Sql
110 })
111}
112
113fn is_merge_mode_last_non_null(options: &HashMap<String, String>) -> bool {
114 options
115 .get(MERGE_MODE_KEY)
116 .map(|mode| mode.eq_ignore_ascii_case("last_non_null"))
117 .unwrap_or(false)
118}
119
120fn encode_insert_plan_request(
121 insert_to: TableName,
122 insert_input_plan: &LogicalPlan,
123) -> Result<api::v1::QueryRequest, Error> {
124 let message = DFLogicalSubstraitConvertor {}
125 .encode(insert_input_plan, DefaultSerializer)
126 .context(SubstraitEncodeLogicalPlanSnafu)?;
127 Ok(api::v1::QueryRequest {
128 query: Some(api::v1::query_request::Query::InsertIntoPlan(
129 api::v1::InsertIntoPlan {
130 table_name: Some(insert_to),
131 logical_plan: message.to_vec(),
132 },
133 )),
134 })
135}
136
137fn format_insert_target_columns(plan: &LogicalPlan) -> String {
138 plan.schema()
139 .fields()
140 .iter()
141 .map(|field| quote_identifier(field.name()).to_string())
142 .collect::<Vec<_>>()
143 .join(", ")
144}
145
146#[derive(Clone)]
147pub struct BatchingTask {
148 pub config: Arc<TaskConfig>,
149 pub state: Arc<RwLock<TaskState>>,
150 execution_lock: Arc<Mutex<()>>,
154}
155
156pub struct TaskArgs<'a> {
158 pub flow_id: FlowId,
159 pub query: &'a str,
160 pub plan: LogicalPlan,
161 pub time_window_expr: Option<TimeWindowExpr>,
162 pub expire_after: Option<i64>,
163 pub sink_table_name: [String; 3],
164 pub source_table_names: Vec<[String; 3]>,
165 pub query_ctx: QueryContextRef,
166 pub catalog_manager: CatalogManagerRef,
167 pub shutdown_rx: oneshot::Receiver<()>,
168 pub batch_opts: Arc<BatchingModeOptions>,
169 pub flow_eval_interval: Option<Duration>,
170 pub eval_schedule: Option<EvalSchedule>,
172}
173
174pub struct PlanInfo {
175 pub plan: LogicalPlan,
176 pub dirty_restore: DirtyRestore,
177 pub coverage: QueryCoverage,
178}
179
180#[derive(Clone)]
181pub enum QueryCoverage {
182 UnfilteredFull,
187 ScopedBaseRepair,
190 FencedRepairChunk { high: BTreeMap<u64, u64> },
194 IncrementalDelta,
196}
197
198impl QueryCoverage {
199 fn is_incremental_delta(&self) -> bool {
202 matches!(self, Self::IncrementalDelta)
203 }
204
205 fn snapshot_seqs(&self) -> HashMap<u64, u64> {
208 match self {
209 Self::FencedRepairChunk { high } => high.iter().map(|(k, v)| (*k, *v)).collect(),
210 _ => HashMap::new(),
211 }
212 }
213}
214
215pub enum DirtyRestore {
216 Scoped(FilterExprInfo),
219 Unscoped(DirtyTimeWindows),
226}
227
228struct ExecuteOnceOutcome {
229 new_query: Option<PlanInfo>,
230 result: Result<Option<(usize, Duration)>, Error>,
236}
237
238impl BatchingTask {
239 #[allow(clippy::too_many_arguments)]
240 pub fn try_new(
241 TaskArgs {
242 flow_id,
243 query,
244 plan,
245 time_window_expr,
246 expire_after,
247 sink_table_name,
248 source_table_names,
249 query_ctx,
250 catalog_manager,
251 shutdown_rx,
252 batch_opts,
253 flow_eval_interval,
254 eval_schedule,
255 }: TaskArgs<'_>,
256 ) -> Result<Self, Error> {
257 let mut state = TaskState::with_dirty_time_windows(
258 query_ctx.clone(),
259 shutdown_rx,
260 DirtyTimeWindows::new(
261 batch_opts.experimental_max_filter_num_per_query,
262 batch_opts.experimental_time_window_merge_threshold,
263 ),
264 );
265 if !batch_opts.experimental_enable_incremental_read {
266 state.disable_incremental();
267 }
268
269 Ok(Self {
270 config: Arc::new(TaskConfig {
271 flow_id,
272 query: query.to_string(),
273 time_window_expr,
274 expire_after,
275 sink_table_name,
276 source_table_names: source_table_names.into_iter().collect(),
277 catalog_manager,
278 output_schema: plan.schema().clone(),
279 query_type: determine_query_type(query, &query_ctx)?,
280 batch_opts,
281 flow_eval_interval,
282 eval_schedule,
283 }),
284 state: Arc::new(RwLock::new(state)),
285 execution_lock: Arc::new(Mutex::new(())),
286 })
287 }
288
289 pub fn last_execution_time_millis(&self) -> Option<i64> {
290 self.state.read().unwrap().last_execution_time_millis()
291 }
292
293 pub fn start_time_millis(&self) -> Option<i64> {
294 self.state.read().unwrap().start_time_millis()
295 }
296
297 fn frontend_extensions(&self) -> HashMap<String, String> {
300 let ctx = self.state.read().unwrap();
301 let all = ctx.query_ctx.extensions();
302 let mut flow_exts = HashMap::new();
303 if let Some(v) = all.get(query::options::FLOW_SCHEDULED_TIME_MILLIS) {
306 flow_exts.insert(
307 query::options::FLOW_SCHEDULED_TIME_MILLIS.to_string(),
308 v.clone(),
309 );
310 }
311 flow_exts
312 }
313
314 pub fn mark_all_windows_as_dirty(&self) -> Result<(), Error> {
318 let now = SystemTime::now();
319 let now = Timestamp::new_second(
320 now.duration_since(UNIX_EPOCH)
321 .expect("Time went backwards")
322 .as_secs() as _,
323 );
324 let lower_bound = self
325 .config
326 .expire_after
327 .map(|e| now.sub_duration(Duration::from_secs(e as _)))
328 .transpose()
329 .map_err(BoxedError::new)
330 .context(ExternalSnafu)?
331 .unwrap_or(Timestamp::new_second(0));
332 debug!(
333 "Flow {} mark range ({:?}, {:?}) as dirty",
334 self.config.flow_id, lower_bound, now
335 );
336 self.state
337 .write()
338 .unwrap()
339 .dirty_time_windows
340 .add_window(lower_bound, Some(now));
341 Ok(())
342 }
343
344 pub async fn check_or_create_sink_table(
346 &self,
347 engine: &QueryEngineRef,
348 frontend_client: &Arc<FrontendClient>,
349 ) -> Result<Option<(usize, Duration)>, Error> {
350 if !self.is_table_exist(&self.config.sink_table_name).await? {
351 let create_table = self.gen_create_table_expr(engine.clone()).await?;
352 info!(
353 "Try creating sink table(if not exists) with expr: {:?}",
354 create_table
355 );
356 self.create_table(frontend_client, create_table).await?;
357 info!(
358 "Sink table {}(if not exists) created",
359 self.config.sink_table_name.join(".")
360 );
361 }
362
363 Ok(None)
364 }
365
366 pub async fn validate_sink_table_schema(&self, engine: &QueryEngineRef) -> Result<(), Error> {
372 let (table, _) = get_table_info_df_schema(
373 self.config.catalog_manager.clone(),
374 self.config.sink_table_name.clone(),
375 )
376 .await?;
377
378 let table_meta = &table.table_info().meta;
379 let merge_mode_last_non_null =
380 is_merge_mode_last_non_null(&table_meta.options.extra_options);
381 let primary_key_indices = table_meta.primary_key_indices.clone();
382 let query_ctx = self.state.read().unwrap().query_ctx.clone();
383
384 gen_plan_with_matching_schema(
385 &self.config.query,
386 query_ctx,
387 engine.clone(),
388 table_meta.schema.clone(),
389 &primary_key_indices,
390 merge_mode_last_non_null,
391 )
392 .await
393 .map(|_| ())
394 }
395
396 async fn is_table_exist(&self, table_name: &[String; 3]) -> Result<bool, Error> {
397 self.config
398 .catalog_manager
399 .table_exists(&table_name[0], &table_name[1], &table_name[2], None)
400 .await
401 .map_err(BoxedError::new)
402 .context(ExternalSnafu)
403 }
404
405 pub(crate) async fn execute_once_serialized(
406 &self,
407 engine: &QueryEngineRef,
408 frontend_client: &Arc<FrontendClient>,
409 max_window_cnt: Option<usize>,
410 ) -> Result<Option<(usize, Duration)>, Error> {
411 let outcome = self
412 .execute_once_serialized_with_outcome(engine, frontend_client, max_window_cnt)
413 .await;
414 outcome.result
415 }
416
417 async fn execute_once_serialized_with_outcome(
420 &self,
421 engine: &QueryEngineRef,
422 frontend_client: &Arc<FrontendClient>,
423 max_window_cnt: Option<usize>,
424 ) -> ExecuteOnceOutcome {
425 let _execution_guard = self.execution_lock.lock().await;
426 self.execute_once_unlocked(engine, frontend_client, max_window_cnt)
427 .await
428 }
429
430 async fn execute_once_unlocked(
432 &self,
433 engine: &QueryEngineRef,
434 frontend_client: &Arc<FrontendClient>,
435 max_window_cnt: Option<usize>,
436 ) -> ExecuteOnceOutcome {
437 let new_query = match self.gen_insert_plan_unlocked(engine, max_window_cnt).await {
438 Ok(new_query) => new_query,
439 Err(err) => {
440 return ExecuteOnceOutcome {
441 new_query: None,
442 result: Err(err),
443 };
444 }
445 };
446
447 if let Some(new_query) = new_query {
448 debug!("Generate new query: {}", new_query.plan);
449 let res = self
450 .execute_logical_plan_unlocked(
451 frontend_client,
452 &new_query.plan,
453 &new_query.dirty_restore,
454 &new_query.coverage,
455 )
456 .await;
457 if res.is_err() {
458 self.handle_executed_query_failure(Some(&new_query));
459 }
460 ExecuteOnceOutcome {
461 new_query: Some(new_query),
462 result: res,
463 }
464 } else {
465 debug!("Generate no query");
466 ExecuteOnceOutcome {
467 new_query: None,
468 result: Ok(None),
469 }
470 }
471 }
472
473 async fn gen_insert_plan_unlocked(
475 &self,
476 engine: &QueryEngineRef,
477 max_window_cnt: Option<usize>,
478 ) -> Result<Option<PlanInfo>, Error> {
479 let (table, df_schema) = get_table_info_df_schema(
480 self.config.catalog_manager.clone(),
481 self.config.sink_table_name.clone(),
482 )
483 .await?;
484
485 let table_meta = &table.table_info().meta;
486 let merge_mode_last_non_null =
487 is_merge_mode_last_non_null(&table_meta.options.extra_options);
488 let primary_key_indices = table_meta.primary_key_indices.clone();
489
490 let new_query = self
491 .gen_query_with_time_window(
492 engine.clone(),
493 &table.table_info().meta.schema,
494 &primary_key_indices,
495 merge_mode_last_non_null,
496 max_window_cnt,
497 )
498 .await?;
499
500 let Some(new_query) = new_query else {
501 return Ok(None);
502 };
503
504 let table_columns = df_schema
507 .columns()
508 .into_iter()
509 .map(|c| c.name)
510 .collect::<BTreeSet<_>>();
511 for column in new_query.plan.schema().columns() {
512 if !table_columns.contains(column.name()) {
513 self.restore_dirty_windows_after_failure(&new_query);
514 return InvalidQuerySnafu {
515 reason: format!(
516 "Column {} not found in sink table with columns {:?}",
517 column, table_columns
518 ),
519 }
520 .fail();
521 }
522 }
523
524 let table_provider = Arc::new(DfTableProviderAdapter::new(table));
525 let table_source = Arc::new(DefaultTableSource::new(table_provider));
526
527 let plan = LogicalPlan::Dml(DmlStatement::new(
529 datafusion_common::TableReference::Full {
530 catalog: self.config.sink_table_name[0].clone().into(),
531 schema: self.config.sink_table_name[1].clone().into(),
532 table: self.config.sink_table_name[2].clone().into(),
533 },
534 table_source,
535 WriteOp::Insert(datafusion_expr::dml::InsertOp::Append),
536 Arc::new(new_query.plan.clone()),
537 ));
538 let insert_into_info = PlanInfo {
539 plan,
540 dirty_restore: new_query.dirty_restore,
541 coverage: new_query.coverage,
542 };
543 let insert_into =
544 match insert_into_info
545 .plan
546 .clone()
547 .recompute_schema()
548 .context(DatafusionSnafu {
549 context: "Failed to recompute schema",
550 }) {
551 Ok(insert_into) => insert_into,
552 Err(err) => {
553 self.restore_dirty_windows_after_failure(&insert_into_info);
554 return Err(err);
555 }
556 };
557
558 Ok(Some(PlanInfo {
559 plan: insert_into,
560 dirty_restore: insert_into_info.dirty_restore,
561 coverage: insert_into_info.coverage,
562 }))
563 }
564
565 pub async fn create_table(
566 &self,
567 frontend_client: &Arc<FrontendClient>,
568 expr: CreateTableExpr,
569 ) -> Result<(), Error> {
570 let catalog = &self.config.sink_table_name[0];
571 let schema = &self.config.sink_table_name[1];
572 frontend_client
573 .create(expr.clone(), catalog, schema)
574 .await?;
575 Ok(())
576 }
577
578 async fn execute_logical_plan_unlocked(
580 &self,
581 frontend_client: &Arc<FrontendClient>,
582 plan: &LogicalPlan,
583 dirty_restore: &DirtyRestore,
584 coverage: &QueryCoverage,
585 ) -> Result<Option<(usize, Duration)>, Error> {
586 let instant = Instant::now();
587 let flow_id = self.config.flow_id;
588
589 debug!(
590 "Executing flow {flow_id}(expire_after={:?} secs) with query {}",
591 self.config.expire_after, &plan
592 );
593
594 let catalog = &self.config.sink_table_name[0];
595 let schema = &self.config.sink_table_name[1];
596
597 let plan = plan
599 .clone()
600 .transform_down_with_subqueries(|p| {
601 if let LogicalPlan::TableScan(mut table_scan) = p {
602 let resolved = table_scan.table_name.resolve(catalog, schema);
603 table_scan.table_name = resolved.into();
604 Ok(Transformed::yes(LogicalPlan::TableScan(table_scan)))
605 } else {
606 Ok(Transformed::no(p))
607 }
608 })
609 .with_context(|_| DatafusionSnafu {
610 context: format!("Failed to fix table ref in logical plan, plan={:?}", plan),
611 })?
612 .data;
613
614 let incremental_plan = if coverage.is_incremental_delta() {
617 self.prepare_plan_for_incremental(&plan).await?
618 } else {
619 None
620 };
621 let incremental_safe = incremental_plan.is_some();
622 if coverage.is_incremental_delta() && !incremental_safe {
623 debug!(
624 "Flow {flow_id} skipped unsafe incremental delta fallback; \
625 restored dirty signal instead of executing an unfiltered full snapshot"
626 );
627 self.restore_dirty_windows(dirty_restore);
628 return Ok(None);
629 }
630 let plan = incremental_plan.unwrap_or_else(|| plan.clone());
631
632 let extensions = self
633 .build_flow_query_extensions(incremental_safe, coverage.is_incremental_delta())
634 .await?;
635 let frontend_extensions = self.frontend_extensions();
636 let extension_refs = extensions
637 .iter()
638 .map(|(key, value)| (*key, value.as_str()))
639 .chain(
640 frontend_extensions
641 .iter()
642 .map(|(key, value)| (key.as_str(), value.as_str())),
643 )
644 .collect::<Vec<_>>();
645 let query_mode = if extensions
646 .iter()
647 .any(|(key, _)| *key == FLOW_INCREMENTAL_MODE)
648 {
649 CheckpointMode::Incremental
650 } else {
651 CheckpointMode::FullSnapshot
652 };
653 Self::record_query_mode(flow_id, query_mode);
654 debug!(
655 "Flow {flow_id} executing batching query with checkpoint_mode={}, extension_count={}",
656 checkpoint_mode_label(query_mode),
657 extensions.len()
658 );
659
660 let mut peer_desc = None;
661 let res = {
662 let _timer = METRIC_FLOW_BATCHING_ENGINE_QUERY_TIME
663 .with_label_values(&[flow_id.to_string().as_str()])
664 .start_timer();
665
666 let req = if let Some((insert_to, insert_input_plan)) =
667 breakup_insert_plan(&plan, catalog, schema)
668 {
669 if query_mode == CheckpointMode::FullSnapshot
670 && matches!(self.config.query_type, QueryType::Sql)
671 && self.config.flow_eval_interval.is_some()
672 && self.config.time_window_expr.is_none()
673 {
674 match df_plan_to_sql(&insert_input_plan) {
684 Ok(select_sql) => {
685 let target_columns = format_insert_target_columns(&insert_input_plan);
686 let sql = format!(
687 "INSERT INTO {} ({}) {}",
688 TableReference::full(
689 insert_to.catalog_name.as_str(),
690 insert_to.schema_name.as_str(),
691 insert_to.table_name.as_str(),
692 )
693 .to_quoted_string(),
694 target_columns,
695 select_sql
696 );
697 api::v1::QueryRequest {
698 query: Some(api::v1::query_request::Query::Sql(sql)),
699 }
700 }
701 Err(err) => {
702 debug!(
703 "Failed to unparse full-snapshot SQL flow {} plan; \
704 falling back to InsertIntoPlan: {:?}",
705 flow_id, err
706 );
707 encode_insert_plan_request(insert_to, &insert_input_plan)?
708 }
709 }
710 } else {
711 encode_insert_plan_request(insert_to, &insert_input_plan)?
712 }
713 } else {
714 let message = DFLogicalSubstraitConvertor {}
715 .encode(&plan, DefaultSerializer)
716 .context(SubstraitEncodeLogicalPlanSnafu)?;
717
718 api::v1::QueryRequest {
719 query: Some(api::v1::query_request::Query::LogicalPlan(message.to_vec())),
720 }
721 };
722
723 let snapshot_seqs = coverage.snapshot_seqs();
724 {
725 let mut state = self.state.write().unwrap();
726 state.record_start_time_if_first();
727 }
728 frontend_client
729 .query_with_terminal_metrics(
730 catalog,
731 schema,
732 req,
733 &extension_refs,
734 &snapshot_seqs,
735 &mut peer_desc,
736 )
737 .await
738 };
739
740 let elapsed = instant.elapsed();
741 let peer_label = peer_desc
742 .as_ref()
743 .map(ToString::to_string)
744 .unwrap_or_else(|| PeerDesc::default().to_string());
745 if let Err(err) = &res {
746 warn!(
747 "Failed to execute Flow {flow_id} on frontend {peer_label}, result: {err:?}, elapsed: {:?} with query: {}",
748 elapsed, &plan
749 );
750 let decision = {
751 let mut state = self.state.write().unwrap();
752 let reason = Self::query_failure_reason(err, coverage);
753 Self::apply_query_failure_to_state(&mut state, elapsed, coverage, reason)
754 };
755 if let Some(decision) = decision {
756 Self::record_checkpoint_decision(flow_id, decision);
757 }
758 }
759
760 if elapsed >= self.config.batch_opts.slow_query_threshold {
762 warn!(
763 "Flow {flow_id} on frontend {peer_label} executed for {:?} before complete, query: {}",
764 elapsed, &plan
765 );
766 let flow_id = flow_id.to_string();
767 METRIC_FLOW_BATCHING_ENGINE_SLOW_QUERY
768 .with_label_values(&[flow_id.as_str(), peer_label.as_str()])
769 .observe(elapsed.as_secs_f64());
770 }
771
772 let res = res?;
773 let (affected_rows, _) = res.output.extract_rows_and_cost();
774 debug!(
775 "Flow {flow_id} executed, affected_rows: {affected_rows:?}, elapsed: {:?}, watermark: {:?}",
776 elapsed,
777 res.region_watermark_map()
778 );
779 METRIC_FLOW_ROWS
780 .with_label_values(&[format!("{}-out-batching", flow_id).as_str()])
781 .inc_by(affected_rows as _);
782 let decision = {
783 let mut state = self.state.write().unwrap();
784 Self::apply_query_result_to_state(&mut state, &res, elapsed, coverage)
785 };
786 Self::record_checkpoint_decision(flow_id, decision);
787
788 Ok(Some((affected_rows, elapsed)))
789 }
790
791 fn restore_dirty_windows(&self, dirty_restore: &DirtyRestore) {
795 match dirty_restore {
796 DirtyRestore::Scoped(filter) => self.restore_scoped_dirty_windows(filter),
797 DirtyRestore::Unscoped(dirty_windows) => self
798 .state
799 .write()
800 .unwrap()
801 .dirty_time_windows
802 .add_dirty_windows(dirty_windows),
803 }
804 }
805
806 fn restore_dirty_windows_after_failure(&self, query: &PlanInfo) {
809 self.restore_dirty_windows(&query.dirty_restore);
810 }
811
812 fn restore_scoped_dirty_windows(&self, filter: &FilterExprInfo) {
815 self.state.write().unwrap().restore_scoped_windows(filter);
816 }
817
818 fn restore_scoped_dirty_windows_on_err<T>(
821 &self,
822 filter: &FilterExprInfo,
823 result: Result<T, Error>,
824 ) -> Result<T, Error> {
825 result.inspect_err(|_| {
826 self.restore_scoped_dirty_windows(filter);
827 })
828 }
829
830 fn restore_unscoped_dirty_windows(&self, dirty_windows: &DirtyTimeWindows) {
833 self.state
834 .write()
835 .unwrap()
836 .dirty_time_windows
837 .add_dirty_windows(dirty_windows);
838 }
839
840 fn restore_unscoped_dirty_windows_on_err<T>(
843 &self,
844 dirty_windows: &DirtyTimeWindows,
845 result: Result<T, Error>,
846 ) -> Result<T, Error> {
847 result.inspect_err(|_| {
848 self.restore_unscoped_dirty_windows(dirty_windows);
849 })
850 }
851
852 fn drain_dirty_windows_signal(&self) -> (bool, DirtyTimeWindows) {
855 let mut state = self.state.write().unwrap();
856 let dirty_windows_to_restore = state.dirty_time_windows.clone();
857 let is_dirty = !dirty_windows_to_restore.is_empty();
858 state.dirty_time_windows.clean();
859 (is_dirty, dirty_windows_to_restore)
860 }
861
862 #[allow(clippy::too_many_arguments)]
863 async fn gen_unfiltered_plan_info(
866 &self,
867 engine: QueryEngineRef,
868 query_ctx: QueryContextRef,
869 sink_table_schema: Arc<Schema>,
870 primary_key_indices: &[usize],
871 allow_partial: bool,
872 dirty_windows_to_restore: DirtyTimeWindows,
873 retention_filter: Option<(&str, Timestamp, &'static str)>,
874 coverage: QueryCoverage,
875 ) -> Result<PlanInfo, Error> {
876 let mut plan = self.restore_unscoped_dirty_windows_on_err(
877 &dirty_windows_to_restore,
878 gen_plan_with_matching_schema(
879 &self.config.query,
880 query_ctx,
881 engine,
882 sink_table_schema,
883 primary_key_indices,
884 allow_partial,
885 )
886 .await,
887 )?;
888
889 if let Some((col_name, lower_bound, context)) = retention_filter {
890 let lower = self.restore_unscoped_dirty_windows_on_err(
891 &dirty_windows_to_restore,
892 to_df_literal(lower_bound),
893 )?;
894 let retention_filter = col(col_name).gt_eq(lit(lower));
895 let mut add_filter = AddFilterRewriter::new(retention_filter);
896 plan = self.restore_unscoped_dirty_windows_on_err(
897 &dirty_windows_to_restore,
898 plan.clone()
899 .rewrite(&mut add_filter)
900 .with_context(|_| DatafusionSnafu {
901 context: format!(
902 "Failed to apply {context} expire_after filter to plan:\n {}\n",
903 plan
904 ),
905 })
906 .map(|rewrite| rewrite.data),
907 )?;
908 }
909
910 Ok(PlanInfo {
911 plan,
912 dirty_restore: DirtyRestore::Unscoped(dirty_windows_to_restore),
913 coverage,
914 })
915 }
916
917 #[allow(clippy::too_many_arguments)]
918 async fn gen_unfiltered_plan_info_if_dirty(
921 &self,
922 engine: QueryEngineRef,
923 query_ctx: QueryContextRef,
924 sink_table_schema: Arc<Schema>,
925 primary_key_indices: &[usize],
926 allow_partial: bool,
927 retention_filter: Option<(&str, Timestamp, &'static str)>,
928 coverage: QueryCoverage,
929 ) -> Result<Option<PlanInfo>, Error> {
930 let (is_dirty, dirty_windows_to_restore) = self.drain_dirty_windows_signal();
931 if !is_dirty {
932 debug!("Flow id={:?}, no new data, not update", self.config.flow_id);
933 return Ok(None);
934 }
935
936 self.gen_unfiltered_plan_info(
937 engine,
938 query_ctx,
939 sink_table_schema,
940 primary_key_indices,
941 allow_partial,
942 dirty_windows_to_restore,
943 retention_filter,
944 coverage,
945 )
946 .await
947 .map(Some)
948 }
949
950 fn handle_executed_query_failure(&self, query: Option<&PlanInfo>) {
951 if let Some(query) = query {
952 self.restore_dirty_windows_after_failure(query);
953 }
954 }
955
956 pub async fn start_executing_loop(
964 &self,
965 engine: QueryEngineRef,
966 frontend_client: Arc<FrontendClient>,
967 ) {
968 if self.config.flow_eval_interval.is_some() {
969 self.start_scheduled_loop(engine, frontend_client).await;
970 } else {
971 self.start_adaptive_loop(engine, frontend_client).await;
972 }
973 }
974
975 async fn start_scheduled_loop(
985 &self,
986 engine: QueryEngineRef,
987 frontend_client: Arc<FrontendClient>,
988 ) {
989 let flow_id_str = self.config.flow_id.to_string();
990
991 let schedule = match &self.config.eval_schedule {
992 Some(s) => s.clone(),
993 None => {
994 let eval_interval_secs = self
995 .config
996 .flow_eval_interval
997 .map(|d| d.as_secs() as i64)
998 .expect("checked by caller");
999
1000 match EvalSchedule::from_config(Some(eval_interval_secs), None) {
1003 Ok(Some(s)) => s,
1004 Ok(None) => {
1005 warn!(
1006 "Flow {}: EVAL INTERVAL set but no schedule parsed; exiting loop",
1007 flow_id_str
1008 );
1009 return;
1010 }
1011 Err(e) => {
1012 warn!(
1013 "Flow {}: Failed to parse eval schedule: {}; exiting loop",
1014 flow_id_str, e
1015 );
1016 return;
1017 }
1018 }
1019 }
1020 };
1021
1022 let mut cursor_secs = schedule.start_secs.saturating_sub(schedule.interval_secs);
1025
1026 info!(
1027 "Flow {}: entering scheduled loop, interval={}s, start={}, anchor={}, policy={:?}, max_runs={}, max_lag={}s",
1028 flow_id_str,
1029 schedule.interval_secs,
1030 schedule.start_secs,
1031 schedule.anchor_secs,
1032 schedule.missed_tick_policy,
1033 schedule.max_runs,
1034 schedule.max_lag_secs,
1035 );
1036
1037 loop {
1038 if self.is_shutdown_signaled() {
1039 break;
1040 }
1041
1042 let wall_now_secs = wall_clock_unix_secs();
1043
1044 let due = match select_due_scheduled_times(&schedule, cursor_secs, wall_now_secs) {
1045 Some(d) => d,
1046 None => {
1047 warn!(
1048 "Flow {}: Invalid schedule (interval <= 0), exiting loop",
1049 flow_id_str
1050 );
1051 return;
1052 }
1053 };
1054
1055 if due.scheduled_times_secs.is_empty() {
1056 if due.skipped > 0 {
1057 warn!(
1058 "Flow {}: all {} due scheduled times skipped by max-lag, advancing cursor to wall-clock ({wall_now_secs}) to avoid re-skipping",
1059 flow_id_str, due.skipped
1060 );
1061 cursor_secs = wall_now_secs;
1062 continue;
1063 }
1064
1065 let next = schedule.next_scheduled_time_after(cursor_secs);
1067 if next <= wall_now_secs {
1068 cursor_secs = wall_now_secs;
1071 continue;
1072 }
1073 let wait_secs = (next - wall_now_secs) as u64;
1074 let wait_dur = Duration::from_secs(wait_secs);
1075 debug!(
1076 "Flow {}: no due scheduled times, sleeping for {}s until next scheduled time at {}",
1077 flow_id_str, wait_secs, next
1078 );
1079 tokio::time::sleep(wait_dur).await;
1080 continue;
1081 }
1082
1083 if due.skipped > 0 {
1084 info!(
1085 "Flow {}: {} due scheduled times, {} skipped (catch-up)",
1086 flow_id_str,
1087 due.scheduled_times_secs.len(),
1088 due.skipped
1089 );
1090 }
1091
1092 for scheduled_time_secs in &due.scheduled_times_secs {
1094 if self.is_shutdown_signaled() {
1095 break;
1096 }
1097
1098 METRIC_FLOW_BATCHING_ENGINE_START_QUERY_CNT
1099 .with_label_values(&[&flow_id_str])
1100 .inc();
1101
1102 let outcome = self
1103 .execute_once_serialized_at_scheduled_time(
1104 &engine,
1105 &frontend_client,
1106 *scheduled_time_secs,
1107 )
1108 .await;
1109
1110 cursor_secs = *scheduled_time_secs;
1112
1113 match outcome.result {
1114 Ok(Some((rows, elapsed))) => {
1115 debug!(
1116 "Flow {}: scheduled time {} completed, rows={}, elapsed={:?}",
1117 flow_id_str, scheduled_time_secs, rows, elapsed
1118 );
1119 }
1120 Ok(None) => {
1121 debug!(
1122 "Flow {}: scheduled time {} produced no query (no dirty signal or no-op)",
1123 flow_id_str, scheduled_time_secs
1124 );
1125 }
1126 Err(err) => {
1127 warn!(
1128 "Flow {}: scheduled time {} failed: {:?}",
1129 flow_id_str, scheduled_time_secs, err
1130 );
1131 METRIC_FLOW_BATCHING_ENGINE_ERROR_CNT
1132 .with_label_values(&[&flow_id_str])
1133 .inc();
1134 }
1138 }
1139 }
1140 }
1141 }
1142
1143 async fn start_adaptive_loop(
1145 &self,
1146 engine: QueryEngineRef,
1147 frontend_client: Arc<FrontendClient>,
1148 ) {
1149 let flow_id_str = self.config.flow_id.to_string();
1150 let mut max_window_cnt = None;
1151 loop {
1152 if self.is_shutdown_signaled() {
1153 break;
1154 }
1155 METRIC_FLOW_BATCHING_ENGINE_START_QUERY_CNT
1156 .with_label_values(&[&flow_id_str])
1157 .inc();
1158
1159 let min_refresh = self.config.batch_opts.experimental_min_refresh_duration;
1160
1161 let outcome = self
1162 .execute_once_serialized_with_outcome(&engine, &frontend_client, max_window_cnt)
1163 .await;
1164
1165 match outcome.result {
1166 Ok(Some(_)) => {
1167 max_window_cnt = max_window_cnt.map(|cnt| {
1168 (cnt + 1).min(self.config.batch_opts.experimental_max_filter_num_per_query)
1169 });
1170
1171 let sleep_until = {
1172 let state = self.state.write().unwrap();
1173
1174 let time_window_size = self
1175 .config
1176 .time_window_expr
1177 .as_ref()
1178 .and_then(|t| *t.time_window_size());
1179
1180 let prefer_short_incremental_cadence = state.checkpoint_mode()
1181 == CheckpointMode::Incremental
1182 && !state.is_incremental_disabled();
1183
1184 state.get_next_start_query_time(
1185 self.config.flow_id,
1186 &time_window_size,
1187 min_refresh,
1188 Some(self.config.batch_opts.query_timeout),
1189 self.config.batch_opts.experimental_max_filter_num_per_query,
1190 prefer_short_incremental_cadence,
1191 )
1192 };
1193
1194 tokio::time::sleep_until(sleep_until).await;
1195 }
1196 Ok(None) => {
1197 debug!(
1198 "Flow id = {:?} found no new data, sleep for {:?} then continue",
1199 self.config.flow_id, min_refresh
1200 );
1201 tokio::time::sleep(min_refresh).await;
1202 continue;
1203 }
1204 Err(err) => {
1205 METRIC_FLOW_BATCHING_ENGINE_ERROR_CNT
1206 .with_label_values(&[&flow_id_str])
1207 .inc();
1208 match outcome.new_query {
1209 Some(query) => {
1210 common_telemetry::error!(err; "Failed to execute query for flow={} with query: {}", self.config.flow_id, query.plan);
1211 max_window_cnt = Some(1);
1212 }
1213 None => {
1214 common_telemetry::error!(err; "Failed to generate query for flow={}", self.config.flow_id)
1215 }
1216 }
1217 tokio::time::sleep(min_refresh).await;
1218 }
1219 }
1220 }
1221 }
1222
1223 fn is_shutdown_signaled(&self) -> bool {
1225 let mut state = self.state.write().unwrap();
1226 match state.shutdown_rx.try_recv() {
1227 Ok(()) | Err(TryRecvError::Closed) => true,
1228 Err(TryRecvError::Empty) => false,
1229 }
1230 }
1231
1232 async fn execute_once_serialized_at_scheduled_time(
1239 &self,
1240 engine: &QueryEngineRef,
1241 frontend_client: &Arc<FrontendClient>,
1242 scheduled_time_secs: i64,
1243 ) -> ExecuteOnceOutcome {
1244 let _execution_guard = self.execution_lock.lock().await;
1245
1246 struct QueryContextRestoreGuard {
1247 state: Arc<RwLock<TaskState>>,
1248 old_ctx: Option<QueryContextRef>,
1249 }
1250
1251 impl Drop for QueryContextRestoreGuard {
1252 fn drop(&mut self) {
1253 let Some(old_ctx) = self.old_ctx.take() else {
1254 return;
1255 };
1256 if let Ok(mut state) = self.state.write() {
1257 state.query_ctx = old_ctx;
1258 }
1259 }
1260 }
1261
1262 let old_ctx = {
1265 let mut state = self.state.write().unwrap();
1266 let old = state.query_ctx.clone();
1267 let mut new_ctx = (*old).clone();
1268 new_ctx.set_extension(
1269 query::options::FLOW_SCHEDULED_TIME_MILLIS,
1270 (scheduled_time_secs.saturating_mul(1000)).to_string(),
1271 );
1272 state.query_ctx = Arc::new(new_ctx);
1273 old
1274 };
1275 let restore_guard = QueryContextRestoreGuard {
1276 state: self.state.clone(),
1277 old_ctx: Some(old_ctx),
1278 };
1279
1280 let outcome = self
1281 .execute_once_unlocked(engine, frontend_client, None)
1282 .await;
1283
1284 drop(restore_guard);
1288
1289 outcome
1290 }
1291
1292 async fn gen_create_table_expr(
1299 &self,
1300 engine: QueryEngineRef,
1301 ) -> Result<CreateTableExpr, Error> {
1302 let query_ctx = self.state.read().unwrap().query_ctx.clone();
1303 let plan =
1304 sql_to_df_plan(query_ctx.clone(), engine.clone(), &self.config.query, true).await?;
1305 create_table_with_expr(&plan, &self.config.sink_table_name, &self.config.query_type)
1306 }
1307
1308 fn should_use_unfiltered_incremental_delta(&self) -> bool {
1311 let state = self.state.read().unwrap();
1312 state.checkpoint_mode() == CheckpointMode::Incremental
1313 && !state.is_incremental_disabled()
1314 && matches!(self.config.query_type, QueryType::Sql)
1315 }
1316
1317 async fn gen_query_with_time_window(
1320 &self,
1321 engine: QueryEngineRef,
1322 sink_table_schema: &Arc<Schema>,
1323 primary_key_indices: &[usize],
1324 allow_partial: bool,
1325 max_window_cnt: Option<usize>,
1326 ) -> Result<Option<PlanInfo>, Error> {
1327 let query_ctx = self.state.read().unwrap().query_ctx.clone();
1328 let start = SystemTime::now();
1329 let since_the_epoch = start
1330 .duration_since(UNIX_EPOCH)
1331 .expect("Time went backwards");
1332 let low_bound = self
1333 .config
1334 .expire_after
1335 .map(|e| since_the_epoch.as_secs() - e as u64)
1336 .unwrap_or(u64::MIN);
1337
1338 let low_bound = Timestamp::new_second(low_bound as i64);
1339
1340 let expire_time_window_bound = self
1341 .config
1342 .time_window_expr
1343 .as_ref()
1344 .map(|expr| expr.eval(low_bound))
1345 .transpose()?;
1346
1347 let (expire_lower_bound, expire_upper_bound) = match (
1348 expire_time_window_bound,
1349 &self.config.query_type,
1350 ) {
1351 (Some((Some(l), Some(u))), QueryType::Sql) => (l, u),
1352 (None, QueryType::Sql) if self.config.flow_eval_interval.is_none() => {
1353 return UnexpectedSnafu {
1354 reason: format!(
1355 "Flow id={} reached execution without a time-window expression or EVAL INTERVAL; create-flow validation should have rejected it",
1356 self.config.flow_id
1357 ),
1358 }
1359 .fail();
1360 }
1361 _ => {
1362 let (_, dirty_windows_to_restore) = self.drain_dirty_windows_signal();
1368
1369 let plan_info = self
1370 .gen_unfiltered_plan_info(
1371 engine,
1372 query_ctx,
1373 sink_table_schema.clone(),
1374 primary_key_indices,
1375 allow_partial,
1376 dirty_windows_to_restore,
1377 None,
1378 QueryCoverage::UnfilteredFull,
1379 )
1380 .await?;
1381
1382 return Ok(Some(plan_info));
1383 }
1384 };
1385
1386 debug!(
1387 "Flow id = {:?}, found time window: precise_lower_bound={:?}, precise_upper_bound={:?} with dirty time windows: {:?}",
1388 self.config.flow_id,
1389 expire_lower_bound,
1390 expire_upper_bound,
1391 self.state.read().unwrap().dirty_time_windows
1392 );
1393 let window_size = expire_upper_bound
1394 .sub(&expire_lower_bound)
1395 .with_context(|| UnexpectedSnafu {
1396 reason: format!(
1397 "Can't get window size from {expire_upper_bound:?} - {expire_lower_bound:?}"
1398 ),
1399 })?;
1400 let col_name = self
1401 .config
1402 .time_window_expr
1403 .as_ref()
1404 .map(|expr| expr.column_name.clone())
1405 .with_context(|| UnexpectedSnafu {
1406 reason: format!(
1407 "Flow id={:?}, Failed to get column name from time window expr",
1408 self.config.flow_id
1409 ),
1410 })?;
1411
1412 if self.should_use_unfiltered_incremental_delta() {
1413 let retention_filter = self
1423 .config
1424 .expire_after
1425 .map(|_| (col_name.as_str(), expire_lower_bound, "incremental"));
1426 return self
1427 .gen_unfiltered_plan_info_if_dirty(
1428 engine,
1429 query_ctx,
1430 sink_table_schema.clone(),
1431 primary_key_indices,
1432 allow_partial,
1433 retention_filter,
1434 QueryCoverage::IncrementalDelta,
1435 )
1436 .await;
1437 }
1438
1439 let (expr, coverage) = {
1440 let mut state = self.state.write().unwrap();
1441 let window_cnt = max_window_cnt
1442 .unwrap_or(self.config.batch_opts.experimental_max_filter_num_per_query);
1443 let expr = state.gen_scoped_filter_exprs(
1444 &col_name,
1445 Some(expire_lower_bound),
1446 window_size,
1447 window_cnt,
1448 self.config.flow_id,
1449 Some(self),
1450 )?;
1451 let repair_high = state
1452 .pending_fenced_repair()
1453 .map(|repair| repair.high().clone());
1454 let coverage = if let Some(high) = repair_high {
1455 QueryCoverage::FencedRepairChunk { high }
1456 } else {
1457 QueryCoverage::ScopedBaseRepair
1458 };
1459 (expr, coverage)
1460 };
1461
1462 let Some(expr) = expr else {
1463 debug!("Flow id={:?}, no new data, not update", self.config.flow_id);
1465 return Ok(None);
1466 };
1467
1468 let filter_sql = expr_to_sql(&expr.expr)
1469 .map(|sql| sql.to_string())
1470 .unwrap_or_else(|err| format!("<failed to format filter expr: {err}>"));
1471
1472 debug!(
1473 "Flow id={:?}, Generated filter expr: {:?}",
1474 self.config.flow_id, filter_sql
1475 );
1476
1477 let mut add_filter = AddFilterRewriter::new(expr.expr.clone());
1478 let mut add_auto_column = ColumnMatcherRewriter::new(
1479 sink_table_schema.clone(),
1480 primary_key_indices.to_vec(),
1481 allow_partial,
1482 );
1483
1484 let plan = self.restore_scoped_dirty_windows_on_err(
1485 &expr,
1486 sql_to_df_plan(query_ctx.clone(), engine.clone(), &self.config.query, false).await,
1487 )?;
1488 let rewrite = self.restore_scoped_dirty_windows_on_err(
1489 &expr,
1490 plan.clone()
1491 .rewrite(&mut add_filter)
1492 .and_then(|p| p.data.rewrite(&mut add_auto_column))
1493 .with_context(|_| DatafusionSnafu {
1494 context: format!("Failed to rewrite plan:\n {}\n", plan),
1495 })
1496 .map(|rewrite| rewrite.data),
1497 )?;
1498 let new_plan = self.restore_scoped_dirty_windows_on_err(
1500 &expr,
1501 apply_df_optimizer(rewrite, &query_ctx).await,
1502 )?;
1503
1504 let info = PlanInfo {
1505 plan: new_plan.clone(),
1506 dirty_restore: DirtyRestore::Scoped(expr),
1507 coverage,
1508 };
1509
1510 Ok(Some(info))
1511 }
1512}
1513
1514#[cfg(test)]
1515mod test;