chore(query): support custom tables for optimizer test case (#18255)

zhang2014 · web-flow · commit 633351de7f86 · 2025-06-26T16:18:49.000+08:00
* chore(query): support custom tables for optimizer test case

* chore(query): support custom tables for optimizer test case
diff --git a/src/query/service/tests/it/sql/planner/optimizer/data/cases/basic/01_cross_join_aggregation.yaml b/src/query/service/tests/it/sql/planner/optimizer/data/cases/basic/01_cross_join_aggregation.yaml
@@ -6,6 +6,9 @@ sql: |
 
 auto_statistics: true
 
+tables:
+  integers: basic/01_integers.sql
+
 good_plan: |
   Result [output: DT.D_YEAR, ITEM.I_BRAND_ID, ITEM.I_BRAND, SUM(...)]
   └── SortWithLimit [limit: 100]
diff --git a/src/query/service/tests/it/sql/planner/optimizer/data/cases/obfuscated/01_multi_join_avg_case_expression.yaml b/src/query/service/tests/it/sql/planner/optimizer/data/cases/obfuscated/01_multi_join_avg_case_expression.yaml
@@ -6,3 +6,9 @@ sql: |
 
 # Reference to external statistics file
 statistics_file: obfuscated/01_multi_join_case_expression_stats.yaml
+
+tables:
+  a0c: obfuscated/01_a0c.sql
+  a1z: obfuscated/01_a1z.sql
+  a2x: obfuscated/01_a2x.sql
+  a5r: obfuscated/01_a5r.sql
diff --git a/src/query/service/tests/it/sql/planner/optimizer/data/cases/obfuscated/01_multi_join_sum_case_expression.yaml b/src/query/service/tests/it/sql/planner/optimizer/data/cases/obfuscated/01_multi_join_sum_case_expression.yaml
@@ -6,3 +6,9 @@ sql: |
 
 # Reference to external statistics file
 statistics_file: obfuscated/01_multi_join_case_expression_stats.yaml
+
+tables:
+  a0c: obfuscated/01_a0c.sql
+  a1z: obfuscated/01_a1z.sql
+  a2x: obfuscated/01_a2x.sql
+  a5r: obfuscated/01_a5r.sql
diff --git a/src/query/service/tests/it/sql/planner/optimizer/data/cases/tpcds/Q01.yaml b/src/query/service/tests/it/sql/planner/optimizer/data/cases/tpcds/Q01.yaml
@@ -58,3 +58,27 @@ good_plan: |
           │                   └── WithClause [CUSTOMER_TOTAL_RETURN] (reference to earlier WITH clause)
           └── JoinFilter [joinKey: (CTR1.CTR_CUSTOMER_SK = CUSTOMER.C_CUSTOMER_SK)]
               └── TableScan [SNOWFLAKE_SAMPLE_DATA.TPCDS_SF10TCL.CUSTOMER] [C_CUSTOMER_SK, C_CUSTOMER_ID] [partitions: 261/261, bytes: 2,328,538,624]
+
+tables:
+  call_center: tpcds/call_center.sql
+  catalog_page: tpcds/catalog_page.sql
+  catalog_returns: tpcds/catalog_returns.sql
+  customer: tpcds/customer.sql
+  customer_address: tpcds/customer_address.sql
+  customer_demographics: tpcds/customer_demographics.sql
+  date_dim: tpcds/date_dim.sql
+  household_demographics: tpcds/household_demographics.sql
+  income_band: tpcds/income_band.sql
+  inventory: tpcds/inventory.sql
+  item: tpcds/item.sql
+  promotion: tpcds/promotion.sql
+  reason: tpcds/reason.sql
+  ship_mode: tpcds/ship_mode.sql
+  store: tpcds/store.sql
+  store_returns: tpcds/store_returns.sql
+  store_sales: tpcds/store_sales.sql
+  time_dim: tpcds/time_dim.sql
+  warehouse: tpcds/warehouse.sql
+  web_page: tpcds/web_page.sql
+  web_returns: tpcds/web_returns.sql
+  web_site: tpcds/web_site.sql
diff --git a/src/query/service/tests/it/sql/planner/optimizer/data/cases/tpcds/Q03.yaml b/src/query/service/tests/it/sql/planner/optimizer/data/cases/tpcds/Q03.yaml
@@ -37,3 +37,27 @@ good_plan: |
                                   └── JoinFilter [join key: (DT.D_DATE_SK = STORE_SALES.SS_SOLD_DATE_SK)]
                                       └── TableScan (STORE_SALES) [partitions: 70,412/72,718, bytes: 1,212,628,258,304]
                                           └── columns: [SS_SOLD_DATE_SK, SS_ITEM_SK, SS_EXT_SALES_PRICE]
+
+tables:
+  call_center: tpcds/call_center.sql
+  catalog_page: tpcds/catalog_page.sql
+  catalog_returns: tpcds/catalog_returns.sql
+  customer: tpcds/customer.sql
+  customer_address: tpcds/customer_address.sql
+  customer_demographics: tpcds/customer_demographics.sql
+  date_dim: tpcds/date_dim.sql
+  household_demographics: tpcds/household_demographics.sql
+  income_band: tpcds/income_band.sql
+  inventory: tpcds/inventory.sql
+  item: tpcds/item.sql
+  promotion: tpcds/promotion.sql
+  reason: tpcds/reason.sql
+  ship_mode: tpcds/ship_mode.sql
+  store: tpcds/store.sql
+  store_returns: tpcds/store_returns.sql
+  store_sales: tpcds/store_sales.sql
+  time_dim: tpcds/time_dim.sql
+  warehouse: tpcds/warehouse.sql
+  web_page: tpcds/web_page.sql
+  web_returns: tpcds/web_returns.sql
+  web_site: tpcds/web_site.sql
diff --git a/src/query/service/tests/it/sql/planner/optimizer/optimizer_test.rs b/src/query/service/tests/it/sql/planner/optimizer/optimizer_test.rs
@@ -19,6 +19,8 @@ use std::path::Path;
 use std::path::PathBuf;
 use std::sync::Arc;
 
+use databend_common_base::base::GlobalUniqName;
+use databend_common_catalog::cluster_info::Cluster;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_catalog::BasicColumnStatistics;
 use databend_common_catalog::TableStatistics;
@@ -29,6 +31,7 @@ use databend_common_expression::types::Number;
 use databend_common_expression::types::NumberScalar;
 use databend_common_expression::types::F64;
 use databend_common_expression::Scalar;
+use databend_common_meta_types::NodeInfo;
 use databend_common_sql::executor::PhysicalPlanBuilder;
 use databend_common_sql::optimize;
 use databend_common_sql::optimizer::ir::SExpr;
@@ -45,6 +48,7 @@ use databend_common_sql::IndexType;
 use databend_common_sql::Metadata;
 use databend_common_sql::MetadataRef;
 use databend_common_storage::Datum;
+use databend_query::clusters::ClusterHelper;
 use databend_query::sessions::QueryContext;
 use databend_query::test_kits::TestFixture;
 use goldenfile::Mint;
@@ -66,7 +70,11 @@ struct TestSpec {
     #[serde(default)]
     statistics_file: Option<String>,
     #[serde(default)]
+    tables: HashMap<String, String>,
+    #[serde(default)]
     auto_statistics: bool,
+    #[serde(default)]
+    node_num: Option<u64>,
     good_plan: Option<String>,
 }
 
@@ -102,6 +110,8 @@ struct TestCase {
     auto_stats: bool,
     stem: String,
     subdir: Option<String>,
+    node_num: Option<u64>,
+    tables: HashMap<String, String>,
 }
 
 struct TestSuite {
@@ -173,24 +183,6 @@ impl TestSuite {
         }
     }
 
-    async fn setup_tables(&self, ctx: &Arc<QueryContext>) -> Result<()> {
-        for file in self.find_files("tables", &["sql"]) {
-            let sql = fs::read_to_string(&file)
-                .map_err(|e| ErrorCode::Internal(format!("Failed to read file: {}", e)))?;
-            for statement in sql.split(';').filter(|s| !s.trim().is_empty()) {
-                match execute_sql(ctx, statement).await {
-                    Ok(_) => {}
-                    Err(e) if e.code() == ErrorCode::TABLE_ALREADY_EXISTS => {
-                        // Ignore table already exists errors
-                        continue;
-                    }
-                    Err(e) => return Err(e),
-                }
-            }
-        }
-        Ok(())
-    }
-
     fn load_cases(&self) -> Result<Vec<TestCase>> {
         let cases_dir = self.base_path.join("cases");
         self.find_files("cases", &["yaml", "yml"])
@@ -218,6 +210,8 @@ impl TestSuite {
         let name = spec.name.clone();
         let sql = spec.sql.clone();
         let auto_stats = spec.auto_statistics;
+        let node_num = spec.node_num;
+        let tables = self.resolve_tables(&spec)?;
         let (table_stats, column_stats) = self.resolve_stats(spec)?;
         let stem = path
             .file_stem()
@@ -235,14 +229,26 @@ impl TestSuite {
         Ok(TestCase {
             name: Box::leak(full_name.into_boxed_str()),
             sql: Box::leak(sql.into_boxed_str()),
+            tables,
             table_stats,
             column_stats,
             auto_stats,
             stem,
             subdir,
+            node_num,
         })
     }
 
+    fn resolve_tables(&self, spec: &TestSpec) -> Result<HashMap<String, String>> {
+        let mut tables = HashMap::with_capacity(spec.tables.len());
+        for (table_name, file_ref) in &spec.tables {
+            let table_define_sql = self.load_table_file(file_ref)?;
+            tables.insert(table_name.clone(), table_define_sql);
+        }
+
+        Ok(tables)
+    }
+
     fn resolve_stats(
         &self,
         mut spec: TestSpec,
@@ -255,6 +261,33 @@ impl TestSuite {
         Ok((spec.table_statistics, spec.column_statistics))
     }
 
+    fn load_table_file(&self, file_ref: &str) -> Result<String> {
+        let stats_files = self.find_files("tables", &["sql"]);
+        let target_stem = Path::new(file_ref)
+            .file_stem()
+            .and_then(|s| s.to_str())
+            .unwrap_or(file_ref);
+
+        stats_files
+            .iter()
+            .find(|path| {
+                path.file_stem()
+                    .and_then(|s| s.to_str())
+                    .map(|s| s == target_stem || s.ends_with(target_stem))
+                    .unwrap_or(false)
+            })
+            .ok_or_else(|| {
+                ErrorCode::Internal(format!(
+                    "Tables file not found: {}, in {:?}",
+                    file_ref, stats_files
+                ))
+            })
+            .and_then(|path| {
+                fs::read_to_string(path)
+                    .map_err(|e| ErrorCode::Internal(format!("Failed to read file: {}", e)))
+            })
+    }
+
     fn load_stats_file(&self, file_ref: &str) -> Result<StatsFile> {
         let stats_files = self.find_files("statistics", &["yaml", "yml"]);
         let target_stem = Path::new(file_ref)
@@ -424,10 +457,6 @@ async fn test_optimizer() -> Result<()> {
 
     let suite = TestSuite::new(base_path.clone(), subdir);
     let fixture = TestFixture::setup().await?;
-    let ctx = fixture.new_query_ctx().await?;
-    ctx.get_settings().set_enable_auto_materialize_cte(0)?;
-
-    suite.setup_tables(&ctx).await?;
 
     let cases = suite.load_cases()?;
     if cases.is_empty() {
@@ -438,17 +467,79 @@ async fn test_optimizer() -> Result<()> {
     let mut root_mint = Mint::new(&results_dir);
     let mut subdir_mints: HashMap<String, Mint> = HashMap::new();
 
+    let local_id = GlobalUniqName::unique();
+    let standalone_cluster = Cluster::create(vec![create_node(&local_id)], local_id);
+
     for case in cases {
         println!("\n========== Testing: {} ==========", case.name);
 
+        let ctx = fixture.new_query_ctx().await?;
+        ctx.get_settings().set_enable_auto_materialize_cte(0)?;
+
+        ctx.set_cluster(standalone_cluster.clone());
+
+        if let Some(nodes) = case.node_num {
+            let mut nodes_info = Vec::with_capacity(nodes as usize);
+            for _ in 0..nodes - 1 {
+                nodes_info.push(create_node(&GlobalUniqName::unique()));
+            }
+
+            let local_id = GlobalUniqName::unique();
+            nodes_info.push(create_node(&local_id));
+            ctx.set_cluster(Cluster::create(nodes_info, local_id))
+        }
+
+        setup_tables(&ctx, &case).await?;
         run_test_case(&ctx, &case, &mut root_mint, &mut subdir_mints, &results_dir).await?;
 
+        clean_tables(&ctx, &case).await?;
+        ctx.set_cluster(standalone_cluster.clone());
         println!("✅ {} test passed!", case.name);
     }
 
     Ok(())
 }
 
+fn create_node(local_id: &str) -> Arc<NodeInfo> {
+    let mut node_info = NodeInfo::create(
+        local_id.to_string(),
+        String::new(),
+        0,
+        String::new(),
+        String::new(),
+        String::new(),
+        String::new(),
+        String::new(),
+    );
+    node_info.cluster_id = "cluster_id".to_string();
+    node_info.warehouse_id = "warehouse_id".to_string();
+    Arc::new(node_info)
+}
+
+async fn setup_tables(ctx: &Arc<QueryContext>, case: &TestCase) -> Result<()> {
+    for sql in case.tables.values() {
+        for statement in sql.split(';').filter(|s| !s.trim().is_empty()) {
+            match execute_sql(ctx, statement).await {
+                Ok(_) => {}
+                Err(e) if e.code() == ErrorCode::TABLE_ALREADY_EXISTS => {
+                    // Ignore table already exists errors
+                    continue;
+                }
+                Err(e) => return Err(e),
+            }
+        }
+    }
+
+    Ok(())
+}
+
+async fn clean_tables(ctx: &Arc<QueryContext>, case: &TestCase) -> Result<()> {
+    for table_name in case.tables.keys() {
+        execute_sql(ctx, &format!("DROP TABLE {}", table_name)).await?;
+    }
+    Ok(())
+}
+
 async fn run_test_case(
     ctx: &Arc<QueryContext>,
     case: &TestCase,
