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 catalog::kvbackend::KvBackendCatalogManager;
22use client::OutputWithMetrics;
23use common_error::ext::BoxedError;
24use common_meta::key::schema_name::SchemaNameKey;
25use common_query::OutputData;
26use common_query::logical_plan::breakup_insert_plan;
27use common_telemetry::tracing::warn;
28use common_telemetry::{debug, info};
29use common_time::{TimeToLive, Timestamp};
30use datafusion::datasource::DefaultTableSource;
31use datafusion::sql::unparser::expr_to_sql;
32use datafusion_common::tree_node::{Transformed, TreeNode};
33use datafusion_common::utils::quote_identifier;
34use datafusion_common::{DFSchemaRef, ScalarValue, TableReference};
35use datafusion_expr::{DmlStatement, LogicalPlan, Projection, WriteOp, col, lit};
36use datatypes::schema::Schema;
37use datatypes::vectors::Helper;
38use futures::TryStreamExt;
39use query::QueryEngineRef;
40use query::options::{
41 FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_SEQUENCE_RANGE,
42 FLOW_RETURN_REGION_SEQ,
43};
44use query::query_engine::DefaultSerializer;
45use session::context::QueryContextRef;
46use snafu::{OptionExt, ResultExt};
47use sql::parsers::utils::is_tql;
48use store_api::mito_engine_options::MERGE_MODE_KEY;
49use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
50use table::TableRef;
51use table::table::adapter::DfTableProviderAdapter;
52use tokio::sync::oneshot::error::TryRecvError;
53use tokio::sync::{Mutex, OwnedMutexGuard, oneshot};
54use tokio::time::Instant;
55
56use crate::batching_mode::BatchingModeOptions;
57use crate::batching_mode::checkpoint::checkpoint_mode_label;
58use crate::batching_mode::eval_schedule::{EvalSchedule, select_due_scheduled_times};
59use crate::batching_mode::frontend_client::{FrontendClient, PeerDesc};
60use crate::batching_mode::state::{
61 CheckpointMode, DirtyTimeWindows, FilterExprInfo, TaskState, to_df_literal,
62};
63use crate::batching_mode::table_creator::{QueryType, create_table_with_expr};
64use crate::batching_mode::time_window::TimeWindowExpr;
65use crate::batching_mode::utils::{
66 AddFilterRewriter, ColumnMatcherRewriter, analyze_incremental_aggregate_plan, df_plan_to_sql,
67 gen_plan_with_matching_schema_and_values, get_table_info_df_schema, sql_to_df_plan,
68};
69use crate::df_optimizer::apply_df_optimizer;
70use crate::error::{
71 DatafusionSnafu, ExternalSnafu, InvalidQuerySnafu, SubstraitEncodeLogicalPlanSnafu,
72 UnexpectedSnafu,
73};
74use crate::metrics::{
75 METRIC_FLOW_BATCHING_ENGINE_ERROR_CNT, METRIC_FLOW_BATCHING_ENGINE_QUERY_TIME,
76 METRIC_FLOW_BATCHING_ENGINE_SLOW_QUERY, METRIC_FLOW_BATCHING_ENGINE_START_QUERY_CNT,
77 METRIC_FLOW_ROWS,
78};
79use crate::{Error, FlowId};
80
81mod ckpt;
82mod inc;
83
84fn wall_clock_unix_secs() -> i64 {
86 SystemTime::now()
87 .duration_since(UNIX_EPOCH)
88 .unwrap_or_default()
89 .as_secs() as i64
90}
91
92fn initial_schedule_cursor(start_secs: i64, interval_secs: i64) -> Result<i64, Error> {
100 let cursor = i128::from(start_secs) - i128::from(interval_secs);
101 i64::try_from(cursor).map_err(|_| {
102 UnexpectedSnafu {
103 reason: format!(
104 "Cannot compute the initial eval schedule cursor one interval before start {start_secs} (interval={interval_secs}): {cursor} does not fit in i64"
105 ),
106 }
107 .build()
108 })
109}
110
111fn sleep_delta_secs(next: i64, wall_now_secs: i64) -> Result<u64, Error> {
119 let delta = i128::from(next) - i128::from(wall_now_secs);
120 u64::try_from(delta).map_err(|_| {
121 UnexpectedSnafu {
122 reason: format!(
123 "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"
124 ),
125 }
126 .build()
127 })
128}
129
130fn scheduled_time_millis(scheduled_time_secs: i64) -> Result<i64, Error> {
137 scheduled_time_secs.checked_mul(1000).ok_or_else(|| {
138 UnexpectedSnafu {
139 reason: format!(
140 "Cannot convert scheduled time {scheduled_time_secs}s to milliseconds: the product exceeds i64"
141 ),
142 }
143 .build()
144 })
145}
146
147#[derive(Clone)]
149pub struct TaskConfig {
150 pub flow_id: FlowId,
151 pub query: String,
152 pub output_schema: DFSchemaRef,
154 pub time_window_expr: Option<TimeWindowExpr>,
155 pub expire_after: Option<i64>,
157 pub sink_table_name: [String; 3],
158 pub source_table_names: HashSet<[String; 3]>,
159 pub catalog_manager: CatalogManagerRef,
160 pub query_type: QueryType,
161 pub batch_opts: Arc<BatchingModeOptions>,
162 pub exact_sequence_range_required: bool,
163 pub flow_eval_interval: Option<Duration>,
164 pub eval_schedule: Option<EvalSchedule>,
166}
167
168fn determine_query_type(query: &str, query_ctx: &QueryContextRef) -> Result<QueryType, Error> {
169 let is_tql = is_tql(query_ctx.sql_dialect(), query)
170 .map_err(BoxedError::new)
171 .context(ExternalSnafu)?;
172 Ok(if is_tql {
173 QueryType::Tql
174 } else {
175 QueryType::Sql
176 })
177}
178
179fn is_merge_mode_last_non_null(options: &HashMap<String, String>) -> bool {
180 options
181 .get(MERGE_MODE_KEY)
182 .map(|mode| mode.eq_ignore_ascii_case("last_non_null"))
183 .unwrap_or(false)
184}
185
186fn encode_insert_plan_request(
187 insert_to: TableName,
188 insert_input_plan: &LogicalPlan,
189) -> Result<api::v1::QueryRequest, Error> {
190 let message = DFLogicalSubstraitConvertor {}
191 .encode(insert_input_plan, DefaultSerializer)
192 .context(SubstraitEncodeLogicalPlanSnafu)?;
193 Ok(api::v1::QueryRequest {
194 query: Some(api::v1::query_request::Query::InsertIntoPlan(
195 api::v1::InsertIntoPlan {
196 table_name: Some(insert_to),
197 logical_plan: message.to_vec(),
198 },
199 )),
200 })
201}
202
203fn recovery_aggregate_input(plan: &LogicalPlan) -> Result<LogicalPlan, Error> {
204 let plan = match plan {
205 LogicalPlan::Projection(projection) => projection.input.as_ref(),
206 _ => plan,
207 };
208 let LogicalPlan::Aggregate(aggregate) = plan else {
209 return UnexpectedSnafu {
210 reason: "Recovery timestamp projection did not find an aggregate".to_string(),
211 }
212 .fail();
213 };
214 Ok(aggregate.input.as_ref().clone())
215}
216
217fn capture_recovery_batch_windows(
218 batch: &common_recordbatch::RecordBatch,
219 time_window_expr: &TimeWindowExpr,
220 windows: &mut BTreeSet<(Timestamp, Timestamp)>,
221) -> Result<(), Error> {
222 if batch.num_columns() != 1 {
223 return UnexpectedSnafu {
224 reason: format!(
225 "Recovery timestamp projection returned {} columns instead of one",
226 batch.num_columns()
227 ),
228 }
229 .fail();
230 }
231 let values = Helper::try_into_vector(batch.column(0).clone())
232 .map_err(BoxedError::new)
233 .context(ExternalSnafu)?;
234 for index in 0..values.len() {
235 let timestamp = values.get(index).as_timestamp().context(UnexpectedSnafu {
236 reason: "Recovery timestamp projection returned a null or non-timestamp value"
237 .to_string(),
238 })?;
239 let (start, end) = time_window_expr.eval(timestamp)?;
240 let window = (
241 start.context(UnexpectedSnafu {
242 reason: "Recovery time-window expression returned no lower bound".to_string(),
243 })?,
244 end.context(UnexpectedSnafu {
245 reason: "Recovery time-window expression returned no upper bound".to_string(),
246 })?,
247 );
248 windows.insert(window);
249 }
250 Ok(())
251}
252
253fn format_insert_target_columns(plan: &LogicalPlan) -> String {
254 plan.schema()
255 .fields()
256 .iter()
257 .map(|field| quote_identifier(field.name()).to_string())
258 .collect::<Vec<_>>()
259 .join(", ")
260}
261
262pub struct BatchingExecutionGuard {
265 _lock: OwnedMutexGuard<()>,
266 restore: Option<(Arc<RwLock<TaskState>>, QueryContextRef)>,
267}
268
269impl BatchingExecutionGuard {
270 fn new(lock: OwnedMutexGuard<()>) -> Self {
271 Self {
272 _lock: lock,
273 restore: None,
274 }
275 }
276
277 fn restore_query_context(&mut self, state: Arc<RwLock<TaskState>>, old_ctx: QueryContextRef) {
278 self.restore = Some((state, old_ctx));
279 }
280}
281
282impl Drop for BatchingExecutionGuard {
283 fn drop(&mut self) {
284 if let Some((state, old_ctx)) = self.restore.take()
285 && let Ok(mut state) = state.write()
286 {
287 state.query_ctx = old_ctx;
288 }
289 }
290}
291
292#[derive(Clone)]
293pub struct BatchingTask {
294 pub config: Arc<TaskConfig>,
295 pub state: Arc<RwLock<TaskState>>,
296 execution_lock: Arc<Mutex<()>>,
300 execution: Option<Arc<dyn crate::BatchingExecution>>,
301}
302
303pub struct TaskArgs<'a> {
305 pub flow_id: FlowId,
306 pub query: &'a str,
307 pub plan: LogicalPlan,
308 pub time_window_expr: Option<TimeWindowExpr>,
309 pub expire_after: Option<i64>,
310 pub sink_table_name: [String; 3],
311 pub source_table_names: Vec<[String; 3]>,
312 pub query_ctx: QueryContextRef,
313 pub catalog_manager: CatalogManagerRef,
314 pub shutdown_rx: oneshot::Receiver<()>,
315 pub batch_opts: Arc<BatchingModeOptions>,
316 pub flow_eval_interval: Option<Duration>,
317 pub eval_schedule: Option<EvalSchedule>,
319}
320
321pub struct PlanInfo {
322 pub plan: LogicalPlan,
323 pub dirty_restore: DirtyRestore,
324 pub coverage: QueryCoverage,
325}
326
327#[derive(Clone)]
328pub enum QueryCoverage {
329 UnfilteredFull,
334 ScopedBaseRepair,
337 FencedRepairChunk { high: BTreeMap<u64, u64> },
341 IncrementalDelta,
343}
344
345impl QueryCoverage {
346 fn is_incremental_delta(&self) -> bool {
349 matches!(self, Self::IncrementalDelta)
350 }
351
352 fn snapshot_seqs(&self) -> HashMap<u64, u64> {
355 match self {
356 Self::FencedRepairChunk { high } => high.iter().map(|(k, v)| (*k, *v)).collect(),
357 _ => HashMap::new(),
358 }
359 }
360}
361
362pub enum DirtyRestore {
363 Scoped(FilterExprInfo),
366 Unscoped(DirtyTimeWindows),
373}
374
375pub struct ExecuteOnceOutcome {
376 pub new_query: Option<PlanInfo>,
377 pub result: Result<Option<(usize, Duration)>, Error>,
383}
384
385impl BatchingTask {
386 #[allow(clippy::too_many_arguments)]
387 pub fn try_new(args: TaskArgs<'_>) -> Result<Self, Error> {
388 Self::try_new_with_exact_sequence_range_required(args, false)
389 }
390
391 pub fn try_new_with_exact_sequence_range_required(
392 TaskArgs {
393 flow_id,
394 query,
395 plan,
396 time_window_expr,
397 expire_after,
398 sink_table_name,
399 source_table_names,
400 query_ctx,
401 catalog_manager,
402 shutdown_rx,
403 batch_opts,
404 flow_eval_interval,
405 eval_schedule,
406 }: TaskArgs<'_>,
407 exact_sequence_range_required: bool,
408 ) -> Result<Self, Error> {
409 let mut state = TaskState::with_dirty_time_windows(
410 query_ctx.clone(),
411 shutdown_rx,
412 DirtyTimeWindows::new(
413 batch_opts.experimental_max_filter_num_per_query,
414 batch_opts.experimental_time_window_merge_threshold,
415 ),
416 );
417 if !batch_opts.experimental_enable_incremental_read {
418 state.disable_incremental();
419 }
420
421 Ok(Self {
422 config: Arc::new(TaskConfig {
423 flow_id,
424 query: query.to_string(),
425 time_window_expr,
426 expire_after,
427 sink_table_name,
428 source_table_names: source_table_names.into_iter().collect(),
429 catalog_manager,
430 output_schema: plan.schema().clone(),
431 query_type: determine_query_type(query, &query_ctx)?,
432 exact_sequence_range_required,
433 batch_opts,
434 flow_eval_interval,
435 eval_schedule,
436 }),
437 state: Arc::new(RwLock::new(state)),
438 execution_lock: Arc::new(Mutex::new(())),
439 execution: None,
440 })
441 }
442
443 pub(crate) fn with_execution(
444 mut self,
445 execution: Option<Arc<dyn crate::BatchingExecution>>,
446 ) -> Self {
447 self.execution = execution;
448 self
449 }
450
451 pub(crate) fn stop_execution(&self) {
452 if let Some(execution) = &self.execution {
453 execution.stop();
454 }
455 }
456
457 pub fn last_execution_time_millis(&self) -> Option<i64> {
458 self.state.read().unwrap().last_execution_time_millis()
459 }
460
461 pub fn start_time_millis(&self) -> Option<i64> {
462 self.state.read().unwrap().start_time_millis()
463 }
464
465 fn frontend_extensions(&self) -> HashMap<String, String> {
468 let ctx = self.state.read().unwrap();
469 let all = ctx.query_ctx.extensions();
470 let mut flow_exts = HashMap::new();
471 if let Some(v) = all.get(query::options::FLOW_SCHEDULED_TIME_MILLIS) {
474 flow_exts.insert(
475 query::options::FLOW_SCHEDULED_TIME_MILLIS.to_string(),
476 v.clone(),
477 );
478 }
479 flow_exts
480 }
481
482 pub fn mark_all_windows_as_dirty(&self) -> Result<(), Error> {
486 let now = SystemTime::now();
487 let now = Timestamp::new_second(
488 now.duration_since(UNIX_EPOCH)
489 .expect("Time went backwards")
490 .as_secs() as _,
491 );
492 let lower_bound = self
493 .config
494 .expire_after
495 .map(|e| now.sub_duration(Duration::from_secs(e as _)))
496 .transpose()
497 .map_err(BoxedError::new)
498 .context(ExternalSnafu)?
499 .unwrap_or(Timestamp::new_second(0));
500 debug!(
501 "Flow {} mark range ({:?}, {:?}) as dirty",
502 self.config.flow_id, lower_bound, now
503 );
504 self.state
505 .write()
506 .unwrap()
507 .dirty_time_windows
508 .add_window(lower_bound, Some(now));
509 Ok(())
510 }
511
512 pub async fn check_or_create_sink_table(
514 &self,
515 engine: &QueryEngineRef,
516 frontend_client: &Arc<FrontendClient>,
517 ) -> Result<TableRef, Error> {
518 if !self.is_table_exist(&self.config.sink_table_name).await? {
519 let create_table = self.gen_create_table_expr(engine.clone()).await?;
520 info!(
521 "Try creating sink table(if not exists) with expr: {:?}",
522 create_table
523 );
524 self.create_table(frontend_client, create_table).await?;
525 info!(
526 "Sink table {}(if not exists) created",
527 self.config.sink_table_name.join(".")
528 );
529 }
530 let (table, _) = get_table_info_df_schema(
531 self.config.catalog_manager.clone(),
532 self.config.sink_table_name.clone(),
533 )
534 .await?;
535 Ok(table)
536 }
537
538 pub fn query_context_snapshot(&self) -> QueryContextRef {
540 let query_ctx = self.state.read().unwrap().query_ctx.clone();
541 Arc::new(query_ctx.fork())
542 }
543
544 pub fn recovery_retention_lower_bound(&self) -> Result<Option<Timestamp>, Error> {
546 let Some(expire_after) = self.config.expire_after else {
547 return Ok(None);
548 };
549 let expire_after = u64::try_from(expire_after).map_err(|_| {
550 UnexpectedSnafu {
551 reason: format!(
552 "Flow {} has negative expire_after {expire_after}",
553 self.config.flow_id
554 ),
555 }
556 .build()
557 })?;
558 let now = Timestamp::new_second(
559 SystemTime::now()
560 .duration_since(UNIX_EPOCH)
561 .map_err(|err| {
562 UnexpectedSnafu {
563 reason: format!("Failed to read recovery wall clock: {err}"),
564 }
565 .build()
566 })?
567 .as_secs() as i64,
568 );
569 let lower = now
570 .sub_duration(Duration::from_secs(expire_after))
571 .map_err(BoxedError::new)
572 .context(ExternalSnafu)?;
573 self.config
574 .time_window_expr
575 .as_ref()
576 .context(UnexpectedSnafu {
577 reason: "Recovery expiry requires a time-window expression".to_string(),
578 })?
579 .eval(lower)?
580 .0
581 .context(UnexpectedSnafu {
582 reason: "Recovery expiry time-window expression returned no lower bound"
583 .to_string(),
584 })
585 .map(Some)
586 }
587
588 pub async fn validate_recovery_retention(
590 &self,
591 retention_lower: Option<Timestamp>,
592 windows: &[(Timestamp, Timestamp)],
593 ) -> Result<(), Error> {
594 for name in &self.config.source_table_names {
595 let table = self
596 .config
597 .catalog_manager
598 .table(&name[0], &name[1], &name[2], None)
599 .await
600 .map_err(BoxedError::new)
601 .context(ExternalSnafu)?
602 .context(UnexpectedSnafu {
603 reason: format!(
604 "Flow {} source table {} is unavailable for recovery retention validation",
605 self.config.flow_id,
606 name.join(".")
607 ),
608 })?;
609 let ttl = if let Some(ttl) = table.table_info().meta.options.ttl {
610 ttl
611 } else {
612 let manager = self
613 .config
614 .catalog_manager
615 .as_any()
616 .downcast_ref::<KvBackendCatalogManager>()
617 .context(UnexpectedSnafu {
618 reason: format!(
619 "Flow {} cannot resolve inherited TTL for source table {} during recovery",
620 self.config.flow_id,
621 name.join(".")
622 ),
623 })?;
624 manager
625 .table_metadata_manager_ref()
626 .schema_manager()
627 .get(SchemaNameKey::new(&name[0], &name[1]))
628 .await
629 .map_err(BoxedError::new)
630 .context(ExternalSnafu)?
631 .context(UnexpectedSnafu {
632 reason: format!(
633 "Flow {} schema {}.{} is unavailable for recovery retention validation",
634 self.config.flow_id, name[0], name[1]
635 ),
636 })?
637 .ttl
638 .map(Into::into)
639 .unwrap_or(TimeToLive::Forever)
640 };
641
642 match ttl {
643 TimeToLive::Forever => {}
644 TimeToLive::Instant => {
645 return UnexpectedSnafu {
646 reason: format!(
647 "Flow {} source table {} has instant TTL and cannot be recovered",
648 self.config.flow_id,
649 name.join(".")
650 ),
651 }
652 .fail();
653 }
654 TimeToLive::Duration(ttl) => {
655 let expire_after = self.config.expire_after.context(UnexpectedSnafu {
656 reason: format!(
657 "Flow {} source table {} has TTL {ttl:?}, but recovery has no expire_after",
658 self.config.flow_id,
659 name.join(".")
660 ),
661 })?;
662 let expire_after = u64::try_from(expire_after).map_err(|_| {
663 UnexpectedSnafu {
664 reason: format!(
665 "Flow {} has negative expire_after {expire_after} during recovery retention validation",
666 self.config.flow_id
667 ),
668 }
669 .build()
670 })?;
671 if ttl <= Duration::from_secs(expire_after) {
672 return UnexpectedSnafu {
673 reason: format!(
674 "Flow {} source table {} TTL {ttl:?} must exceed expire_after {expire_after}s for recovery",
675 self.config.flow_id,
676 name.join(".")
677 ),
678 }
679 .fail();
680 }
681 let mut oldest = retention_lower.context(UnexpectedSnafu {
682 reason: format!(
683 "Flow {} source table {} has finite TTL {ttl:?}, but recovery has no retained lower bound",
684 self.config.flow_id,
685 name.join(".")
686 ),
687 })?;
688 for (start, _) in windows {
689 oldest = oldest.min(*start);
690 }
691 let cutoff = Timestamp::current_millis()
692 .sub_duration(ttl)
693 .map_err(BoxedError::new)
694 .context(ExternalSnafu)?;
695 if oldest <= cutoff {
696 return UnexpectedSnafu {
697 reason: format!(
698 "Flow {} recovery requires source table {} data from {oldest:?}, older than TTL {ttl:?} cutoff {cutoff:?}",
699 self.config.flow_id,
700 name.join(".")
701 ),
702 }
703 .fail();
704 }
705 }
706 }
707 }
708 Ok(())
709 }
710
711 pub async fn capture_recovery_windows(
715 &self,
716 engine: &QueryEngineRef,
717 frontend_client: &FrontendClient,
718 lower: &BTreeMap<u64, u64>,
719 ) -> Result<(BTreeMap<u64, u64>, Vec<(Timestamp, Timestamp)>), Error> {
720 let retention_lower = self.recovery_retention_lower_bound()?;
721 self.capture_recovery_windows_since(engine, frontend_client, lower, retention_lower)
722 .await
723 }
724
725 pub async fn capture_recovery_windows_since(
731 &self,
732 engine: &QueryEngineRef,
733 frontend_client: &FrontendClient,
734 lower: &BTreeMap<u64, u64>,
735 retention_lower: Option<Timestamp>,
736 ) -> Result<(BTreeMap<u64, u64>, Vec<(Timestamp, Timestamp)>), Error> {
737 if lower.is_empty() {
738 return UnexpectedSnafu {
739 reason: format!(
740 "Flow {} recovery window capture requires nonempty lower sequence bounds",
741 self.config.flow_id
742 ),
743 }
744 .fail();
745 }
746 if !self.sequence_range_capable().await? {
747 return UnexpectedSnafu {
748 reason: format!(
749 "Flow {} recovery window capture requires sequence_range-capable sources",
750 self.config.flow_id
751 ),
752 }
753 .fail();
754 }
755
756 let query_ctx = self.query_context_snapshot();
757 let plan = sql_to_df_plan(query_ctx, engine.clone(), &self.config.query, false).await?;
758 let input = recovery_aggregate_input(&plan)?;
759 let Some(analysis) = analyze_incremental_aggregate_plan(&plan)? else {
760 return UnexpectedSnafu {
761 reason: format!(
762 "Flow {} recovery window capture requires a supported aggregate input",
763 self.config.flow_id
764 ),
765 }
766 .fail();
767 };
768 if !analysis.unsupported_exprs.is_empty() {
769 return UnexpectedSnafu {
770 reason: format!(
771 "Flow {} recovery window capture has unsupported aggregate expressions: {:?}",
772 self.config.flow_id, analysis.unsupported_exprs
773 ),
774 }
775 .fail();
776 }
777 let time_window_expr = self
778 .config
779 .time_window_expr
780 .as_ref()
781 .context(UnexpectedSnafu {
782 reason: format!(
783 "Flow {} recovery window capture requires a time-window expression",
784 self.config.flow_id
785 ),
786 })?;
787 let input = if let Some(retention_lower) = retention_lower {
788 let mut add_filter = AddFilterRewriter::new(
789 col(&time_window_expr.column_name).gt_eq(lit(to_df_literal(retention_lower)?)),
790 );
791 input
792 .rewrite(&mut add_filter)
793 .with_context(|_| DatafusionSnafu {
794 context: "Failed to apply recovery expire_after filter".to_string(),
795 })?
796 .data
797 } else {
798 input
799 };
800 let timestamp_plan = LogicalPlan::Projection(
801 Projection::try_new(vec![col(&time_window_expr.column_name)], Arc::new(input))
802 .context(DatafusionSnafu {
803 context: "Failed to project recovery source timestamps".to_string(),
804 })?,
805 );
806 let catalog = &self.config.sink_table_name[0];
807 let schema = &self.config.sink_table_name[1];
808 let timestamp_plan = timestamp_plan
809 .clone()
810 .transform_down_with_subqueries(|p| {
811 if let LogicalPlan::TableScan(mut table_scan) = p {
812 let resolved = table_scan.table_name.resolve(catalog, schema);
813 table_scan.table_name = resolved.into();
814 Ok(Transformed::yes(LogicalPlan::TableScan(table_scan)))
815 } else {
816 Ok(Transformed::no(p))
817 }
818 })
819 .with_context(|_| DatafusionSnafu {
820 context: format!(
821 "Failed to fix table ref in recovery timestamp plan, plan={:?}",
822 timestamp_plan
823 ),
824 })?
825 .data;
826 let message = DFLogicalSubstraitConvertor {}
827 .encode(×tamp_plan, DefaultSerializer)
828 .context(SubstraitEncodeLogicalPlanSnafu)?;
829 let request = api::v1::QueryRequest {
830 query: Some(api::v1::query_request::Query::LogicalPlan(message.to_vec())),
831 };
832 let lower_json = serde_json::to_string(lower).map_err(|err| {
833 UnexpectedSnafu {
834 reason: format!("Failed to serialize recovery lower sequence bounds: {err}"),
835 }
836 .build()
837 })?;
838 let extensions = [
839 (FLOW_RETURN_REGION_SEQ, "true"),
840 (FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_SEQUENCE_RANGE),
841 (FLOW_INCREMENTAL_AFTER_SEQS, lower_json.as_str()),
842 ];
843 let mut peer_desc = None;
844 let result = frontend_client
845 .query_with_terminal_metrics(
846 catalog,
847 schema,
848 request,
849 &extensions,
850 &HashMap::new(),
851 &mut peer_desc,
852 )
853 .await?;
854 let mut aligned_windows = BTreeSet::new();
855 let metrics = result.metrics.clone();
856 match result.output.data {
857 OutputData::AffectedRows(_) => {
858 return UnexpectedSnafu {
859 reason: "Recovery timestamp projection unexpectedly returned affected rows"
860 .to_string(),
861 }
862 .fail();
863 }
864 OutputData::RecordBatches(batches) => {
865 for batch in batches.iter() {
866 capture_recovery_batch_windows(batch, time_window_expr, &mut aligned_windows)?;
867 }
868 }
869 OutputData::Stream(mut stream) => {
870 while let Some(batch) = stream
871 .try_next()
872 .await
873 .map_err(BoxedError::new)
874 .context(ExternalSnafu)?
875 {
876 capture_recovery_batch_windows(&batch, time_window_expr, &mut aligned_windows)?;
877 }
878 }
879 }
880 if !metrics.is_ready() {
881 return UnexpectedSnafu {
882 reason: "Recovery timestamp projection ended without terminal metrics".to_string(),
883 }
884 .fail();
885 }
886 let participating = metrics.participating_regions().context(UnexpectedSnafu {
887 reason: "Recovery timestamp projection has no participating-region proof".to_string(),
888 })?;
889 let high = metrics.region_watermark_map().context(UnexpectedSnafu {
890 reason: "Recovery timestamp projection has no terminal watermark proof".to_string(),
891 })?;
892 if participating.len() != lower.len()
893 || high.len() != lower.len()
894 || participating
895 .iter()
896 .any(|region| !lower.contains_key(region))
897 || lower
898 .iter()
899 .any(|(region, low)| high.get(region).is_none_or(|watermark| watermark < low))
900 || high.keys().any(|region| !participating.contains(region))
901 {
902 return UnexpectedSnafu {
903 reason: format!(
904 "Recovery timestamp projection returned incomplete or regressing terminal proof: lower={lower:?}, participating={participating:?}, high={high:?}"
905 ),
906 }
907 .fail();
908 }
909 Ok((
910 high.into_iter().collect(),
911 aligned_windows.into_iter().collect(),
912 ))
913 }
914
915 pub async fn validate_sink_table_schema(&self, engine: &QueryEngineRef) -> Result<(), Error> {
917 self.validate_sink_table_schema_with_values(engine, &BTreeMap::new())
918 .await
919 }
920
921 pub async fn validate_sink_table_schema_with_values(
923 &self,
924 engine: &QueryEngineRef,
925 values: &BTreeMap<String, ScalarValue>,
926 ) -> Result<(), Error> {
927 let (table, _) = get_table_info_df_schema(
928 self.config.catalog_manager.clone(),
929 self.config.sink_table_name.clone(),
930 )
931 .await?;
932
933 let table_meta = &table.table_info().meta;
934 let merge_mode_last_non_null =
935 is_merge_mode_last_non_null(&table_meta.options.extra_options);
936 let primary_key_indices = table_meta.primary_key_indices.clone();
937
938 gen_plan_with_matching_schema_and_values(
939 &self.config.query,
940 self.query_context_snapshot(),
941 engine.clone(),
942 table_meta.schema.clone(),
943 &primary_key_indices,
944 merge_mode_last_non_null,
945 values,
946 )
947 .await
948 .map(|_| ())
949 }
950
951 async fn is_table_exist(&self, table_name: &[String; 3]) -> Result<bool, Error> {
952 self.config
953 .catalog_manager
954 .table_exists(&table_name[0], &table_name[1], &table_name[2], None)
955 .await
956 .map_err(BoxedError::new)
957 .context(ExternalSnafu)
958 }
959
960 pub(crate) async fn execute_once_serialized(
961 &self,
962 engine: &QueryEngineRef,
963 frontend_client: &Arc<FrontendClient>,
964 max_window_cnt: Option<usize>,
965 ) -> Result<Option<(usize, Duration)>, Error> {
966 let outcome = self
967 .execute_once_serialized_with_outcome(engine, frontend_client, max_window_cnt)
968 .await;
969 outcome.result
970 }
971
972 async fn execute_once_serialized_with_outcome(
975 &self,
976 engine: &QueryEngineRef,
977 frontend_client: &Arc<FrontendClient>,
978 max_window_cnt: Option<usize>,
979 ) -> ExecuteOnceOutcome {
980 let guard = BatchingExecutionGuard::new(self.execution_lock.clone().lock_owned().await);
981 self.execute_once_with_guard(guard, engine, frontend_client, max_window_cnt)
982 .await
983 }
984
985 async fn execute_once_with_guard(
986 &self,
987 guard: BatchingExecutionGuard,
988 engine: &QueryEngineRef,
989 frontend_client: &Arc<FrontendClient>,
990 max_window_cnt: Option<usize>,
991 ) -> ExecuteOnceOutcome {
992 if let Some(execution) = &self.execution {
993 return execution
994 .clone()
995 .execute_once(guard, self, engine, frontend_client, max_window_cnt)
996 .await;
997 }
998 self.execute_once_default_unlocked(engine, frontend_client, max_window_cnt)
999 .await
1000 }
1001
1002 async fn execute_once_default_unlocked(
1004 &self,
1005 engine: &QueryEngineRef,
1006 frontend_client: &Arc<FrontendClient>,
1007 max_window_cnt: Option<usize>,
1008 ) -> ExecuteOnceOutcome {
1009 let new_query = match self.gen_insert_plan_unlocked(engine, max_window_cnt).await {
1010 Ok(new_query) => new_query,
1011 Err(err) => {
1012 return ExecuteOnceOutcome {
1013 new_query: None,
1014 result: Err(err),
1015 };
1016 }
1017 };
1018
1019 if let Some(new_query) = new_query {
1020 debug!("Generate new query: {}", new_query.plan);
1021 let res = self
1022 .execute_logical_plan_unlocked(
1023 engine,
1024 frontend_client,
1025 &new_query.plan,
1026 &new_query.dirty_restore,
1027 &new_query.coverage,
1028 )
1029 .await;
1030 if res.is_err() {
1031 self.handle_executed_query_failure(Some(&new_query));
1032 }
1033 ExecuteOnceOutcome {
1034 new_query: Some(new_query),
1035 result: res,
1036 }
1037 } else {
1038 debug!("Generate no query");
1039 ExecuteOnceOutcome {
1040 new_query: None,
1041 result: Ok(None),
1042 }
1043 }
1044 }
1045
1046 async fn gen_insert_plan_unlocked(
1048 &self,
1049 engine: &QueryEngineRef,
1050 max_window_cnt: Option<usize>,
1051 ) -> Result<Option<PlanInfo>, Error> {
1052 self.gen_insert_plan_with_values_unlocked(engine, max_window_cnt, &BTreeMap::new(), false)
1053 .await
1054 }
1055
1056 pub async fn gen_insert_plan_with_values_unlocked(
1059 &self,
1060 engine: &QueryEngineRef,
1061 max_window_cnt: Option<usize>,
1062 values: &BTreeMap<String, ScalarValue>,
1063 force_full_snapshot: bool,
1064 ) -> Result<Option<PlanInfo>, Error> {
1065 let (table, df_schema) = get_table_info_df_schema(
1066 self.config.catalog_manager.clone(),
1067 self.config.sink_table_name.clone(),
1068 )
1069 .await?;
1070
1071 let table_meta = &table.table_info().meta;
1072 let merge_mode_last_non_null =
1073 is_merge_mode_last_non_null(&table_meta.options.extra_options);
1074 let primary_key_indices = table_meta.primary_key_indices.clone();
1075
1076 let new_query = self
1077 .gen_query_with_time_window_with_values(
1078 engine.clone(),
1079 &table.table_info().meta.schema,
1080 &primary_key_indices,
1081 merge_mode_last_non_null,
1082 max_window_cnt,
1083 values,
1084 force_full_snapshot,
1085 )
1086 .await?;
1087
1088 let Some(new_query) = new_query else {
1089 return Ok(None);
1090 };
1091
1092 let table_columns = df_schema
1095 .columns()
1096 .into_iter()
1097 .map(|c| c.name)
1098 .collect::<BTreeSet<_>>();
1099 for column in new_query.plan.schema().columns() {
1100 if !table_columns.contains(column.name()) {
1101 self.restore_dirty_windows_after_failure(&new_query);
1102 return InvalidQuerySnafu {
1103 reason: format!(
1104 "Column {} not found in sink table with columns {:?}",
1105 column, table_columns
1106 ),
1107 }
1108 .fail();
1109 }
1110 }
1111
1112 let table_provider = Arc::new(DfTableProviderAdapter::new(table));
1113 let table_source = Arc::new(DefaultTableSource::new(table_provider));
1114
1115 let plan = LogicalPlan::Dml(DmlStatement::new(
1117 datafusion_common::TableReference::Full {
1118 catalog: self.config.sink_table_name[0].clone().into(),
1119 schema: self.config.sink_table_name[1].clone().into(),
1120 table: self.config.sink_table_name[2].clone().into(),
1121 },
1122 table_source,
1123 WriteOp::Insert(datafusion_expr::dml::InsertOp::Append),
1124 Arc::new(new_query.plan.clone()),
1125 ));
1126 let insert_into_info = PlanInfo {
1127 plan,
1128 dirty_restore: new_query.dirty_restore,
1129 coverage: new_query.coverage,
1130 };
1131 let insert_into =
1132 match insert_into_info
1133 .plan
1134 .clone()
1135 .recompute_schema()
1136 .context(DatafusionSnafu {
1137 context: "Failed to recompute schema",
1138 }) {
1139 Ok(insert_into) => insert_into,
1140 Err(err) => {
1141 self.restore_dirty_windows_after_failure(&insert_into_info);
1142 return Err(err);
1143 }
1144 };
1145
1146 Ok(Some(PlanInfo {
1147 plan: insert_into,
1148 dirty_restore: insert_into_info.dirty_restore,
1149 coverage: insert_into_info.coverage,
1150 }))
1151 }
1152
1153 pub async fn create_table(
1154 &self,
1155 frontend_client: &Arc<FrontendClient>,
1156 expr: CreateTableExpr,
1157 ) -> Result<(), Error> {
1158 let catalog = &self.config.sink_table_name[0];
1159 let schema = &self.config.sink_table_name[1];
1160 frontend_client
1161 .create(expr.clone(), catalog, schema)
1162 .await?;
1163 Ok(())
1164 }
1165
1166 async fn execute_logical_plan_unlocked(
1168 &self,
1169 engine: &QueryEngineRef,
1170 frontend_client: &Arc<FrontendClient>,
1171 plan: &LogicalPlan,
1172 dirty_restore: &DirtyRestore,
1173 coverage: &QueryCoverage,
1174 ) -> Result<Option<(usize, Duration)>, Error> {
1175 let flow_id = self.config.flow_id;
1176 let Some((res, elapsed)) = self
1177 .execute_plan_unlocked(engine, frontend_client, plan, dirty_restore, coverage)
1178 .await?
1179 else {
1180 return Ok(None);
1181 };
1182
1183 if let Err(err) = &res {
1184 let decision = {
1185 let mut state = self.state.write().unwrap();
1186 let reason = Self::query_failure_reason(err, coverage);
1187 Self::apply_query_failure_to_state(&mut state, elapsed, coverage, reason)
1188 };
1189 if let Some(decision) = decision {
1190 Self::record_checkpoint_decision(flow_id, decision);
1191 }
1192 }
1193
1194 let res = res?;
1195 let (affected_rows, _) = res.output.extract_rows_and_cost();
1196 let decision = {
1197 let mut state = self.state.write().unwrap();
1198 Self::apply_query_result_to_state(&mut state, &res, elapsed, coverage)
1199 };
1200 Self::record_checkpoint_decision(flow_id, decision);
1201 Ok(Some((affected_rows, elapsed)))
1202 }
1203
1204 pub async fn execute_plan_unlocked(
1208 &self,
1209 engine: &QueryEngineRef,
1210 frontend_client: &Arc<FrontendClient>,
1211 plan: &LogicalPlan,
1212 dirty_restore: &DirtyRestore,
1213 coverage: &QueryCoverage,
1214 ) -> Result<Option<(Result<OutputWithMetrics, Error>, Duration)>, Error> {
1215 let instant = Instant::now();
1216 let flow_id = self.config.flow_id;
1217
1218 debug!(
1219 "Executing flow {flow_id}(expire_after={:?} secs) with query {}",
1220 self.config.expire_after, &plan
1221 );
1222
1223 let catalog = &self.config.sink_table_name[0];
1224 let schema = &self.config.sink_table_name[1];
1225
1226 let plan = plan
1228 .clone()
1229 .transform_down_with_subqueries(|p| {
1230 if let LogicalPlan::TableScan(mut table_scan) = p {
1231 let resolved = table_scan.table_name.resolve(catalog, schema);
1232 table_scan.table_name = resolved.into();
1233 Ok(Transformed::yes(LogicalPlan::TableScan(table_scan)))
1234 } else {
1235 Ok(Transformed::no(p))
1236 }
1237 })
1238 .with_context(|_| DatafusionSnafu {
1239 context: format!("Failed to fix table ref in logical plan, plan={:?}", plan),
1240 })?
1241 .data;
1242
1243 let incremental_plan = if coverage.is_incremental_delta() {
1246 self.prepare_plan_for_incremental(engine, &plan).await?
1247 } else {
1248 None
1249 };
1250 let incremental_safe = incremental_plan.is_some();
1251 if coverage.is_incremental_delta() && !incremental_safe {
1252 debug!(
1253 "Flow {flow_id} skipped unsafe incremental delta fallback; \
1254 restored dirty signal instead of executing an unfiltered full snapshot"
1255 );
1256 self.restore_dirty_windows(dirty_restore);
1257 return Ok(None);
1258 }
1259 let plan = incremental_plan.unwrap_or_else(|| plan.clone());
1260 let plan = match &self.execution {
1261 Some(execution) => execution.rewrite_plan(self, plan)?,
1262 None => plan,
1263 };
1264
1265 let extensions = self
1266 .build_flow_query_extensions(incremental_safe, coverage.is_incremental_delta())
1267 .await?;
1268 let frontend_extensions = self.frontend_extensions();
1269 let extension_refs = extensions
1270 .iter()
1271 .map(|(key, value)| (*key, value.as_str()))
1272 .chain(
1273 frontend_extensions
1274 .iter()
1275 .map(|(key, value)| (key.as_str(), value.as_str())),
1276 )
1277 .collect::<Vec<_>>();
1278 let query_mode = if extensions
1279 .iter()
1280 .any(|(key, _)| *key == FLOW_INCREMENTAL_MODE)
1281 {
1282 CheckpointMode::Incremental
1283 } else {
1284 CheckpointMode::FullSnapshot
1285 };
1286 Self::record_query_mode(flow_id, query_mode);
1287 debug!(
1288 "Flow {flow_id} executing batching query with checkpoint_mode={}, extension_count={}",
1289 checkpoint_mode_label(query_mode),
1290 extensions.len()
1291 );
1292
1293 let mut peer_desc = None;
1294 let res = {
1295 let _timer = METRIC_FLOW_BATCHING_ENGINE_QUERY_TIME
1296 .with_label_values(&[flow_id.to_string().as_str()])
1297 .start_timer();
1298
1299 let req = if let Some((insert_to, insert_input_plan)) =
1300 breakup_insert_plan(&plan, catalog, schema)
1301 {
1302 if query_mode == CheckpointMode::FullSnapshot
1303 && matches!(self.config.query_type, QueryType::Sql)
1304 && self.config.flow_eval_interval.is_some()
1305 && self.config.time_window_expr.is_none()
1306 {
1307 match df_plan_to_sql(&insert_input_plan) {
1317 Ok(select_sql) => {
1318 let target_columns = format_insert_target_columns(&insert_input_plan);
1319 let sql = format!(
1320 "INSERT INTO {} ({}) {}",
1321 TableReference::full(
1322 insert_to.catalog_name.as_str(),
1323 insert_to.schema_name.as_str(),
1324 insert_to.table_name.as_str(),
1325 )
1326 .to_quoted_string(),
1327 target_columns,
1328 select_sql
1329 );
1330 api::v1::QueryRequest {
1331 query: Some(api::v1::query_request::Query::Sql(sql)),
1332 }
1333 }
1334 Err(err) => {
1335 debug!(
1336 "Failed to unparse full-snapshot SQL flow {} plan; \
1337 falling back to InsertIntoPlan: {:?}",
1338 flow_id, err
1339 );
1340 encode_insert_plan_request(insert_to, &insert_input_plan)?
1341 }
1342 }
1343 } else {
1344 encode_insert_plan_request(insert_to, &insert_input_plan)?
1345 }
1346 } else {
1347 let message = DFLogicalSubstraitConvertor {}
1348 .encode(&plan, DefaultSerializer)
1349 .context(SubstraitEncodeLogicalPlanSnafu)?;
1350
1351 api::v1::QueryRequest {
1352 query: Some(api::v1::query_request::Query::LogicalPlan(message.to_vec())),
1353 }
1354 };
1355
1356 let snapshot_seqs = coverage.snapshot_seqs();
1357 {
1358 let mut state = self.state.write().unwrap();
1359 state.record_start_time_if_first();
1360 }
1361 frontend_client
1362 .query_with_terminal_metrics(
1363 catalog,
1364 schema,
1365 req,
1366 &extension_refs,
1367 &snapshot_seqs,
1368 &mut peer_desc,
1369 )
1370 .await
1371 };
1372
1373 let elapsed = instant.elapsed();
1374 let peer_label = peer_desc
1375 .as_ref()
1376 .map(ToString::to_string)
1377 .unwrap_or_else(|| PeerDesc::default().to_string());
1378 if let Err(err) = &res {
1379 warn!(
1383 "Failed to execute Flow {flow_id} on frontend {peer_label}, result: {err:?}, elapsed: {:?} with query: {}",
1384 elapsed, &plan
1385 );
1386 }
1387
1388 if elapsed >= self.config.batch_opts.slow_query_threshold {
1390 warn!(
1391 "Flow {flow_id} on frontend {peer_label} executed for {:?} before complete, query: {}",
1392 elapsed, &plan
1393 );
1394 let flow_id = flow_id.to_string();
1395 METRIC_FLOW_BATCHING_ENGINE_SLOW_QUERY
1396 .with_label_values(&[flow_id.as_str(), peer_label.as_str()])
1397 .observe(elapsed.as_secs_f64());
1398 }
1399
1400 match res {
1401 Ok(res) => {
1402 let (affected_rows, _) = res.output.extract_rows_and_cost();
1403 debug!(
1404 "Flow {flow_id} executed, affected_rows: {affected_rows:?}, elapsed: {:?}, watermark: {:?}",
1405 elapsed,
1406 res.region_watermark_map()
1407 );
1408 METRIC_FLOW_ROWS
1409 .with_label_values(&[format!("{}-out-batching", flow_id).as_str()])
1410 .inc_by(affected_rows as _);
1411 Ok(Some((Ok(res), elapsed)))
1412 }
1413 Err(err) => Ok(Some((Err(err), elapsed))),
1414 }
1415 }
1416
1417 pub fn restore_dirty_windows(&self, dirty_restore: &DirtyRestore) {
1421 match dirty_restore {
1422 DirtyRestore::Scoped(filter) => self.restore_scoped_dirty_windows(filter),
1423 DirtyRestore::Unscoped(dirty_windows) => self
1424 .state
1425 .write()
1426 .unwrap()
1427 .dirty_time_windows
1428 .add_dirty_windows(dirty_windows),
1429 }
1430 }
1431
1432 fn restore_dirty_windows_after_failure(&self, query: &PlanInfo) {
1435 self.restore_dirty_windows(&query.dirty_restore);
1436 }
1437
1438 fn restore_scoped_dirty_windows(&self, filter: &FilterExprInfo) {
1441 self.state.write().unwrap().restore_scoped_windows(filter);
1442 }
1443
1444 fn restore_scoped_dirty_windows_on_err<T>(
1447 &self,
1448 filter: &FilterExprInfo,
1449 result: Result<T, Error>,
1450 ) -> Result<T, Error> {
1451 result.inspect_err(|_| {
1452 self.restore_scoped_dirty_windows(filter);
1453 })
1454 }
1455
1456 fn restore_unscoped_dirty_windows(&self, dirty_windows: &DirtyTimeWindows) {
1459 self.state
1460 .write()
1461 .unwrap()
1462 .dirty_time_windows
1463 .add_dirty_windows(dirty_windows);
1464 }
1465
1466 fn restore_unscoped_dirty_windows_on_err<T>(
1469 &self,
1470 dirty_windows: &DirtyTimeWindows,
1471 result: Result<T, Error>,
1472 ) -> Result<T, Error> {
1473 result.inspect_err(|_| {
1474 self.restore_unscoped_dirty_windows(dirty_windows);
1475 })
1476 }
1477
1478 fn drain_dirty_windows_signal(&self) -> (bool, DirtyTimeWindows) {
1481 let mut state = self.state.write().unwrap();
1482 let dirty_windows_to_restore = state.dirty_time_windows.clone();
1483 let is_dirty = !dirty_windows_to_restore.is_empty();
1484 state.dirty_time_windows.clean();
1485 (is_dirty, dirty_windows_to_restore)
1486 }
1487
1488 #[allow(clippy::too_many_arguments)]
1489 async fn gen_unfiltered_plan_info(
1492 &self,
1493 engine: QueryEngineRef,
1494 query_ctx: QueryContextRef,
1495 sink_table_schema: Arc<Schema>,
1496 primary_key_indices: &[usize],
1497 allow_partial: bool,
1498 values: &BTreeMap<String, ScalarValue>,
1499 dirty_windows_to_restore: DirtyTimeWindows,
1500 retention_filter: Option<(&str, Timestamp, &'static str)>,
1501 coverage: QueryCoverage,
1502 ) -> Result<PlanInfo, Error> {
1503 let mut plan = self.restore_unscoped_dirty_windows_on_err(
1504 &dirty_windows_to_restore,
1505 gen_plan_with_matching_schema_and_values(
1506 &self.config.query,
1507 query_ctx,
1508 engine,
1509 sink_table_schema,
1510 primary_key_indices,
1511 allow_partial,
1512 values,
1513 )
1514 .await,
1515 )?;
1516
1517 if let Some((col_name, lower_bound, context)) = retention_filter {
1518 let lower = self.restore_unscoped_dirty_windows_on_err(
1519 &dirty_windows_to_restore,
1520 to_df_literal(lower_bound),
1521 )?;
1522 let retention_filter = col(col_name).gt_eq(lit(lower));
1523 let mut add_filter = AddFilterRewriter::new(retention_filter);
1524 plan = self.restore_unscoped_dirty_windows_on_err(
1525 &dirty_windows_to_restore,
1526 plan.clone()
1527 .rewrite(&mut add_filter)
1528 .with_context(|_| DatafusionSnafu {
1529 context: format!(
1530 "Failed to apply {context} expire_after filter to plan:\n {}\n",
1531 plan
1532 ),
1533 })
1534 .map(|rewrite| rewrite.data),
1535 )?;
1536 }
1537
1538 Ok(PlanInfo {
1539 plan,
1540 dirty_restore: DirtyRestore::Unscoped(dirty_windows_to_restore),
1541 coverage,
1542 })
1543 }
1544
1545 #[allow(clippy::too_many_arguments)]
1546 async fn gen_unfiltered_plan_info_if_dirty(
1549 &self,
1550 engine: QueryEngineRef,
1551 query_ctx: QueryContextRef,
1552 sink_table_schema: Arc<Schema>,
1553 primary_key_indices: &[usize],
1554 allow_partial: bool,
1555 values: &BTreeMap<String, ScalarValue>,
1556 retention_filter: Option<(&str, Timestamp, &'static str)>,
1557 coverage: QueryCoverage,
1558 ) -> Result<Option<PlanInfo>, Error> {
1559 let (is_dirty, dirty_windows_to_restore) = self.drain_dirty_windows_signal();
1560 if !is_dirty {
1561 debug!("Flow id={:?}, no new data, not update", self.config.flow_id);
1562 return Ok(None);
1563 }
1564
1565 self.gen_unfiltered_plan_info(
1566 engine,
1567 query_ctx,
1568 sink_table_schema,
1569 primary_key_indices,
1570 allow_partial,
1571 values,
1572 dirty_windows_to_restore,
1573 retention_filter,
1574 coverage,
1575 )
1576 .await
1577 .map(Some)
1578 }
1579
1580 fn handle_executed_query_failure(&self, query: Option<&PlanInfo>) {
1581 if let Some(query) = query {
1582 self.restore_dirty_windows_after_failure(query);
1583 }
1584 }
1585
1586 pub async fn start_executing_loop(
1594 &self,
1595 engine: QueryEngineRef,
1596 frontend_client: Arc<FrontendClient>,
1597 ) {
1598 if self.config.flow_eval_interval.is_some() {
1599 self.start_scheduled_loop(engine, frontend_client).await;
1600 } else {
1601 self.start_adaptive_loop(engine, frontend_client).await;
1602 }
1603 }
1604
1605 async fn start_scheduled_loop(
1615 &self,
1616 engine: QueryEngineRef,
1617 frontend_client: Arc<FrontendClient>,
1618 ) {
1619 let flow_id_str = self.config.flow_id.to_string();
1620
1621 let schedule = match &self.config.eval_schedule {
1622 Some(s) => s.clone(),
1623 None => {
1624 let eval_interval_secs = self
1625 .config
1626 .flow_eval_interval
1627 .map(|d| d.as_secs() as i64)
1628 .expect("checked by caller");
1629
1630 match EvalSchedule::from_config(Some(eval_interval_secs), None) {
1633 Ok(Some(s)) => s,
1634 Ok(None) => {
1635 warn!(
1636 "Flow {}: EVAL INTERVAL set but no schedule parsed; exiting loop",
1637 flow_id_str
1638 );
1639 return;
1640 }
1641 Err(e) => {
1642 warn!(
1643 "Flow {}: Failed to parse eval schedule: {}; exiting loop",
1644 flow_id_str, e
1645 );
1646 return;
1647 }
1648 }
1649 }
1650 };
1651
1652 let mut cursor_secs =
1657 match initial_schedule_cursor(schedule.start_secs, schedule.interval_secs) {
1658 Ok(cursor) => cursor,
1659 Err(e) => {
1660 warn!(
1661 "Flow {}: invalid eval schedule, exiting loop: {}",
1662 flow_id_str, e
1663 );
1664 return;
1665 }
1666 };
1667
1668 info!(
1669 "Flow {}: entering scheduled loop, interval={}s, start={}, anchor={}, policy={:?}, max_runs={}, max_lag={}s",
1670 flow_id_str,
1671 schedule.interval_secs,
1672 schedule.start_secs,
1673 schedule.anchor_secs,
1674 schedule.missed_tick_policy,
1675 schedule.max_runs,
1676 schedule.max_lag_secs,
1677 );
1678
1679 loop {
1680 if self.is_shutdown_signaled() {
1681 break;
1682 }
1683
1684 let wall_now_secs = wall_clock_unix_secs();
1685
1686 let due = match select_due_scheduled_times(&schedule, cursor_secs, wall_now_secs) {
1687 Ok(d) => d,
1688 Err(e) => {
1689 warn!(
1690 "Flow {}: invalid eval schedule, exiting loop: {}",
1691 flow_id_str, e
1692 );
1693 return;
1694 }
1695 };
1696
1697 if due.scheduled_times_secs.is_empty() {
1698 if due.skipped > 0 {
1699 warn!(
1700 "Flow {}: all {} due scheduled times skipped by max-lag, advancing cursor to wall-clock ({wall_now_secs}) to avoid re-skipping",
1701 flow_id_str, due.skipped
1702 );
1703 cursor_secs = wall_now_secs;
1704 continue;
1705 }
1706
1707 let next = match schedule.next_scheduled_time_after(cursor_secs) {
1709 Ok(next) => next,
1710 Err(e) => {
1711 warn!(
1712 "Flow {}: cannot advance eval schedule past cursor {cursor_secs}: {e}; exiting loop",
1713 flow_id_str
1714 );
1715 return;
1716 }
1717 };
1718 if next <= wall_now_secs {
1719 cursor_secs = wall_now_secs;
1722 continue;
1723 }
1724 let wait_secs = match sleep_delta_secs(next, wall_now_secs) {
1725 Ok(wait_secs) => wait_secs,
1726 Err(e) => {
1727 warn!(
1728 "Flow {}: cannot sleep until next scheduled time {}: {e}; exiting loop",
1729 flow_id_str, next
1730 );
1731 return;
1732 }
1733 };
1734 let wait_dur = Duration::from_secs(wait_secs);
1735 debug!(
1736 "Flow {}: no due scheduled times, sleeping for {}s until next scheduled time at {}",
1737 flow_id_str, wait_secs, next
1738 );
1739 tokio::time::sleep(wait_dur).await;
1740 continue;
1741 }
1742
1743 if due.skipped > 0 {
1744 info!(
1745 "Flow {}: {} due scheduled times, {} skipped (catch-up)",
1746 flow_id_str,
1747 due.scheduled_times_secs.len(),
1748 due.skipped
1749 );
1750 }
1751
1752 for scheduled_time_secs in &due.scheduled_times_secs {
1754 if self.is_shutdown_signaled() {
1755 break;
1756 }
1757
1758 METRIC_FLOW_BATCHING_ENGINE_START_QUERY_CNT
1759 .with_label_values(&[&flow_id_str])
1760 .inc();
1761
1762 let outcome = self
1763 .execute_once_serialized_at_scheduled_time(
1764 &engine,
1765 &frontend_client,
1766 *scheduled_time_secs,
1767 )
1768 .await;
1769
1770 cursor_secs = *scheduled_time_secs;
1772
1773 match outcome.result {
1774 Ok(Some((rows, elapsed))) => {
1775 debug!(
1776 "Flow {}: scheduled time {} completed, rows={}, elapsed={:?}",
1777 flow_id_str, scheduled_time_secs, rows, elapsed
1778 );
1779 }
1780 Ok(None) => {
1781 debug!(
1782 "Flow {}: scheduled time {} produced no query (no dirty signal or no-op)",
1783 flow_id_str, scheduled_time_secs
1784 );
1785 }
1786 Err(err) => {
1787 warn!(
1788 "Flow {}: scheduled time {} failed: {:?}",
1789 flow_id_str, scheduled_time_secs, err
1790 );
1791 METRIC_FLOW_BATCHING_ENGINE_ERROR_CNT
1792 .with_label_values(&[&flow_id_str])
1793 .inc();
1794 }
1798 }
1799 }
1800 }
1801 }
1802
1803 async fn start_adaptive_loop(
1805 &self,
1806 engine: QueryEngineRef,
1807 frontend_client: Arc<FrontendClient>,
1808 ) {
1809 let flow_id_str = self.config.flow_id.to_string();
1810 let mut max_window_cnt = None;
1811 loop {
1812 if self.is_shutdown_signaled() {
1813 break;
1814 }
1815 METRIC_FLOW_BATCHING_ENGINE_START_QUERY_CNT
1816 .with_label_values(&[&flow_id_str])
1817 .inc();
1818
1819 let min_refresh = self.config.batch_opts.experimental_min_refresh_duration;
1820
1821 let outcome = self
1822 .execute_once_serialized_with_outcome(&engine, &frontend_client, max_window_cnt)
1823 .await;
1824
1825 match outcome.result {
1826 Ok(Some(_)) => {
1827 max_window_cnt = max_window_cnt.map(|cnt| {
1828 (cnt + 1).min(self.config.batch_opts.experimental_max_filter_num_per_query)
1829 });
1830
1831 let sleep_until = {
1832 let state = self.state.write().unwrap();
1833
1834 let time_window_size = self
1835 .config
1836 .time_window_expr
1837 .as_ref()
1838 .and_then(|t| *t.time_window_size());
1839
1840 let prefer_short_incremental_cadence = state.checkpoint_mode()
1841 == CheckpointMode::Incremental
1842 && !state.is_incremental_disabled();
1843
1844 state.get_next_start_query_time(
1845 self.config.flow_id,
1846 &time_window_size,
1847 min_refresh,
1848 Some(self.config.batch_opts.query_timeout),
1849 self.config.batch_opts.experimental_max_filter_num_per_query,
1850 prefer_short_incremental_cadence,
1851 )
1852 };
1853
1854 tokio::time::sleep_until(sleep_until).await;
1855 }
1856 Ok(None) => {
1857 debug!(
1858 "Flow id = {:?} found no new data, sleep for {:?} then continue",
1859 self.config.flow_id, min_refresh
1860 );
1861 tokio::time::sleep(min_refresh).await;
1862 continue;
1863 }
1864 Err(err) => {
1865 METRIC_FLOW_BATCHING_ENGINE_ERROR_CNT
1866 .with_label_values(&[&flow_id_str])
1867 .inc();
1868 match outcome.new_query {
1869 Some(query) => {
1870 common_telemetry::error!(err; "Failed to execute query for flow={} with query: {}", self.config.flow_id, query.plan);
1871 max_window_cnt = Some(1);
1872 }
1873 None => {
1874 common_telemetry::error!(err; "Failed to generate query for flow={}", self.config.flow_id)
1875 }
1876 }
1877 tokio::time::sleep(min_refresh).await;
1878 }
1879 }
1880 }
1881 }
1882
1883 fn is_shutdown_signaled(&self) -> bool {
1885 let mut state = self.state.write().unwrap();
1886 match state.shutdown_rx.try_recv() {
1887 Ok(()) | Err(TryRecvError::Closed) => true,
1888 Err(TryRecvError::Empty) => false,
1889 }
1890 }
1891
1892 async fn execute_once_serialized_at_scheduled_time(
1899 &self,
1900 engine: &QueryEngineRef,
1901 frontend_client: &Arc<FrontendClient>,
1902 scheduled_time_secs: i64,
1903 ) -> ExecuteOnceOutcome {
1904 let mut guard = BatchingExecutionGuard::new(self.execution_lock.clone().lock_owned().await);
1905
1906 let scheduled_time_millis = match scheduled_time_millis(scheduled_time_secs) {
1910 Ok(millis) => millis,
1911 Err(e) => {
1912 return ExecuteOnceOutcome {
1913 new_query: None,
1914 result: Err(e),
1915 };
1916 }
1917 };
1918
1919 let old_ctx = {
1922 let mut state = self.state.write().unwrap();
1923 let old = state.query_ctx.clone();
1924 let mut new_ctx = (*old).clone();
1925 new_ctx.set_extension(
1926 query::options::FLOW_SCHEDULED_TIME_MILLIS,
1927 scheduled_time_millis.to_string(),
1928 );
1929 state.query_ctx = Arc::new(new_ctx);
1930 old
1931 };
1932 guard.restore_query_context(self.state.clone(), old_ctx);
1933
1934 self.execute_once_with_guard(guard, engine, frontend_client, None)
1937 .await
1938 }
1939
1940 async fn gen_create_table_expr(
1947 &self,
1948 engine: QueryEngineRef,
1949 ) -> Result<CreateTableExpr, Error> {
1950 let query_ctx = self.state.read().unwrap().query_ctx.clone();
1951 let plan =
1952 sql_to_df_plan(query_ctx.clone(), engine.clone(), &self.config.query, true).await?;
1953 create_table_with_expr(&plan, &self.config.sink_table_name, &self.config.query_type)
1954 }
1955
1956 fn should_use_unfiltered_incremental_delta(&self) -> bool {
1959 let state = self.state.read().unwrap();
1960 state.checkpoint_mode() == CheckpointMode::Incremental
1961 && !state.is_incremental_disabled()
1962 && matches!(self.config.query_type, QueryType::Sql)
1963 }
1964
1965 async fn gen_query_with_time_window(
1967 &self,
1968 engine: QueryEngineRef,
1969 sink_table_schema: &Arc<Schema>,
1970 primary_key_indices: &[usize],
1971 allow_partial: bool,
1972 max_window_cnt: Option<usize>,
1973 ) -> Result<Option<PlanInfo>, Error> {
1974 self.gen_query_with_time_window_with_values(
1975 engine,
1976 sink_table_schema,
1977 primary_key_indices,
1978 allow_partial,
1979 max_window_cnt,
1980 &BTreeMap::new(),
1981 false,
1982 )
1983 .await
1984 }
1985
1986 #[allow(clippy::too_many_arguments)]
1988 async fn gen_query_with_time_window_with_values(
1989 &self,
1990 engine: QueryEngineRef,
1991 sink_table_schema: &Arc<Schema>,
1992 primary_key_indices: &[usize],
1993 allow_partial: bool,
1994 max_window_cnt: Option<usize>,
1995 values: &BTreeMap<String, ScalarValue>,
1996 force_full_snapshot: bool,
1997 ) -> Result<Option<PlanInfo>, Error> {
1998 let query_ctx = self.state.read().unwrap().query_ctx.clone();
1999 let start = SystemTime::now();
2000 let since_the_epoch = start
2001 .duration_since(UNIX_EPOCH)
2002 .expect("Time went backwards");
2003 let low_bound = self
2004 .config
2005 .expire_after
2006 .map(|e| since_the_epoch.as_secs() - e as u64)
2007 .unwrap_or(u64::MIN);
2008
2009 let low_bound = Timestamp::new_second(low_bound as i64);
2010
2011 let expire_time_window_bound = self
2012 .config
2013 .time_window_expr
2014 .as_ref()
2015 .map(|expr| expr.eval(low_bound))
2016 .transpose()?;
2017
2018 let (expire_lower_bound, expire_upper_bound) = match (
2019 expire_time_window_bound,
2020 &self.config.query_type,
2021 ) {
2022 (Some((Some(l), Some(u))), QueryType::Sql) => (l, u),
2023 (None, QueryType::Sql) if self.config.flow_eval_interval.is_none() => {
2024 return UnexpectedSnafu {
2025 reason: format!(
2026 "Flow id={} reached execution without a time-window expression or EVAL INTERVAL; create-flow validation should have rejected it",
2027 self.config.flow_id
2028 ),
2029 }
2030 .fail();
2031 }
2032 _ => {
2033 let (_, dirty_windows_to_restore) = self.drain_dirty_windows_signal();
2039
2040 let plan_info = self
2041 .gen_unfiltered_plan_info(
2042 engine,
2043 query_ctx,
2044 sink_table_schema.clone(),
2045 primary_key_indices,
2046 allow_partial,
2047 values,
2048 dirty_windows_to_restore,
2049 None,
2050 QueryCoverage::UnfilteredFull,
2051 )
2052 .await?;
2053
2054 return Ok(Some(plan_info));
2055 }
2056 };
2057
2058 debug!(
2059 "Flow id = {:?}, found time window: precise_lower_bound={:?}, precise_upper_bound={:?} with dirty time windows: {:?}",
2060 self.config.flow_id,
2061 expire_lower_bound,
2062 expire_upper_bound,
2063 self.state.read().unwrap().dirty_time_windows
2064 );
2065 let window_size = expire_upper_bound
2066 .sub(&expire_lower_bound)
2067 .with_context(|| UnexpectedSnafu {
2068 reason: format!(
2069 "Can't get window size from {expire_upper_bound:?} - {expire_lower_bound:?}"
2070 ),
2071 })?;
2072 let col_name = self
2073 .config
2074 .time_window_expr
2075 .as_ref()
2076 .map(|expr| expr.column_name.clone())
2077 .with_context(|| UnexpectedSnafu {
2078 reason: format!(
2079 "Flow id={:?}, Failed to get column name from time window expr",
2080 self.config.flow_id
2081 ),
2082 })?;
2083
2084 if force_full_snapshot {
2085 let retention_filter = self.config.expire_after.map(|_| {
2089 (
2090 col_name.as_str(),
2091 expire_lower_bound,
2092 "forced full snapshot",
2093 )
2094 });
2095 let (_, dirty_windows_to_restore) = self.drain_dirty_windows_signal();
2096 return self
2097 .gen_unfiltered_plan_info(
2098 engine,
2099 query_ctx,
2100 sink_table_schema.clone(),
2101 primary_key_indices,
2102 allow_partial,
2103 values,
2104 dirty_windows_to_restore,
2105 retention_filter,
2106 QueryCoverage::UnfilteredFull,
2107 )
2108 .await
2109 .map(Some);
2110 }
2111
2112 if self.should_use_unfiltered_incremental_delta() {
2113 let retention_filter = self
2123 .config
2124 .expire_after
2125 .map(|_| (col_name.as_str(), expire_lower_bound, "incremental"));
2126 return self
2127 .gen_unfiltered_plan_info_if_dirty(
2128 engine,
2129 query_ctx,
2130 sink_table_schema.clone(),
2131 primary_key_indices,
2132 allow_partial,
2133 values,
2134 retention_filter,
2135 QueryCoverage::IncrementalDelta,
2136 )
2137 .await;
2138 }
2139
2140 let (expr, coverage) = {
2141 let mut state = self.state.write().unwrap();
2142 let window_cnt = max_window_cnt
2143 .unwrap_or(self.config.batch_opts.experimental_max_filter_num_per_query);
2144 let expr = state.gen_scoped_filter_exprs(
2145 &col_name,
2146 Some(expire_lower_bound),
2147 window_size,
2148 window_cnt,
2149 self.config.flow_id,
2150 Some(self),
2151 )?;
2152 let repair_high = state
2153 .pending_fenced_repair()
2154 .map(|repair| repair.high().clone());
2155 let coverage = if let Some(high) = repair_high {
2156 QueryCoverage::FencedRepairChunk { high }
2157 } else {
2158 QueryCoverage::ScopedBaseRepair
2159 };
2160 (expr, coverage)
2161 };
2162
2163 let Some(expr) = expr else {
2164 debug!("Flow id={:?}, no new data, not update", self.config.flow_id);
2166 return Ok(None);
2167 };
2168
2169 let filter_sql = expr_to_sql(&expr.expr)
2170 .map(|sql| sql.to_string())
2171 .unwrap_or_else(|err| format!("<failed to format filter expr: {err}>"));
2172
2173 debug!(
2174 "Flow id={:?}, Generated filter expr: {:?}",
2175 self.config.flow_id, filter_sql
2176 );
2177
2178 let mut add_filter = AddFilterRewriter::new(expr.expr.clone());
2179 let mut add_auto_column = ColumnMatcherRewriter::new_with_values(
2180 sink_table_schema.clone(),
2181 primary_key_indices.to_vec(),
2182 allow_partial,
2183 values.clone(),
2184 );
2185
2186 let plan = self.restore_scoped_dirty_windows_on_err(
2187 &expr,
2188 sql_to_df_plan(query_ctx.clone(), engine.clone(), &self.config.query, false).await,
2189 )?;
2190 let rewrite = self.restore_scoped_dirty_windows_on_err(
2191 &expr,
2192 plan.clone()
2193 .rewrite(&mut add_filter)
2194 .and_then(|p| p.data.rewrite(&mut add_auto_column))
2195 .with_context(|_| DatafusionSnafu {
2196 context: format!("Failed to rewrite plan:\n {}\n", plan),
2197 })
2198 .map(|rewrite| rewrite.data),
2199 )?;
2200 let new_plan = self.restore_scoped_dirty_windows_on_err(
2202 &expr,
2203 apply_df_optimizer(rewrite, &query_ctx).await,
2204 )?;
2205
2206 let info = PlanInfo {
2207 plan: new_plan.clone(),
2208 dirty_restore: DirtyRestore::Scoped(expr),
2209 coverage,
2210 };
2211
2212 Ok(Some(info))
2213 }
2214}
2215
2216#[cfg(test)]
2217mod test;