rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-12 16:32:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(flow): remove redundant merge state and test scaffolding

Browse Source 
Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-09-08 14:05:06 +08:00
parent d5ecdd3361
commit 65ebaba42a
6 changed files with 221 additions and 405 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/flow/src/batching_mode/engine.rs
+17 -27
View File
@@ -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;
src/flow/src/batching_mode/task.rs
+4 -4
View File
@@ -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",
)
})
src/flow/src/batching_mode/task/inc.rs
+1 -4
View File
@@ -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<TableId, Error> {
let table = self
src/flow/src/batching_mode/utils.rs
+7 -29
View File
@@ -98,7 +98,6 @@ pub enum IncrementalAggregateMergeOp {
BitAnd,
BitOr,
BitXor,
AvgDeltaMerge,
StateDeltaMerge {
function_name: &'static str,
params: Vec<Expr>,
@@ -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<String, String>,
literal_columns: HashSet<String>,
output_field_names: Vec<String>,
}
@@ -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::<Vec<_>>();
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<Expr, Error> {
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(),
}
src/flow/src/batching_mode/utils/test.rs
+118 -218
View File
@@ -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::<Vec<_>>();
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::<Vec<_>>();
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}");
}
src/query/src/datafusion.rs
+74 -123
View File
@@ -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<common_base::AffectedRows> {
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<common_base::AffectedRows> {
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<common_base::AffectedRows> {
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<common_base::AffectedRows> {
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<common_base::AffectedRows> {
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<common_base::AffectedRows> {
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<common_base::AffectedRows> {
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<common_base::AffectedRows> {
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::<arrow::array::Date32Array>()
.unwrap()
.value(0),
0
);
assert_eq!(
batch
.column(1)
.as_any()
.downcast_ref::<arrow::array::Date32Array>()
.unwrap()
.value(0),
2
);
assert_eq!(
batch
.column(2)
.as_any()
.downcast_ref::<arrow::array::Decimal128Array>()
.unwrap()
.value(0),
30000
);
assert_eq!(
batch
.column(3)
.as_any()
.downcast_ref::<arrow::array::DurationMillisecondArray>()
.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::<arrow::array::Decimal128Array>()
.unwrap()
.value(0),
60000
);
assert_eq!(
batch
.column(1)
.as_any()
.downcast_ref::<arrow::array::DurationMillisecondArray>()
.unwrap()
.value(0),
60
);
assert!(
batch
.column(2)
.as_any()
.downcast_ref::<arrow::array::BooleanArray>()
.unwrap()
.value(0)
);
assert!(
batch
.column(3)
.as_any()
.downcast_ref::<arrow::array::BooleanArray>()
.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
);
}
}
}