mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-08-18 20:18:30 +00:00
fix(query): avoid unsafe count wildcard rewrites (#8522)
* fix(query): avoid unsafe count wildcard rewrites Signed-off-by: discord9 <discord9@163.com> * fix(query): preserve outer count alias Signed-off-by: discord9 <discord9@163.com> * fix(query): address review comments on count wildcard rewrite - Remove the has_projection check: the row count is correct regardless of whether a projection exists (per review). - Explain why checking the first input is equivalent to checking all inputs (a plan with zero inputs falls back to count(1)). - Rename qa_ prefixed tests to follow the module convention. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(tql): update tql-cte expectations for count wildcard rewrite The QP-026 count-wildcard fix rewrites count(*) -> count(time_index), so the EXPLAIN output for the filtered/final CTE aggregates names the time-index column. Aligns tql-cte.result with the actual output (CI failure). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <discord9@163.com> Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
GitHub
@@ -85,20 +85,23 @@ impl CountWildcardToTimeIndexRule {
|
||||
plan.visit(&mut finder).unwrap();
|
||||
let col = finder.into_column();
|
||||
|
||||
// check if the time index is a valid column as for current plan
|
||||
// The resolved time index must be present and non-nullable in the
|
||||
// immediate input schema. Schema-changing nodes can otherwise expose
|
||||
// a nullable field with the same name as the source time index, and
|
||||
// `count(<col>)` would then count fewer rows than `count(*)`.
|
||||
if let Some(col) = &col {
|
||||
let mut is_valid = false;
|
||||
// if more than one input, we give up and just use `count(1)`
|
||||
if plan.inputs().len() > 1 {
|
||||
return None;
|
||||
}
|
||||
for input in plan.inputs() {
|
||||
if input.schema().has_column(col) {
|
||||
is_valid = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if !is_valid {
|
||||
// The guard above guarantees exactly one input here, so checking
|
||||
// the first input is equivalent to checking all inputs as the rule
|
||||
// used to: a plan with zero inputs also falls back to `count(1)`.
|
||||
let input = plan.inputs().first().copied()?;
|
||||
let Ok((_, field)) = input.schema().qualified_field_from_column(col) else {
|
||||
return None;
|
||||
};
|
||||
if field.is_nullable() {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
@@ -147,7 +150,8 @@ impl TreeNodeVisitor<'_> for TimeIndexFinder {
|
||||
|
||||
fn f_down(&mut self, node: &Self::Node) -> DataFusionResult<TreeNodeRecursion> {
|
||||
if let LogicalPlan::SubqueryAlias(subquery_alias) = node {
|
||||
self.table_alias = Some(subquery_alias.alias.clone());
|
||||
self.table_alias
|
||||
.get_or_insert_with(|| subquery_alias.alias.clone());
|
||||
}
|
||||
|
||||
if let LogicalPlan::TableScan(table_scan) = &node
|
||||
@@ -199,17 +203,20 @@ mod test {
|
||||
use common_catalog::consts::DEFAULT_CATALOG_NAME;
|
||||
use common_error::ext::{BoxedError, ErrorExt, StackError};
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_recordbatch::SendableRecordBatchStream;
|
||||
use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
|
||||
use datafusion::functions_aggregate::count::count_all;
|
||||
use datafusion::functions_aggregate::min_max::max;
|
||||
use datafusion_common::Column;
|
||||
use datafusion_expr::LogicalPlanBuilder;
|
||||
use datafusion_sql::TableReference;
|
||||
use datatypes::data_type::ConcreteDataType;
|
||||
use datatypes::schema::{ColumnSchema, SchemaBuilder};
|
||||
use datatypes::schema::{ColumnSchema, Schema, SchemaBuilder};
|
||||
use datatypes::vectors::{Int64Vector, TimestampMillisecondVector, VectorRef};
|
||||
use store_api::data_source::DataSource;
|
||||
use store_api::storage::ScanRequest;
|
||||
use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder, TableType};
|
||||
use table::table::numbers::NumbersTable;
|
||||
use table::test_util::MemTable;
|
||||
use table::{Table, TableRef};
|
||||
|
||||
use super::*;
|
||||
@@ -319,6 +326,216 @@ mod test {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn count_wildcard_shape_matrix() {
|
||||
let config = datafusion::config::ConfigOptions::default();
|
||||
|
||||
let direct = CountWildcardToTimeIndexRule
|
||||
.analyze(count_star(source_plan("source")), &config)
|
||||
.unwrap();
|
||||
assert_count_argument_column(&direct, "source", "ts");
|
||||
|
||||
let simple_alias = count_star(
|
||||
LogicalPlanBuilder::from(source_plan("source"))
|
||||
.alias("projected")
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
let simple_alias = CountWildcardToTimeIndexRule
|
||||
.analyze(simple_alias, &config)
|
||||
.unwrap();
|
||||
assert_count_argument_column(&simple_alias, "projected", "ts");
|
||||
|
||||
let nested_alias = count_star(
|
||||
LogicalPlanBuilder::from(source_plan("source"))
|
||||
.alias("inner")
|
||||
.unwrap()
|
||||
.alias("outer")
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
let nested_alias = CountWildcardToTimeIndexRule
|
||||
.analyze(nested_alias, &config)
|
||||
.unwrap();
|
||||
assert_count_argument_column(&nested_alias, "outer", "ts");
|
||||
|
||||
let nested_rename = count_star(
|
||||
LogicalPlanBuilder::from(source_plan("source"))
|
||||
.project(vec![col("ts").alias("renamed")])
|
||||
.unwrap()
|
||||
.alias("projected")
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
let nested_rename = CountWildcardToTimeIndexRule
|
||||
.analyze(nested_rename, &config)
|
||||
.unwrap();
|
||||
assert_count_argument_literal_one(&nested_rename);
|
||||
|
||||
let nested_rename_with_payload_reorder = count_star(
|
||||
LogicalPlanBuilder::from(source_plan("source"))
|
||||
.project(vec![col("payload"), col("ts").alias("renamed")])
|
||||
.unwrap()
|
||||
.alias("projected")
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
let nested_rename_with_payload_reorder = CountWildcardToTimeIndexRule
|
||||
.analyze(nested_rename_with_payload_reorder, &config)
|
||||
.unwrap();
|
||||
assert_count_argument_literal_one(&nested_rename_with_payload_reorder);
|
||||
|
||||
let multi_input = count_star(
|
||||
LogicalPlanBuilder::from(source_plan("left"))
|
||||
.cross_join(source_plan("right"))
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
let multi_input = CountWildcardToTimeIndexRule
|
||||
.analyze(multi_input, &config)
|
||||
.unwrap();
|
||||
assert_count_argument_literal_one(&multi_input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn projection_name_collision_falls_back_to_literal_one() {
|
||||
let before = count_star(
|
||||
LogicalPlanBuilder::from(source_plan("source"))
|
||||
.project(vec![col("payload").alias("ts")])
|
||||
.unwrap()
|
||||
.alias("projected")
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
let aggregate = aggregate_plan(&before);
|
||||
let field = aggregate
|
||||
.input
|
||||
.schema()
|
||||
.qualified_field_with_name(Some(&TableReference::bare("projected")), "ts")
|
||||
.unwrap();
|
||||
assert!(field.1.is_nullable());
|
||||
|
||||
let after = CountWildcardToTimeIndexRule
|
||||
.analyze(before, &datafusion::config::ConfigOptions::default())
|
||||
.unwrap();
|
||||
assert_count_argument_literal_one(&after);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inner_aggregate_nullable_time_index_name_falls_back_to_literal_one() {
|
||||
let before = count_star(
|
||||
LogicalPlanBuilder::from(source_plan("source"))
|
||||
.aggregate(Vec::<Expr>::new(), vec![max(col("payload")).alias("ts")])
|
||||
.unwrap()
|
||||
.alias("aggregated")
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
let aggregate = aggregate_plan(&before);
|
||||
let field = aggregate
|
||||
.input
|
||||
.schema()
|
||||
.qualified_field_with_name(Some(&TableReference::bare("aggregated")), "ts")
|
||||
.unwrap();
|
||||
assert!(field.1.is_nullable());
|
||||
|
||||
let after = CountWildcardToTimeIndexRule
|
||||
.analyze(before, &datafusion::config::ConfigOptions::default())
|
||||
.unwrap();
|
||||
assert_count_argument_literal_one(&after);
|
||||
}
|
||||
|
||||
fn source_plan(table_name: &str) -> LogicalPlan {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
ColumnSchema::new(
|
||||
"ts",
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
false,
|
||||
)
|
||||
.with_time_index(true),
|
||||
ColumnSchema::new("payload", ConcreteDataType::int64_datatype(), true),
|
||||
]));
|
||||
let columns: Vec<VectorRef> = vec![
|
||||
Arc::new(TimestampMillisecondVector::from_slice([1, 2, 3])),
|
||||
Arc::new(Int64Vector::from(vec![Some(10), None, Some(30)])),
|
||||
];
|
||||
let table = MemTable::table(
|
||||
table_name,
|
||||
RecordBatch::new(schema, columns).expect("test record batch must be valid"),
|
||||
);
|
||||
let source = Arc::new(DefaultTableSource::new(Arc::new(
|
||||
DfTableProviderAdapter::new(table),
|
||||
)));
|
||||
LogicalPlanBuilder::scan_with_filters(table_name, source, None, vec![])
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn count_star(input: LogicalPlan) -> LogicalPlan {
|
||||
LogicalPlanBuilder::from(input)
|
||||
.aggregate(Vec::<Expr>::new(), vec![count_all()])
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn count_aggregate(plan: &LogicalPlan) -> &AggregateFunction {
|
||||
let LogicalPlan::Aggregate(aggregate) = plan else {
|
||||
panic!("expected aggregate plan, got {plan:?}");
|
||||
};
|
||||
assert_eq!(1, aggregate.aggr_expr.len());
|
||||
let expr = unwrap_aliases(&aggregate.aggr_expr[0]);
|
||||
let Expr::AggregateFunction(count) = expr else {
|
||||
panic!("expected count aggregate, got {:?}", aggregate.aggr_expr[0]);
|
||||
};
|
||||
assert_eq!("count", count.func.name());
|
||||
count
|
||||
}
|
||||
|
||||
fn unwrap_aliases(expr: &Expr) -> &Expr {
|
||||
match expr {
|
||||
Expr::Alias(alias) => unwrap_aliases(alias.expr.as_ref()),
|
||||
expr => expr,
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_count_argument_column(plan: &LogicalPlan, relation: &str, name: &str) {
|
||||
let count = count_aggregate(plan);
|
||||
let [Expr::Column(column)] = count.params.args.as_slice() else {
|
||||
panic!(
|
||||
"expected one column count argument, got {:?}",
|
||||
count.params.args
|
||||
);
|
||||
};
|
||||
assert_eq!(Some(TableReference::bare(relation)), column.relation);
|
||||
assert_eq!(name, column.name);
|
||||
}
|
||||
|
||||
fn assert_count_argument_literal_one(plan: &LogicalPlan) {
|
||||
let count = count_aggregate(plan);
|
||||
assert!(matches!(
|
||||
count.params.args.as_slice(),
|
||||
[Expr::Literal(ScalarValue::Int64(Some(1)), _)]
|
||||
));
|
||||
}
|
||||
|
||||
fn aggregate_plan(plan: &LogicalPlan) -> &datafusion_expr::logical_plan::Aggregate {
|
||||
let LogicalPlan::Aggregate(aggregate) = plan else {
|
||||
panic!("expected aggregate plan, got {plan:?}");
|
||||
};
|
||||
aggregate
|
||||
}
|
||||
|
||||
fn build_time_index_table(table_name: &str, schema_name: &str, catalog_name: &str) -> TableRef {
|
||||
let column_schemas = vec![
|
||||
ColumnSchema::new(
|
||||
|
||||
@@ -155,7 +155,7 @@ SELECT count(*) FROM filtered;
|
||||
+---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------+
|
||||
| logical_plan | MergeScan [is_placeholder=false, remote_input=[ |
|
||||
| | Projection: count(Int64(1)) AS count(*) |
|
||||
| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] |
|
||||
| | Aggregate: groupBy=[[]], aggr=[[count(filtered.ts) AS count(Int64(1))]] |
|
||||
| | SubqueryAlias: filtered |
|
||||
| | Projection: tql_data.ts, tql_data.val |
|
||||
| | SubqueryAlias: tql_data |
|
||||
@@ -642,7 +642,7 @@ SELECT count(*) as high_values FROM final;
|
||||
+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|
||||
| logical_plan | MergeScan [is_placeholder=false, remote_input=[ |
|
||||
| | Projection: count(Int64(1)) AS count(*) AS high_values |
|
||||
| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] |
|
||||
| | Aggregate: groupBy=[[]], aggr=[[count(final.ts) AS count(Int64(1))]] |
|
||||
| | SubqueryAlias: final |
|
||||
| | Projection: processed.ts AS ts, processed.percent AS percent |
|
||||
| | Projection: processed.ts, processed.percent |
|
||||
|
||||
Reference in New Issue
Block a user