diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index fada4e18d5..5ebf8797a0 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/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, - ) -> Result, 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, diff --git a/src/flow/src/batching_mode/utils.rs b/src/flow/src/batching_mode/utils.rs index 454abf1182..53d3fe5a6f 100644 --- a/src/flow/src/batching_mode/utils.rs +++ b/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::>(); - 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>, + ordinary_values: &BTreeMap, ) -> Result { 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() diff --git a/src/flow/src/batching_mode/utils/test.rs b/src/flow/src/batching_mode/utils/test.rs index 5479874f26..e5213bb8f7 100644 --- a/src/flow/src/batching_mode/utils/test.rs +++ b/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()