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

refactor(flow): delegate durable batching execution to extensions

Browse Source 
Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-09-07 19:29:45 +08:00
parent c7b0280473
commit 5f91fa9589
14 changed files with 608 additions and 863 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/flow/src/batching_mode.rs
+5 -5
View File
@@ -28,15 +28,15 @@ pub enum IncrementalMode {
SequenceRange,
}
mod checkpoint;
pub(crate) mod batching_execution;
pub(crate) mod checkpoint;
pub(crate) mod engine;
mod eval_schedule;
pub(crate) mod frontend_client;
pub(crate) mod persistence;
mod state;
pub(crate) mod state;
mod table_creator;
mod task;
mod time_window;
pub(crate) mod task;
pub(crate) mod time_window;
pub(crate) mod utils;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
src/flow/src/batching_mode/batching_execution.rs
+46
View File
@@ -0,0 +1,46 @@
// 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 execution collaborator for batching tasks.
use std::sync::Arc;
use query::QueryEngineRef;
use table::TableRef;
use crate::Result;
use crate::batching_mode::frontend_client::FrontendClient;
use crate::batching_mode::task::{BatchingTask, ExecuteOnceOutcome};
#[async_trait::async_trait]
pub trait BatchingExecution: Send + Sync + 'static {
async fn execute_once(
&self,
task: &BatchingTask,
engine: &QueryEngineRef,
frontend: &Arc<FrontendClient>,
max_window_cnt: Option<usize>,
) -> ExecuteOnceOutcome;
}
#[async_trait::async_trait]
pub trait BatchingExecutionFactory: Send + Sync + 'static {
async fn create(
&self,
task: &BatchingTask,
sink: TableRef,
engine: &QueryEngineRef,
frontend: &Arc<FrontendClient>,
) -> Result<Option<Arc<dyn BatchingExecution>>>;
}
src/flow/src/batching_mode/checkpoint.rs
+4 -11
View File
@@ -87,13 +87,6 @@ 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
@@ -121,7 +114,7 @@ impl FlowCheckpointDecision {
checkpoint_mode_label(CheckpointMode::FullSnapshot)
}
Self::AdvancedIncremental { .. } => checkpoint_mode_label(CheckpointMode::Incremental),
Self::CompletedFullRepair { .. } | Self::ContinuedFencedRepair { .. } => {
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)
@@ -134,9 +127,9 @@ impl FlowCheckpointDecision {
pub(super) fn decision_label(self) -> &'static str {
match self {
Self::AdvancedFromFullSnapshot { .. }
| Self::AdvancedIncremental { .. }
| Self::CompletedFullRepair { .. } => CHECKPOINT_DECISION_ADVANCE,
Self::AdvancedFromFullSnapshot { .. } | Self::AdvancedIncremental { .. } => {
CHECKPOINT_DECISION_ADVANCE
}
Self::ContinuedFencedRepair { .. } => CHECKPOINT_DECISION_CONTINUE_REPAIR,
Self::FallbackToFullSnapshot { .. } => CHECKPOINT_DECISION_FALLBACK,
}
src/flow/src/batching_mode/engine.rs
+242 -186
View File
@@ -49,7 +49,6 @@ 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};
@@ -76,7 +75,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>,
execution_factory: Option<Arc<dyn crate::BatchingExecutionFactory>>,
}
#[derive(Default)]
@@ -133,7 +132,7 @@ impl BatchingEngine {
catalog_manager: CatalogManagerRef,
batch_opts: BatchingModeOptions,
) -> Self {
Self::new_with_persistence(
Self::new_with_execution(
frontend_client,
query_engine,
flow_metadata_manager,
@@ -144,14 +143,14 @@ impl BatchingEngine {
)
}
pub fn new_with_persistence(
pub fn new_with_execution(
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>,
execution_factory: Option<Arc<dyn crate::BatchingExecutionFactory>>,
) -> Self {
Self {
runtime: Default::default(),
@@ -161,7 +160,7 @@ impl BatchingEngine {
catalog_manager,
query_engine,
batch_opts: Arc::new(batch_opts),
persistence_factory,
execution_factory,
}
}
@@ -745,10 +744,6 @@ impl BatchingEngine {
exact_sequence_range_required,
)?;
let task_inner = task.clone();
let engine = self.query_engine.clone();
let frontend = self.frontend_client.clone();
if task.config.exact_sequence_range_required {
ensure!(
task.sequence_range_capable().await?,
@@ -760,64 +755,19 @@ impl BatchingEngine {
);
}
// 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 only when persistence is actually created. Flows
// without a created collaborator keep the existing creation-time validation.
let engine = self.query_engine.clone();
let frontend = self.frontend_client.clone();
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.exact_sequence_range_required
|| (task.config.batch_opts.experimental_enable_incremental_read
&& task
.sequence_range_capable()
.await
.is_ok_and(|capable| capable))
{
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?
let execution = if let Some(factory) = &self.execution_factory {
factory.create(&task, table, &engine, &frontend).await?
} else {
None
};
if persistence.is_none() {
if execution.is_none() {
task.validate_sink_table_schema(&engine).await?;
}
task.set_persistence(persistence).await?;
let task = task.with_execution(execution);
let task_inner = task.clone();
let (start_tx, start_rx) = oneshot::channel();
@@ -1151,7 +1101,7 @@ impl FlowEngine for BatchingEngine {
mod tests {
use api::v1::flow::{DirtyWindowRequest, TimeRange};
use catalog::RegisterTableRequest;
use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager};
use catalog::memory::MemoryCatalogManager;
use common_meta::key::TableMetadataManager;
use common_meta::key::flow::FlowMetadataManager;
use common_meta::key::table_route::TableRouteValue;
@@ -1163,11 +1113,10 @@ mod tests {
use datatypes::vectors::{TimestampMillisecondVector, UInt32Vector, VectorRef};
use query::options::QueryOptions;
use session::context::QueryContext;
use tokio::sync::Notify;
use super::*;
use crate::batching_mode::persistence::{
BatchingAttempt, BatchingPersistence, Factory, FactoryPlugin, RestoreOutcome,
};
use crate::ExecuteOnceOutcome;
use crate::test_utils::create_test_query_engine;
struct DropNotify(Option<oneshot::Sender<()>>);
@@ -1180,50 +1129,13 @@ mod tests {
}
}
struct TestPersistenceFactory {
create_persistence: bool,
}
struct TestPersistence;
#[async_trait::async_trait]
impl BatchingPersistence for TestPersistence {
async fn restore(&self) -> crate::Result<RestoreOutcome> {
Ok(RestoreOutcome::TrustedCheckpoint(BTreeMap::new()))
}
async fn begin_attempt(&self) -> crate::Result<BatchingAttempt> {
Ok(BatchingAttempt::default())
}
async fn persist(
&self,
_attempt: BatchingAttempt,
_validated_checkpoints: BTreeMap<u64, u64>,
) -> crate::Result<()> {
Ok(())
}
}
#[async_trait::async_trait]
impl Factory for TestPersistenceFactory {
async fn create(
&self,
_context: PersistenceContext,
) -> crate::Result<Option<Arc<dyn BatchingPersistence>>> {
Ok(self
.create_persistence
.then_some(Arc::new(TestPersistence) as Arc<dyn BatchingPersistence>))
}
}
async fn new_test_engine() -> BatchingEngine {
let kv_backend = Arc::new(MemoryKvBackend::new());
let table_meta = Arc::new(TableMetadataManager::new(kv_backend.clone()));
table_meta.init().await.unwrap();
let flow_meta = Arc::new(FlowMetadataManager::new(kv_backend));
let catalog_manager = new_memory_catalog_manager().unwrap();
let query_engine = create_test_query_engine();
let catalog_manager = query_engine.engine_state().catalog_manager().clone();
let (frontend_client, _handler) =
FrontendClient::from_empty_grpc_handler(QueryOptions::default());
@@ -1237,8 +1149,8 @@ mod tests {
)
}
async fn new_test_engine_with_persistence(
persistence_factory: Option<FactoryPlugin>,
async fn new_test_engine_with_execution(
execution_factory: Option<Arc<dyn crate::BatchingExecutionFactory>>,
) -> BatchingEngine {
let kv_backend = Arc::new(MemoryKvBackend::new());
let table_meta = Arc::new(TableMetadataManager::new(kv_backend.clone()));
@@ -1249,14 +1161,14 @@ mod tests {
let (frontend_client, _handler) =
FrontendClient::from_empty_grpc_handler(QueryOptions::default());
let engine = BatchingEngine::new_with_persistence(
let engine = BatchingEngine::new_with_execution(
Arc::new(frontend_client),
query_engine,
flow_meta,
table_meta,
catalog_manager,
BatchingModeOptions::default(),
persistence_factory,
execution_factory,
);
engine
.table_meta
@@ -1270,8 +1182,8 @@ mod tests {
engine
}
fn register_sink_with_schema(engine: &BatchingEngine, name: &str, extended: bool) {
let mut columns = vec![
fn register_sink_with_schema(engine: &BatchingEngine, name: &str) {
let schema = Arc::new(Schema::new(vec![
ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), false),
ColumnSchema::new(
"ts",
@@ -1279,35 +1191,53 @@ mod tests {
false,
)
.with_time_index(true),
];
let mut vectors: Vec<VectorRef> = vec![
Arc::new(UInt32Vector::from_slice([1_u32])),
Arc::new(TimestampMillisecondVector::from_slice([0_i64])),
];
if extended {
columns.push(ColumnSchema::new(
"checkpoint",
ConcreteDataType::uint32_datatype(),
false,
));
vectors.push(Arc::new(UInt32Vector::from_slice([1_u32])));
}
let schema = Arc::new(Schema::new(columns));
let recordbatch = RecordBatch::new(schema, vectors).unwrap();
let table = table::test_util::MemTable::table(name, recordbatch);
let request = RegisterTableRequest {
catalog: "greptime".to_string(),
schema: "public".to_string(),
table_name: name.to_string(),
table_id: 9000,
table,
};
]));
let recordbatch = RecordBatch::new(
schema,
vec![
Arc::new(UInt32Vector::from_slice([1_u32])) as VectorRef,
Arc::new(TimestampMillisecondVector::from_slice([0_i64])) as VectorRef,
],
)
.unwrap();
engine
.catalog_manager
.as_any()
.downcast_ref::<MemoryCatalogManager>()
.unwrap()
.register_table_sync(request)
.register_table_sync(RegisterTableRequest {
catalog: "greptime".to_string(),
schema: "public".to_string(),
table_name: name.to_string(),
table_id: 9000,
table: table::test_util::MemTable::table(name, recordbatch),
})
.unwrap();
}
fn register_number_only_sink(engine: &BatchingEngine, name: &str) {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
"number",
ConcreteDataType::uint32_datatype(),
false,
)]));
let recordbatch = RecordBatch::new(
schema,
vec![Arc::new(UInt32Vector::from_slice([1_u32])) as VectorRef],
)
.unwrap();
engine
.catalog_manager
.as_any()
.downcast_ref::<MemoryCatalogManager>()
.unwrap()
.register_table_sync(RegisterTableRequest {
catalog: "greptime".to_string(),
schema: "public".to_string(),
table_name: name.to_string(),
table_id: 9001,
table: table::test_util::MemTable::table(name, recordbatch),
})
.unwrap();
}
@@ -1332,61 +1262,192 @@ mod tests {
}
}
#[tokio::test]
async fn test_installed_persistence_factory_none_still_validates_sink_schema() {
let engine = new_test_engine_with_persistence(Some(FactoryPlugin(Arc::new(
TestPersistenceFactory {
create_persistence: false,
},
))))
.await;
register_sink_with_schema(&engine, "sink_factory_none", true);
struct TestExecution {
manual_calls: std::sync::atomic::AtomicUsize,
}
let result = engine
.create_flow_inner(flow_create_args(1, "sink_factory_none"))
.await;
#[async_trait::async_trait]
impl crate::BatchingExecution for TestExecution {
async fn execute_once(
&self,
task: &BatchingTask,
_engine: &QueryEngineRef,
_frontend: &Arc<FrontendClient>,
_max_window_cnt: Option<usize>,
) -> ExecuteOnceOutcome {
if task
.state
.read()
.unwrap()
.query_ctx
.extension(query::options::FLOW_SCHEDULED_TIME_MILLIS)
.is_none()
{
self.manual_calls
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
}
ExecuteOnceOutcome {
new_query: None,
result: Ok(None),
}
}
}
assert!(
result.is_err(),
"ordinary sink validation must reject mismatch"
);
assert!(!engine.flow_exist_inner(1).await);
struct TestExecutionFactory {
entered: Option<Arc<Notify>>,
release: Option<Arc<Notify>>,
result: std::sync::Mutex<Option<crate::Result<Option<Arc<dyn crate::BatchingExecution>>>>>,
}
#[async_trait::async_trait]
impl crate::BatchingExecutionFactory for TestExecutionFactory {
async fn create(
&self,
_task: &BatchingTask,
_sink: table::TableRef,
_engine: &QueryEngineRef,
_frontend: &Arc<FrontendClient>,
) -> crate::Result<Option<Arc<dyn crate::BatchingExecution>>> {
if let Some(entered) = &self.entered {
entered.notify_one();
}
if let Some(release) = &self.release {
release.notified().await;
}
self.result
.lock()
.unwrap()
.take()
.expect("execution factory should only be called once")
}
}
#[tokio::test]
async fn test_matching_persistence_factory_allows_extended_sink_schema() {
let engine = new_test_engine_with_persistence(Some(FactoryPlugin(Arc::new(
TestPersistenceFactory {
create_persistence: true,
},
))))
.await;
register_sink_with_schema(&engine, "sink_factory_some", true);
async fn test_execution_factory_finishes_before_task_publication() {
const GATE_TIMEOUT: Duration = Duration::from_secs(1);
let result = engine
.create_flow_inner(flow_create_args(2, "sink_factory_some"))
let execution = Arc::new(TestExecution {
manual_calls: Default::default(),
});
let entered = Arc::new(Notify::new());
let release = Arc::new(Notify::new());
let factory = Arc::new(TestExecutionFactory {
entered: Some(entered.clone()),
release: Some(release.clone()),
result: std::sync::Mutex::new(Some(Ok(Some(execution.clone())))),
});
let engine = Arc::new(new_test_engine_with_execution(Some(factory)).await);
register_sink_with_schema(&engine, "factory_sink");
let entered_wait = entered.notified();
let mut args = flow_create_args(6, "factory_sink");
args.eval_interval = Some(86_400);
let mut create = tokio::spawn({
let engine = engine.clone();
async move { engine.create_flow_inner(args).await }
});
if tokio::time::timeout(GATE_TIMEOUT, entered_wait)
.await
.is_err()
{
release.notify_one();
if tokio::time::timeout(GATE_TIMEOUT, &mut create)
.await
.is_err()
{
create.abort();
}
panic!("execution factory should be entered");
}
let unpublished = {
let runtime = engine.runtime.read().await;
!runtime.tasks.contains_key(&6) && !runtime.shutdown_txs.contains_key(&6)
};
release.notify_one();
assert!(unpublished);
let created = tokio::time::timeout(GATE_TIMEOUT, &mut create)
.await
.expect("flow creation should finish after factory release");
created.unwrap().unwrap();
assert!(engine.flow_exist_inner(6).await);
let calls_before_flush = execution
.manual_calls
.load(std::sync::atomic::Ordering::SeqCst);
assert_eq!(engine.flush_flow_inner(6).await.unwrap(), 0);
assert_eq!(
execution
.manual_calls
.load(std::sync::atomic::Ordering::SeqCst),
calls_before_flush + 1
);
engine.remove_flow_inner(6).await.unwrap();
}
#[tokio::test]
async fn test_execution_factory_result_controls_sink_validation_and_publication() {
let no_factory_engine = new_test_engine_with_execution(None).await;
register_number_only_sink(&no_factory_engine, "no_factory_sink");
assert!(
no_factory_engine
.create_flow_inner(flow_create_args(7, "no_factory_sink"))
.await
.is_err()
);
assert!(!no_factory_engine.flow_exist_inner(7).await);
let declined_engine =
new_test_engine_with_execution(Some(Arc::new(TestExecutionFactory {
entered: None,
release: None,
result: std::sync::Mutex::new(Some(Ok(None))),
})))
.await;
register_number_only_sink(&declined_engine, "declined_factory_sink");
assert!(
declined_engine
.create_flow_inner(flow_create_args(8, "declined_factory_sink"))
.await
.is_err()
);
assert!(!declined_engine.flow_exist_inner(8).await);
let accepted_engine =
new_test_engine_with_execution(Some(Arc::new(TestExecutionFactory {
entered: None,
release: None,
result: std::sync::Mutex::new(Some(Ok(Some(Arc::new(TestExecution {
manual_calls: Default::default(),
}))))),
})))
.await;
register_number_only_sink(&accepted_engine, "accepted_factory_sink");
accepted_engine
.create_flow_inner(flow_create_args(9, "accepted_factory_sink"))
.await
.unwrap();
assert!(accepted_engine.flow_exist_inner(9).await);
accepted_engine.remove_flow_inner(9).await.unwrap();
assert_eq!(Some(2), result);
assert!(engine.flow_exist_inner(2).await);
engine.remove_flow_inner(2).await.unwrap();
}
#[tokio::test]
async fn test_no_persistence_factory_still_validates_sink_schema() {
let engine = new_test_engine_with_persistence(None).await;
register_sink_with_schema(&engine, "sink_no_factory", true);
let result = engine
.create_flow_inner(flow_create_args(3, "sink_no_factory"))
.await;
let error_engine = new_test_engine_with_execution(Some(Arc::new(TestExecutionFactory {
entered: None,
release: None,
result: std::sync::Mutex::new(Some(
UnexpectedSnafu {
reason: "test execution factory failure".to_string(),
}
.fail(),
)),
})))
.await;
register_sink_with_schema(&error_engine, "error_factory_sink");
assert!(
result.is_err(),
"ordinary sink validation must reject mismatch"
error_engine
.create_flow_inner(flow_create_args(10, "error_factory_sink"))
.await
.is_err()
);
assert!(!engine.flow_exist_inner(3).await);
assert!(!error_engine.flow_exist_inner(10).await);
}
#[tokio::test]
@@ -1429,13 +1490,8 @@ mod tests {
#[tokio::test]
async fn test_forged_query_context_does_not_enable_exact_sequence_range() {
let engine = new_test_engine_with_persistence(Some(FactoryPlugin(Arc::new(
TestPersistenceFactory {
create_persistence: true,
},
))))
.await;
register_sink_with_schema(&engine, "forged_query_context", true);
let engine = new_test_engine_with_execution(None).await;
register_sink_with_schema(&engine, "forged_query_context");
let mut args = flow_create_args(4, "forged_query_context");
let mut query_ctx = QueryContext::arc().as_ref().clone();
query_ctx.set_extension("__old_forged_required_extension", "true");
@@ -1449,7 +1505,7 @@ mod tests {
#[tokio::test]
async fn test_exact_sequence_range_capability_is_checked_before_task_startup() {
let engine = new_test_engine_with_persistence(None).await;
let engine = new_test_engine_with_execution(None).await;
let mut args = flow_create_args(5, "exact_requires_capability");
args.flow_options.insert(
FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY.to_string(),
src/flow/src/batching_mode/frontend_client.rs
+2 -2
View File
@@ -463,7 +463,7 @@ impl FrontendClient {
/// Execute a flow query and return terminal metrics. `snapshot_seqs` are
/// optional read upper bounds used only by snapshot-fenced repair chunks.
pub(crate) async fn query_with_terminal_metrics(
pub async fn query_with_terminal_metrics(
&self,
catalog: &str,
schema: &str,
@@ -666,7 +666,7 @@ fn wrap_standalone_output_with_terminal_metrics(
/// Describe a peer of frontend
#[derive(Debug, Default, Clone)]
pub(crate) enum PeerDesc {
pub enum PeerDesc {
/// The query failed before a frontend peer was selected.
#[default]
Unknown,
src/flow/src/batching_mode/persistence.rs
-141
View File
@@ -1,141 +0,0 @@
// 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
+5 -73
View File
@@ -35,19 +35,6 @@ use crate::metrics::{
};
use crate::{Error, FlowId};
#[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,
@@ -64,7 +51,6 @@ 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>,
@@ -97,7 +83,6 @@ 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,
@@ -140,6 +125,11 @@ impl TaskState {
self.start_time_millis
}
/// Pending dirty work, without permitting mutation of the task state.
pub fn dirty_time_windows(&self) -> &DirtyTimeWindows {
&self.dirty_time_windows
}
pub fn checkpoint_mode(&self) -> CheckpointMode {
self.checkpoint_mode
}
@@ -148,64 +138,6 @@ 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> {
src/flow/src/batching_mode/task.rs
+88 -167
View File
@@ -28,7 +28,7 @@ use datafusion::datasource::DefaultTableSource;
use datafusion::sql::unparser::expr_to_sql;
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_common::utils::quote_identifier;
use datafusion_common::{DFSchemaRef, TableReference};
use datafusion_common::{DFSchemaRef, ScalarValue, TableReference};
use datafusion_expr::{DmlStatement, LogicalPlan, WriteOp, col, lit};
use datatypes::schema::Schema;
use query::QueryEngineRef;
@@ -49,15 +49,13 @@ 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, TaskStateCheckpointSnapshot,
to_df_literal,
CheckpointMode, DirtyTimeWindows, FilterExprInfo, TaskState, 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,
AddFilterRewriter, ColumnMatcherRewriter, df_plan_to_sql,
gen_plan_with_matching_schema_and_values, get_table_info_df_schema, sql_to_df_plan,
};
use crate::df_optimizer::apply_df_optimizer;
@@ -211,8 +209,7 @@ 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>>>>,
execution: Option<Arc<dyn crate::BatchingExecution>>,
}
/// Arguments for creating batching task
@@ -286,17 +283,16 @@ 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 {
new_query: Option<PlanInfo>,
pub struct ExecuteOnceOutcome {
pub new_query: Option<PlanInfo>,
/// Execution result of the generated insert plan.
///
/// `Ok(Some((affected_rows, elapsed)))` means a query was executed.
/// `Ok(None)` means no query was generated because there was no dirty signal.
/// `Err(_)` means plan generation or execution failed.
result: Result<Option<(usize, Duration)>, Error>,
pub result: Result<Option<(usize, Duration)>, Error>,
}
impl BatchingTask {
@@ -353,31 +349,16 @@ impl BatchingTask {
}),
state: Arc::new(RwLock::new(state)),
execution_lock: Arc::new(Mutex::new(())),
persistence: Arc::new(RwLock::new(None)),
execution: 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(crate) fn with_execution(
mut self,
execution: Option<Arc<dyn crate::BatchingExecution>>,
) -> Self {
self.execution = execution;
self
}
pub fn last_execution_time_millis(&self) -> Option<i64> {
@@ -473,33 +454,53 @@ impl BatchingTask {
pub async fn validate_sink_table_schema(
&self,
engine: &QueryEngineRef,
) -> Result<Arc<Schema>, Error> {
self.validate_sink_table_schema_with_values(engine, &BTreeMap::new())
.await
}
pub async fn validate_sink_table_schema_with_values(
&self,
engine: &QueryEngineRef,
values: &BTreeMap<String, ScalarValue>,
) -> 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)
self.validate_sink_table_schema_with_table_and_values(engine, table, values)
.await
}
async fn validate_sink_table_schema_with_table(
pub(crate) async fn validate_sink_table_schema_with_table(
&self,
engine: &QueryEngineRef,
table: TableRef,
) -> Result<Arc<Schema>, Error> {
self.validate_sink_table_schema_with_table_and_values(engine, table, &BTreeMap::new())
.await
}
async fn validate_sink_table_schema_with_table_and_values(
&self,
engine: &QueryEngineRef,
table: TableRef,
values: &BTreeMap<String, ScalarValue>,
) -> 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(
gen_plan_with_matching_schema_and_values(
&self.config.query,
query_ctx,
engine.clone(),
table_meta.schema.clone(),
&primary_key_indices,
merge_mode_last_non_null,
Some(values),
)
.await
.map(|_| table_meta.schema.clone())
@@ -511,15 +512,16 @@ impl BatchingTask {
engine: &QueryEngineRef,
max_window_cnt: Option<usize>,
) -> Result<Option<PlanInfo>, Error> {
self.gen_insert_plan_unlocked_with_attempt(engine, max_window_cnt, None)
self.gen_insert_plan_with_values_unlocked(engine, max_window_cnt, &BTreeMap::new(), false)
.await
}
async fn gen_insert_plan_unlocked_with_attempt(
pub async fn gen_insert_plan_with_values_unlocked(
&self,
engine: &QueryEngineRef,
max_window_cnt: Option<usize>,
attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>,
values: &BTreeMap<String, ScalarValue>,
force_full_snapshot: bool,
) -> Result<Option<PlanInfo>, Error> {
let (table, df_schema) = get_table_info_df_schema(
self.config.catalog_manager.clone(),
@@ -533,13 +535,14 @@ impl BatchingTask {
let primary_key_indices = table_meta.primary_key_indices.clone();
let new_query = self
.gen_query_with_time_window_with_attempt(
.gen_query_with_time_window_with_values(
engine.clone(),
&table.table_info().meta.schema,
&primary_key_indices,
merge_mode_last_non_null,
max_window_cnt,
attempt,
values,
force_full_snapshot,
)
.await?;
@@ -621,54 +624,6 @@ impl BatchingTask {
Ok(())
}
/// 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,
engine: &QueryEngineRef,
@@ -676,7 +631,6 @@ impl BatchingTask {
plan: &LogicalPlan,
dirty_restore: &DirtyRestore,
coverage: &QueryCoverage,
attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>,
) -> Result<Option<(usize, Duration)>, Error> {
let flow_id = self.config.flow_id;
let Some((res, elapsed)) = self
@@ -690,13 +644,7 @@ impl BatchingTask {
let decision = {
let mut state = self.state.write().unwrap();
let reason = Self::query_failure_reason(err, coverage);
Self::apply_query_failure_to_state(
&mut state,
elapsed,
coverage,
reason,
attempt.is_some(),
)
Self::apply_query_failure_to_state(&mut state, elapsed, coverage, reason)
};
if let Some(decision) = decision {
Self::record_checkpoint_decision(flow_id, decision);
@@ -705,42 +653,16 @@ impl BatchingTask {
let res = res?;
let (affected_rows, _) = res.output.extract_rows_and_cost();
let (decision, checkpoint_txn) = {
let decision = {
let mut state = self.state.write().unwrap();
let snapshot = state.checkpoint_snapshot();
let repair_required = snapshot.full_repair_required;
let decision = Self::apply_query_result_to_state(
&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)))
Self::apply_query_result_to_state(&mut state, &res, elapsed, coverage)
};
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)))
}
async fn execute_plan_unlocked(
pub async fn execute_plan_unlocked(
&self,
engine: &QueryEngineRef,
frontend_client: &Arc<FrontendClient>,
@@ -946,10 +868,10 @@ impl BatchingTask {
/// Restore dirty windows consumed by a failed query so they are retried on
/// the next execution.
///
fn restore_dirty_windows(&self, dirty_restore: &DirtyRestore) {
pub fn restore_dirty_windows(&self, dirty_restore: &DirtyRestore) {
match dirty_restore {
DirtyRestore::Scoped(filter) => self.restore_scoped_dirty_windows(filter),
DirtyRestore::Unscoped(dirty_windows) | DirtyRestore::FullRepair(dirty_windows) => self
DirtyRestore::Unscoped(dirty_windows) => self
.state
.write()
.unwrap()
@@ -1027,7 +949,7 @@ impl BatchingTask {
dirty_windows_to_restore: DirtyTimeWindows,
retention_filter: Option<(&str, Timestamp, &'static str)>,
coverage: QueryCoverage,
attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>,
values: &BTreeMap<String, ScalarValue>,
) -> Result<PlanInfo, Error> {
let mut plan = self.restore_unscoped_dirty_windows_on_err(
&dirty_windows_to_restore,
@@ -1038,7 +960,7 @@ impl BatchingTask {
sink_table_schema,
primary_key_indices,
allow_partial,
attempt.map(|a| &a.ordinary_values),
Some(values),
)
.await,
)?;
@@ -1066,11 +988,7 @@ impl BatchingTask {
Ok(PlanInfo {
plan,
dirty_restore: if self.state.read().unwrap().full_repair_required() {
DirtyRestore::FullRepair(dirty_windows_to_restore)
} else {
DirtyRestore::Unscoped(dirty_windows_to_restore)
},
dirty_restore: DirtyRestore::Unscoped(dirty_windows_to_restore),
coverage,
})
}
@@ -1087,7 +1005,7 @@ impl BatchingTask {
allow_partial: bool,
retention_filter: Option<(&str, Timestamp, &'static str)>,
coverage: QueryCoverage,
attempt: Option<&crate::batching_mode::persistence::BatchingAttempt>,
values: &BTreeMap<String, ScalarValue>,
) -> Result<Option<PlanInfo>, Error> {
let (is_dirty, dirty_windows_to_restore) = self.drain_dirty_windows_signal();
if !is_dirty {
@@ -1104,7 +1022,7 @@ impl BatchingTask {
dirty_windows_to_restore,
retention_filter,
coverage,
attempt,
values,
)
.await
.map(Some)
@@ -1441,23 +1359,22 @@ impl BatchingTask {
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())
if let Some(execution) = &self.execution {
return execution
.execute_once(self, engine, frontend_client, max_window_cnt)
.await;
}
self.execute_once_default_unlocked(engine, frontend_client, max_window_cnt)
.await
{
}
async fn execute_once_default_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 {
@@ -1479,7 +1396,6 @@ impl BatchingTask {
&new_query.plan,
&new_query.dirty_restore,
&new_query.coverage,
attempt.as_ref(),
)
.await;
match res {
@@ -1614,25 +1530,27 @@ impl BatchingTask {
allow_partial: bool,
max_window_cnt: Option<usize>,
) -> Result<Option<PlanInfo>, Error> {
self.gen_query_with_time_window_with_attempt(
self.gen_query_with_time_window_with_values(
engine,
sink_table_schema,
primary_key_indices,
allow_partial,
max_window_cnt,
None,
&BTreeMap::new(),
false,
)
.await
}
async fn gen_query_with_time_window_with_attempt(
async fn gen_query_with_time_window_with_values(
&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>,
values: &BTreeMap<String, ScalarValue>,
force_full_snapshot: bool,
) -> Result<Option<PlanInfo>, Error> {
let query_ctx = self.state.read().unwrap().query_ctx.clone();
let start = SystemTime::now();
@@ -1654,15 +1572,20 @@ impl BatchingTask {
.map(|expr| expr.eval(low_bound))
.transpose()?;
let full_repair = self.state.read().unwrap().full_repair_required();
if full_repair {
if force_full_snapshot {
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"))
.map(|lower| {
(
expr.column_name.as_str(),
lower,
"forced full snapshot retention",
)
})
})
});
return self
@@ -1675,7 +1598,7 @@ impl BatchingTask {
detached,
retention_filter,
QueryCoverage::UnfilteredFull,
attempt,
values,
)
.await
.map(Some);
@@ -1713,7 +1636,7 @@ impl BatchingTask {
dirty_windows_to_restore,
None,
QueryCoverage::UnfilteredFull,
attempt,
values,
)
.await?;
@@ -1770,7 +1693,7 @@ impl BatchingTask {
allow_partial,
retention_filter,
QueryCoverage::IncrementalDelta,
attempt,
values,
)
.await;
}
@@ -1818,9 +1741,7 @@ impl BatchingTask {
sink_table_schema.clone(),
primary_key_indices.to_vec(),
allow_partial,
attempt
.map(|a| a.ordinary_values.clone())
.unwrap_or_default(),
values.clone(),
);
let plan = self.restore_scoped_dirty_windows_on_err(
src/flow/src/batching_mode/task/ckpt.rs
-18
View File
@@ -58,7 +58,6 @@ impl BatchingTask {
elapsed: Duration,
coverage: &QueryCoverage,
reason: FlowQueryFallbackReason,
persistence_backed: bool,
) -> Option<FlowCheckpointDecision> {
state.after_query_exec(elapsed, false);
let checkpoint_mode = state.checkpoint_mode();
@@ -80,9 +79,6 @@ impl BatchingTask {
if checkpoint_mode == CheckpointMode::Incremental {
state.mark_full_snapshot();
if persistence_backed && matches!(coverage, QueryCoverage::IncrementalDelta) {
state.request_full_repair();
}
}
Some(FlowCheckpointDecision::FallbackToFullSnapshot {
previous_mode: checkpoint_mode,
@@ -100,7 +96,6 @@ impl BatchingTask {
res: &OutputWithMetrics,
elapsed: Duration,
coverage: &QueryCoverage,
started_full_repair: bool,
) -> FlowCheckpointDecision {
state.after_query_exec(elapsed, true);
let checkpoint_mode = state.checkpoint_mode();
@@ -201,11 +196,6 @@ 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,
@@ -276,14 +266,6 @@ 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
+1 -1
View File
@@ -73,7 +73,7 @@ impl BatchingTask {
/// 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) -> Result<bool, Error> {
pub async fn sequence_range_capable(&self) -> Result<bool, Error> {
for name in &self.config.source_table_names {
let table = match self
.config
src/flow/src/batching_mode/task/test.rs
+200 -248
View File
@@ -13,10 +13,7 @@
// limitations under the License.
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::pin::Pin;
use std::task::{Context, Poll};
use bytes::Bytes;
use catalog::memory::MemoryCatalogManager;
use catalog::{DeregisterTableRequest, RegisterTableRequest};
use client::OutputWithMetrics;
@@ -25,12 +22,12 @@ use common_error::ext::BoxedError;
use common_error::mock::MockError;
use common_error::status_code::StatusCode;
use common_query::Output;
use common_recordbatch::RecordBatch;
use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry};
use common_recordbatch::{RecordBatch, RecordBatchStream};
use datatypes::data_type::ConcreteDataType as CDT;
use datatypes::schema::ColumnSchema;
use datatypes::vectors::{
TimestampMillisecondVector, TimestampNanosecondVector, UInt32Vector, VectorRef,
StringVector, TimestampMillisecondVector, TimestampNanosecondVector, UInt32Vector, VectorRef,
};
use pretty_assertions::assert_eq;
use query::options::{
@@ -49,7 +46,6 @@ use crate::batching_mode::checkpoint::{
FlowCheckpointDecision, FlowQueryFallbackReason,
};
use crate::batching_mode::eval_schedule::{FlowMissedTickPolicy, FlowScheduleConfig};
use crate::batching_mode::persistence::{BatchingAttempt, BatchingPersistence, RestoreOutcome};
use crate::batching_mode::state::CheckpointMode;
use crate::batching_mode::time_window::find_time_window_expr;
use crate::test_utils::create_test_query_engine;
@@ -61,6 +57,68 @@ fn incremental_batch_opts() -> Arc<BatchingModeOptions> {
})
}
struct CountingExecution {
calls: std::sync::atomic::AtomicUsize,
active: std::sync::atomic::AtomicUsize,
max_active: std::sync::atomic::AtomicUsize,
}
#[async_trait::async_trait]
impl crate::BatchingExecution for CountingExecution {
async fn execute_once(
&self,
_task: &BatchingTask,
_engine: &QueryEngineRef,
_frontend: &Arc<FrontendClient>,
_max_window_cnt: Option<usize>,
) -> ExecuteOnceOutcome {
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let active = self
.active
.fetch_add(1, std::sync::atomic::Ordering::SeqCst)
+ 1;
self.max_active
.fetch_max(active, std::sync::atomic::Ordering::SeqCst);
tokio::task::yield_now().await;
self.active
.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
ExecuteOnceOutcome {
new_query: None,
result: Ok(None),
}
}
}
#[tokio::test]
async fn test_execution_delegate_dispatch_is_serialized() {
let TestTaskParts {
task, query_engine, ..
} = new_test_task_engine_and_plan_with_query("SELECT number, ts FROM numbers_with_ts", "sink")
.await;
let execution = Arc::new(CountingExecution {
calls: Default::default(),
active: Default::default(),
max_active: Default::default(),
});
let task = task.with_execution(Some(execution.clone()));
let (frontend, _handler) = FrontendClient::from_empty_grpc_handler(QueryOptions::default());
let frontend = Arc::new(frontend);
let first = task.execute_once_serialized(&query_engine, &frontend, None);
let second = task.execute_once_serialized(&query_engine, &frontend, None);
let (first, second) = tokio::join!(first, second);
assert_eq!(first.unwrap(), None);
assert_eq!(second.unwrap(), None);
assert_eq!(execution.calls.load(std::sync::atomic::Ordering::SeqCst), 2);
assert_eq!(
execution
.max_active
.load(std::sync::atomic::Ordering::SeqCst),
1,
"the existing execution_lock must span delegate execution"
);
}
async fn new_test_task_and_plan_with_missing_sink() -> (BatchingTask, LogicalPlan) {
new_test_task_engine_and_plan_with_query(
"SELECT number, ts FROM numbers_with_ts",
@@ -554,6 +612,35 @@ fn register_twe_sink(query_engine: &QueryEngineRef, table_name: &str, table_id:
memory_catalog.register_table_sync(request).unwrap();
}
fn register_twe_sink_with_metadata(query_engine: &QueryEngineRef, table_name: &str, table_id: u32) {
let schema = Arc::new(Schema::new(vec![
ColumnSchema::new("number", CDT::uint32_datatype(), false),
ColumnSchema::new("time_window", CDT::timestamp_millisecond_datatype(), false)
.with_time_index(true),
ColumnSchema::new("metadata", CDT::string_datatype(), false),
]));
let columns: Vec<VectorRef> = vec![
Arc::new(UInt32Vector::from_slice([1_u32])),
Arc::new(TimestampMillisecondVector::from_slice([0_i64])),
Arc::new(StringVector::from_slice(&["existing"])),
];
let recordbatch = RecordBatch::new(schema, columns).unwrap();
let table = MemTable::table(table_name, recordbatch);
let request = RegisterTableRequest {
catalog: DEFAULT_CATALOG_NAME.to_string(),
schema: DEFAULT_SCHEMA_NAME.to_string(),
table_name: table_name.to_string(),
table_id,
table,
};
let catalog_manager = query_engine.engine_state().catalog_manager();
let memory_catalog = catalog_manager
.as_any()
.downcast_ref::<MemoryCatalogManager>()
.unwrap();
memory_catalog.register_table_sync(request).unwrap();
}
fn register_scheduled_now_sink(query_engine: &QueryEngineRef, table_name: &str, table_id: u32) {
let schema = Arc::new(Schema::new(vec![
ColumnSchema::new("ts", CDT::timestamp_nanosecond_datatype(), false).with_time_index(true),
@@ -944,7 +1031,6 @@ fn test_apply_query_result_to_state_advances_full_snapshot_to_incremental() {
&result,
std::time::Duration::from_millis(1),
&QueryCoverage::UnfilteredFull,
false,
);
assert_eq!(
@@ -976,7 +1062,6 @@ fn test_apply_query_result_to_state_stays_full_snapshot_when_incremental_disable
&result,
std::time::Duration::from_millis(1),
&QueryCoverage::UnfilteredFull,
false,
);
// Should NOT claim advancement to incremental; should fallback with correct reason.
@@ -1008,7 +1093,6 @@ fn test_apply_query_result_to_state_rejects_unproved_watermark() {
&result,
std::time::Duration::from_millis(1),
&QueryCoverage::UnfilteredFull,
false,
);
assert_eq!(
@@ -1034,7 +1118,6 @@ fn test_apply_query_result_to_state_reports_missing_watermark() {
&result,
std::time::Duration::from_millis(1),
&QueryCoverage::UnfilteredFull,
false,
);
assert_eq!(
@@ -1065,7 +1148,6 @@ fn test_apply_query_result_to_state_advances_incremental_subset() {
&result,
std::time::Duration::from_millis(1),
&QueryCoverage::IncrementalDelta,
false,
);
assert_eq!(
@@ -1101,7 +1183,6 @@ fn test_scoped_base_repair_with_dirty_backlog_starts_fenced_repair_from_full_sna
&result,
std::time::Duration::from_millis(1),
&QueryCoverage::ScopedBaseRepair,
false,
);
assert_eq!(
@@ -1161,7 +1242,6 @@ fn test_fenced_repair_chunk_with_pending_windows_stays_full_snapshot() {
&output_with_region_watermarks([(1_u64, Some(10_u64)), (2_u64, Some(20_u64))]),
std::time::Duration::from_millis(1),
&QueryCoverage::FencedRepairChunk { high },
false,
);
assert_eq!(
@@ -1218,7 +1298,6 @@ fn test_continued_fenced_repair_uses_pending_snapshot_not_later_live_dirty() {
&output_with_region_watermarks([(1_u64, Some(10_u64)), (2_u64, Some(20_u64))]),
std::time::Duration::from_millis(1),
&QueryCoverage::FencedRepairChunk { high },
false,
);
assert_eq!(
decision,
@@ -1256,7 +1335,6 @@ fn test_final_fenced_repair_chunk_advances_to_high() {
&output_with_region_watermarks([(1_u64, Some(10_u64)), (2_u64, Some(20_u64))]),
std::time::Duration::from_millis(1),
&QueryCoverage::FencedRepairChunk { high: high.clone() },
false,
);
assert_eq!(
@@ -1331,7 +1409,6 @@ fn test_fenced_repair_chunk_watermark_mismatch_restores_pending_but_consumes_inf
&output_with_region_watermarks([(1_u64, Some(11_u64)), (2_u64, Some(20_u64))]),
std::time::Duration::from_millis(1),
&QueryCoverage::FencedRepairChunk { high },
false,
);
assert_eq!(
@@ -1375,7 +1452,6 @@ async fn test_fenced_repair_mismatch_next_plan_is_scoped_base_repair() {
&output_with_region_watermarks([(1_u64, Some(11_u64)), (2_u64, Some(20_u64))]),
std::time::Duration::from_millis(1),
&QueryCoverage::FencedRepairChunk { high },
false,
);
assert_eq!(
decision,
@@ -1422,7 +1498,6 @@ fn test_apply_query_failure_to_state_falls_back_from_incremental() {
std::time::Duration::from_millis(1),
&QueryCoverage::IncrementalDelta,
FlowQueryFallbackReason::IncrementalQueryFailure,
false,
);
assert_eq!(
@@ -1450,7 +1525,6 @@ fn test_apply_query_failure_to_state_records_full_snapshot_failure() {
std::time::Duration::from_millis(1),
&QueryCoverage::UnfilteredFull,
FlowQueryFallbackReason::QueryFailure,
false,
);
assert_eq!(
@@ -1485,11 +1559,11 @@ fn test_query_failure_reason_distinguishes_fenced_repair_stale_fence() {
let generic_err = flow_error_with_status(StatusCode::Unexpected);
assert_eq!(
BatchingTask::query_failure_reason(&generic_err, &QueryCoverage::ScopedBaseRepair),
FlowQueryFallbackReason::QueryFailure
FlowQueryFallbackReason::QueryFailure,
);
assert_eq!(
BatchingTask::query_failure_reason(&generic_err, &QueryCoverage::IncrementalDelta),
FlowQueryFallbackReason::IncrementalQueryFailure
FlowQueryFallbackReason::IncrementalQueryFailure,
);
}
@@ -1542,7 +1616,6 @@ async fn test_fenced_repair_stale_fence_next_plan_is_scoped_base_repair() {
std::time::Duration::from_millis(1),
&QueryCoverage::FencedRepairChunk { high },
FlowQueryFallbackReason::SnapshotFenceExpired,
false,
);
assert_eq!(
decision,
@@ -1605,7 +1678,6 @@ fn test_fenced_repair_transient_non_stale_failure_retries_same_high() {
std::time::Duration::from_millis(1),
&QueryCoverage::FencedRepairChunk { high: high.clone() },
FlowQueryFallbackReason::QueryFailure,
false,
);
assert_eq!(
@@ -2047,6 +2119,44 @@ async fn test_full_snapshot_seeding_applies_expire_after_retention_filter() {
);
}
#[tokio::test]
async fn test_forced_full_snapshot_keeps_retention_without_dirty_windows() {
let TestTaskParts {
mut task,
query_engine,
..
} = new_time_window_test_task_with_query(
"SELECT max(number) AS number, date_bin(INTERVAL '5 second', ts) AS time_window FROM numbers_with_ts GROUP BY time_window",
)
.await;
Arc::get_mut(&mut task.config)
.expect("test task config should be uniquely owned")
.expire_after = Some(expire_after_for_retention_filter_test());
let sink_schema = aggregate_time_window_sink_schema();
let plan = task
.gen_query_with_time_window_with_values(
query_engine,
&sink_schema,
&[],
false,
None,
&BTreeMap::new(),
true,
)
.await
.unwrap()
.expect("forced full snapshot must not require a dirty signal");
assert!(matches!(plan.coverage, QueryCoverage::UnfilteredFull));
assert!(task.state.read().unwrap().dirty_time_windows.is_empty());
let plan_text = plan.plan.to_string();
assert!(
plan_text.contains("Filter: ts >= TimestampMillisecond("),
"forced full snapshot retains the configured retention predicate: {plan_text}"
);
}
#[tokio::test]
async fn test_incremental_plan_does_not_add_dirty_window_filter() {
let TestTaskParts {
@@ -2165,7 +2275,6 @@ async fn test_successful_incremental_checkpoint_fallback_consumes_unscoped_dirty
&result,
std::time::Duration::from_millis(1),
&plan_info.coverage,
false,
)
};
assert_eq!(
@@ -2423,7 +2532,6 @@ async fn test_unsafe_incremental_plan_skip_restores_dirty_without_query() {
&dml_plan,
&dirty_restore,
&QueryCoverage::IncrementalDelta,
None,
)
.await
.unwrap();
@@ -2575,229 +2683,6 @@ async fn test_auto_created_sql_aggregate_sink_reaches_incremental_safe() {
);
}
struct TestIncrementalPersistence;
#[async_trait::async_trait]
impl BatchingPersistence for TestIncrementalPersistence {
async fn restore(&self) -> crate::Result<RestoreOutcome> {
Ok(RestoreOutcome::TrustedCheckpoint(BTreeMap::from([
(1_u64, 10_u64),
(2_u64, 20_u64),
])))
}
async fn begin_attempt(&self) -> crate::Result<BatchingAttempt> {
Ok(BatchingAttempt::default())
}
async fn persist(
&self,
_attempt: BatchingAttempt,
_validated_checkpoints: BTreeMap<u64, u64>,
) -> crate::Result<()> {
Ok(())
}
}
#[derive(Debug, PartialEq, Eq)]
struct ModeledSink {
a: u64,
b: u64,
}
struct StatefulRepairHandler {
sink: Arc<std::sync::Mutex<ModeledSink>>,
query_engine: QueryEngineRef,
}
struct RepairOutputStream {
schema: Arc<Schema>,
metrics: RecordBatchMetrics,
}
impl futures::Stream for RepairOutputStream {
type Item = common_recordbatch::error::Result<RecordBatch>;
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Poll::Ready(None)
}
}
impl RecordBatchStream for RepairOutputStream {
fn name(&self) -> &str {
"StatefulRepairHandler"
}
fn schema(&self) -> datatypes::schema::SchemaRef {
self.schema.clone()
}
fn output_ordering(&self) -> Option<&[common_recordbatch::OrderOption]> {
None
}
fn metrics(&self) -> Option<RecordBatchMetrics> {
Some(self.metrics.clone())
}
}
#[async_trait::async_trait]
impl crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError
for StatefulRepairHandler
{
async fn do_query(
&self,
query: api::v1::greptime_request::Request,
ctx: QueryContextRef,
) -> std::result::Result<Output, BoxedError> {
let api::v1::greptime_request::Request::Query(api::v1::QueryRequest {
query: Some(api::v1::query_request::Query::InsertIntoPlan(insert)),
}) = query
else {
panic!("expected an InsertIntoPlan request");
};
let incremental_mode = ctx.extension(FLOW_INCREMENTAL_MODE);
if incremental_mode.is_some() {
assert_eq!(incremental_mode, Some(FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY));
let after_seqs = ctx
.extension(FLOW_INCREMENTAL_AFTER_SEQS)
.expect("incremental mode must carry sequence bounds");
assert_eq!(
serde_json::from_str::<BTreeMap<u64, u64>>(after_seqs).unwrap(),
BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)])
);
// This models the boundary failure: the incremental A+B delta is
// sent to the sink, but only A is applied before the sink errors.
let mut sink = self.sink.lock().unwrap();
assert_eq!((sink.a, sink.b), (10, 20));
sink.a = 11;
return Err(BoxedError::new(MockError::new(StatusCode::Unexpected)));
}
assert!(
ctx.extension(FLOW_INCREMENTAL_AFTER_SEQS).is_none(),
"full repair must not carry incremental sequence bounds"
);
let decoder = self
.query_engine
.engine_context(ctx.clone())
.new_plan_decoder()
.unwrap();
let catalog = Arc::new(
catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx(
self.query_engine.engine_state().catalog_manager().clone(),
ctx.clone(),
),
);
let logical_plan = decoder
.decode(Bytes::from(insert.logical_plan), catalog, false)
.await
.unwrap();
assert!(
!logical_plan.to_string().contains("Filter"),
"full repair must not include a dirty-window predicate: {logical_plan}"
);
// The fixture's source timestamps share one window; this models the
// unfiltered A+B sink repair explicitly rather than asserting source
// rows for distinct windows.
let mut sink = self.sink.lock().unwrap();
assert_eq!(sink.a, 11);
assert_eq!(sink.b, 20);
sink.b = 21;
let metrics = RecordBatchMetrics {
region_watermarks: vec![
RegionWatermarkEntry {
region_id: 1,
watermark: Some(11),
},
RegionWatermarkEntry {
region_id: 2,
watermark: Some(21),
},
],
..Default::default()
};
Ok(Output::new_with_stream(Box::pin(RepairOutputStream {
schema: aggregate_time_window_sink_schema(),
metrics,
})))
}
}
#[tokio::test]
async fn test_persistence_full_repair_retries_unfiltered_snapshot_after_partial_sink_failure() {
let TestTaskParts {
mut task,
query_engine,
..
} = new_time_window_test_task_with_query(
"SELECT max(number) AS number, date_bin(INTERVAL '5 second', ts) AS time_window FROM numbers_with_ts GROUP BY time_window",
)
.await;
let sink_table = "incremental_failure_repair_sink";
Arc::get_mut(&mut task.config)
.expect("test task config should be uniquely owned")
.sink_table_name[2] = sink_table.to_string();
register_twe_sink(&query_engine, sink_table, 9103);
task.set_persistence(Some(Arc::new(TestIncrementalPersistence)))
.await
.unwrap();
{
let mut state = task.state.write().unwrap();
// Only A is dirty. B is intentionally absent from the dirty signal;
// the later full repair must nevertheless repair the modeled A+B sink.
state
.dirty_time_windows
.add_window(Timestamp::new_second(10), Some(Timestamp::new_second(15)));
}
let sink = Arc::new(std::sync::Mutex::new(ModeledSink { a: 10, b: 20 }));
let handler: Arc<dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError> =
Arc::new(StatefulRepairHandler {
sink: sink.clone(),
query_engine: query_engine.clone(),
});
let frontend_client = Arc::new(FrontendClient::from_grpc_handler(
Arc::downgrade(&handler),
QueryOptions::default(),
));
let first = task
.execute_once_serialized(&query_engine, &frontend_client, Some(2))
.await;
assert!(
first.is_err(),
"the partial incremental sink write must fail"
);
assert_eq!(*sink.lock().unwrap(), ModeledSink { a: 11, b: 20 });
{
let state = task.state.read().unwrap();
assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot);
assert_eq!(
state.checkpoints(),
&BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)])
);
assert!(state.full_repair_required());
}
let second = task
.execute_once_serialized(&query_engine, &frontend_client, Some(2))
.await
.expect("the unfiltered full repair should succeed");
assert!(second.is_some());
assert_eq!(*sink.lock().unwrap(), ModeledSink { a: 11, b: 21 });
let state = task.state.read().unwrap();
assert_eq!(state.checkpoint_mode(), CheckpointMode::Incremental);
assert_eq!(
state.checkpoints(),
&BTreeMap::from([(1_u64, 11_u64), (2_u64, 21_u64)])
);
assert!(!state.full_repair_required());
}
#[tokio::test]
async fn test_unscoped_failure_restores_consumed_dirty_signal() {
assert_unscoped_failure_restore(dirty_marker(), DirtyTimeWindows::default(), 1, 0).await;
@@ -2876,6 +2761,73 @@ async fn test_scoped_plan_generation_failure_restores_consumed_dirty_windows() {
);
}
#[tokio::test]
async fn test_validate_sink_table_schema_with_values_accepts_typed_metadata() {
let sink_table = "metadata_validation_sink";
let TestTaskParts {
mut task,
query_engine,
..
} = 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;
Arc::get_mut(&mut task.config)
.expect("test task config should be uniquely owned")
.sink_table_name[2] = sink_table.to_string();
register_twe_sink_with_metadata(&query_engine, sink_table, 9104);
let values = BTreeMap::from([(
"metadata".to_string(),
datafusion_common::ScalarValue::Utf8(Some("typed".to_string())),
)]);
assert!(
task.validate_sink_table_schema_with_values(&query_engine, &values)
.await
.is_ok()
);
assert!(task.state.read().unwrap().dirty_time_windows.is_empty());
}
#[tokio::test]
async fn test_validate_sink_table_schema_with_values_rejects_output_mismatch_without_dirty_change()
{
let sink_table = "metadata_validation_mismatch_sink";
let TestTaskParts {
mut task,
query_engine,
..
} = 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;
Arc::get_mut(&mut task.config)
.expect("test task config should be uniquely owned")
.sink_table_name[2] = sink_table.to_string();
register_number_only_sink(&query_engine, sink_table);
task.state
.write()
.unwrap()
.dirty_time_windows
.add_window(Timestamp::new_second(10), Some(Timestamp::new_second(15)));
let values = BTreeMap::from([(
"metadata".to_string(),
datafusion_common::ScalarValue::Utf8(Some("typed".to_string())),
)]);
assert!(
task.validate_sink_table_schema_with_values(&query_engine, &values)
.await
.is_err()
);
let state = task.state.read().unwrap();
assert_eq!(state.dirty_time_windows.len(), 1);
assert_eq!(
state.dirty_time_windows.window_size(),
std::time::Duration::from_secs(5)
);
}
#[tokio::test]
async fn test_insert_plan_matching_failure_restores_consumed_dirty_marker() {
let sink_table = "partial_sink";
src/flow/src/batching_mode/utils.rs
+2 -2
View File
@@ -1345,7 +1345,7 @@ impl ColumnMatcherRewriter {
return self.modify_project_exprs_with_partial(exprs);
}
// Ordinary values are persistence-owned columns, not flow outputs. Remove them from the
// Ordinary values are execution-owned columns, not flow outputs. Remove them from the
// effective sink sequence while deciding whether the existing auto-column rules apply.
// This keeps those columns from hiding an auto-created update_at column that precedes them.
let effective_sink_columns = self
@@ -1428,7 +1428,7 @@ impl ColumnMatcherRewriter {
&effective_sink_columns,
)?;
// Put persistence-owned values back at their physical sink positions only after matching
// Put execution-owned values back at their physical sink positions only after matching
// flow expressions against the effective sequence.
let mut exprs = exprs;
for (idx, column) in self.schema.column_schemas().iter().enumerate() {
src/flow/src/lib.rs
+10 -5
View File
@@ -43,16 +43,21 @@ 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 use batching_mode::batching_execution::{BatchingExecution, BatchingExecutionFactory};
pub use batching_mode::frontend_client::{
FrontendClient, GrpcQueryHandlerWithBoxedError, PeerDesc,
};
pub use batching_mode::task::{
BatchingTask, DirtyRestore, ExecuteOnceOutcome, PlanInfo, QueryCoverage, TaskArgs,
};
pub use batching_mode::time_window::{TimeWindowExpr, find_time_window_expr};
pub use batching_mode::utils::sql_to_df_plan;
pub use batching_mode::{BatchingModeOptions, IncrementalMode};
pub(crate) use engine::{CreateFlowArgs, FlowId, TableName};
pub use error::{Error, Result};
pub use server::{
FlownodeBuilder, FlownodeInstance, FlownodeServer, FlownodeServiceBuilder, FrontendInvoker,
};
pub use transform::register_function_to_query_engine;
pub use crate::adapter::FlownodeOptions;
src/flow/src/server.rs
+3 -4
View File
@@ -57,7 +57,6 @@ 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,
@@ -66,7 +65,7 @@ use crate::heartbeat::HeartbeatTask;
use crate::metrics::{METRIC_FLOW_PROCESSING_TIME, METRIC_FLOW_ROWS};
use crate::transform::register_function_to_query_engine;
use crate::utils::{SizeReportSender, StateReportHandler};
use crate::{Error, FlownodeOptions, FrontendClient, StreamingEngine};
use crate::{BatchingExecutionFactory, Error, FlownodeOptions, FrontendClient, StreamingEngine};
pub const FLOW_NODE_SERVER_NAME: &str = "FLOW_NODE_SERVER";
/// wrapping flow node manager to avoid orphan rule with Arc<...>
@@ -405,14 +404,14 @@ impl FlownodeBuilder {
self.build_manager(query_engine_factory.query_engine())
.await?,
);
let batching = Arc::new(BatchingEngine::new_with_persistence(
let batching = Arc::new(BatchingEngine::new_with_execution(
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.plugins.get::<FactoryPlugin>(),
self.plugins.get::<Arc<dyn BatchingExecutionFactory>>(),
));
let dual = Arc::new(FlowDualEngine::new(
manager.clone(),