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

fix: resolve optimization issue for extended query (#7824)

Browse Source 
* fix: resolve optimization issue for extended query

* fix: type cast from subquery

* chore: update error information in sqlness

* chore: switch to released pgwire

* refactor: remove optimize function completely

* chore: add more tests

* test: attempt to fix the fuzz issue

* fix: try to resolve the test issue
This commit is contained in:
Ning Sun
2026-03-20 03:58:39 +00:00
committed by GitHub
parent f034255fe6
commit d14817bfa6
9 changed files with 94 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+3 -3
View File
@@ -7301,7 +7301,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667"
dependencies = [
"cfg-if",
"windows-targets 0.48.5",
"windows-targets 0.52.6",
]
[[package]]
@@ -9620,9 +9620,9 @@ dependencies = [
[[package]]
name = "pgwire"
version = "0.38.0"
version = "0.38.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "89d5e5a60d3f6e40c91f6a2a7f8d09665e636272bd5611977253559b6651aabb"
checksum = "f2a798d130b8975a566c2cf6d8955746e1f09a9ee2c3ff2e6020a2c6528c5bd1"
dependencies = [
"async-trait",
"base64 0.22.1",
src/catalog/src/system_schema/information_schema/region_peers.rs
+1 -1
View File
@@ -267,7 +267,7 @@ impl InformationSchemaRegionPeersBuilder {
];
if !predicates.eval(&row) {
return;
continue;
}
self.table_catalogs.push(Some(table_catalog));
src/query/src/datafusion.rs
+11 -45
View File
@@ -354,25 +354,6 @@ impl DatafusionQueryEngine {
Ok(physical_plan)
}
#[tracing::instrument(skip_all)]
pub fn optimize(
&self,
context: &QueryEngineContext,
plan: &LogicalPlan,
) -> Result<LogicalPlan> {
let _timer = metrics::OPTIMIZE_LOGICAL_ELAPSED.start_timer();
// Optimized by extension rules
let optimized_plan = self
.state
.optimize_by_extension_rules(plan.clone(), context)?;
// Optimized by datafusion optimizer
let optimized_plan = self.state.session_state().optimize(&optimized_plan)?;
Ok(optimized_plan)
}
#[tracing::instrument(skip_all)]
fn optimize_physical_plan(
&self,
@@ -444,32 +425,17 @@ impl QueryEngine for DatafusionQueryEngine {
async fn describe(
&self,
plan: LogicalPlan,
query_ctx: QueryContextRef,
_query_ctx: QueryContextRef,
) -> Result<DescribeResult> {
let ctx = self.engine_context(query_ctx);
if let Ok(optimised_plan) = self.optimize(&ctx, &plan) {
let schema = optimised_plan
.schema()
.clone()
.try_into()
.context(ConvertSchemaSnafu)?;
Ok(DescribeResult {
schema,
logical_plan: optimised_plan,
})
} else {
// Table's like those in information_schema cannot be optimized when
// it contains parameters. So we fallback to original plans.
let schema = plan
.schema()
.clone()
.try_into()
.context(ConvertSchemaSnafu)?;
Ok(DescribeResult {
schema,
logical_plan: plan,
})
}
let schema = plan
.schema()
.clone()
.try_into()
.context(ConvertSchemaSnafu)?;
Ok(DescribeResult {
schema,
logical_plan: plan,
})
}
async fn execute(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<Output> {
@@ -924,7 +890,7 @@ mod tests {
)
);
assert_eq!(
"Limit: skip=0, fetch=20\n Aggregate: groupBy=[[]], aggr=[[sum(CAST(numbers.number AS UInt64))]]\n TableScan: numbers projection=[number]",
"Limit: skip=0, fetch=20\n Projection: sum(numbers.number)\n Aggregate: groupBy=[[]], aggr=[[sum(numbers.number)]]\n TableScan: numbers",
format!("{}", logical_plan.display_indent())
);
}
src/query/src/planner.rs
+40 -2
View File
@@ -28,6 +28,7 @@ use datafusion::execution::context::SessionState;
use datafusion::sql::planner::PlannerContext;
use datafusion_common::ToDFSchema;
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
use datafusion_expr::expr::{Exists, InSubquery};
use datafusion_expr::{
Analyze, Explain, ExplainFormat, Expr as DfExpr, LogicalPlan, LogicalPlanBuilder, PlanType,
ToStringifiedPlan, col,
@@ -424,9 +425,20 @@ impl DfLogicalPlanner {
let mut placeholder_types = HashMap::new();
let mut casted_placeholders = HashSet::new();
Self::extract_from_plan(plan, &mut placeholder_types, &mut casted_placeholders)?;
Ok(placeholder_types)
}
fn extract_from_plan(
plan: &LogicalPlan,
placeholder_types: &mut HashMap<String, Option<DataType>>,
casted_placeholders: &mut HashSet<String>,
) -> Result<()> {
plan.apply(|node| {
for expr in node.expressions() {
let _ = expr.apply(|e| {
// Handle casted placeholders
if let DfExpr::Cast(cast) = e
&& let DfExpr::Placeholder(ph) = &*cast.expr
{
@@ -434,6 +446,7 @@ impl DfLogicalPlanner {
casted_placeholders.insert(ph.id.clone());
}
// Handle bare (non-casted) placeholders
if let DfExpr::Placeholder(ph) = e
&& !casted_placeholders.contains(&ph.id)
&& !placeholder_types.contains_key(&ph.id)
@@ -441,13 +454,26 @@ impl DfLogicalPlanner {
placeholder_types.insert(ph.id.clone(), None);
}
// Recurse into subquery plans embedded in expressions
match e {
DfExpr::Exists(Exists { subquery, .. })
| DfExpr::InSubquery(InSubquery { subquery, .. })
| DfExpr::ScalarSubquery(subquery) => {
Self::extract_from_plan(
&subquery.subquery,
placeholder_types,
casted_placeholders,
)?;
}
_ => {}
}
Ok(TreeNodeRecursion::Continue)
});
}
Ok(TreeNodeRecursion::Continue)
})?;
Ok(placeholder_types)
Ok(())
}
/// Gets inferred parameter types from a logical plan.
@@ -619,4 +645,16 @@ mod tests {
assert_eq!(type_2, &Some(DataType::Utf8));
assert_eq!(type_3, &Some(DataType::Int32));
}
#[tokio::test]
async fn test_get_inferred_parameter_types_subquery() {
let plan = parse_sql_to_plan(
r#"SELECT * FROM test WHERE id = (SELECT id FROM test CROSS JOIN (SELECT parse_ident($1::TEXT) AS parts) p LIMIT 1)"#,
).await;
let types = DfLogicalPlanner::get_inferred_parameter_types(&plan).unwrap();
assert_eq!(types.len(), 1);
let type_1 = types.get("$1").unwrap();
assert_eq!(type_1, &Some(DataType::Utf8));
}
}
src/servers/Cargo.toml
+1 -1
View File
@@ -89,7 +89,7 @@ operator.workspace = true
otel-arrow-rust.workspace = true
parking_lot.workspace = true
pg_interval = { version = "0.5.2", package = "pg_interval_2" }
pgwire = { version = "0.38", default-features = false, features = [
pgwire = { version = "0.38.1", default-features = false, features = [
"server-api-ring",
"pg-ext-types",
] }
tests-fuzz/src/utils/partition.rs
+1 -1
View File
@@ -36,7 +36,7 @@ pub struct PartitionCount {
}
pub async fn count_partitions(db: &MySqlPool, datanode_id: u64) -> Result<PartitionCount> {
let sql = "select count(1) as count from information_schema.region_peers where peer_id == ?";
let sql = "select count(1) as count from information_schema.region_peers where peer_id = ?";
sqlx::query_as::<_, PartitionCount>(sql)
.bind(datanode_id)
.fetch_one(db)
tests-fuzz/targets/migration/fuzz_migrate_mito_regions.rs
+8 -3
View File
@@ -261,13 +261,18 @@ async fn migrate_regions(ctx: &FuzzContext, migrations: &[Migration]) -> Result<
{
let output = procedure_state(&greptime, &procedure_id).await;
info!("Checking procedure: {procedure_id}, output: {output}");
(fetch_partition(&greptime, region_id).await.unwrap(), output)
(fetch_partition(&greptime, region_id).await.ok(), output)
}
})
},
|(partition, output)| {
info!("Region: {region_id}, datanode: {}", partition.datanode_id);
partition.datanode_id == migration.to_peer && output.contains("Done")
if let Some(partition) = partition {
info!("Region: {region_id}, datanode: {}", partition.datanode_id);
partition.datanode_id == migration.to_peer && output.contains("Done")
} else {
info!("Region: {region_id}, partition not found yet");
false
}
},
Duration::from_secs(5),
)
tests/cases/standalone/common/prepare/mysql_prepare.result
+20 -1
View File
@@ -42,7 +42,7 @@ affected_rows: 0
-- SQLNESS PROTOCOL MYSQL
EXECUTE stmt USING 'a';
Failed to execute query, err: MySqlError { ERROR 1815 (HY000): (EngineExecuteQuery): Cast error: Cannot cast string 'a' to value of Int32 type }
Failed to execute query, err: MySqlError { ERROR 1210 (HY000): (InvalidArguments): Invalid request parameter: Unable to convert a to datatype Int32(Int32Type) }
-- SQLNESS PROTOCOL MYSQL
DEALLOCATE stmt;
@@ -124,6 +124,25 @@ DEALLOCATE stmt;
affected_rows: 0
-- SQLNESS PROTOCOL MYSQL
PREPARE stmt FROM 'SELECT table_name, table_schema FROM information_schema.tables WHERE table_name = ?';
affected_rows: 0
-- SQLNESS PROTOCOL MYSQL
EXECUTE stmt USING 'cake';
+------------+--------------+
| table_name | table_schema |
+------------+--------------+
| cake | public |
+------------+--------------+
-- SQLNESS PROTOCOL MYSQL
DEALLOCATE stmt;
affected_rows: 0
-- SQLNESS PROTOCOL MYSQL
DROP TABLE cake;
tests/cases/standalone/common/prepare/mysql_prepare.sql
+9
View File
@@ -72,5 +72,14 @@ EXECUTE stmt USING 'happy', 42, 0;
-- SQLNESS PROTOCOL MYSQL
DEALLOCATE stmt;
-- SQLNESS PROTOCOL MYSQL
PREPARE stmt FROM 'SELECT table_name, table_schema FROM information_schema.tables WHERE table_name = ?';
-- SQLNESS PROTOCOL MYSQL
EXECUTE stmt USING 'cake';
-- SQLNESS PROTOCOL MYSQL
DEALLOCATE stmt;
-- SQLNESS PROTOCOL MYSQL
DROP TABLE cake;