refactor(flow): consolidate internal planning and dirty drains

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

src/flow/src/batching_mode/task.rs

@@ -469,16 +469,6 @@ impl BatchingTask {
             self.config.sink_table_name.clone(),
         )
         .await?;
-        self.validate_sink_table_schema_with_table_and_values(engine, table, values)
-            .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);
@@ -491,7 +481,7 @@ impl BatchingTask {
             table_meta.schema.clone(),
             &primary_key_indices,
             merge_mode_last_non_null,
-            Some(values),
+            values,
         )
         .await
         .map(|_| table_meta.schema.clone())
@@ -920,10 +910,8 @@ impl BatchingTask {
     /// Consume the live dirty signal for an unscoped query while keeping a copy
     /// that can be restored if planning or execution fails.
     fn drain_dirty_windows_signal(&self) -> (bool, DirtyTimeWindows) {
-        let mut state = self.state.write().unwrap();
-        let dirty_windows_to_restore = state.dirty_time_windows.clone();
+        let dirty_windows_to_restore = self.state.write().unwrap().dirty_time_windows.detach();
         let is_dirty = !dirty_windows_to_restore.is_empty();
-        state.dirty_time_windows.clean();
         (is_dirty, dirty_windows_to_restore)
     }
 
@@ -951,7 +939,7 @@ impl BatchingTask {
                 sink_table_schema,
                 primary_key_indices,
                 allow_partial,
-                Some(values),
+                values,
             )
             .await,
         )?;
@@ -1513,6 +1501,7 @@ impl BatchingTask {
 
     /// Generate the next plan and classify its coverage so checkpoint handling
     /// knows whether it is full-query, scoped repair, fenced repair, or delta.
+    #[cfg(test)]
     async fn gen_query_with_time_window(
         &self,
         engine: QueryEngineRef,

src/flow/src/batching_mode/utils.rs

@@ -790,8 +790,7 @@ pub async fn rewrite_incremental_aggregate_with_sink_merge(
         .iter()
         .map(|c| (&c.output_field_name, c))
         .collect::<HashMap<_, _>>();
-    let mut projection_exprs = Vec::with_capacity(analysis.output_field_names.len());
-    let mut group_exprs = Vec::new();
+    let mut ordinary_exprs = Vec::with_capacity(analysis.output_field_names.len());
     let mut state_aggr_exprs = Vec::new();
     for output_field_name in &analysis.output_field_names {
         if group_key_names.contains(output_field_name)
@@ -799,8 +798,7 @@ pub async fn rewrite_incremental_aggregate_with_sink_merge(
         {
             let expr =
                 qualified_col(delta_alias, output_field_name.clone()).alias(output_field_name);
-            projection_exprs.push(expr.clone());
-            group_exprs.push(expr);
+            ordinary_exprs.push(expr);
         } else if let Some(merge_col) = merge_columns.get(output_field_name) {
             if matches!(
                 &merge_col.merge_op,
@@ -809,8 +807,7 @@ pub async fn rewrite_incremental_aggregate_with_sink_merge(
                 state_aggr_exprs.push(build_state_delta_merge_expr(engine, merge_col)?);
             } else {
                 let expr = build_left_join_merge_expr(delta_alias, sink_alias, merge_col)?;
-                projection_exprs.push(expr.clone());
-                group_exprs.push(expr);
+                ordinary_exprs.push(expr);
             }
         } else {
             return InvalidQuerySnafu {
@@ -824,7 +821,7 @@ pub async fn rewrite_incremental_aggregate_with_sink_merge(
 
     if state_merge {
         let aggregated = LogicalPlanBuilder::from(joined)
-            .aggregate(group_exprs, state_aggr_exprs)
+            .aggregate(ordinary_exprs, state_aggr_exprs)
             .with_context(|_| DatafusionSnafu {
                 context: "Failed to aggregate state delta merge plan".to_string(),
             })?
@@ -849,7 +846,7 @@ pub async fn rewrite_incremental_aggregate_with_sink_merge(
             })
     } else {
         LogicalPlanBuilder::from(joined)
-            .project(projection_exprs)
+            .project(ordinary_exprs)
             .with_context(|_| DatafusionSnafu {
                 context: "Failed to build projection merge plan for incremental sink merge"
                     .to_string(),
@@ -1073,6 +1070,7 @@ pub async fn sql_to_df_plan(
 
 /// Generate a plan that matches the schema of the sink table
 /// from given sql by alias and adding auto columns
+#[cfg(test)]
 pub(crate) async fn gen_plan_with_matching_schema(
     sql: &str,
     query_ctx: QueryContextRef,
@@ -1088,7 +1086,7 @@ pub(crate) async fn gen_plan_with_matching_schema(
         sink_table_schema,
         primary_key_indices,
         allow_partial,
-        None,
+        &BTreeMap::new(),
     )
     .await
 }
@@ -1100,7 +1098,7 @@ pub(crate) async fn gen_plan_with_matching_schema_and_values(
     sink_table_schema: SchemaRef,
     primary_key_indices: &[usize],
     allow_partial: bool,
-    ordinary_values: Option<&BTreeMap<String, ScalarValue>>,
+    ordinary_values: &BTreeMap<String, ScalarValue>,
 ) -> Result<LogicalPlan, Error> {
     let plan = sql_to_df_plan(query_ctx.clone(), engine.clone(), sql, false).await?;
 
@@ -1108,7 +1106,7 @@ pub(crate) async fn gen_plan_with_matching_schema_and_values(
         sink_table_schema,
         primary_key_indices.to_vec(),
         allow_partial,
-        ordinary_values.cloned().unwrap_or_default(),
+        ordinary_values.clone(),
     );
     let plan = plan
         .clone()

src/flow/src/batching_mode/utils/test.rs

@@ -960,7 +960,7 @@ async fn test_gen_plan_with_matching_schema_injects_attempt_columns_in_sink_orde
         sink_schema,
         &[],
         false,
-        Some(&values),
+        &values,
     )
     .await
     .unwrap();
@@ -1002,7 +1002,7 @@ async fn test_gen_plan_with_matching_schema_rejects_arbitrary_missing_attempt_co
         sink_schema,
         &[],
         false,
-        Some(&BTreeMap::new()),
+        &BTreeMap::new(),
     )
     .await
     .unwrap_err()
@@ -2485,7 +2485,7 @@ async fn test_gen_plan_with_matching_schema_rejects_unknown_attempt_column() {
             sink_schema,
             primary_key_indices,
             allow_partial,
-            Some(&values),
+            &values,
         )
         .await
         .unwrap_err()
@@ -2520,7 +2520,7 @@ async fn test_gen_plan_with_matching_schema_rejects_wrong_attempt_column_type()
         sink_schema,
         &[0],
         false,
-        Some(&values),
+        &values,
     )
     .await
     .unwrap_err()
@@ -2554,7 +2554,7 @@ async fn test_gen_plan_with_matching_schema_matches_positional_alias_and_injects
         sink_schema,
         &[0],
         false,
-        Some(&values),
+        &values,
     )
     .await
     .unwrap();
@@ -2594,7 +2594,7 @@ async fn test_gen_plan_with_matching_schema_injects_ordinary_columns_after_auto_
         sink_schema,
         &[],
         false,
-        Some(&ordinary_values),
+        &ordinary_values,
     )
     .await
     .unwrap();
@@ -2686,7 +2686,7 @@ async fn test_gen_plan_with_matching_schema_rejects_attempt_output_collision() {
             sink_schema,
             &[0],
             allow_partial,
-            Some(&values),
+            &values,
         )
         .await
         .unwrap_err()