diff --git a/src/flow/src/batching_mode/engine.rs b/src/flow/src/batching_mode/engine.rs index a89f8ab63d..458df50a4a 100644 --- a/src/flow/src/batching_mode/engine.rs +++ b/src/flow/src/batching_mode/engine.rs @@ -1448,17 +1448,16 @@ mod tests { async fn test_flow_option_parser_matrix() { let engine = new_test_engine().await; let cases = [ - (None, false, false), - (Some("true"), true, false), - (Some("false"), false, false), + (None, Ok((false, false))), + (Some("true"), Ok((true, false))), + (Some("false"), Ok((false, false))), ( Some(FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_SEQUENCE_RANGE), - true, - true, + Ok((true, true)), ), - (Some("malformed"), false, false), + (Some("malformed"), Err(())), ]; - for (value, enabled, required) in cases { + for (value, expected) in cases { let options = value .map(|value| { HashMap::from([( @@ -1467,36 +1466,27 @@ mod tests { )]) }) .unwrap_or_default(); - let required_from_validated_sentinel = options + let exact_sequence_range_required = options .get(FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY) .is_some_and(|value| { value == FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_SEQUENCE_RANGE }); - match engine.batch_opts_for_flow_options(&options, required_from_validated_sentinel) { - Ok(opts) => { + match ( + engine.batch_opts_for_flow_options(&options, exact_sequence_range_required), + expected, + ) { + (Ok(opts), Ok((enabled, required))) => { assert_eq!(opts.experimental_enable_incremental_read, enabled); - assert_eq!(required_from_validated_sentinel, required); + assert_eq!(exact_sequence_range_required, required); } - Err(_) => assert!(!enabled && !required), + (Err(_), Err(())) => {} + (result, expected) => panic!( + "unexpected parser result for {value:?}: {result:?}, expected {expected:?}" + ), } } } - #[tokio::test] - async fn test_forged_query_context_does_not_enable_exact_sequence_range() { - 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"); - args.query_ctx = Some(query_ctx); - - engine.create_flow_inner(args).await.unwrap(); - let task = engine.runtime.read().await.tasks.get(&4).cloned().unwrap(); - assert!(!task.config.exact_sequence_range_required); - engine.remove_flow_inner(4).await.unwrap(); - } - #[tokio::test] async fn test_exact_sequence_range_capability_is_checked_before_task_startup() { let engine = new_test_engine_with_execution(None).await; diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index 4f5e858b1d..34b8423e29 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -1568,13 +1568,13 @@ impl BatchingTask { 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) + expire_time_window_bound + .as_ref() + .and_then(|(lower, _)| lower.as_ref()) .map(|lower| { ( expr.column_name.as_str(), - lower, + lower.clone(), "forced full snapshot retention", ) }) diff --git a/src/flow/src/batching_mode/task/inc.rs b/src/flow/src/batching_mode/task/inc.rs index a87a14b215..61dd9934ac 100644 --- a/src/flow/src/batching_mode/task/inc.rs +++ b/src/flow/src/batching_mode/task/inc.rs @@ -21,7 +21,7 @@ use datafusion_expr::{DmlStatement, LogicalPlan}; use query::QueryEngineRef; use query::options::{ FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY, - FLOW_SINK_TABLE_ID, + FLOW_INCREMENTAL_MODE_SEQUENCE_RANGE, FLOW_SINK_TABLE_ID, }; use snafu::ResultExt; use store_api::mito_engine_options::PRESERVE_ROW_SEQUENCE; @@ -37,9 +37,6 @@ use crate::batching_mode::utils::{ }; use crate::error::{ExternalSnafu, UnexpectedSnafu}; -// Kept local until the query-side extension enum exposes the exact scan mode. -const FLOW_INCREMENTAL_MODE_SEQUENCE_RANGE: &str = "sequence_range"; - impl BatchingTask { async fn sink_table_id(&self) -> Result { let table = self diff --git a/src/flow/src/batching_mode/utils.rs b/src/flow/src/batching_mode/utils.rs index 01beb82714..0bb79bd196 100644 --- a/src/flow/src/batching_mode/utils.rs +++ b/src/flow/src/batching_mode/utils.rs @@ -98,7 +98,6 @@ pub enum IncrementalAggregateMergeOp { BitAnd, BitOr, BitXor, - AvgDeltaMerge, StateDeltaMerge { function_name: &'static str, params: Vec, @@ -225,8 +224,6 @@ struct OutputProjectionInfo { has_top_level_projection: bool, /// Aggregate expression name and projected output field, in projection order. aggregate_outputs: Vec<(String, String)>, - /// Original single-instance resolver mapping, retained for compatibility. - output_aliases: HashMap, literal_columns: HashSet, output_field_names: Vec, } @@ -260,7 +257,6 @@ fn collect_output_projection_info(plan: &LogicalPlan) -> OutputProjectionInfo { ..Default::default() }; - let mut output_aliases = HashMap::new(); if let LogicalPlan::Projection(projection) = plan { for expr in &projection.expr { match expr { @@ -271,9 +267,6 @@ fn collect_output_projection_info(plan: &LogicalPlan) -> OutputProjectionInfo { // other output wrappers. let alias_name = alias.name.clone(); if let Expr::Column(column) = alias.expr.as_ref() { - output_aliases - .entry(column.name.clone()) - .or_insert_with(|| alias_name.clone()); projection_info .aggregate_outputs .push((column.name.clone(), alias_name)); @@ -281,9 +274,6 @@ fn collect_output_projection_info(plan: &LogicalPlan) -> OutputProjectionInfo { && inner_alias.name.eq_ignore_ascii_case("count(*)") && let Expr::Column(column) = inner_alias.expr.as_ref() { - output_aliases - .entry(column.name.clone()) - .or_insert_with(|| alias_name.clone()); projection_info .aggregate_outputs .push((column.name.clone(), alias_name)); @@ -316,7 +306,6 @@ fn collect_output_projection_info(plan: &LogicalPlan) -> OutputProjectionInfo { .insert(AUTO_CREATED_PLACEHOLDER_TS_COL.to_string()); } - projection_info.output_aliases = output_aliases; projection_info } @@ -374,7 +363,7 @@ fn merge_op_for_aggregate_expr( "bit_xor" => Ok(IncrementalAggregateMergeOp::BitXor), // Preserve state-family parameters; value coercion is handled by the aggregate. "avg_state" if aggr_func.params.args.len() == 1 => { - Ok(IncrementalAggregateMergeOp::AvgDeltaMerge) + state_delta_merge("__avg_state_delta_merge", vec![]) } "hll" if aggr_func.params.args.len() == 1 => state_delta_merge("__hll_delta_merge", vec![]), "stddev_pop_state" if aggr_func.params.args.len() == 1 => { @@ -400,7 +389,7 @@ fn merge_op_for_aggregate_expr( if aggr_func.params.args.len() == 1 && is_type(&aggr_func.params.args[0], ArrowDataType::Binary) => { - Ok(IncrementalAggregateMergeOp::AvgDeltaMerge) + state_delta_merge("__avg_state_delta_merge", vec![]) } _ => Err(aggr_expr.to_string()), } @@ -419,19 +408,12 @@ fn resolve_aggregate_output_fields( // one aggregate input field for identical expressions. let raw_name = aggr_expr.qualified_name().1; if projection_info.has_top_level_projection { - let outputs = projection_info + projection_info .aggregate_outputs .iter() .filter(|(input_name, _)| input_name == &raw_name) .cloned() - .collect::>(); - if outputs.len() > 1 { - outputs - } else if let Some(alias) = projection_info.output_aliases.get(&raw_name) { - vec![(raw_name, alias.clone())] - } else { - outputs - } + .collect() } else if output_field_name_set.contains(&raw_name) { vec![(raw_name.clone(), raw_name)] } else { @@ -692,8 +674,7 @@ pub async fn rewrite_incremental_aggregate_with_sink_merge( let state_merge = analysis.merge_columns.iter().any(|column| { matches!( &column.merge_op, - IncrementalAggregateMergeOp::AvgDeltaMerge - | IncrementalAggregateMergeOp::StateDeltaMerge { .. } + IncrementalAggregateMergeOp::StateDeltaMerge { .. } ) }); let mut selected_columns = analysis.group_key_names.clone(); @@ -840,8 +821,7 @@ pub async fn rewrite_incremental_aggregate_with_sink_merge( } else if let Some(merge_col) = merge_columns.get(output_field_name) { if matches!( &merge_col.merge_op, - IncrementalAggregateMergeOp::AvgDeltaMerge - | IncrementalAggregateMergeOp::StateDeltaMerge { .. } + IncrementalAggregateMergeOp::StateDeltaMerge { .. } ) { state_aggr_exprs.push(build_state_delta_merge_expr(engine, merge_col)?); } else { @@ -903,7 +883,6 @@ fn build_state_delta_merge_expr( merge_col: &IncrementalAggregateMergeColumn, ) -> Result { let (function_name, params) = match &merge_col.merge_op { - IncrementalAggregateMergeOp::AvgDeltaMerge => ("__avg_state_delta_merge", vec![]), IncrementalAggregateMergeOp::StateDeltaMerge { function_name, params, @@ -1000,8 +979,7 @@ fn build_left_join_merge_expr( .with_context(|_| DatafusionSnafu { context: "Failed to build BIT_XOR merge expression".to_string(), })?, - IncrementalAggregateMergeOp::AvgDeltaMerge - | IncrementalAggregateMergeOp::StateDeltaMerge { .. } => { + IncrementalAggregateMergeOp::StateDeltaMerge { .. } => { return InvalidQuerySnafu { reason: "state aggregate must be built with its delta UDAF".to_string(), } diff --git a/src/flow/src/batching_mode/utils/test.rs b/src/flow/src/batching_mode/utils/test.rs index 7bc9302986..7aa48865a1 100644 --- a/src/flow/src/batching_mode/utils/test.rs +++ b/src/flow/src/batching_mode/utils/test.rs @@ -23,7 +23,7 @@ use common_time::Timestamp; use datafusion_common::tree_node::TreeNode as _; use datafusion_expr::GroupingSet; use datatypes::arrow::array::{Array, AsArray}; -use datatypes::arrow::datatypes::{Float64Type, Int64Type, UInt64Type}; +use datatypes::arrow::datatypes::{DataType as ArrowDataType, Float64Type, Int64Type, UInt64Type}; use datatypes::prelude::{ConcreteDataType, MutableVector, Scalar, ScalarVectorBuilder, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::timestamp::TimestampMillisecond; @@ -1010,34 +1010,6 @@ async fn test_gen_plan_with_matching_schema_rejects_arbitrary_missing_attempt_co assert!(err.contains("missing sink columns"), "{err}"); } -#[tokio::test] -async fn test_gen_plan_with_matching_schema_no_attempt_still_rejects_missing_column() { - let query_engine = create_test_query_engine(); - let ctx = QueryContext::arc(); - let sink_schema = Arc::new(Schema::new(vec![ - ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true), - ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new("state", ConcreteDataType::uint32_datatype(), true), - ])); - assert!( - gen_plan_with_matching_schema( - "SELECT number, ts FROM numbers_with_ts", - ctx, - query_engine, - sink_schema, - &[], - false, - ) - .await - .is_err() - ); -} - #[tokio::test] async fn test_gen_plan_with_matching_schema_allow_partial_fills_nullable_columns() { let query_engine = create_test_query_engine(); @@ -1809,10 +1781,13 @@ async fn test_analyze_incremental_aggregate_plan_supports_avg_state() { ); assert_eq!(analysis.merge_columns.len(), 1); assert_eq!(analysis.merge_columns[0].output_field_name, "avg_num"); - assert_eq!( - analysis.merge_columns[0].merge_op, - IncrementalAggregateMergeOp::AvgDeltaMerge - ); + assert!(matches!( + &analysis.merge_columns[0].merge_op, + IncrementalAggregateMergeOp::StateDeltaMerge { + function_name: "__avg_state_delta_merge", + params, + } if params.is_empty() + )); } #[tokio::test] @@ -1829,10 +1804,13 @@ async fn test_analyze_incremental_aggregate_plan_supports_avg_merge() { analysis.unsupported_exprs ); assert_eq!(analysis.merge_columns.len(), 1); - assert_eq!( - analysis.merge_columns[0].merge_op, - IncrementalAggregateMergeOp::AvgDeltaMerge - ); + assert!(matches!( + &analysis.merge_columns[0].merge_op, + IncrementalAggregateMergeOp::StateDeltaMerge { + function_name: "__avg_state_delta_merge", + params, + } if params.is_empty() + )); } #[tokio::test] @@ -1844,12 +1822,15 @@ async fn test_analyze_incremental_aggregate_plan_supports_duplicate_avg_projecti assert!(analysis.unsupported_exprs.is_empty()); assert_eq!(analysis.merge_columns.len(), 2); - assert!( - analysis - .merge_columns - .iter() - .all(|column| { column.merge_op == IncrementalAggregateMergeOp::AvgDeltaMerge }) - ); + assert!(analysis.merge_columns.iter().all(|column| { + matches!( + &column.merge_op, + IncrementalAggregateMergeOp::StateDeltaMerge { + function_name: "__avg_state_delta_merge", + params, + } if params.is_empty() + ) + })); assert!( analysis .merge_columns @@ -1875,7 +1856,13 @@ async fn test_analyze_incremental_aggregate_plan_supports_avg_with_native_aggreg assert_eq!(analysis.merge_columns.len(), 2); assert!(analysis.merge_columns.iter().any(|column| { column.output_field_name == "avg_num" - && column.merge_op == IncrementalAggregateMergeOp::AvgDeltaMerge + && matches!( + &column.merge_op, + IncrementalAggregateMergeOp::StateDeltaMerge { + function_name: "__avg_state_delta_merge", + params, + } if params.is_empty() + ) })); assert!(analysis.merge_columns.iter().any(|column| { column.output_field_name == "total" && column.merge_op == IncrementalAggregateMergeOp::Sum @@ -2568,32 +2555,36 @@ async fn test_gen_plan_with_matching_schema_last_non_null_rejects_extra_flow_col #[tokio::test] async fn test_gen_plan_with_matching_schema_rejects_unknown_attempt_column() { - let query_engine = create_test_query_engine(); - let ctx = QueryContext::arc(); - let sink_schema = Arc::new(Schema::new(vec![ - ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true), - ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, + for allow_partial in [false, true] { + let query_engine = create_test_query_engine(); + let ctx = QueryContext::arc(); + let sink_schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ])); + let values = + BTreeMap::from([(String::from("unknown_attempt"), ScalarValue::Int64(Some(1)))]); + let primary_key_indices: &[usize] = if allow_partial { &[0] } else { &[] }; + let err = gen_plan_with_matching_schema_and_values( + "SELECT number, ts FROM numbers_with_ts", + ctx, + query_engine, + sink_schema, + primary_key_indices, + allow_partial, + Some(&values), ) - .with_time_index(true), - ])); - let values = BTreeMap::from([(String::from("unknown_attempt"), ScalarValue::Int64(Some(1)))]); - let err = gen_plan_with_matching_schema_and_values( - "SELECT number, ts FROM numbers_with_ts", - ctx, - query_engine, - sink_schema, - &[0], - false, - Some(&values), - ) - .await - .unwrap_err() - .to_string(); - assert!(err.contains("unknown_attempt"), "{err}"); - assert!(err.contains("does not exist in sink schema"), "{err}"); + .await + .unwrap_err() + .to_string(); + assert!(err.contains("unknown_attempt"), "{err}"); + assert!(err.contains("does not exist in sink schema"), "{err}"); + } } #[tokio::test] @@ -2672,59 +2663,24 @@ async fn test_gen_plan_with_matching_schema_matches_positional_alias_and_injects async fn test_gen_plan_with_matching_schema_injects_ordinary_columns_after_auto_update_at() { let query_engine = create_test_query_engine(); let ctx = QueryContext::arc(); - let mut sink_columns = (0..16) - .map(|idx| { - ColumnSchema::new( - format!("state_{idx}"), - ConcreteDataType::int32_datatype(), - true, - ) - }) - .collect::>(); - sink_columns.push(ColumnSchema::new( - "update_at", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - )); - sink_columns.extend([ + let sink_schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("state_0", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new("state_1", ConcreteDataType::int32_datatype(), true), ColumnSchema::new( - "__ee_checkpoint_epoch", - ConcreteDataType::uint32_datatype(), + "update_at", + ConcreteDataType::timestamp_millisecond_datatype(), false, ), - ColumnSchema::new( - "__ee_checkpoint_sequence", - ConcreteDataType::uint32_datatype(), - false, - ), - ColumnSchema::new( - "__ee_checkpoint_region", - ConcreteDataType::uint32_datatype(), - false, - ), - ]); - let sink_schema = Arc::new(Schema::new(sink_columns)); + ColumnSchema::new("metadata_a", ConcreteDataType::uint32_datatype(), false), + ColumnSchema::new("metadata_b", ConcreteDataType::uint32_datatype(), false), + ])); let ordinary_values = BTreeMap::from([ - ( - "__ee_checkpoint_epoch".to_string(), - ScalarValue::UInt32(Some(1)), - ), - ( - "__ee_checkpoint_sequence".to_string(), - ScalarValue::UInt32(Some(2)), - ), - ( - "__ee_checkpoint_region".to_string(), - ScalarValue::UInt32(Some(3)), - ), + ("metadata_a".to_string(), ScalarValue::UInt32(Some(1))), + ("metadata_b".to_string(), ScalarValue::UInt32(Some(2))), ]); - let flow_exprs = (0..16) - .map(|idx| format!("number AS state_{idx}")) - .collect::>(); - let sql = format!("SELECT {} FROM numbers_with_ts", flow_exprs.join(", ")); let plan = gen_plan_with_matching_schema_and_values( - &sql, + "SELECT number AS state_0, number AS state_1 FROM numbers_with_ts", ctx, query_engine, sink_schema, @@ -2746,26 +2702,28 @@ async fn test_gen_plan_with_matching_schema_injects_ordinary_columns_after_auto_ vec![ "state_0", "state_1", - "state_2", - "state_3", - "state_4", - "state_5", - "state_6", - "state_7", - "state_8", - "state_9", - "state_10", - "state_11", - "state_12", - "state_13", - "state_14", - "state_15", "update_at", - "__ee_checkpoint_epoch", - "__ee_checkpoint_sequence", - "__ee_checkpoint_region", + "metadata_a", + "metadata_b", ] ); + assert_eq!(plan.schema().field(3).data_type(), &ArrowDataType::UInt32); + assert_eq!(plan.schema().field(4).data_type(), &ArrowDataType::UInt32); + let LogicalPlan::Projection(projection) = plan else { + panic!("expected projection plan"); + }; + assert!(matches!( + &projection.expr[3], + Expr::Alias(alias) + if alias.name == "metadata_a" + && matches!(alias.expr.as_ref(), Expr::Literal(ScalarValue::UInt32(Some(1)), _)) + )); + assert!(matches!( + &projection.expr[4], + Expr::Alias(alias) + if alias.name == "metadata_b" + && matches!(alias.expr.as_ref(), Expr::Literal(ScalarValue::UInt32(Some(2)), _)) + )); } #[tokio::test] @@ -2797,65 +2755,37 @@ async fn test_gen_plan_with_matching_schema_rejects_no_attempt_strict_mismatch() assert!(err.contains("attempt"), "{err}"); } -#[tokio::test] -async fn test_gen_plan_with_matching_schema_rejects_unknown_attempt_column_in_partial_mode() { - let query_engine = create_test_query_engine(); - let ctx = QueryContext::arc(); - let sink_schema = Arc::new(Schema::new(vec![ - ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true), - ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new("optional", ConcreteDataType::uint32_datatype(), true), - ])); - let values = BTreeMap::from([(String::from("unknown_attempt"), ScalarValue::Int64(Some(1)))]); - let err = gen_plan_with_matching_schema_and_values( - "SELECT number, ts FROM numbers_with_ts", - ctx, - query_engine, - sink_schema, - &[0], - true, - Some(&values), - ) - .await - .unwrap_err() - .to_string(); - assert!(err.contains("unknown_attempt"), "{err}"); - assert!(err.contains("does not exist in sink schema"), "{err}"); -} - #[tokio::test] async fn test_gen_plan_with_matching_schema_rejects_attempt_output_collision() { - let query_engine = create_test_query_engine(); - let ctx = QueryContext::arc(); - let sink_schema = Arc::new(Schema::new(vec![ - ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true), - ColumnSchema::new("attempt", ConcreteDataType::uint32_datatype(), true), - ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, + for allow_partial in [false, true] { + let query_engine = create_test_query_engine(); + let ctx = QueryContext::arc(); + let sink_schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true), + ColumnSchema::new("attempt", ConcreteDataType::uint32_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ])); + let values = BTreeMap::from([(String::from("attempt"), ScalarValue::UInt32(Some(1)))]); + let err = gen_plan_with_matching_schema_and_values( + "SELECT number, number AS attempt, ts FROM numbers_with_ts", + ctx, + query_engine, + sink_schema, + &[0], + allow_partial, + Some(&values), ) - .with_time_index(true), - ])); - let values = BTreeMap::from([(String::from("attempt"), ScalarValue::UInt32(Some(1)))]); - let err = gen_plan_with_matching_schema_and_values( - "SELECT number, number AS attempt, ts FROM numbers_with_ts", - ctx, - query_engine, - sink_schema, - &[0], - false, - Some(&values), - ) - .await - .unwrap_err() - .to_string(); - assert!(err.contains("collides with a flow output"), "{err}"); + .await + .unwrap_err() + .to_string(); + assert!(err.contains("attempt"), "{err}"); + assert!(err.contains("collides with a flow output"), "{err}"); + } } #[tokio::test] @@ -2885,33 +2815,3 @@ async fn test_gen_plan_with_matching_schema_rejects_duplicate_original_outputs() assert!(diagnostic.contains("duplicate column"), "{diagnostic}"); assert!(diagnostic.contains("number"), "{diagnostic}"); } - -#[tokio::test] -async fn test_gen_plan_with_matching_schema_rejects_attempt_output_collision_in_partial_mode() { - let query_engine = create_test_query_engine(); - let ctx = QueryContext::arc(); - let sink_schema = Arc::new(Schema::new(vec![ - ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true), - ColumnSchema::new("attempt", ConcreteDataType::uint32_datatype(), true), - ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ])); - let values = BTreeMap::from([(String::from("attempt"), ScalarValue::UInt32(Some(1)))]); - let err = gen_plan_with_matching_schema_and_values( - "SELECT number, number AS attempt, ts FROM numbers_with_ts", - ctx, - query_engine, - sink_schema, - &[0], - true, - Some(&values), - ) - .await - .unwrap_err() - .to_string(); - assert!(err.contains("collides with a flow output"), "{err}"); -} diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 89dfef4ccb..9e1f5d9cbb 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -826,7 +826,7 @@ mod tests { use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr}; use datafusion::prelude::{col, lit}; - use datafusion_common::{JoinType, NullEquality}; + use datafusion_common::{JoinType, NullEquality, ScalarValue}; use datafusion_physical_expr::expressions::Column; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, SchemaRef}; @@ -1446,7 +1446,7 @@ mod tests { _request: table::requests::DeleteRequest, _ctx: session::context::QueryContextRef, ) -> common_query::error::Result { - Ok(0) + unimplemented!("unexpected delete") } async fn flush( @@ -1454,7 +1454,7 @@ mod tests { _request: table::requests::FlushTableRequest, _ctx: session::context::QueryContextRef, ) -> common_query::error::Result { - Ok(0) + unimplemented!("unexpected flush") } async fn compact( @@ -1462,7 +1462,7 @@ mod tests { _request: table::requests::CompactTableRequest, _ctx: session::context::QueryContextRef, ) -> common_query::error::Result { - Ok(0) + unimplemented!("unexpected compact") } async fn build_index( @@ -1470,7 +1470,7 @@ mod tests { _request: table::requests::BuildIndexTableRequest, _ctx: session::context::QueryContextRef, ) -> common_query::error::Result { - Ok(0) + unimplemented!("unexpected build_index") } async fn flush_region( @@ -1478,7 +1478,7 @@ mod tests { _region_id: store_api::storage::RegionId, _ctx: session::context::QueryContextRef, ) -> common_query::error::Result { - Ok(0) + unimplemented!("unexpected flush_region") } async fn compact_region( @@ -1486,7 +1486,7 @@ mod tests { _region_id: store_api::storage::RegionId, _ctx: session::context::QueryContextRef, ) -> common_query::error::Result { - Ok(0) + unimplemented!("unexpected compact_region") } async fn discard_unflushed_data( @@ -1494,7 +1494,7 @@ mod tests { _region_id: store_api::storage::RegionId, _ctx: session::context::QueryContextRef, ) -> common_query::error::Result { - Ok(0) + unimplemented!("unexpected discard_unflushed_data") } async fn discard_unflushed_data_by_table( @@ -1502,7 +1502,7 @@ mod tests { _table_name: table::table_name::TableName, _ctx: session::context::QueryContextRef, ) -> common_query::error::Result { - Ok(0) + unimplemented!("unexpected discard_unflushed_data_by_table") } } @@ -1675,18 +1675,17 @@ mod tests { assert!(vector.is_null(0), "{name}"); } assert!(!request.columns_values["ts"].is_null(0)); - assert_eq!( - request.columns_values["marker"].data_type(), - ConcreteDataType::uint8_datatype() - ); - assert_eq!( - request.columns_values["payload"].data_type(), - ConcreteDataType::binary_datatype() - ); - assert_eq!( - request.columns_values["epoch"].data_type(), - ConcreteDataType::uint64_datatype() - ); + for (name, data_type) in [ + ("marker", ConcreteDataType::uint8_datatype()), + ("payload", ConcreteDataType::binary_datatype()), + ("epoch", ConcreteDataType::uint64_datatype()), + ] { + assert_eq!( + request.columns_values[name].data_type(), + data_type, + "{name}" + ); + } let batches = run_sql( &engine, @@ -1695,58 +1694,33 @@ mod tests { .await; let batch = &batches[0]; assert_eq!(batch.num_rows(), 1); - assert_eq!( - batch.schema.column_schemas()[0].data_type, - ConcreteDataType::date_datatype() - ); - assert_eq!( - batch.schema.column_schemas()[1].data_type, - ConcreteDataType::date_datatype() - ); - assert_eq!( - batch.schema.column_schemas()[2].data_type, - ConcreteDataType::decimal128_datatype(30, 2) - ); - assert_eq!( - batch.schema.column_schemas()[3].data_type, - ConcreteDataType::duration_millisecond_datatype() - ); - assert_eq!( - batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 0 - ); - assert_eq!( - batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 2 - ); - assert_eq!( - batch - .column(2) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 30000 - ); - assert_eq!( - batch - .column(3) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 30 - ); + for (index, (data_type, value)) in [ + ( + ConcreteDataType::date_datatype(), + ScalarValue::Date32(Some(0)), + ), + ( + ConcreteDataType::date_datatype(), + ScalarValue::Date32(Some(2)), + ), + ( + ConcreteDataType::decimal128_datatype(30, 2), + ScalarValue::Decimal128(Some(30000), 30, 2), + ), + ( + ConcreteDataType::duration_millisecond_datatype(), + ScalarValue::DurationMillisecond(Some(30)), + ), + ] + .into_iter() + .enumerate() + { + assert_eq!(batch.schema.column_schemas()[index].data_type, data_type); + assert_eq!( + ScalarValue::try_from_array(batch.column(index).as_ref(), 0).unwrap(), + value + ); + } let batches = run_sql( &engine, @@ -1760,55 +1734,32 @@ mod tests { .await; let batch = &batches[0]; assert_eq!(batch.num_rows(), 1); - assert_eq!( - batch.schema.column_schemas()[0].data_type, - ConcreteDataType::decimal128_datatype(31, 2) - ); - assert_eq!( - batch.schema.column_schemas()[1].data_type, - ConcreteDataType::duration_millisecond_datatype() - ); - assert_eq!( - batch.schema.column_schemas()[2].data_type, - ConcreteDataType::boolean_datatype() - ); - assert_eq!( - batch.schema.column_schemas()[3].data_type, - ConcreteDataType::boolean_datatype() - ); - assert_eq!( - batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 60000 - ); - assert_eq!( - batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 60 - ); - assert!( - batch - .column(2) - .as_any() - .downcast_ref::() - .unwrap() - .value(0) - ); - assert!( - batch - .column(3) - .as_any() - .downcast_ref::() - .unwrap() - .value(0) - ); + for (index, (data_type, value)) in [ + ( + ConcreteDataType::decimal128_datatype(31, 2), + ScalarValue::Decimal128(Some(60000), 31, 2), + ), + ( + ConcreteDataType::duration_millisecond_datatype(), + ScalarValue::DurationMillisecond(Some(60)), + ), + ( + ConcreteDataType::boolean_datatype(), + ScalarValue::Boolean(Some(true)), + ), + ( + ConcreteDataType::boolean_datatype(), + ScalarValue::Boolean(Some(true)), + ), + ] + .into_iter() + .enumerate() + { + assert_eq!(batch.schema.column_schemas()[index].data_type, data_type); + assert_eq!( + ScalarValue::try_from_array(batch.column(index).as_ref(), 0).unwrap(), + value + ); + } } }