From e3dad9401c17ad515f8462c8dfe4dd9a58b34135 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 21 Sep 2026 19:47:46 +0800 Subject: [PATCH] feat(flow): allow execution hook to rewrite completed plans Signed-off-by: discord9 --- src/flow/src/batching_mode/batching_execution.rs | 6 ++++++ src/flow/src/batching_mode/task.rs | 4 ++++ 2 files changed, 10 insertions(+) diff --git a/src/flow/src/batching_mode/batching_execution.rs b/src/flow/src/batching_mode/batching_execution.rs index 0a699a227be..298522b8137 100644 --- a/src/flow/src/batching_mode/batching_execution.rs +++ b/src/flow/src/batching_mode/batching_execution.rs @@ -16,6 +16,7 @@ use std::sync::Arc; +use datafusion_expr::LogicalPlan; use query::QueryEngineRef; use table::TableRef; @@ -37,6 +38,11 @@ pub trait BatchingExecution: Send + Sync + 'static { max_window_cnt: Option, ) -> ExecuteOnceOutcome; + /// Rewrite the completed query plan after incremental merging and before execution. + fn rewrite_plan(&self, _task: &BatchingTask, plan: LogicalPlan) -> Result { + Ok(plan) + } + /// Retire this execution instance, rejecting new work and requesting that any /// retained local work stop. This is not an acknowledgement of remote quiescence. fn stop(&self) {} diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index c0ee6d1c1b5..346d39d1d46 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -828,6 +828,10 @@ impl BatchingTask { return Ok(None); } let plan = incremental_plan.unwrap_or_else(|| plan.clone()); + let plan = match &self.execution { + Some(execution) => execution.rewrite_plan(self, plan)?, + None => plan, + }; let extensions = self .build_flow_query_extensions(incremental_safe, coverage.is_incremental_delta())