rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-13 00:42:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(flow): add neutral batching persistence seam

Browse Source 
Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-09-04 18:57:31 +08:00
parent a932433d21
commit c068019ebc
12 changed files with 1097 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/flow/src/batching_mode.rs
+9
View File
@@ -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;
src/flow/src/batching_mode/checkpoint.rs
+13 -8
View File
@@ -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,
}
src/flow/src/batching_mode/engine.rs
+83 -5
View File
@@ -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<BatchingModeOptions>,
persistence_factory: Option<crate::FactoryPlugin>,
}
#[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<FrontendClient>,
query_engine: QueryEngineRef,
flow_metadata_manager: FlowMetadataManagerRef,
table_meta: TableMetadataManagerRef,
catalog_manager: CatalogManagerRef,
batch_opts: BatchingModeOptions,
persistence_factory: Option<crate::FactoryPlugin>,
) -> 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<String, String>,
) -> 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();
src/flow/src/batching_mode/persistence.rs
+141
View File
@@ -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<Output>;
}
/// 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<String>,
pub columns: Vec<BatchingMetadataColumn>,
pub ordered_primary_key_indices: Vec<usize>,
pub time_index: Option<usize>,
}
#[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<u64, u64>),
FullRepair,
}
/// One serialized execution attempt.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct BatchingAttempt {
pub ordinary_values: BTreeMap<String, ScalarValue>,
}
#[async_trait::async_trait]
pub trait BatchingPersistence: Send + Sync + 'static {
async fn restore(&self) -> Result<RestoreOutcome>;
async fn begin_attempt(&self) -> Result<BatchingAttempt>;
async fn persist(
&self,
attempt: BatchingAttempt,
validated_checkpoints: BTreeMap<u64, u64>,
) -> 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<dyn BatchingQueryExecutor>,
}
/// 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<Option<Arc<dyn BatchingPersistence>>>;
}
/// Typed plugin wrapper around a persistence factory.
#[derive(Clone)]
pub struct FactoryPlugin(pub Arc<dyn Factory>);
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<crate::FrontendClient>,
}
impl FrontendBatchingQueryExecutor {
pub(crate) fn new(client: Arc<crate::FrontendClient>) -> Self {
Self { client }
}
}
#[async_trait::async_trait]
impl BatchingQueryExecutor for FrontendBatchingQueryExecutor {
async fn execute_sql(&self, catalog: &str, schema: &str, sql: &str) -> Result<Output> {
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!(),
})
}
}
src/flow/src/batching_mode/state.rs
+87 -3
View File
@@ -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<u64, u64>,
pub(crate) pending_fenced_repair: Option<FencedRepair>,
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<i64>,
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<FencedRepair>,
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<u64, u64>,
@@ -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<i64> {
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<datafusion_common::Scala
}
#[derive(Debug, Clone)]
enum ExecState {
pub(crate) enum ExecState {
Idle,
Executing,
}
src/flow/src/batching_mode/task.rs
+302 -102
View File
@@ -38,6 +38,7 @@ use snafu::{OptionExt, ResultExt};
use sql::parsers::utils::is_tql;
use store_api::mito_engine_options::MERGE_MODE_KEY;
use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
use table::TableRef;
use table::table::adapter::DfTableProviderAdapter;
use tokio::sync::oneshot::error::TryRecvError;
use tokio::sync::{Mutex, oneshot};
@@ -47,14 +48,16 @@ use crate::batching_mode::BatchingModeOptions;
use crate::batching_mode::checkpoint::checkpoint_mode_label;
use crate::batching_mode::eval_schedule::{EvalSchedule, select_due_scheduled_times};
use crate::batching_mode::frontend_client::{FrontendClient, PeerDesc};
use crate::batching_mode::persistence::RestoreOutcome;
use crate::batching_mode::state::{
CheckpointMode, DirtyTimeWindows, FilterExprInfo, TaskState, to_df_literal,
CheckpointMode, DirtyTimeWindows, FilterExprInfo, TaskState, TaskStateCheckpointSnapshot,
to_df_literal,
};
use crate::batching_mode::table_creator::{QueryType, create_table_with_expr};
use crate::batching_mode::time_window::TimeWindowExpr;
use crate::batching_mode::utils::{
AddFilterRewriter, ColumnMatcherRewriter, df_plan_to_sql, gen_plan_with_matching_schema,
get_table_info_df_schema, sql_to_df_plan,
gen_plan_with_matching_schema_and_values, get_table_info_df_schema, sql_to_df_plan,
};
use crate::df_optimizer::apply_df_optimizer;
use crate::error::{
@@ -206,6 +209,8 @@ pub struct BatchingTask {
/// window restoration for this flow. Without this, a manual flush and the
/// background loop can process the same checkpoint range concurrently.
execution_lock: Arc<Mutex<()>>,
persistence:
Arc<RwLock<Option<Arc<dyn crate::batching_mode::persistence::BatchingPersistence>>>>,
}
/// 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<Arc<dyn crate::batching_mode::persistence::BatchingPersistence>>,
) -> 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<i64> {
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<FrontendClient>,
) -> Result<Option<(usize, Duration)>, Error> {
if !self.is_table_exist(&self.config.sink_table_name).await? {
) -> Result<TableRef, Error> {
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<Arc<Schema>, 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<Arc<Schema>, 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<bool, Error> {
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<FrontendClient>,
max_window_cnt: Option<usize>,
) -> Result<Option<(usize, Duration)>, 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<FrontendClient>,
max_window_cnt: Option<usize>,
) -> 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<FrontendClient>,
max_window_cnt: Option<usize>,
) -> 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<usize>,
) -> Result<Option<PlanInfo>, 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<usize>,
attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>,
) -> Result<Option<PlanInfo>, 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<FrontendClient>,
plan: &LogicalPlan,
dirty_restore: &DirtyRestore,
coverage: &QueryCoverage,
attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>,
) -> Result<Option<(usize, Duration)>, 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<PlanInfo, Error> {
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<Option<PlanInfo>, 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<FrontendClient>,
max_window_cnt: Option<usize>,
) -> Result<Option<(usize, Duration)>, 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<FrontendClient>,
max_window_cnt: Option<usize>,
) -> 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<FrontendClient>,
max_window_cnt: Option<usize>,
) -> 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<usize>,
) -> Result<Option<PlanInfo>, 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<Schema>,
primary_key_indices: &[usize],
allow_partial: bool,
max_window_cnt: Option<usize>,
attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>,
) -> Result<Option<PlanInfo>, 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(
src/flow/src/batching_mode/task/ckpt.rs
+27
View File
@@ -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,
src/flow/src/batching_mode/task/inc.rs
+52
View File
@@ -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.
///
src/flow/src/batching_mode/task/test.rs
+314 -1
View File
@@ -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<RestoreOutcome> {
Ok(RestoreOutcome::TrustedCheckpoint(BTreeMap::new()))
}
async fn begin_attempt(&self) -> Result<BatchingAttempt> {
Ok(BatchingAttempt {
ordinary_values: BTreeMap::new(),
})
}
async fn persist(
&self,
_attempt: BatchingAttempt,
_checkpoints: BTreeMap<u64, u64>,
) -> 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<BlockingPersistence> {
Arc::new(BlockingPersistence {
persists: AtomicUsize::new(0),
fail: AtomicBool::new(fail),
started: Notify::new(),
release: Notify::new(),
})
}
fn install_persistence(task: &BatchingTask, persistence: Arc<BlockingPersistence>) {
*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<RestoreOutcome> {
self.restores.fetch_add(1, Ordering::SeqCst);
Ok(RestoreOutcome::TrustedCheckpoint(BTreeMap::from([(1, 2)])))
}
async fn begin_attempt(&self) -> Result<BatchingAttempt> {
self.begins.fetch_add(1, Ordering::SeqCst);
Ok(BatchingAttempt {
ordinary_values: BTreeMap::new(),
})
}
async fn persist(
&self,
_attempt: BatchingAttempt,
checkpoints: BTreeMap<u64, u64>,
) -> 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<RestoreOutcome> {
Ok(RestoreOutcome::FullRepair)
}
async fn begin_attempt(&self) -> Result<BatchingAttempt> {
Ok(BatchingAttempt {
ordinary_values: BTreeMap::new(),
})
}
async fn persist(
&self,
_attempt: BatchingAttempt,
_checkpoints: BTreeMap<u64, u64>,
) -> 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());
}
src/flow/src/batching_mode/utils.rs
+53 -4
View File
@@ -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<LogicalPlan, Error> {
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<String, ScalarValue>>,
) -> Result<LogicalPlan, Error> {
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<usize>,
pub allow_partial: bool,
pub ordinary_values: BTreeMap<String, ScalarValue>,
}
impl ColumnMatcherRewriter {
pub fn new(schema: SchemaRef, primary_key_indices: Vec<usize>, 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<usize>,
allow_partial: bool,
ordinary_values: BTreeMap<String, ScalarValue>,
) -> Self {
Self {
schema,
is_rewritten: false,
primary_key_indices,
allow_partial,
ordinary_values,
}
}
@@ -1114,12 +1148,27 @@ impl ColumnMatcherRewriter {
mut exprs: Vec<Expr>,
input_schema: &DFSchema,
) -> DfResult<Vec<Expr>> {
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()
src/flow/src/lib.rs
+5
View File
@@ -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::{
src/flow/src/server.rs
+11 -1
View File
@@ -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<StateReportHandler>,
frontend_client: Arc<FrontendClient>,
batching_persistence_factory: Option<FactoryPlugin>,
}
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(),