From 8b47fd53bdcdded141491326ea6ccf81b6f549d2 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 22 Jul 2026 18:43:24 +0800 Subject: [PATCH] fix: qualify scalar-subquery tables in persisted views (#8581) fix(query): traverse scalar subqueries in views Signed-off-by: discord9 --- src/query/src/plan.rs | 146 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 133 insertions(+), 13 deletions(-) diff --git a/src/query/src/plan.rs b/src/query/src/plan.rs index d3eafd84c0..b4e5e664b6 100644 --- a/src/query/src/plan.rs +++ b/src/query/src/plan.rs @@ -16,7 +16,7 @@ use std::collections::HashSet; use datafusion::datasource::DefaultTableSource; use datafusion_common::TableReference; -use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; +use datafusion_common::tree_node::{Transformed, TreeNodeRewriter}; use datafusion_expr::{Expr, LogicalPlan}; use session::context::QueryContextRef; pub use table::metadata::TableType; @@ -116,7 +116,7 @@ pub fn extract_and_rewrite_full_table_names( query_ctx: QueryContextRef, ) -> Result<(HashSet, LogicalPlan)> { let mut extractor = TableNamesExtractAndRewriter::new(query_ctx); - let plan = plan.rewrite(&mut extractor)?; + let plan = plan.rewrite_with_subqueries(&mut extractor)?; Ok((extractor.table_names, plan.data)) } @@ -142,23 +142,26 @@ pub(crate) mod tests { use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use common_catalog::consts::DEFAULT_CATALOG_NAME; use datafusion::logical_expr::builder::LogicalTableSource; - use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, col, lit}; + use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, col, lit, scalar_subquery}; use session::context::QueryContextBuilder; use super::*; - fn mock_plan() -> LogicalPlan { + fn mock_table_source() -> Arc { let schema = Schema::new(vec![ Field::new("id", DataType::Int32, true), Field::new("name", DataType::Utf8, true), Field::new("ts", DataType::Timestamp(TimeUnit::Millisecond, None), true), ]); - let table_source = LogicalTableSource::new(SchemaRef::new(schema)); + Arc::new(LogicalTableSource::new(SchemaRef::new(schema))) + } + + fn mock_plan() -> LogicalPlan { + let table_source = mock_table_source(); let projection = None; - let builder = - LogicalPlanBuilder::scan("devices", Arc::new(table_source), projection).unwrap(); + let builder = LogicalPlanBuilder::scan("devices", table_source, projection).unwrap(); builder .filter(col("id").gt(lit(500))) @@ -167,6 +170,55 @@ pub(crate) mod tests { .unwrap() } + fn scalar_subquery_plan(table_name: TableReference) -> LogicalPlan { + let subquery = LogicalPlanBuilder::scan(table_name, mock_table_source(), None) + .unwrap() + .project(vec![col("id")]) + .unwrap() + .build() + .unwrap(); + + LogicalPlanBuilder::empty(false) + .project(vec![scalar_subquery(Arc::new(subquery))]) + .unwrap() + .build() + .unwrap() + } + + fn assert_dependencies(actual: &HashSet, expected: &[(&str, &str, &str)]) { + let expected = expected + .iter() + .map(|(catalog, schema, table)| TableName::new(*catalog, *schema, *table)) + .collect::>(); + assert_eq!(&expected, actual); + } + + fn assert_nested_scalar_subquery_table_name(plan: &LogicalPlan, expected: TableReference) { + let LogicalPlan::Projection(projection) = plan else { + panic!("expected scalar-subquery projection, got {plan:?}"); + }; + let [Expr::ScalarSubquery(subquery)] = projection.expr.as_slice() else { + panic!( + "expected one scalar-subquery expression, got {:?}", + projection.expr + ); + }; + let LogicalPlan::Projection(projection) = subquery.subquery.as_ref() else { + panic!( + "expected scalar-subquery projection, got {:?}", + subquery.subquery + ); + }; + let LogicalPlan::TableScan(scan) = projection.input.as_ref() else { + panic!( + "expected scalar-subquery table scan, got {:?}", + projection.input + ); + }; + + assert_eq!(expected, scan.table_name); + } + #[test] fn test_extract_full_table_names() { let ctx = QueryContextBuilder::default() @@ -176,16 +228,84 @@ pub(crate) mod tests { let (table_names, plan) = extract_and_rewrite_full_table_names(mock_plan(), Arc::new(ctx)).unwrap(); - assert_eq!(1, table_names.len()); - assert!(table_names.contains(&TableName::new( - DEFAULT_CATALOG_NAME.to_string(), - "test".to_string(), - "devices".to_string() - ))); + assert_dependencies(&table_names, &[(DEFAULT_CATALOG_NAME, "test", "devices")]); assert_eq!( "Filter: devices.id > Int32(500)\n TableScan: greptime.test.devices", plan.to_string() ); } + + #[test] + fn test_extract_full_table_names_from_scalar_subquery_bare_table_scan() { + let ctx = QueryContextBuilder::default() + .current_catalog("qp031_catalog".to_string()) + .current_schema("qp031_current_schema".to_string()) + .build(); + + let (table_names, plan) = extract_and_rewrite_full_table_names( + scalar_subquery_plan(TableReference::bare("lookup")), + Arc::new(ctx), + ) + .unwrap(); + + assert_dependencies( + &table_names, + &[("qp031_catalog", "qp031_current_schema", "lookup")], + ); + assert_nested_scalar_subquery_table_name( + &plan, + TableReference::full("qp031_catalog", "qp031_current_schema", "lookup"), + ); + } + + #[test] + fn test_extract_full_table_names_from_scalar_subquery_partial_table_scan() { + let ctx = QueryContextBuilder::default() + .current_catalog("qp031_catalog".to_string()) + .current_schema("qp031_current_schema".to_string()) + .build(); + + let (table_names, plan) = extract_and_rewrite_full_table_names( + scalar_subquery_plan(TableReference::partial("qp031_lookup_schema", "lookup")), + Arc::new(ctx), + ) + .unwrap(); + + assert_dependencies( + &table_names, + &[("qp031_catalog", "qp031_lookup_schema", "lookup")], + ); + assert_nested_scalar_subquery_table_name( + &plan, + TableReference::full("qp031_catalog", "qp031_lookup_schema", "lookup"), + ); + } + + #[test] + fn test_extract_full_table_names_from_scalar_subquery_full_table_scan() { + let ctx = QueryContextBuilder::default() + .current_catalog("qp031_catalog".to_string()) + .current_schema("qp031_current_schema".to_string()) + .build(); + + let (table_names, plan) = extract_and_rewrite_full_table_names( + scalar_subquery_plan(TableReference::full( + "qp031_external_catalog", + "qp031_external_schema", + "lookup", + )), + Arc::new(ctx), + ) + .unwrap(); + + assert_dependencies( + &table_names, + &[("qp031_external_catalog", "qp031_external_schema", "lookup")], + ); + assert_nested_scalar_subquery_table_name( + &plan, + TableReference::full("qp031_external_catalog", "qp031_external_schema", "lookup"), + ); + } }