diff --git a/src/flow/src/batching_mode.rs b/src/flow/src/batching_mode.rs index 3981446721..60e7969156 100644 --- a/src/flow/src/batching_mode.rs +++ b/src/flow/src/batching_mode.rs @@ -20,10 +20,19 @@ use common_grpc::channel_manager::ClientTlsOption; use serde::{Deserialize, Serialize}; use session::ReadPreference; +/// Runtime-only mode used for incremental source scans. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum IncrementalMode { + #[default] + MemtableOnly, + SequenceRange, +} + mod checkpoint; pub(crate) mod engine; mod eval_schedule; pub(crate) mod frontend_client; +pub(crate) mod persistence; mod state; mod table_creator; mod task; diff --git a/src/flow/src/batching_mode/checkpoint.rs b/src/flow/src/batching_mode/checkpoint.rs index 11be951a9a..f650f4cd5a 100644 --- a/src/flow/src/batching_mode/checkpoint.rs +++ b/src/flow/src/batching_mode/checkpoint.rs @@ -87,6 +87,13 @@ pub(super) enum FlowCheckpointDecision { participating_regions: usize, watermarks: usize, }, + /// The full repair requested by persistence completed and proved a new + /// checkpoint. This is distinct from an ordinary full snapshot because it + /// is the only decision that clears the repair request. + CompletedFullRepair { + participating_regions: usize, + watermarks: usize, + }, /// FullSnapshot stayed in full snapshot mode because a scoped base repair /// found additional dirty windows that may be concurrent with the returned /// high watermark. These windows must be repaired under the fixed high @@ -114,11 +121,9 @@ impl FlowCheckpointDecision { checkpoint_mode_label(CheckpointMode::FullSnapshot) } Self::AdvancedIncremental { .. } => checkpoint_mode_label(CheckpointMode::Incremental), - // Fenced repair is intentionally a FullSnapshot sub-state, not a - // third top-level checkpoint mode, so metrics keep the - // `full_snapshot` mode label while the decision label carries - // `continue_repair`. - Self::ContinuedFencedRepair { .. } => { + Self::CompletedFullRepair { .. } | Self::ContinuedFencedRepair { .. } => { + // Fenced repair and completion of a requested full repair are + // FullSnapshot sub-states, not third top-level modes. checkpoint_mode_label(CheckpointMode::FullSnapshot) } Self::FallbackToFullSnapshot { previous_mode, .. } => { @@ -129,9 +134,9 @@ impl FlowCheckpointDecision { pub(super) fn decision_label(self) -> &'static str { match self { - Self::AdvancedFromFullSnapshot { .. } | Self::AdvancedIncremental { .. } => { - CHECKPOINT_DECISION_ADVANCE - } + Self::AdvancedFromFullSnapshot { .. } + | Self::AdvancedIncremental { .. } + | Self::CompletedFullRepair { .. } => CHECKPOINT_DECISION_ADVANCE, Self::ContinuedFencedRepair { .. } => CHECKPOINT_DECISION_CONTINUE_REPAIR, Self::FallbackToFullSnapshot { .. } => CHECKPOINT_DECISION_FALLBACK, } diff --git a/src/flow/src/batching_mode/engine.rs b/src/flow/src/batching_mode/engine.rs index d05eec9251..943092d0fd 100644 --- a/src/flow/src/batching_mode/engine.rs +++ b/src/flow/src/batching_mode/engine.rs @@ -46,6 +46,7 @@ use tokio::sync::{RwLock, oneshot}; use crate::batching_mode::BatchingModeOptions; use crate::batching_mode::eval_schedule::EvalSchedule; use crate::batching_mode::frontend_client::FrontendClient; +use crate::batching_mode::persistence::{FrontendBatchingQueryExecutor, PersistenceContext}; use crate::batching_mode::state::DirtyTimeWindows; use crate::batching_mode::task::{BatchingTask, TaskArgs}; use crate::batching_mode::time_window::{TimeWindowExpr, find_time_window_expr}; @@ -72,6 +73,7 @@ pub struct BatchingEngine { /// Batching mode options for control how batching mode query works /// pub(crate) batch_opts: Arc, + persistence_factory: Option, } #[derive(Default)] @@ -127,6 +129,26 @@ impl BatchingEngine { table_meta: TableMetadataManagerRef, catalog_manager: CatalogManagerRef, batch_opts: BatchingModeOptions, + ) -> Self { + Self::new_with_persistence( + frontend_client, + query_engine, + flow_metadata_manager, + table_meta, + catalog_manager, + batch_opts, + None, + ) + } + + pub fn new_with_persistence( + frontend_client: Arc, + query_engine: QueryEngineRef, + flow_metadata_manager: FlowMetadataManagerRef, + table_meta: TableMetadataManagerRef, + catalog_manager: CatalogManagerRef, + batch_opts: BatchingModeOptions, + persistence_factory: Option, ) -> Self { Self { runtime: Default::default(), @@ -136,6 +158,7 @@ impl BatchingEngine { catalog_manager, query_engine, batch_opts: Arc::new(batch_opts), + persistence_factory, } } @@ -491,6 +514,14 @@ impl BatchingEngine { .is_some_and(|value| value.eq_ignore_ascii_case("true")) } + fn table_options_enable_merge_mode_last_non_null( + extra_options: &HashMap, + ) -> bool { + extra_options + .get(store_api::mito_engine_options::MERGE_MODE_KEY) + .is_some_and(|value| value.eq_ignore_ascii_case("last_non_null")) + } + /// SQL flows without a usable time-window expression can only run as an /// explicit full-query flow, so require `EVAL INTERVAL` at creation time. fn ensure_sql_flow_has_twe_or_eval_interval( @@ -703,11 +734,58 @@ impl BatchingEngine { let engine = self.query_engine.clone(); let frontend = self.frontend_client.clone(); - // Create sink table if needed, then validate an existing/created sink schema before - // spawning the background task. This catches user-created sink schema mismatches at - // CREATE FLOW time instead of surfacing them later in the execution loop. - task.check_or_create_sink_table(&engine, &frontend).await?; - task.validate_sink_table_schema(&engine).await?; + // Create the sink before configuring persistence. A persistence-backed sink may + // contain ordinary metadata columns supplied by `begin_attempt`, so strict plan/schema + // validation is deferred to execution for that path. OSS flows without a factory keep + // the existing creation-time validation. + let table = task.check_or_create_sink_table(&engine, &frontend).await?; + + let persistence = if let Some(factory) = &self.persistence_factory { + let table_info = table.table_info(); + let meta = &table_info.meta; + let effective_mode = if task.config.batch_opts.experimental_enable_incremental_read + && task.sequence_range_capable().await + { + crate::IncrementalMode::SequenceRange + } else { + crate::IncrementalMode::MemtableOnly + }; + let context = PersistenceContext { + flow_id, + sink: crate::batching_mode::persistence::SinkLayout { + table_id: table_info.table_id(), + table_name: task.config.sink_table_name.clone(), + engine: meta.engine.clone(), + append: Self::table_options_enable_append_mode(&meta.options.extra_options), + merge_mode: if Self::table_options_enable_merge_mode_last_non_null( + &meta.options.extra_options, + ) { + Some("last_non_null".to_string()) + } else { + None + }, + columns: meta + .schema + .column_schemas() + .iter() + .map(|column| crate::BatchingMetadataColumn { + name: column.name.clone(), + data_type: column.data_type.clone(), + nullable: column.is_nullable(), + }) + .collect(), + ordered_primary_key_indices: meta.primary_key_indices.clone(), + time_index: meta.schema.timestamp_index(), + }, + executor: Arc::new(FrontendBatchingQueryExecutor::new(frontend.clone())), + incremental_mode: effective_mode, + }; + factory.create(context).await? + } else { + task.validate_sink_table_schema(&engine).await?; + None + }; + task.set_persistence(persistence).await?; let (start_tx, start_rx) = oneshot::channel(); diff --git a/src/flow/src/batching_mode/persistence.rs b/src/flow/src/batching_mode/persistence.rs new file mode 100644 index 0000000000..1241e6e324 --- /dev/null +++ b/src/flow/src/batching_mode/persistence.rs @@ -0,0 +1,141 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Optional collaborator boundary for batching state. +//! +//! The batching engine only deals in the small, typed values described here. It +//! does not interpret or persist any of them. + +use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; + +use common_error::ext::BoxedError; +use common_query::Output; +use datafusion_common::ScalarValue; +use datatypes::prelude::ConcreteDataType; +use store_api::storage::TableId; + +use crate::Result; +use crate::batching_mode::IncrementalMode; + +/// SQL execution exposed to a persistence implementation for reading and +/// writing internal state. The batching engine does not interpret this SQL; +/// the persistence implementation owns the SQL and its schema. +#[async_trait::async_trait] +pub trait BatchingQueryExecutor: Send + Sync + 'static { + async fn execute_sql(&self, catalog: &str, schema: &str, sql: &str) -> Result; +} + +/// Describes the validated sink table available to a persistence collaborator. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SinkLayout { + pub table_id: TableId, + pub table_name: [String; 3], + pub engine: String, + pub append: bool, + pub merge_mode: Option, + pub columns: Vec, + pub ordered_primary_key_indices: Vec, + pub time_index: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BatchingMetadataColumn { + pub name: String, + pub data_type: ConcreteDataType, + pub nullable: bool, +} + +/// Result of restoring a task's opaque state. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RestoreOutcome { + TrustedCheckpoint(BTreeMap), + FullRepair, +} + +/// One serialized execution attempt. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct BatchingAttempt { + pub ordinary_values: BTreeMap, +} + +#[async_trait::async_trait] +pub trait BatchingPersistence: Send + Sync + 'static { + async fn restore(&self) -> Result; + async fn begin_attempt(&self) -> Result; + async fn persist( + &self, + attempt: BatchingAttempt, + validated_checkpoints: BTreeMap, + ) -> Result<()>; +} + +/// Context supplied to the typed persistence factory for one batching flow. +#[derive(Clone)] +pub struct PersistenceContext { + pub flow_id: crate::FlowId, + pub incremental_mode: IncrementalMode, + pub sink: SinkLayout, + pub executor: Arc, +} + +/// Typed factory implementation used by the batching engine. +#[async_trait::async_trait] +pub trait Factory: Send + Sync + 'static { + async fn create( + &self, + context: PersistenceContext, + ) -> Result>>; +} + +/// Typed plugin wrapper around a persistence factory. +#[derive(Clone)] +pub struct FactoryPlugin(pub Arc); + +impl std::ops::Deref for FactoryPlugin { + type Target = dyn Factory; + + fn deref(&self) -> &Self::Target { + self.0.as_ref() + } +} + +/// Query executor backed by the existing flownode frontend client. +pub(crate) struct FrontendBatchingQueryExecutor { + client: Arc, +} + +impl FrontendBatchingQueryExecutor { + pub(crate) fn new(client: Arc) -> Self { + Self { client } + } +} + +#[async_trait::async_trait] +impl BatchingQueryExecutor for FrontendBatchingQueryExecutor { + async fn execute_sql(&self, catalog: &str, schema: &str, sql: &str) -> Result { + let mut peer = None; + let request = api::v1::QueryRequest { + query: Some(api::v1::query_request::Query::Sql(sql.to_string())), + }; + self.client + .query_with_terminal_metrics(catalog, schema, request, &[], &HashMap::new(), &mut peer) + .await + .map(|output| output.into_output()) + .map_err(|err| crate::Error::External { + source: BoxedError::new(err), + location: snafu::location!(), + }) + } +} diff --git a/src/flow/src/batching_mode/state.rs b/src/flow/src/batching_mode/state.rs index 59c4f0fe10..84d4d7b022 100644 --- a/src/flow/src/batching_mode/state.rs +++ b/src/flow/src/batching_mode/state.rs @@ -35,8 +35,19 @@ use crate::metrics::{ }; use crate::{Error, FlowId}; -/// The state of the [`BatchingTask`]. -#[derive(Debug)] +#[derive(Debug, Clone)] +pub(crate) struct TaskStateCheckpointSnapshot { + pub(crate) checkpoint_mode: CheckpointMode, + pub(crate) checkpoints: BTreeMap, + pub(crate) pending_fenced_repair: Option, + pub(crate) full_repair_required: bool, + pub(crate) dirty_time_windows: DirtyTimeWindows, + pub(crate) last_update_time: Instant, + pub(crate) last_query_duration: Duration, + pub(crate) last_exec_time_millis: Option, + pub(crate) exec_state: ExecState, +} + pub struct TaskState { /// Query context pub(crate) query_ctx: QueryContextRef, @@ -53,6 +64,7 @@ pub struct TaskState { pub(crate) dirty_time_windows: DirtyTimeWindows, checkpoint_mode: CheckpointMode, pending_fenced_repair: Option, + full_repair_required: bool, /// Region id -> last consumed watermark sequence. Incremental scans use /// this as the next lower sequence bound for each source region. checkpoints: BTreeMap, @@ -85,6 +97,7 @@ impl TaskState { dirty_time_windows, checkpoint_mode: CheckpointMode::FullSnapshot, pending_fenced_repair: None, + full_repair_required: false, checkpoints: Default::default(), incremental_disabled: false, exec_state: ExecState::Idle, @@ -114,6 +127,10 @@ impl TaskState { } } + pub(crate) fn last_query_duration(&self) -> Duration { + self.last_query_duration + } + pub fn last_execution_time_millis(&self) -> Option { self.last_exec_time_millis } @@ -131,6 +148,64 @@ impl TaskState { &self.checkpoints } + pub(crate) fn checkpoint_snapshot(&self) -> TaskStateCheckpointSnapshot { + TaskStateCheckpointSnapshot { + checkpoint_mode: self.checkpoint_mode, + checkpoints: self.checkpoints.clone(), + pending_fenced_repair: self.pending_fenced_repair.clone(), + full_repair_required: self.full_repair_required, + dirty_time_windows: self.dirty_time_windows.clone(), + last_update_time: self.last_update_time, + last_query_duration: self.last_query_duration, + last_exec_time_millis: self.last_exec_time_millis, + exec_state: self.exec_state.clone(), + } + } + + pub(crate) fn restore_checkpoint_snapshot(&mut self, snapshot: TaskStateCheckpointSnapshot) { + let live_dirty = self.dirty_time_windows.clone(); + self.checkpoint_mode = snapshot.checkpoint_mode; + self.checkpoints = snapshot.checkpoints; + self.pending_fenced_repair = snapshot.pending_fenced_repair; + self.full_repair_required = snapshot.full_repair_required; + self.dirty_time_windows = snapshot.dirty_time_windows; + self.dirty_time_windows.add_dirty_windows(&live_dirty); + self.last_update_time = snapshot.last_update_time; + self.last_query_duration = snapshot.last_query_duration; + self.last_exec_time_millis = snapshot.last_exec_time_millis; + self.exec_state = snapshot.exec_state; + } + + pub(crate) fn request_full_repair(&mut self) { + self.full_repair_required = true; + } + + pub(crate) fn full_repair_required(&self) -> bool { + self.full_repair_required + } + + pub(crate) fn commit_checkpoint_candidate( + &mut self, + candidate: &TaskStateCheckpointSnapshot, + successful_full_repair: bool, + ) { + let live_dirty = self.dirty_time_windows.clone(); + self.checkpoint_mode = candidate.checkpoint_mode; + self.checkpoints = candidate.checkpoints.clone(); + self.pending_fenced_repair = candidate.pending_fenced_repair.clone(); + self.full_repair_required = if successful_full_repair { + false + } else { + candidate.full_repair_required + }; + self.dirty_time_windows = candidate.dirty_time_windows.clone(); + self.dirty_time_windows.add_dirty_windows(&live_dirty); + self.last_update_time = candidate.last_update_time; + self.last_query_duration = candidate.last_query_duration; + self.last_exec_time_millis = candidate.last_exec_time_millis; + self.exec_state = candidate.exec_state.clone(); + } + /// Returns the in-progress fenced repair, if the task is repairing dirty /// windows under a frozen full-snapshot high watermark. pub fn pending_fenced_repair(&self) -> Option<&FencedRepair> { @@ -551,6 +626,15 @@ impl DirtyTimeWindows { self.windows.clear(); } + pub(crate) fn detach(&mut self) -> Self { + let mut detached = Self::new( + self.max_filter_num_per_query, + self.time_window_merge_threshold, + ); + std::mem::swap(&mut detached.windows, &mut self.windows); + detached + } + /// Set windows to be dirty, only useful for full aggr without time window /// to mark some new data is inserted pub fn set_dirty(&mut self) { @@ -907,7 +991,7 @@ pub(crate) fn to_df_literal(value: Timestamp) -> Result>, + persistence: + Arc>>>, } /// Arguments for creating batching task @@ -267,6 +272,7 @@ impl QueryCoverage { } } +#[derive(Clone)] pub enum DirtyRestore { /// The query was scoped to dirty time ranges; restore those ranges if the /// run fails. @@ -278,6 +284,7 @@ pub enum DirtyRestore { /// TODO(discord9): Full-query runs only need a dirty bool flag. Refactor /// the unscoped path to stop reusing `DirtyTimeWindows` for this signal. Unscoped(DirtyTimeWindows), + FullRepair(DirtyTimeWindows), } struct ExecuteOnceOutcome { @@ -338,9 +345,33 @@ impl BatchingTask { }), state: Arc::new(RwLock::new(state)), execution_lock: Arc::new(Mutex::new(())), + persistence: Arc::new(RwLock::new(None)), }) } + pub(crate) async fn set_persistence( + &self, + persistence: Option>, + ) -> Result<(), Error> { + if let Some(persistence) = &persistence { + let outcome = persistence.restore().await?; + let mut state = self.state.write().unwrap(); + match outcome { + RestoreOutcome::TrustedCheckpoint(checkpoints) => { + if !checkpoints.is_empty() { + state.advance_checkpoints(checkpoints.into_iter().collect()); + } + } + RestoreOutcome::FullRepair => { + state.mark_full_snapshot(); + state.request_full_repair(); + } + } + } + *self.persistence.write().unwrap() = persistence; + Ok(()) + } + pub fn last_execution_time_millis(&self) -> Option { self.state.read().unwrap().last_execution_time_millis() } @@ -396,26 +427,34 @@ impl BatchingTask { Ok(()) } - /// Create sink table if not exists + /// Create the sink table if needed, then return it. pub async fn check_or_create_sink_table( &self, engine: &QueryEngineRef, frontend_client: &Arc, - ) -> Result, Error> { - if !self.is_table_exist(&self.config.sink_table_name).await? { + ) -> Result { + if !self + .config + .catalog_manager + .table_exists( + &self.config.sink_table_name[0], + &self.config.sink_table_name[1], + &self.config.sink_table_name[2], + None, + ) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)? + { let create_table = self.gen_create_table_expr(engine.clone()).await?; - info!( - "Try creating sink table(if not exists) with expr: {:?}", - create_table - ); self.create_table(frontend_client, create_table).await?; - info!( - "Sink table {}(if not exists) created", - self.config.sink_table_name.join(".") - ); } - - Ok(None) + let (table, _) = get_table_info_df_schema( + self.config.catalog_manager.clone(), + self.config.sink_table_name.clone(), + ) + .await?; + Ok(table) } /// Validates that the sink table schema can accept this flow's output. @@ -423,19 +462,29 @@ impl BatchingTask { /// This is a dry-run of the same schema matching logic used by insert-plan /// generation, but without adding dirty-window filters or executing the query. It is used /// during CREATE FLOW to catch existing sink table mismatches early. - pub async fn validate_sink_table_schema(&self, engine: &QueryEngineRef) -> Result<(), Error> { + pub async fn validate_sink_table_schema( + &self, + engine: &QueryEngineRef, + ) -> Result, Error> { let (table, _) = get_table_info_df_schema( self.config.catalog_manager.clone(), self.config.sink_table_name.clone(), ) .await?; + self.validate_sink_table_schema_with_table(engine, table) + .await + } + async fn validate_sink_table_schema_with_table( + &self, + engine: &QueryEngineRef, + table: TableRef, + ) -> Result, Error> { let table_meta = &table.table_info().meta; let merge_mode_last_non_null = is_merge_mode_last_non_null(&table_meta.options.extra_options); let primary_key_indices = table_meta.primary_key_indices.clone(); let query_ctx = self.state.read().unwrap().query_ctx.clone(); - gen_plan_with_matching_schema( &self.config.query, query_ctx, @@ -445,84 +494,7 @@ impl BatchingTask { merge_mode_last_non_null, ) .await - .map(|_| ()) - } - - async fn is_table_exist(&self, table_name: &[String; 3]) -> Result { - self.config - .catalog_manager - .table_exists(&table_name[0], &table_name[1], &table_name[2], None) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu) - } - - pub(crate) async fn execute_once_serialized( - &self, - engine: &QueryEngineRef, - frontend_client: &Arc, - max_window_cnt: Option, - ) -> Result, Error> { - let outcome = self - .execute_once_serialized_with_outcome(engine, frontend_client, max_window_cnt) - .await; - outcome.result - } - - /// Executes one flow evaluation under `execution_lock` and keeps the - /// generated query context for the background loop's error logging/backoff. - async fn execute_once_serialized_with_outcome( - &self, - engine: &QueryEngineRef, - frontend_client: &Arc, - max_window_cnt: Option, - ) -> ExecuteOnceOutcome { - let _execution_guard = self.execution_lock.lock().await; - self.execute_once_unlocked(engine, frontend_client, max_window_cnt) - .await - } - - /// Executes one flow evaluation. Caller must hold `execution_lock`. - async fn execute_once_unlocked( - &self, - engine: &QueryEngineRef, - frontend_client: &Arc, - max_window_cnt: Option, - ) -> ExecuteOnceOutcome { - let new_query = match self.gen_insert_plan_unlocked(engine, max_window_cnt).await { - Ok(new_query) => new_query, - Err(err) => { - return ExecuteOnceOutcome { - new_query: None, - result: Err(err), - }; - } - }; - - if let Some(new_query) = new_query { - debug!("Generate new query: {}", new_query.plan); - let res = self - .execute_logical_plan_unlocked( - frontend_client, - &new_query.plan, - &new_query.dirty_restore, - &new_query.coverage, - ) - .await; - if res.is_err() { - self.handle_executed_query_failure(Some(&new_query)); - } - ExecuteOnceOutcome { - new_query: Some(new_query), - result: res, - } - } else { - debug!("Generate no query"); - ExecuteOnceOutcome { - new_query: None, - result: Ok(None), - } - } + .map(|_| table_meta.schema.clone()) } /// Generates the insert plan. Caller must reach this through the serialized path. @@ -530,6 +502,16 @@ impl BatchingTask { &self, engine: &QueryEngineRef, max_window_cnt: Option, + ) -> Result, Error> { + self.gen_insert_plan_unlocked_with_attempt(engine, max_window_cnt, None) + .await + } + + async fn gen_insert_plan_unlocked_with_attempt( + &self, + engine: &QueryEngineRef, + max_window_cnt: Option, + attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, ) -> Result, Error> { let (table, df_schema) = get_table_info_df_schema( self.config.catalog_manager.clone(), @@ -543,12 +525,13 @@ impl BatchingTask { let primary_key_indices = table_meta.primary_key_indices.clone(); let new_query = self - .gen_query_with_time_window( + .gen_query_with_time_window_with_attempt( engine.clone(), &table.table_info().meta.schema, &primary_key_indices, merge_mode_last_non_null, max_window_cnt, + attempt, ) .await?; @@ -631,12 +614,60 @@ impl BatchingTask { } /// Executes the insert plan. Caller must reach this through the serialized path. + async fn persist_checkpoint_candidate( + &self, + snapshot: TaskStateCheckpointSnapshot, + candidate: TaskStateCheckpointSnapshot, + attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, + dirty_restore: DirtyRestore, + ) -> Result<(), Error> { + let persistence = self.persistence.read().unwrap().as_ref().cloned(); + let Some(persistence) = persistence else { + self.state.write().unwrap().commit_checkpoint_candidate( + &candidate, + matches!(dirty_restore, DirtyRestore::FullRepair(_)), + ); + return Ok(()); + }; + let Some(attempt) = attempt else { + self.state.write().unwrap().commit_checkpoint_candidate( + &candidate, + matches!(dirty_restore, DirtyRestore::FullRepair(_)), + ); + return Ok(()); + }; + + let successful_full_repair = matches!(dirty_restore, DirtyRestore::FullRepair(_)); + if let Err(source) = persistence + .persist(attempt.clone(), candidate.checkpoints.clone()) + .await + { + let mut state = self.state.write().unwrap(); + state.restore_checkpoint_snapshot(snapshot); + match dirty_restore { + DirtyRestore::Scoped(filter) => state.restore_scoped_windows(&filter), + DirtyRestore::Unscoped(windows) | DirtyRestore::FullRepair(windows) => { + state.dirty_time_windows.add_dirty_windows(&windows) + } + } + state.mark_full_snapshot(); + state.request_full_repair(); + return Err(source); + } + self.state + .write() + .unwrap() + .commit_checkpoint_candidate(&candidate, successful_full_repair); + Ok(()) + } + async fn execute_logical_plan_unlocked( &self, frontend_client: &Arc, plan: &LogicalPlan, dirty_restore: &DirtyRestore, coverage: &QueryCoverage, + attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, ) -> Result, Error> { let instant = Instant::now(); let flow_id = self.config.flow_id; @@ -834,10 +865,36 @@ impl BatchingTask { METRIC_FLOW_ROWS .with_label_values(&[format!("{}-out-batching", flow_id).as_str()]) .inc_by(affected_rows as _); - let decision = { + let (decision, checkpoint_txn) = { let mut state = self.state.write().unwrap(); - Self::apply_query_result_to_state(&mut state, &res, elapsed, coverage) + let snapshot = state.checkpoint_snapshot(); + let repair_required = snapshot.full_repair_required; + let decision = Self::apply_query_result_to_state_with_repair( + &mut state, + &res, + elapsed, + coverage, + repair_required, + ); + let eligible = matches!( + decision, + crate::batching_mode::checkpoint::FlowCheckpointDecision::AdvancedFromFullSnapshot { .. } + | crate::batching_mode::checkpoint::FlowCheckpointDecision::AdvancedIncremental { .. } + | crate::batching_mode::checkpoint::FlowCheckpointDecision::CompletedFullRepair { .. } + ); + if eligible { + let candidate = state.checkpoint_snapshot(); + state.restore_checkpoint_snapshot(snapshot.clone()); + Some((snapshot, candidate)) + } else { + None + } + .map_or((decision, None), |txn| (decision, Some(txn))) }; + if let Some((snapshot, candidate)) = checkpoint_txn { + self.persist_checkpoint_candidate(snapshot, candidate, attempt, dirty_restore.clone()) + .await?; + } Self::record_checkpoint_decision(flow_id, decision); Ok(Some((affected_rows, elapsed))) @@ -849,7 +906,7 @@ impl BatchingTask { fn restore_dirty_windows(&self, dirty_restore: &DirtyRestore) { match dirty_restore { DirtyRestore::Scoped(filter) => self.restore_scoped_dirty_windows(filter), - DirtyRestore::Unscoped(dirty_windows) => self + DirtyRestore::Unscoped(dirty_windows) | DirtyRestore::FullRepair(dirty_windows) => self .state .write() .unwrap() @@ -927,16 +984,18 @@ impl BatchingTask { dirty_windows_to_restore: DirtyTimeWindows, retention_filter: Option<(&str, Timestamp, &'static str)>, coverage: QueryCoverage, + attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, ) -> Result { let mut plan = self.restore_unscoped_dirty_windows_on_err( &dirty_windows_to_restore, - gen_plan_with_matching_schema( + gen_plan_with_matching_schema_and_values( &self.config.query, query_ctx, engine, sink_table_schema, primary_key_indices, allow_partial, + attempt.map(|a| &a.ordinary_values), ) .await, )?; @@ -964,7 +1023,11 @@ impl BatchingTask { Ok(PlanInfo { plan, - dirty_restore: DirtyRestore::Unscoped(dirty_windows_to_restore), + dirty_restore: if self.state.read().unwrap().full_repair_required() { + DirtyRestore::FullRepair(dirty_windows_to_restore) + } else { + DirtyRestore::Unscoped(dirty_windows_to_restore) + }, coverage, }) } @@ -981,6 +1044,7 @@ impl BatchingTask { allow_partial: bool, retention_filter: Option<(&str, Timestamp, &'static str)>, coverage: QueryCoverage, + attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, ) -> Result, Error> { let (is_dirty, dirty_windows_to_restore) = self.drain_dirty_windows_signal(); if !is_dirty { @@ -997,6 +1061,7 @@ impl BatchingTask { dirty_windows_to_restore, retention_filter, coverage, + attempt, ) .await .map(Some) @@ -1305,6 +1370,89 @@ impl BatchingTask { } } + pub(crate) async fn execute_once_serialized( + &self, + engine: &QueryEngineRef, + frontend_client: &Arc, + max_window_cnt: Option, + ) -> Result, Error> { + self.execute_once_serialized_with_outcome(engine, frontend_client, max_window_cnt) + .await + .result + } + + async fn execute_once_serialized_with_outcome( + &self, + engine: &QueryEngineRef, + frontend_client: &Arc, + max_window_cnt: Option, + ) -> ExecuteOnceOutcome { + let _execution_guard = self.execution_lock.lock().await; + self.execute_once_unlocked(engine, frontend_client, max_window_cnt) + .await + } + + async fn execute_once_unlocked( + &self, + engine: &QueryEngineRef, + frontend_client: &Arc, + max_window_cnt: Option, + ) -> ExecuteOnceOutcome { + let persistence = self.persistence.read().unwrap().as_ref().cloned(); + let attempt = match persistence { + Some(state) => match state.begin_attempt().await { + Ok(attempt) => Some(attempt), + Err(source) => { + return ExecuteOnceOutcome { + new_query: None, + result: Err(source), + }; + } + }, + None => None, + }; + let new_query = match self + .gen_insert_plan_unlocked_with_attempt(engine, max_window_cnt, attempt.as_ref()) + .await + { + Ok(new_query) => new_query, + Err(err) => { + return ExecuteOnceOutcome { + new_query: None, + result: Err(err), + }; + } + }; + let Some(new_query) = new_query else { + return ExecuteOnceOutcome { + new_query: None, + result: Ok(None), + }; + }; + let res = self + .execute_logical_plan_unlocked( + frontend_client, + &new_query.plan, + &new_query.dirty_restore, + &new_query.coverage, + attempt.as_ref(), + ) + .await; + match res { + Ok(result) => ExecuteOnceOutcome { + new_query: Some(new_query), + result: Ok(result), + }, + Err(err) => { + self.handle_executed_query_failure(Some(&new_query)); + ExecuteOnceOutcome { + new_query: Some(new_query), + result: Err(err), + } + } + } + } + /// Check whether the shutdown signal has been received. fn is_shutdown_signaled(&self) -> bool { let mut state = self.state.write().unwrap(); @@ -1421,6 +1569,26 @@ impl BatchingTask { primary_key_indices: &[usize], allow_partial: bool, max_window_cnt: Option, + ) -> Result, Error> { + self.gen_query_with_time_window_with_attempt( + engine, + sink_table_schema, + primary_key_indices, + allow_partial, + max_window_cnt, + None, + ) + .await + } + + async fn gen_query_with_time_window_with_attempt( + &self, + engine: QueryEngineRef, + sink_table_schema: &Arc, + primary_key_indices: &[usize], + allow_partial: bool, + max_window_cnt: Option, + attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>, ) -> Result, Error> { let query_ctx = self.state.read().unwrap().query_ctx.clone(); let start = SystemTime::now(); @@ -1442,6 +1610,33 @@ impl BatchingTask { .map(|expr| expr.eval(low_bound)) .transpose()?; + let full_repair = self.state.read().unwrap().full_repair_required(); + if full_repair { + let detached = self.state.write().unwrap().dirty_time_windows.detach(); + let retention_filter = self.config.expire_after.and_then(|_| { + self.config.time_window_expr.as_ref().and_then(|expr| { + expr.eval(low_bound) + .ok() + .and_then(|(lower, _)| lower) + .map(|lower| (expr.column_name.as_str(), lower, "full repair retention")) + }) + }); + return self + .gen_unfiltered_plan_info( + engine, + query_ctx, + sink_table_schema.clone(), + primary_key_indices, + allow_partial, + detached, + retention_filter, + QueryCoverage::UnfilteredFull, + attempt, + ) + .await + .map(Some); + } + let (expire_lower_bound, expire_upper_bound) = match ( expire_time_window_bound, &self.config.query_type, @@ -1474,6 +1669,7 @@ impl BatchingTask { dirty_windows_to_restore, None, QueryCoverage::UnfilteredFull, + attempt, ) .await?; @@ -1530,6 +1726,7 @@ impl BatchingTask { allow_partial, retention_filter, QueryCoverage::IncrementalDelta, + attempt, ) .await; } @@ -1573,10 +1770,13 @@ impl BatchingTask { ); let mut add_filter = AddFilterRewriter::new(expr.expr.clone()); - let mut add_auto_column = ColumnMatcherRewriter::new( + let mut add_auto_column = ColumnMatcherRewriter::new_with_values( sink_table_schema.clone(), primary_key_indices.to_vec(), allow_partial, + attempt + .map(|a| a.ordinary_values.clone()) + .unwrap_or_default(), ); let plan = self.restore_scoped_dirty_windows_on_err( diff --git a/src/flow/src/batching_mode/task/ckpt.rs b/src/flow/src/batching_mode/task/ckpt.rs index 960e541c73..c9e8716481 100644 --- a/src/flow/src/batching_mode/task/ckpt.rs +++ b/src/flow/src/batching_mode/task/ckpt.rs @@ -138,6 +138,20 @@ impl BatchingTask { res: &OutputWithMetrics, elapsed: Duration, coverage: &QueryCoverage, + ) -> FlowCheckpointDecision { + Self::apply_query_result_to_state_with_repair(state, res, elapsed, coverage, false) + } + + /// Apply checkpoint transitions while retaining whether this attempt + /// started with a persisted full-repair request. That fact is deliberately + /// captured by the caller before execution; a repair request raised later + /// must not turn an ordinary full snapshot into a completed repair. + pub(super) fn apply_query_result_to_state_with_repair( + state: &mut TaskState, + res: &OutputWithMetrics, + elapsed: Duration, + coverage: &QueryCoverage, + started_full_repair: bool, ) -> FlowCheckpointDecision { state.after_query_exec(elapsed, true); let checkpoint_mode = state.checkpoint_mode(); @@ -238,6 +252,11 @@ impl BatchingTask { previous_mode: CheckpointMode::FullSnapshot, reason: FlowQueryFallbackReason::IncrementalDisabled, } + } else if started_full_repair { + FlowCheckpointDecision::CompletedFullRepair { + participating_regions: participating_region_count, + watermarks: watermark_count, + } } else { FlowCheckpointDecision::AdvancedFromFullSnapshot { participating_regions: participating_region_count, @@ -308,6 +327,14 @@ impl BatchingTask { "Flow {flow_id} switched to incremental mode after full snapshot, participating_regions={participating_regions}, watermarks={watermarks}" ); } + FlowCheckpointDecision::CompletedFullRepair { + participating_regions, + watermarks, + } => { + info!( + "Flow {flow_id} completed full repair, participating_regions={participating_regions}, watermarks={watermarks}" + ); + } FlowCheckpointDecision::AdvancedIncremental { participating_regions, watermarks, diff --git a/src/flow/src/batching_mode/task/inc.rs b/src/flow/src/batching_mode/task/inc.rs index fbea3abd25..d019d7380f 100644 --- a/src/flow/src/batching_mode/task/inc.rs +++ b/src/flow/src/batching_mode/task/inc.rs @@ -23,6 +23,7 @@ use query::options::{ FLOW_SINK_TABLE_ID, }; use snafu::ResultExt; +use store_api::mito_engine_options::PRESERVE_ROW_SEQUENCE; use table::metadata::TableId; use crate::Error; @@ -61,6 +62,57 @@ impl BatchingTask { Ok(table.table_info().table_id()) } + /// Whether every source table provably supports exact sequence-range + /// scans: it must be the canonical mito engine and declare the + /// `preserve_row_sequence` capability (a mito region option enforced to + /// require append-only mode). When the capability is unknown — a source + /// table that cannot be resolved, is not the mito engine, or lacks the + /// option — this returns `false` so the caller keeps the historical + /// `memtable_only` mode instead of upgrading. + pub(crate) async fn sequence_range_capable(&self) -> bool { + for name in &self.config.source_table_names { + let table = match self + .config + .catalog_manager + .table(&name[0], &name[1], &name[2], None) + .await + { + Ok(Some(table)) => table, + Ok(None) => { + debug!( + "Flow {} source table {} not found; retaining memtable_only incremental mode", + self.config.flow_id, + name.join(".") + ); + return false; + } + Err(err) => { + warn!( + "Flow {} failed to resolve source table {} for sequence_range capability check; \ + retaining memtable_only incremental mode: {:?}", + self.config.flow_id, + name.join("."), + err + ); + return false; + } + }; + + let info = table.table_info(); + let preserves = info.meta.engine == "mito" + && info + .meta + .options + .extra_options + .get(PRESERVE_ROW_SEQUENCE) + .is_some_and(|value| value.eq_ignore_ascii_case("true")); + if !preserves { + return false; + } + } + !self.config.source_table_names.is_empty() + } + /// For incremental-mode SQL queries, attempt to prepare an executable plan /// that is safe for incremental scan extensions. /// diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index 24cdbf883e..2d0064be88 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/src/flow/src/batching_mode/task/test.rs @@ -13,6 +13,8 @@ // limitations under the License. use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use catalog::RegisterTableRequest; use catalog::memory::MemoryCatalogManager; @@ -37,14 +39,18 @@ use query::options::{ use session::context::QueryContext; use snafu::ResultExt; use table::test_util::MemTable; +use table::{Table, TableRef}; +use tokio::sync::Notify; use super::*; +use crate::Result; use crate::batching_mode::checkpoint::{ CHECKPOINT_DECISION_ADVANCE, CHECKPOINT_DECISION_FALLBACK, CHECKPOINT_REASON_NONE, FlowCheckpointDecision, FlowQueryFallbackReason, }; use crate::batching_mode::eval_schedule::{FlowMissedTickPolicy, FlowScheduleConfig}; -use crate::batching_mode::state::CheckpointMode; +use crate::batching_mode::persistence::{BatchingAttempt, BatchingPersistence, RestoreOutcome}; +use crate::batching_mode::state::{CheckpointMode, TaskStateCheckpointSnapshot}; use crate::batching_mode::time_window::find_time_window_expr; use crate::test_utils::create_test_query_engine; @@ -2423,6 +2429,7 @@ async fn test_unsafe_incremental_plan_skip_restores_dirty_without_query() { &dml_plan, &dirty_restore, &QueryCoverage::IncrementalDelta, + None, ) .await .unwrap(); @@ -2680,3 +2687,309 @@ async fn test_insert_plan_matching_failure_restores_consumed_dirty_marker() { std::time::Duration::from_secs(5) ); } + +struct TestPersistence { + restores: AtomicUsize, + begins: AtomicUsize, + persists: AtomicUsize, + fail_persist: AtomicBool, +} + +struct BlockingPersistence { + persists: AtomicUsize, + fail: AtomicBool, + started: Notify, + release: Notify, +} + +#[async_trait::async_trait] +impl BatchingPersistence for BlockingPersistence { + async fn restore(&self) -> Result { + Ok(RestoreOutcome::TrustedCheckpoint(BTreeMap::new())) + } + + async fn begin_attempt(&self) -> Result { + Ok(BatchingAttempt { + ordinary_values: BTreeMap::new(), + }) + } + + async fn persist( + &self, + _attempt: BatchingAttempt, + _checkpoints: BTreeMap, + ) -> Result<()> { + self.persists.fetch_add(1, Ordering::SeqCst); + self.started.notify_one(); + self.release.notified().await; + if self.fail.load(Ordering::SeqCst) { + Err(crate::Error::External { + source: BoxedError::new(MockError::new(StatusCode::Internal)), + location: snafu::location!(), + }) + } else { + Ok(()) + } + } +} + +fn blocking_persistence(fail: bool) -> Arc { + Arc::new(BlockingPersistence { + persists: AtomicUsize::new(0), + fail: AtomicBool::new(fail), + started: Notify::new(), + release: Notify::new(), + }) +} + +fn install_persistence(task: &BatchingTask, persistence: Arc) { + *task.persistence.write().unwrap() = Some(persistence); +} + +fn candidate_transaction_states( + task: &BatchingTask, +) -> (TaskStateCheckpointSnapshot, TaskStateCheckpointSnapshot) { + let mut state = task.state.write().unwrap(); + state.advance_checkpoints(HashMap::from([(1, 10)])); + state.request_full_repair(); + let snapshot = state.checkpoint_snapshot(); + let mut candidate = snapshot.clone(); + candidate.checkpoint_mode = CheckpointMode::Incremental; + candidate.checkpoints = BTreeMap::from([(1, 20)]); + candidate.last_query_duration = Duration::from_millis(42); + candidate.last_exec_time_millis = Some(42); + state.restore_checkpoint_snapshot(snapshot.clone()); + (snapshot, candidate) +} + +fn test_attempt() -> BatchingAttempt { + BatchingAttempt { + ordinary_values: BTreeMap::new(), + } +} + +fn assert_candidate_state(task: &BatchingTask) { + let state = task.state.read().unwrap(); + assert_eq!(state.checkpoints(), &BTreeMap::from([(1, 20)])); + assert_eq!(state.last_query_duration(), Duration::from_millis(42)); + assert!(state.last_execution_time_millis().is_some()); +} + +#[tokio::test] +async fn test_checkpoint_persist_success_commits_candidate_after_release() { + let task = new_test_task_engine_and_plan_with_query( + "SELECT number, ts FROM numbers_with_ts", + "missing_sink", + ) + .await + .task; + let persistence = blocking_persistence(false); + install_persistence(&task, persistence.clone()); + let (snapshot, candidate) = candidate_transaction_states(&task); + let task_for_persist = task.clone(); + let attempt = test_attempt(); + let persist = tokio::spawn(async move { + task_for_persist + .persist_checkpoint_candidate( + snapshot.clone(), + candidate, + Some(&attempt), + DirtyRestore::Unscoped(DirtyTimeWindows::default()), + ) + .await + }); + persistence.started.notified().await; + { + let state = task.state.read().unwrap(); + assert_eq!(state.checkpoints(), &BTreeMap::from([(1, 10)])); + } + persistence.release.notify_one(); + persist.await.unwrap().unwrap(); + assert_candidate_state(&task); + assert_eq!(persistence.persists.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_checkpoint_persist_error_restores_pre_state_and_unions_dirty_once() { + let task = new_test_task_engine_and_plan_with_query( + "SELECT number, ts FROM numbers_with_ts", + "missing_sink", + ) + .await + .task; + let persistence = blocking_persistence(true); + install_persistence(&task, persistence.clone()); + let (snapshot, candidate) = candidate_transaction_states(&task); + let detached = dirty_range(1, 2); + let live = dirty_range(3, 4); + let task_for_persist = task.clone(); + let attempt = test_attempt(); + let persist = tokio::spawn(async move { + task_for_persist + .persist_checkpoint_candidate( + snapshot, + candidate, + Some(&attempt), + DirtyRestore::Unscoped(detached), + ) + .await + }); + persistence.started.notified().await; + task.state + .write() + .unwrap() + .dirty_time_windows + .add_dirty_windows(&live); + persistence.release.notify_one(); + assert!(persist.await.unwrap().is_err()); + let state = task.state.read().unwrap(); + assert_eq!(state.checkpoints(), &BTreeMap::from([(1, 10)])); + assert_eq!(state.dirty_time_windows.len(), 2); + assert_eq!(persistence.persists.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_full_repair_checkpoint_persist_handles_success_and_error() { + for fail in [false, true] { + let task = new_test_task_engine_and_plan_with_query( + "SELECT number, ts FROM numbers_with_ts", + "missing_sink", + ) + .await + .task; + let persistence = blocking_persistence(fail); + install_persistence(&task, persistence.clone()); + let (snapshot, candidate) = candidate_transaction_states(&task); + let dirty = dirty_range(1, 2); + let task_for_persist = task.clone(); + let attempt = test_attempt(); + let persist = tokio::spawn(async move { + task_for_persist + .persist_checkpoint_candidate( + snapshot, + candidate, + Some(&attempt), + DirtyRestore::FullRepair(dirty), + ) + .await + }); + persistence.started.notified().await; + task.state + .write() + .unwrap() + .dirty_time_windows + .add_dirty_windows(&dirty_range(1, 2)); + persistence.release.notify_one(); + let result = persist.await.unwrap(); + assert_eq!(result.is_err(), fail); + let state = task.state.read().unwrap(); + assert_eq!(state.dirty_time_windows.len(), 1); + assert_eq!(state.full_repair_required(), fail); + assert_eq!(persistence.persists.load(Ordering::SeqCst), 1); + } +} + +#[async_trait::async_trait] +impl BatchingPersistence for TestPersistence { + async fn restore(&self) -> Result { + self.restores.fetch_add(1, Ordering::SeqCst); + Ok(RestoreOutcome::TrustedCheckpoint(BTreeMap::from([(1, 2)]))) + } + async fn begin_attempt(&self) -> Result { + self.begins.fetch_add(1, Ordering::SeqCst); + Ok(BatchingAttempt { + ordinary_values: BTreeMap::new(), + }) + } + async fn persist( + &self, + _attempt: BatchingAttempt, + checkpoints: BTreeMap, + ) -> Result<()> { + assert_eq!(checkpoints, BTreeMap::from([(1, 2)])); + self.persists.fetch_add(1, Ordering::SeqCst); + if self.fail_persist.load(Ordering::SeqCst) { + Err(crate::Error::External { + source: BoxedError::new(MockError::new(StatusCode::Internal)), + location: snafu::location!(), + }) + } else { + Ok(()) + } + } +} + +#[tokio::test] +async fn test_persistence_restore_is_wired() { + let parts = new_test_task_engine_and_plan_with_query( + "SELECT number, ts FROM numbers_with_ts", + "missing_sink", + ) + .await; + let state = Arc::new(TestPersistence { + restores: AtomicUsize::new(0), + begins: AtomicUsize::new(0), + persists: AtomicUsize::new(0), + fail_persist: AtomicBool::new(false), + }); + let config = state.clone(); + parts.task.set_persistence(Some(config)).await.unwrap(); + assert_eq!(state.restores.load(Ordering::SeqCst), 1); + assert_eq!( + parts.task.state.read().unwrap().checkpoints(), + &BTreeMap::from([(1, 2)]) + ); +} + +#[tokio::test] +async fn test_full_repair_restore_is_sticky_and_unfiltered() { + let parts = new_time_window_test_task_with_query( + "SELECT number, date_bin(INTERVAL '5 second', ts) AS time_window FROM numbers_with_ts GROUP BY time_window, number", + ) + .await; + // Use a persistence implementation whose restore requests a full repair. + struct FullRepairPersistence; + #[async_trait::async_trait] + impl BatchingPersistence for FullRepairPersistence { + async fn restore(&self) -> Result { + Ok(RestoreOutcome::FullRepair) + } + async fn begin_attempt(&self) -> Result { + Ok(BatchingAttempt { + ordinary_values: BTreeMap::new(), + }) + } + async fn persist( + &self, + _attempt: BatchingAttempt, + _checkpoints: BTreeMap, + ) -> Result<()> { + Ok(()) + } + } + let config = Arc::new(FullRepairPersistence); + let sink = aggregate_time_window_sink_schema(); + parts + .task + .state + .write() + .unwrap() + .dirty_time_windows + .add_window(Timestamp::new_second(10), Some(Timestamp::new_second(15))); + parts.task.set_persistence(Some(config)).await.unwrap(); + let plan = parts + .task + .gen_query_with_time_window(parts.query_engine, &sink, &[], false, Some(1)) + .await + .unwrap() + .expect("full repair should always produce a plan"); + + assert!(matches!(plan.coverage, QueryCoverage::UnfilteredFull)); + assert!(matches!(plan.dirty_restore, DirtyRestore::FullRepair(_))); + let plan_text = plan.plan.to_string(); + assert!(!plan_text.contains("Filter:")); + assert!(!plan_text.contains("TimestampMillisecond(")); + // Full repair detaches the consumed dirty ownership from the live signal. + assert_eq!(parts.task.state.read().unwrap().dirty_time_windows.len(), 0); + assert!(parts.task.state.read().unwrap().full_repair_required()); +} diff --git a/src/flow/src/batching_mode/utils.rs b/src/flow/src/batching_mode/utils.rs index a707c2d9eb..51b5961a40 100644 --- a/src/flow/src/batching_mode/utils.rs +++ b/src/flow/src/batching_mode/utils.rs @@ -14,7 +14,7 @@ //! some utils for helping with batching mode -use std::collections::{BTreeSet, HashMap, HashSet}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::sync::Arc; use catalog::CatalogManagerRef; @@ -36,6 +36,7 @@ use datafusion_expr::{ Distinct, ExprSchemable, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, Projection, and, binary_expr, bitwise_and, bitwise_or, bitwise_xor, is_null, or, when, }; +use datatypes::data_type::DataType; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, SchemaRef}; use query::QueryEngineRef; @@ -950,13 +951,35 @@ pub(crate) async fn gen_plan_with_matching_schema( sink_table_schema: SchemaRef, primary_key_indices: &[usize], allow_partial: bool, +) -> Result { + gen_plan_with_matching_schema_and_values( + sql, + query_ctx, + engine, + sink_table_schema, + primary_key_indices, + allow_partial, + None, + ) + .await +} + +pub(crate) async fn gen_plan_with_matching_schema_and_values( + sql: &str, + query_ctx: QueryContextRef, + engine: QueryEngineRef, + sink_table_schema: SchemaRef, + primary_key_indices: &[usize], + allow_partial: bool, + ordinary_values: Option<&BTreeMap>, ) -> Result { let plan = sql_to_df_plan(query_ctx.clone(), engine.clone(), sql, false).await?; - let mut add_auto_column = ColumnMatcherRewriter::new( + let mut add_auto_column = ColumnMatcherRewriter::new_with_values( sink_table_schema, primary_key_indices.to_vec(), allow_partial, + ordinary_values.cloned().unwrap_or_default(), ); let plan = plan .clone() @@ -1096,15 +1119,26 @@ pub struct ColumnMatcherRewriter { pub is_rewritten: bool, pub primary_key_indices: Vec, pub allow_partial: bool, + pub ordinary_values: BTreeMap, } impl ColumnMatcherRewriter { pub fn new(schema: SchemaRef, primary_key_indices: Vec, allow_partial: bool) -> Self { + Self::new_with_values(schema, primary_key_indices, allow_partial, BTreeMap::new()) + } + + pub fn new_with_values( + schema: SchemaRef, + primary_key_indices: Vec, + allow_partial: bool, + ordinary_values: BTreeMap, + ) -> Self { Self { schema, is_rewritten: false, primary_key_indices, allow_partial, + ordinary_values, } } @@ -1114,12 +1148,27 @@ impl ColumnMatcherRewriter { mut exprs: Vec, input_schema: &DFSchema, ) -> DfResult> { + let original_exprs = exprs.clone(); + for column in self.schema.column_schemas() { + if let Some(value) = self.ordinary_values.get(&column.name) { + if value.data_type() != column.data_type.as_arrow_type() { + return Err(DataFusionError::Plan(format!( + "Configured batching metadata column {} has incompatible type", + column.name + ))); + } + if !exprs + .iter() + .any(|expr| expr.qualified_name().1 == column.name) + { + exprs.push(datafusion_expr::lit(value.clone()).alias(column.name.clone())); + } + } + } if self.allow_partial { return self.modify_project_exprs_with_partial(exprs); } - let original_exprs = exprs.clone(); - let all_names = self .schema .column_schemas() diff --git a/src/flow/src/lib.rs b/src/flow/src/lib.rs index e55e7bab76..94ff463a57 100644 --- a/src/flow/src/lib.rs +++ b/src/flow/src/lib.rs @@ -43,7 +43,12 @@ mod test_utils; pub use adapter::flownode_impl::FlowDualEngineRef; pub use adapter::{FlowConfig, FlowStreamingEngineRef, StreamingEngine}; +pub use batching_mode::IncrementalMode; pub use batching_mode::frontend_client::{FrontendClient, GrpcQueryHandlerWithBoxedError}; +pub use batching_mode::persistence::{ + BatchingAttempt, BatchingMetadataColumn, BatchingPersistence, BatchingQueryExecutor, Factory, + FactoryPlugin, PersistenceContext, RestoreOutcome, SinkLayout, +}; pub(crate) use engine::{CreateFlowArgs, FlowId, TableName}; pub use error::{Error, Result}; pub use server::{ diff --git a/src/flow/src/server.rs b/src/flow/src/server.rs index 9bda516f24..e3742bb96a 100644 --- a/src/flow/src/server.rs +++ b/src/flow/src/server.rs @@ -57,6 +57,7 @@ use tonic::{Request, Response, Status}; use crate::adapter::flownode_impl::{FlowDualEngine, FlowDualEngineRef}; use crate::adapter::{FlowStreamingEngineRef, create_worker}; use crate::batching_mode::engine::BatchingEngine; +use crate::batching_mode::persistence::FactoryPlugin; use crate::error::{ CacheRequiredSnafu, DatafusionSnafu, ExternalSnafu, ListFlowsSnafu, ParseAddrSnafu, ShutdownServerSnafu, StartServerSnafu, UnexpectedSnafu, to_status_with_last_err, @@ -326,6 +327,7 @@ pub struct FlownodeBuilder { /// receive a oneshot sender to send state size report state_report_handler: Option, frontend_client: Arc, + batching_persistence_factory: Option, } impl FlownodeBuilder { @@ -347,9 +349,16 @@ impl FlownodeBuilder { heartbeat_task: None, state_report_handler: None, frontend_client, + batching_persistence_factory: None, } } + /// Inject the optional batching persistence collaborator. + pub fn with_batching_persistence_factory(mut self, factory: FactoryPlugin) -> Self { + self.batching_persistence_factory = Some(factory); + self + } + pub fn with_heartbeat_task(self, heartbeat_task: HeartbeatTask) -> Self { let (sender, receiver) = SizeReportSender::new(); Self { @@ -404,13 +413,14 @@ impl FlownodeBuilder { self.build_manager(query_engine_factory.query_engine()) .await?, ); - let batching = Arc::new(BatchingEngine::new( + let batching = Arc::new(BatchingEngine::new_with_persistence( self.frontend_client.clone(), query_engine_factory.query_engine(), self.flow_metadata_manager.clone(), self.table_meta.clone(), self.catalog_manager.clone(), self.opts.flow.batching_mode.clone(), + self.batching_persistence_factory.clone(), )); let dual = Arc::new(FlowDualEngine::new( manager.clone(),