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
82fn initial_schedule_cursor(start_secs: i64, interval_secs: i64) -> Result<i64, Error> {
90 let cursor = i128::from(start_secs) - i128::from(interval_secs);
91 i64::try_from(cursor).map_err(|_| {
92 UnexpectedSnafu {
93 reason: format!(
94 "Cannot compute the initial eval schedule cursor one interval before start {start_secs} (interval={interval_secs}): {cursor} does not fit in i64"
95 ),
96 }
97 .build()
98 })
99}
100
101fn sleep_delta_secs(next: i64, wall_now_secs: i64) -> Result<u64, Error> {
109 let delta = i128::from(next) - i128::from(wall_now_secs);
110 u64::try_from(delta).map_err(|_| {
111 UnexpectedSnafu {
112 reason: format!(
113 "Cannot sleep until the next scheduled time {next}: the delta from wall clock {wall_now_secs} is {delta} seconds, which does not fit in u64"
114 ),
115 }
116 .build()
117 })
118}
119
120fn scheduled_time_millis(scheduled_time_secs: i64) -> Result<i64, Error> {
127 scheduled_time_secs.checked_mul(1000).ok_or_else(|| {
128 UnexpectedSnafu {
129 reason: format!(
130 "Cannot convert scheduled time {scheduled_time_secs}s to milliseconds: the product exceeds i64"
131 ),
132 }
133 .build()
134 })
135}
136
137#[derive(Clone)]
139pub struct TaskConfig {
140 pub flow_id: FlowId,
141 pub query: String,
142 pub output_schema: DFSchemaRef,
144 pub time_window_expr: Option<TimeWindowExpr>,
145 pub expire_after: Option<i64>,
147 pub sink_table_name: [String; 3],
148 pub source_table_names: HashSet<[String; 3]>,
149 pub catalog_manager: CatalogManagerRef,
150 pub query_type: QueryType,
151 pub batch_opts: Arc<BatchingModeOptions>,
152 pub flow_eval_interval: Option<Duration>,
153 pub eval_schedule: Option<EvalSchedule>,
155}
156
157fn determine_query_type(query: &str, query_ctx: &QueryContextRef) -> Result<QueryType, Error> {
158 let is_tql = is_tql(query_ctx.sql_dialect(), query)
159 .map_err(BoxedError::new)
160 .context(ExternalSnafu)?;
161 Ok(if is_tql {
162 QueryType::Tql
163 } else {
164 QueryType::Sql
165 })
166}
167
168fn is_merge_mode_last_non_null(options: &HashMap<String, String>) -> bool {
169 options
170 .get(MERGE_MODE_KEY)
171 .map(|mode| mode.eq_ignore_ascii_case("last_non_null"))
172 .unwrap_or(false)
173}
174
175fn encode_insert_plan_request(
176 insert_to: TableName,
177 insert_input_plan: &LogicalPlan,
178) -> Result<api::v1::QueryRequest, Error> {
179 let message = DFLogicalSubstraitConvertor {}
180 .encode(insert_input_plan, DefaultSerializer)
181 .context(SubstraitEncodeLogicalPlanSnafu)?;
182 Ok(api::v1::QueryRequest {
183 query: Some(api::v1::query_request::Query::InsertIntoPlan(
184 api::v1::InsertIntoPlan {
185 table_name: Some(insert_to),
186 logical_plan: message.to_vec(),
187 },
188 )),
189 })
190}
191
192fn format_insert_target_columns(plan: &LogicalPlan) -> String {
193 plan.schema()
194 .fields()
195 .iter()
196 .map(|field| quote_identifier(field.name()).to_string())
197 .collect::<Vec<_>>()
198 .join(", ")
199}
200
201#[derive(Clone)]
202pub struct BatchingTask {
203 pub config: Arc<TaskConfig>,
204 pub state: Arc<RwLock<TaskState>>,
205 execution_lock: Arc<Mutex<()>>,
209}
210
211pub struct TaskArgs<'a> {
213 pub flow_id: FlowId,
214 pub query: &'a str,
215 pub plan: LogicalPlan,
216 pub time_window_expr: Option<TimeWindowExpr>,
217 pub expire_after: Option<i64>,
218 pub sink_table_name: [String; 3],
219 pub source_table_names: Vec<[String; 3]>,
220 pub query_ctx: QueryContextRef,
221 pub catalog_manager: CatalogManagerRef,
222 pub shutdown_rx: oneshot::Receiver<()>,
223 pub batch_opts: Arc<BatchingModeOptions>,
224 pub flow_eval_interval: Option<Duration>,
225 pub eval_schedule: Option<EvalSchedule>,
227}
228
229pub struct PlanInfo {
230 pub plan: LogicalPlan,
231 pub dirty_restore: DirtyRestore,
232 pub coverage: QueryCoverage,
233}
234
235#[derive(Clone)]
236pub enum QueryCoverage {
237 UnfilteredFull,
242 ScopedBaseRepair,
245 FencedRepairChunk { high: BTreeMap<u64, u64> },
249 IncrementalDelta,
251}
252
253impl QueryCoverage {
254 fn is_incremental_delta(&self) -> bool {
257 matches!(self, Self::IncrementalDelta)
258 }
259
260 fn snapshot_seqs(&self) -> HashMap<u64, u64> {
263 match self {
264 Self::FencedRepairChunk { high } => high.iter().map(|(k, v)| (*k, *v)).collect(),
265 _ => HashMap::new(),
266 }
267 }
268}
269
270pub enum DirtyRestore {
271 Scoped(FilterExprInfo),
274 Unscoped(DirtyTimeWindows),
281}
282
283struct ExecuteOnceOutcome {
284 new_query: Option<PlanInfo>,
285 result: Result<Option<(usize, Duration)>, Error>,
291}
292
293impl BatchingTask {
294 #[allow(clippy::too_many_arguments)]
295 pub fn try_new(
296 TaskArgs {
297 flow_id,
298 query,
299 plan,
300 time_window_expr,
301 expire_after,
302 sink_table_name,
303 source_table_names,
304 query_ctx,
305 catalog_manager,
306 shutdown_rx,
307 batch_opts,
308 flow_eval_interval,
309 eval_schedule,
310 }: TaskArgs<'_>,
311 ) -> Result<Self, Error> {
312 let mut state = TaskState::with_dirty_time_windows(
313 query_ctx.clone(),
314 shutdown_rx,
315 DirtyTimeWindows::new(
316 batch_opts.experimental_max_filter_num_per_query,
317 batch_opts.experimental_time_window_merge_threshold,
318 ),
319 );
320 if !batch_opts.experimental_enable_incremental_read {
321 state.disable_incremental();
322 }
323
324 Ok(Self {
325 config: Arc::new(TaskConfig {
326 flow_id,
327 query: query.to_string(),
328 time_window_expr,
329 expire_after,
330 sink_table_name,
331 source_table_names: source_table_names.into_iter().collect(),
332 catalog_manager,
333 output_schema: plan.schema().clone(),
334 query_type: determine_query_type(query, &query_ctx)?,
335 batch_opts,
336 flow_eval_interval,
337 eval_schedule,
338 }),
339 state: Arc::new(RwLock::new(state)),
340 execution_lock: Arc::new(Mutex::new(())),
341 })
342 }
343
344 pub fn last_execution_time_millis(&self) -> Option<i64> {
345 self.state.read().unwrap().last_execution_time_millis()
346 }
347
348 pub fn start_time_millis(&self) -> Option<i64> {
349 self.state.read().unwrap().start_time_millis()
350 }
351
352 fn frontend_extensions(&self) -> HashMap<String, String> {
355 let ctx = self.state.read().unwrap();
356 let all = ctx.query_ctx.extensions();
357 let mut flow_exts = HashMap::new();
358 if let Some(v) = all.get(query::options::FLOW_SCHEDULED_TIME_MILLIS) {
361 flow_exts.insert(
362 query::options::FLOW_SCHEDULED_TIME_MILLIS.to_string(),
363 v.clone(),
364 );
365 }
366 flow_exts
367 }
368
369 pub fn mark_all_windows_as_dirty(&self) -> Result<(), Error> {
373 let now = SystemTime::now();
374 let now = Timestamp::new_second(
375 now.duration_since(UNIX_EPOCH)
376 .expect("Time went backwards")
377 .as_secs() as _,
378 );
379 let lower_bound = self
380 .config
381 .expire_after
382 .map(|e| now.sub_duration(Duration::from_secs(e as _)))
383 .transpose()
384 .map_err(BoxedError::new)
385 .context(ExternalSnafu)?
386 .unwrap_or(Timestamp::new_second(0));
387 debug!(
388 "Flow {} mark range ({:?}, {:?}) as dirty",
389 self.config.flow_id, lower_bound, now
390 );
391 self.state
392 .write()
393 .unwrap()
394 .dirty_time_windows
395 .add_window(lower_bound, Some(now));
396 Ok(())
397 }
398
399 pub async fn check_or_create_sink_table(
401 &self,
402 engine: &QueryEngineRef,
403 frontend_client: &Arc<FrontendClient>,
404 ) -> Result<Option<(usize, Duration)>, Error> {
405 if !self.is_table_exist(&self.config.sink_table_name).await? {
406 let create_table = self.gen_create_table_expr(engine.clone()).await?;
407 info!(
408 "Try creating sink table(if not exists) with expr: {:?}",
409 create_table
410 );
411 self.create_table(frontend_client, create_table).await?;
412 info!(
413 "Sink table {}(if not exists) created",
414 self.config.sink_table_name.join(".")
415 );
416 }
417
418 Ok(None)
419 }
420
421 pub async fn validate_sink_table_schema(&self, engine: &QueryEngineRef) -> Result<(), Error> {
427 let (table, _) = get_table_info_df_schema(
428 self.config.catalog_manager.clone(),
429 self.config.sink_table_name.clone(),
430 )
431 .await?;
432
433 let table_meta = &table.table_info().meta;
434 let merge_mode_last_non_null =
435 is_merge_mode_last_non_null(&table_meta.options.extra_options);
436 let primary_key_indices = table_meta.primary_key_indices.clone();
437 let query_ctx = self.state.read().unwrap().query_ctx.clone();
438
439 gen_plan_with_matching_schema(
440 &self.config.query,
441 query_ctx,
442 engine.clone(),
443 table_meta.schema.clone(),
444 &primary_key_indices,
445 merge_mode_last_non_null,
446 )
447 .await
448 .map(|_| ())
449 }
450
451 async fn is_table_exist(&self, table_name: &[String; 3]) -> Result<bool, Error> {
452 self.config
453 .catalog_manager
454 .table_exists(&table_name[0], &table_name[1], &table_name[2], None)
455 .await
456 .map_err(BoxedError::new)
457 .context(ExternalSnafu)
458 }
459
460 pub(crate) async fn execute_once_serialized(
461 &self,
462 engine: &QueryEngineRef,
463 frontend_client: &Arc<FrontendClient>,
464 max_window_cnt: Option<usize>,
465 ) -> Result<Option<(usize, Duration)>, Error> {
466 let outcome = self
467 .execute_once_serialized_with_outcome(engine, frontend_client, max_window_cnt)
468 .await;
469 outcome.result
470 }
471
472 async fn execute_once_serialized_with_outcome(
475 &self,
476 engine: &QueryEngineRef,
477 frontend_client: &Arc<FrontendClient>,
478 max_window_cnt: Option<usize>,
479 ) -> ExecuteOnceOutcome {
480 let _execution_guard = self.execution_lock.lock().await;
481 self.execute_once_unlocked(engine, frontend_client, max_window_cnt)
482 .await
483 }
484
485 async fn execute_once_unlocked(
487 &self,
488 engine: &QueryEngineRef,
489 frontend_client: &Arc<FrontendClient>,
490 max_window_cnt: Option<usize>,
491 ) -> ExecuteOnceOutcome {
492 let new_query = match self.gen_insert_plan_unlocked(engine, max_window_cnt).await {
493 Ok(new_query) => new_query,
494 Err(err) => {
495 return ExecuteOnceOutcome {
496 new_query: None,
497 result: Err(err),
498 };
499 }
500 };
501
502 if let Some(new_query) = new_query {
503 debug!("Generate new query: {}", new_query.plan);
504 let res = self
505 .execute_logical_plan_unlocked(
506 frontend_client,
507 &new_query.plan,
508 &new_query.dirty_restore,
509 &new_query.coverage,
510 )
511 .await;
512 if res.is_err() {
513 self.handle_executed_query_failure(Some(&new_query));
514 }
515 ExecuteOnceOutcome {
516 new_query: Some(new_query),
517 result: res,
518 }
519 } else {
520 debug!("Generate no query");
521 ExecuteOnceOutcome {
522 new_query: None,
523 result: Ok(None),
524 }
525 }
526 }
527
528 async fn gen_insert_plan_unlocked(
530 &self,
531 engine: &QueryEngineRef,
532 max_window_cnt: Option<usize>,
533 ) -> Result<Option<PlanInfo>, Error> {
534 let (table, df_schema) = get_table_info_df_schema(
535 self.config.catalog_manager.clone(),
536 self.config.sink_table_name.clone(),
537 )
538 .await?;
539
540 let table_meta = &table.table_info().meta;
541 let merge_mode_last_non_null =
542 is_merge_mode_last_non_null(&table_meta.options.extra_options);
543 let primary_key_indices = table_meta.primary_key_indices.clone();
544
545 let new_query = self
546 .gen_query_with_time_window(
547 engine.clone(),
548 &table.table_info().meta.schema,
549 &primary_key_indices,
550 merge_mode_last_non_null,
551 max_window_cnt,
552 )
553 .await?;
554
555 let Some(new_query) = new_query else {
556 return Ok(None);
557 };
558
559 let table_columns = df_schema
562 .columns()
563 .into_iter()
564 .map(|c| c.name)
565 .collect::<BTreeSet<_>>();
566 for column in new_query.plan.schema().columns() {
567 if !table_columns.contains(column.name()) {
568 self.restore_dirty_windows_after_failure(&new_query);
569 return InvalidQuerySnafu {
570 reason: format!(
571 "Column {} not found in sink table with columns {:?}",
572 column, table_columns
573 ),
574 }
575 .fail();
576 }
577 }
578
579 let table_provider = Arc::new(DfTableProviderAdapter::new(table));
580 let table_source = Arc::new(DefaultTableSource::new(table_provider));
581
582 let plan = LogicalPlan::Dml(DmlStatement::new(
584 datafusion_common::TableReference::Full {
585 catalog: self.config.sink_table_name[0].clone().into(),
586 schema: self.config.sink_table_name[1].clone().into(),
587 table: self.config.sink_table_name[2].clone().into(),
588 },
589 table_source,
590 WriteOp::Insert(datafusion_expr::dml::InsertOp::Append),
591 Arc::new(new_query.plan.clone()),
592 ));
593 let insert_into_info = PlanInfo {
594 plan,
595 dirty_restore: new_query.dirty_restore,
596 coverage: new_query.coverage,
597 };
598 let insert_into =
599 match insert_into_info
600 .plan
601 .clone()
602 .recompute_schema()
603 .context(DatafusionSnafu {
604 context: "Failed to recompute schema",
605 }) {
606 Ok(insert_into) => insert_into,
607 Err(err) => {
608 self.restore_dirty_windows_after_failure(&insert_into_info);
609 return Err(err);
610 }
611 };
612
613 Ok(Some(PlanInfo {
614 plan: insert_into,
615 dirty_restore: insert_into_info.dirty_restore,
616 coverage: insert_into_info.coverage,
617 }))
618 }
619
620 pub async fn create_table(
621 &self,
622 frontend_client: &Arc<FrontendClient>,
623 expr: CreateTableExpr,
624 ) -> Result<(), Error> {
625 let catalog = &self.config.sink_table_name[0];
626 let schema = &self.config.sink_table_name[1];
627 frontend_client
628 .create(expr.clone(), catalog, schema)
629 .await?;
630 Ok(())
631 }
632
633 async fn execute_logical_plan_unlocked(
635 &self,
636 frontend_client: &Arc<FrontendClient>,
637 plan: &LogicalPlan,
638 dirty_restore: &DirtyRestore,
639 coverage: &QueryCoverage,
640 ) -> Result<Option<(usize, Duration)>, Error> {
641 let instant = Instant::now();
642 let flow_id = self.config.flow_id;
643
644 debug!(
645 "Executing flow {flow_id}(expire_after={:?} secs) with query {}",
646 self.config.expire_after, &plan
647 );
648
649 let catalog = &self.config.sink_table_name[0];
650 let schema = &self.config.sink_table_name[1];
651
652 let plan = plan
654 .clone()
655 .transform_down_with_subqueries(|p| {
656 if let LogicalPlan::TableScan(mut table_scan) = p {
657 let resolved = table_scan.table_name.resolve(catalog, schema);
658 table_scan.table_name = resolved.into();
659 Ok(Transformed::yes(LogicalPlan::TableScan(table_scan)))
660 } else {
661 Ok(Transformed::no(p))
662 }
663 })
664 .with_context(|_| DatafusionSnafu {
665 context: format!("Failed to fix table ref in logical plan, plan={:?}", plan),
666 })?
667 .data;
668
669 let incremental_plan = if coverage.is_incremental_delta() {
672 self.prepare_plan_for_incremental(&plan).await?
673 } else {
674 None
675 };
676 let incremental_safe = incremental_plan.is_some();
677 if coverage.is_incremental_delta() && !incremental_safe {
678 debug!(
679 "Flow {flow_id} skipped unsafe incremental delta fallback; \
680 restored dirty signal instead of executing an unfiltered full snapshot"
681 );
682 self.restore_dirty_windows(dirty_restore);
683 return Ok(None);
684 }
685 let plan = incremental_plan.unwrap_or_else(|| plan.clone());
686
687 let extensions = self
688 .build_flow_query_extensions(incremental_safe, coverage.is_incremental_delta())
689 .await?;
690 let frontend_extensions = self.frontend_extensions();
691 let extension_refs = extensions
692 .iter()
693 .map(|(key, value)| (*key, value.as_str()))
694 .chain(
695 frontend_extensions
696 .iter()
697 .map(|(key, value)| (key.as_str(), value.as_str())),
698 )
699 .collect::<Vec<_>>();
700 let query_mode = if extensions
701 .iter()
702 .any(|(key, _)| *key == FLOW_INCREMENTAL_MODE)
703 {
704 CheckpointMode::Incremental
705 } else {
706 CheckpointMode::FullSnapshot
707 };
708 Self::record_query_mode(flow_id, query_mode);
709 debug!(
710 "Flow {flow_id} executing batching query with checkpoint_mode={}, extension_count={}",
711 checkpoint_mode_label(query_mode),
712 extensions.len()
713 );
714
715 let mut peer_desc = None;
716 let res = {
717 let _timer = METRIC_FLOW_BATCHING_ENGINE_QUERY_TIME
718 .with_label_values(&[flow_id.to_string().as_str()])
719 .start_timer();
720
721 let req = if let Some((insert_to, insert_input_plan)) =
722 breakup_insert_plan(&plan, catalog, schema)
723 {
724 if query_mode == CheckpointMode::FullSnapshot
725 && matches!(self.config.query_type, QueryType::Sql)
726 && self.config.flow_eval_interval.is_some()
727 && self.config.time_window_expr.is_none()
728 {
729 match df_plan_to_sql(&insert_input_plan) {
739 Ok(select_sql) => {
740 let target_columns = format_insert_target_columns(&insert_input_plan);
741 let sql = format!(
742 "INSERT INTO {} ({}) {}",
743 TableReference::full(
744 insert_to.catalog_name.as_str(),
745 insert_to.schema_name.as_str(),
746 insert_to.table_name.as_str(),
747 )
748 .to_quoted_string(),
749 target_columns,
750 select_sql
751 );
752 api::v1::QueryRequest {
753 query: Some(api::v1::query_request::Query::Sql(sql)),
754 }
755 }
756 Err(err) => {
757 debug!(
758 "Failed to unparse full-snapshot SQL flow {} plan; \
759 falling back to InsertIntoPlan: {:?}",
760 flow_id, err
761 );
762 encode_insert_plan_request(insert_to, &insert_input_plan)?
763 }
764 }
765 } else {
766 encode_insert_plan_request(insert_to, &insert_input_plan)?
767 }
768 } else {
769 let message = DFLogicalSubstraitConvertor {}
770 .encode(&plan, DefaultSerializer)
771 .context(SubstraitEncodeLogicalPlanSnafu)?;
772
773 api::v1::QueryRequest {
774 query: Some(api::v1::query_request::Query::LogicalPlan(message.to_vec())),
775 }
776 };
777
778 let snapshot_seqs = coverage.snapshot_seqs();
779 {
780 let mut state = self.state.write().unwrap();
781 state.record_start_time_if_first();
782 }
783 frontend_client
784 .query_with_terminal_metrics(
785 catalog,
786 schema,
787 req,
788 &extension_refs,
789 &snapshot_seqs,
790 &mut peer_desc,
791 )
792 .await
793 };
794
795 let elapsed = instant.elapsed();
796 let peer_label = peer_desc
797 .as_ref()
798 .map(ToString::to_string)
799 .unwrap_or_else(|| PeerDesc::default().to_string());
800 if let Err(err) = &res {
801 warn!(
802 "Failed to execute Flow {flow_id} on frontend {peer_label}, result: {err:?}, elapsed: {:?} with query: {}",
803 elapsed, &plan
804 );
805 let decision = {
806 let mut state = self.state.write().unwrap();
807 let reason = Self::query_failure_reason(err, coverage);
808 Self::apply_query_failure_to_state(&mut state, elapsed, coverage, reason)
809 };
810 if let Some(decision) = decision {
811 Self::record_checkpoint_decision(flow_id, decision);
812 }
813 }
814
815 if elapsed >= self.config.batch_opts.slow_query_threshold {
817 warn!(
818 "Flow {flow_id} on frontend {peer_label} executed for {:?} before complete, query: {}",
819 elapsed, &plan
820 );
821 let flow_id = flow_id.to_string();
822 METRIC_FLOW_BATCHING_ENGINE_SLOW_QUERY
823 .with_label_values(&[flow_id.as_str(), peer_label.as_str()])
824 .observe(elapsed.as_secs_f64());
825 }
826
827 let res = res?;
828 let (affected_rows, _) = res.output.extract_rows_and_cost();
829 debug!(
830 "Flow {flow_id} executed, affected_rows: {affected_rows:?}, elapsed: {:?}, watermark: {:?}",
831 elapsed,
832 res.region_watermark_map()
833 );
834 METRIC_FLOW_ROWS
835 .with_label_values(&[format!("{}-out-batching", flow_id).as_str()])
836 .inc_by(affected_rows as _);
837 let decision = {
838 let mut state = self.state.write().unwrap();
839 Self::apply_query_result_to_state(&mut state, &res, elapsed, coverage)
840 };
841 Self::record_checkpoint_decision(flow_id, decision);
842
843 Ok(Some((affected_rows, elapsed)))
844 }
845
846 fn restore_dirty_windows(&self, dirty_restore: &DirtyRestore) {
850 match dirty_restore {
851 DirtyRestore::Scoped(filter) => self.restore_scoped_dirty_windows(filter),
852 DirtyRestore::Unscoped(dirty_windows) => self
853 .state
854 .write()
855 .unwrap()
856 .dirty_time_windows
857 .add_dirty_windows(dirty_windows),
858 }
859 }
860
861 fn restore_dirty_windows_after_failure(&self, query: &PlanInfo) {
864 self.restore_dirty_windows(&query.dirty_restore);
865 }
866
867 fn restore_scoped_dirty_windows(&self, filter: &FilterExprInfo) {
870 self.state.write().unwrap().restore_scoped_windows(filter);
871 }
872
873 fn restore_scoped_dirty_windows_on_err<T>(
876 &self,
877 filter: &FilterExprInfo,
878 result: Result<T, Error>,
879 ) -> Result<T, Error> {
880 result.inspect_err(|_| {
881 self.restore_scoped_dirty_windows(filter);
882 })
883 }
884
885 fn restore_unscoped_dirty_windows(&self, dirty_windows: &DirtyTimeWindows) {
888 self.state
889 .write()
890 .unwrap()
891 .dirty_time_windows
892 .add_dirty_windows(dirty_windows);
893 }
894
895 fn restore_unscoped_dirty_windows_on_err<T>(
898 &self,
899 dirty_windows: &DirtyTimeWindows,
900 result: Result<T, Error>,
901 ) -> Result<T, Error> {
902 result.inspect_err(|_| {
903 self.restore_unscoped_dirty_windows(dirty_windows);
904 })
905 }
906
907 fn drain_dirty_windows_signal(&self) -> (bool, DirtyTimeWindows) {
910 let mut state = self.state.write().unwrap();
911 let dirty_windows_to_restore = state.dirty_time_windows.clone();
912 let is_dirty = !dirty_windows_to_restore.is_empty();
913 state.dirty_time_windows.clean();
914 (is_dirty, dirty_windows_to_restore)
915 }
916
917 #[allow(clippy::too_many_arguments)]
918 async fn gen_unfiltered_plan_info(
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 dirty_windows_to_restore: DirtyTimeWindows,
928 retention_filter: Option<(&str, Timestamp, &'static str)>,
929 coverage: QueryCoverage,
930 ) -> Result<PlanInfo, Error> {
931 let mut plan = self.restore_unscoped_dirty_windows_on_err(
932 &dirty_windows_to_restore,
933 gen_plan_with_matching_schema(
934 &self.config.query,
935 query_ctx,
936 engine,
937 sink_table_schema,
938 primary_key_indices,
939 allow_partial,
940 )
941 .await,
942 )?;
943
944 if let Some((col_name, lower_bound, context)) = retention_filter {
945 let lower = self.restore_unscoped_dirty_windows_on_err(
946 &dirty_windows_to_restore,
947 to_df_literal(lower_bound),
948 )?;
949 let retention_filter = col(col_name).gt_eq(lit(lower));
950 let mut add_filter = AddFilterRewriter::new(retention_filter);
951 plan = self.restore_unscoped_dirty_windows_on_err(
952 &dirty_windows_to_restore,
953 plan.clone()
954 .rewrite(&mut add_filter)
955 .with_context(|_| DatafusionSnafu {
956 context: format!(
957 "Failed to apply {context} expire_after filter to plan:\n {}\n",
958 plan
959 ),
960 })
961 .map(|rewrite| rewrite.data),
962 )?;
963 }
964
965 Ok(PlanInfo {
966 plan,
967 dirty_restore: DirtyRestore::Unscoped(dirty_windows_to_restore),
968 coverage,
969 })
970 }
971
972 #[allow(clippy::too_many_arguments)]
973 async fn gen_unfiltered_plan_info_if_dirty(
976 &self,
977 engine: QueryEngineRef,
978 query_ctx: QueryContextRef,
979 sink_table_schema: Arc<Schema>,
980 primary_key_indices: &[usize],
981 allow_partial: bool,
982 retention_filter: Option<(&str, Timestamp, &'static str)>,
983 coverage: QueryCoverage,
984 ) -> Result<Option<PlanInfo>, Error> {
985 let (is_dirty, dirty_windows_to_restore) = self.drain_dirty_windows_signal();
986 if !is_dirty {
987 debug!("Flow id={:?}, no new data, not update", self.config.flow_id);
988 return Ok(None);
989 }
990
991 self.gen_unfiltered_plan_info(
992 engine,
993 query_ctx,
994 sink_table_schema,
995 primary_key_indices,
996 allow_partial,
997 dirty_windows_to_restore,
998 retention_filter,
999 coverage,
1000 )
1001 .await
1002 .map(Some)
1003 }
1004
1005 fn handle_executed_query_failure(&self, query: Option<&PlanInfo>) {
1006 if let Some(query) = query {
1007 self.restore_dirty_windows_after_failure(query);
1008 }
1009 }
1010
1011 pub async fn start_executing_loop(
1019 &self,
1020 engine: QueryEngineRef,
1021 frontend_client: Arc<FrontendClient>,
1022 ) {
1023 if self.config.flow_eval_interval.is_some() {
1024 self.start_scheduled_loop(engine, frontend_client).await;
1025 } else {
1026 self.start_adaptive_loop(engine, frontend_client).await;
1027 }
1028 }
1029
1030 async fn start_scheduled_loop(
1040 &self,
1041 engine: QueryEngineRef,
1042 frontend_client: Arc<FrontendClient>,
1043 ) {
1044 let flow_id_str = self.config.flow_id.to_string();
1045
1046 let schedule = match &self.config.eval_schedule {
1047 Some(s) => s.clone(),
1048 None => {
1049 let eval_interval_secs = self
1050 .config
1051 .flow_eval_interval
1052 .map(|d| d.as_secs() as i64)
1053 .expect("checked by caller");
1054
1055 match EvalSchedule::from_config(Some(eval_interval_secs), None) {
1058 Ok(Some(s)) => s,
1059 Ok(None) => {
1060 warn!(
1061 "Flow {}: EVAL INTERVAL set but no schedule parsed; exiting loop",
1062 flow_id_str
1063 );
1064 return;
1065 }
1066 Err(e) => {
1067 warn!(
1068 "Flow {}: Failed to parse eval schedule: {}; exiting loop",
1069 flow_id_str, e
1070 );
1071 return;
1072 }
1073 }
1074 }
1075 };
1076
1077 let mut cursor_secs =
1082 match initial_schedule_cursor(schedule.start_secs, schedule.interval_secs) {
1083 Ok(cursor) => cursor,
1084 Err(e) => {
1085 warn!(
1086 "Flow {}: invalid eval schedule, exiting loop: {}",
1087 flow_id_str, e
1088 );
1089 return;
1090 }
1091 };
1092
1093 info!(
1094 "Flow {}: entering scheduled loop, interval={}s, start={}, anchor={}, policy={:?}, max_runs={}, max_lag={}s",
1095 flow_id_str,
1096 schedule.interval_secs,
1097 schedule.start_secs,
1098 schedule.anchor_secs,
1099 schedule.missed_tick_policy,
1100 schedule.max_runs,
1101 schedule.max_lag_secs,
1102 );
1103
1104 loop {
1105 if self.is_shutdown_signaled() {
1106 break;
1107 }
1108
1109 let wall_now_secs = wall_clock_unix_secs();
1110
1111 let due = match select_due_scheduled_times(&schedule, cursor_secs, wall_now_secs) {
1112 Ok(d) => d,
1113 Err(e) => {
1114 warn!(
1115 "Flow {}: invalid eval schedule, exiting loop: {}",
1116 flow_id_str, e
1117 );
1118 return;
1119 }
1120 };
1121
1122 if due.scheduled_times_secs.is_empty() {
1123 if due.skipped > 0 {
1124 warn!(
1125 "Flow {}: all {} due scheduled times skipped by max-lag, advancing cursor to wall-clock ({wall_now_secs}) to avoid re-skipping",
1126 flow_id_str, due.skipped
1127 );
1128 cursor_secs = wall_now_secs;
1129 continue;
1130 }
1131
1132 let next = match schedule.next_scheduled_time_after(cursor_secs) {
1134 Ok(next) => next,
1135 Err(e) => {
1136 warn!(
1137 "Flow {}: cannot advance eval schedule past cursor {cursor_secs}: {e}; exiting loop",
1138 flow_id_str
1139 );
1140 return;
1141 }
1142 };
1143 if next <= wall_now_secs {
1144 cursor_secs = wall_now_secs;
1147 continue;
1148 }
1149 let wait_secs = match sleep_delta_secs(next, wall_now_secs) {
1150 Ok(wait_secs) => wait_secs,
1151 Err(e) => {
1152 warn!(
1153 "Flow {}: cannot sleep until next scheduled time {}: {e}; exiting loop",
1154 flow_id_str, next
1155 );
1156 return;
1157 }
1158 };
1159 let wait_dur = Duration::from_secs(wait_secs);
1160 debug!(
1161 "Flow {}: no due scheduled times, sleeping for {}s until next scheduled time at {}",
1162 flow_id_str, wait_secs, next
1163 );
1164 tokio::time::sleep(wait_dur).await;
1165 continue;
1166 }
1167
1168 if due.skipped > 0 {
1169 info!(
1170 "Flow {}: {} due scheduled times, {} skipped (catch-up)",
1171 flow_id_str,
1172 due.scheduled_times_secs.len(),
1173 due.skipped
1174 );
1175 }
1176
1177 for scheduled_time_secs in &due.scheduled_times_secs {
1179 if self.is_shutdown_signaled() {
1180 break;
1181 }
1182
1183 METRIC_FLOW_BATCHING_ENGINE_START_QUERY_CNT
1184 .with_label_values(&[&flow_id_str])
1185 .inc();
1186
1187 let outcome = self
1188 .execute_once_serialized_at_scheduled_time(
1189 &engine,
1190 &frontend_client,
1191 *scheduled_time_secs,
1192 )
1193 .await;
1194
1195 cursor_secs = *scheduled_time_secs;
1197
1198 match outcome.result {
1199 Ok(Some((rows, elapsed))) => {
1200 debug!(
1201 "Flow {}: scheduled time {} completed, rows={}, elapsed={:?}",
1202 flow_id_str, scheduled_time_secs, rows, elapsed
1203 );
1204 }
1205 Ok(None) => {
1206 debug!(
1207 "Flow {}: scheduled time {} produced no query (no dirty signal or no-op)",
1208 flow_id_str, scheduled_time_secs
1209 );
1210 }
1211 Err(err) => {
1212 warn!(
1213 "Flow {}: scheduled time {} failed: {:?}",
1214 flow_id_str, scheduled_time_secs, err
1215 );
1216 METRIC_FLOW_BATCHING_ENGINE_ERROR_CNT
1217 .with_label_values(&[&flow_id_str])
1218 .inc();
1219 }
1223 }
1224 }
1225 }
1226 }
1227
1228 async fn start_adaptive_loop(
1230 &self,
1231 engine: QueryEngineRef,
1232 frontend_client: Arc<FrontendClient>,
1233 ) {
1234 let flow_id_str = self.config.flow_id.to_string();
1235 let mut max_window_cnt = None;
1236 loop {
1237 if self.is_shutdown_signaled() {
1238 break;
1239 }
1240 METRIC_FLOW_BATCHING_ENGINE_START_QUERY_CNT
1241 .with_label_values(&[&flow_id_str])
1242 .inc();
1243
1244 let min_refresh = self.config.batch_opts.experimental_min_refresh_duration;
1245
1246 let outcome = self
1247 .execute_once_serialized_with_outcome(&engine, &frontend_client, max_window_cnt)
1248 .await;
1249
1250 match outcome.result {
1251 Ok(Some(_)) => {
1252 max_window_cnt = max_window_cnt.map(|cnt| {
1253 (cnt + 1).min(self.config.batch_opts.experimental_max_filter_num_per_query)
1254 });
1255
1256 let sleep_until = {
1257 let state = self.state.write().unwrap();
1258
1259 let time_window_size = self
1260 .config
1261 .time_window_expr
1262 .as_ref()
1263 .and_then(|t| *t.time_window_size());
1264
1265 let prefer_short_incremental_cadence = state.checkpoint_mode()
1266 == CheckpointMode::Incremental
1267 && !state.is_incremental_disabled();
1268
1269 state.get_next_start_query_time(
1270 self.config.flow_id,
1271 &time_window_size,
1272 min_refresh,
1273 Some(self.config.batch_opts.query_timeout),
1274 self.config.batch_opts.experimental_max_filter_num_per_query,
1275 prefer_short_incremental_cadence,
1276 )
1277 };
1278
1279 tokio::time::sleep_until(sleep_until).await;
1280 }
1281 Ok(None) => {
1282 debug!(
1283 "Flow id = {:?} found no new data, sleep for {:?} then continue",
1284 self.config.flow_id, min_refresh
1285 );
1286 tokio::time::sleep(min_refresh).await;
1287 continue;
1288 }
1289 Err(err) => {
1290 METRIC_FLOW_BATCHING_ENGINE_ERROR_CNT
1291 .with_label_values(&[&flow_id_str])
1292 .inc();
1293 match outcome.new_query {
1294 Some(query) => {
1295 common_telemetry::error!(err; "Failed to execute query for flow={} with query: {}", self.config.flow_id, query.plan);
1296 max_window_cnt = Some(1);
1297 }
1298 None => {
1299 common_telemetry::error!(err; "Failed to generate query for flow={}", self.config.flow_id)
1300 }
1301 }
1302 tokio::time::sleep(min_refresh).await;
1303 }
1304 }
1305 }
1306 }
1307
1308 fn is_shutdown_signaled(&self) -> bool {
1310 let mut state = self.state.write().unwrap();
1311 match state.shutdown_rx.try_recv() {
1312 Ok(()) | Err(TryRecvError::Closed) => true,
1313 Err(TryRecvError::Empty) => false,
1314 }
1315 }
1316
1317 async fn execute_once_serialized_at_scheduled_time(
1324 &self,
1325 engine: &QueryEngineRef,
1326 frontend_client: &Arc<FrontendClient>,
1327 scheduled_time_secs: i64,
1328 ) -> ExecuteOnceOutcome {
1329 let _execution_guard = self.execution_lock.lock().await;
1330
1331 struct QueryContextRestoreGuard {
1332 state: Arc<RwLock<TaskState>>,
1333 old_ctx: Option<QueryContextRef>,
1334 }
1335
1336 impl Drop for QueryContextRestoreGuard {
1337 fn drop(&mut self) {
1338 let Some(old_ctx) = self.old_ctx.take() else {
1339 return;
1340 };
1341 if let Ok(mut state) = self.state.write() {
1342 state.query_ctx = old_ctx;
1343 }
1344 }
1345 }
1346
1347 let scheduled_time_millis = match scheduled_time_millis(scheduled_time_secs) {
1351 Ok(millis) => millis,
1352 Err(e) => {
1353 return ExecuteOnceOutcome {
1354 new_query: None,
1355 result: Err(e),
1356 };
1357 }
1358 };
1359
1360 let old_ctx = {
1363 let mut state = self.state.write().unwrap();
1364 let old = state.query_ctx.clone();
1365 let mut new_ctx = (*old).clone();
1366 new_ctx.set_extension(
1367 query::options::FLOW_SCHEDULED_TIME_MILLIS,
1368 scheduled_time_millis.to_string(),
1369 );
1370 state.query_ctx = Arc::new(new_ctx);
1371 old
1372 };
1373 let restore_guard = QueryContextRestoreGuard {
1374 state: self.state.clone(),
1375 old_ctx: Some(old_ctx),
1376 };
1377
1378 let outcome = self
1379 .execute_once_unlocked(engine, frontend_client, None)
1380 .await;
1381
1382 drop(restore_guard);
1386
1387 outcome
1388 }
1389
1390 async fn gen_create_table_expr(
1397 &self,
1398 engine: QueryEngineRef,
1399 ) -> Result<CreateTableExpr, Error> {
1400 let query_ctx = self.state.read().unwrap().query_ctx.clone();
1401 let plan =
1402 sql_to_df_plan(query_ctx.clone(), engine.clone(), &self.config.query, true).await?;
1403 create_table_with_expr(&plan, &self.config.sink_table_name, &self.config.query_type)
1404 }
1405
1406 fn should_use_unfiltered_incremental_delta(&self) -> bool {
1409 let state = self.state.read().unwrap();
1410 state.checkpoint_mode() == CheckpointMode::Incremental
1411 && !state.is_incremental_disabled()
1412 && matches!(self.config.query_type, QueryType::Sql)
1413 }
1414
1415 async fn gen_query_with_time_window(
1418 &self,
1419 engine: QueryEngineRef,
1420 sink_table_schema: &Arc<Schema>,
1421 primary_key_indices: &[usize],
1422 allow_partial: bool,
1423 max_window_cnt: Option<usize>,
1424 ) -> Result<Option<PlanInfo>, Error> {
1425 let query_ctx = self.state.read().unwrap().query_ctx.clone();
1426 let start = SystemTime::now();
1427 let since_the_epoch = start
1428 .duration_since(UNIX_EPOCH)
1429 .expect("Time went backwards");
1430 let low_bound = self
1431 .config
1432 .expire_after
1433 .map(|e| since_the_epoch.as_secs() - e as u64)
1434 .unwrap_or(u64::MIN);
1435
1436 let low_bound = Timestamp::new_second(low_bound as i64);
1437
1438 let expire_time_window_bound = self
1439 .config
1440 .time_window_expr
1441 .as_ref()
1442 .map(|expr| expr.eval(low_bound))
1443 .transpose()?;
1444
1445 let (expire_lower_bound, expire_upper_bound) = match (
1446 expire_time_window_bound,
1447 &self.config.query_type,
1448 ) {
1449 (Some((Some(l), Some(u))), QueryType::Sql) => (l, u),
1450 (None, QueryType::Sql) if self.config.flow_eval_interval.is_none() => {
1451 return UnexpectedSnafu {
1452 reason: format!(
1453 "Flow id={} reached execution without a time-window expression or EVAL INTERVAL; create-flow validation should have rejected it",
1454 self.config.flow_id
1455 ),
1456 }
1457 .fail();
1458 }
1459 _ => {
1460 let (_, dirty_windows_to_restore) = self.drain_dirty_windows_signal();
1466
1467 let plan_info = self
1468 .gen_unfiltered_plan_info(
1469 engine,
1470 query_ctx,
1471 sink_table_schema.clone(),
1472 primary_key_indices,
1473 allow_partial,
1474 dirty_windows_to_restore,
1475 None,
1476 QueryCoverage::UnfilteredFull,
1477 )
1478 .await?;
1479
1480 return Ok(Some(plan_info));
1481 }
1482 };
1483
1484 debug!(
1485 "Flow id = {:?}, found time window: precise_lower_bound={:?}, precise_upper_bound={:?} with dirty time windows: {:?}",
1486 self.config.flow_id,
1487 expire_lower_bound,
1488 expire_upper_bound,
1489 self.state.read().unwrap().dirty_time_windows
1490 );
1491 let window_size = expire_upper_bound
1492 .sub(&expire_lower_bound)
1493 .with_context(|| UnexpectedSnafu {
1494 reason: format!(
1495 "Can't get window size from {expire_upper_bound:?} - {expire_lower_bound:?}"
1496 ),
1497 })?;
1498 let col_name = self
1499 .config
1500 .time_window_expr
1501 .as_ref()
1502 .map(|expr| expr.column_name.clone())
1503 .with_context(|| UnexpectedSnafu {
1504 reason: format!(
1505 "Flow id={:?}, Failed to get column name from time window expr",
1506 self.config.flow_id
1507 ),
1508 })?;
1509
1510 if self.should_use_unfiltered_incremental_delta() {
1511 let retention_filter = self
1521 .config
1522 .expire_after
1523 .map(|_| (col_name.as_str(), expire_lower_bound, "incremental"));
1524 return self
1525 .gen_unfiltered_plan_info_if_dirty(
1526 engine,
1527 query_ctx,
1528 sink_table_schema.clone(),
1529 primary_key_indices,
1530 allow_partial,
1531 retention_filter,
1532 QueryCoverage::IncrementalDelta,
1533 )
1534 .await;
1535 }
1536
1537 let (expr, coverage) = {
1538 let mut state = self.state.write().unwrap();
1539 let window_cnt = max_window_cnt
1540 .unwrap_or(self.config.batch_opts.experimental_max_filter_num_per_query);
1541 let expr = state.gen_scoped_filter_exprs(
1542 &col_name,
1543 Some(expire_lower_bound),
1544 window_size,
1545 window_cnt,
1546 self.config.flow_id,
1547 Some(self),
1548 )?;
1549 let repair_high = state
1550 .pending_fenced_repair()
1551 .map(|repair| repair.high().clone());
1552 let coverage = if let Some(high) = repair_high {
1553 QueryCoverage::FencedRepairChunk { high }
1554 } else {
1555 QueryCoverage::ScopedBaseRepair
1556 };
1557 (expr, coverage)
1558 };
1559
1560 let Some(expr) = expr else {
1561 debug!("Flow id={:?}, no new data, not update", self.config.flow_id);
1563 return Ok(None);
1564 };
1565
1566 let filter_sql = expr_to_sql(&expr.expr)
1567 .map(|sql| sql.to_string())
1568 .unwrap_or_else(|err| format!("<failed to format filter expr: {err}>"));
1569
1570 debug!(
1571 "Flow id={:?}, Generated filter expr: {:?}",
1572 self.config.flow_id, filter_sql
1573 );
1574
1575 let mut add_filter = AddFilterRewriter::new(expr.expr.clone());
1576 let mut add_auto_column = ColumnMatcherRewriter::new(
1577 sink_table_schema.clone(),
1578 primary_key_indices.to_vec(),
1579 allow_partial,
1580 );
1581
1582 let plan = self.restore_scoped_dirty_windows_on_err(
1583 &expr,
1584 sql_to_df_plan(query_ctx.clone(), engine.clone(), &self.config.query, false).await,
1585 )?;
1586 let rewrite = self.restore_scoped_dirty_windows_on_err(
1587 &expr,
1588 plan.clone()
1589 .rewrite(&mut add_filter)
1590 .and_then(|p| p.data.rewrite(&mut add_auto_column))
1591 .with_context(|_| DatafusionSnafu {
1592 context: format!("Failed to rewrite plan:\n {}\n", plan),
1593 })
1594 .map(|rewrite| rewrite.data),
1595 )?;
1596 let new_plan = self.restore_scoped_dirty_windows_on_err(
1598 &expr,
1599 apply_df_optimizer(rewrite, &query_ctx).await,
1600 )?;
1601
1602 let info = PlanInfo {
1603 plan: new_plan.clone(),
1604 dirty_restore: DirtyRestore::Scoped(expr),
1605 coverage,
1606 };
1607
1608 Ok(Some(info))
1609 }
1610}
1611
1612#[cfg(test)]
1613mod test;