feat(flow): allow execution hook to rewrite completed plans

Signed-off-by: discord9 <discord9@163.com>

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<usize>,
     ) -> ExecuteOnceOutcome;
 
+    /// Rewrite the completed query plan after incremental merging and before execution.
+    fn rewrite_plan(&self, _task: &BatchingTask, plan: LogicalPlan) -> Result<LogicalPlan> {
+        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) {}

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())