fix: qualify scalar-subquery tables in persisted views (#8581)

fix(query): traverse scalar subqueries in views

Signed-off-by: discord9 <discord9@163.com>

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<TableName>, 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<LogicalTableSource> {
         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<TableName>, expected: &[(&str, &str, &str)]) {
+        let expected = expected
+            .iter()
+            .map(|(catalog, schema, table)| TableName::new(*catalog, *schema, *table))
+            .collect::<HashSet<_>>();
+        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"),
+        );
+    }
 }