feat(sqlsmith): enable sub-query generation

src/tests/sqlsmith/src/lib.rs

@@ -87,6 +87,11 @@ struct SqlGenerator<'a, R: Rng> {
     /// We might not read from all tables.
     bound_relations: Vec<Table>,
 
+    /// Relations in generated in 'From' statement shall be added to bound_relations and
+    /// bound_columns after all the relations are generated because they are not useable among
+    /// the relations
+    parallel_relations: Vec<Table>,
+
     /// Columns bound in generated query.
     /// May not contain all columns from Self::bound_relations.
     /// e.g. GROUP BY clause will constrain bound_columns.
@@ -110,6 +115,7 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
             relation_id: 0,
             is_distinct_allowed,
             bound_relations: vec![],
+            parallel_relations: vec![],
             bound_columns: vec![],
             is_mview: false,
         }
@@ -123,15 +129,22 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
             relation_id: 0,
             is_distinct_allowed: false,
             bound_relations: vec![],
+            parallel_relations: vec![],
             bound_columns: vec![],
             is_mview: true,
         }
     }
 
     fn add_relation_to_context(&mut self, table: Table) {
-        let mut bound_columns = table.get_qualified_columns();
-        self.bound_columns.append(&mut bound_columns);
-        self.bound_relations.push(table);
+        self.parallel_relations.push(table);
+    }
+
+    fn merge_parallel_to_relation(&mut self) {
+        for rel in &self.parallel_relations {
+            let mut bound_columns = rel.get_qualified_columns();
+            self.bound_columns.append(&mut bound_columns);
+        }
+        self.bound_relations.append(&mut self.parallel_relations);
     }
 
     fn gen_stmt(&mut self) -> Statement {
@@ -204,7 +217,7 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
     }
 
     /// Generates a query with local context.
-    /// Used by `WITH`, (and perhaps subquery should use this too)
+    /// Used by `WITH`, subquery
     fn gen_local_query(&mut self) -> (Query, Vec<Column>) {
         let old_ctxt = self.new_local_context();
         let t = self.gen_query();
@@ -350,7 +363,9 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
                 .expect("with tables should not be empty");
             vec![create_table_with_joins_from_table(with_table)]
         } else {
-            vec![self.gen_from_relation()]
+            let rel = self.gen_from_relation();
+            self.merge_parallel_to_relation();
+            vec![rel]
         };
 
         if self.is_mview {
@@ -365,6 +380,8 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
                 from.push(self.gen_from_relation());
             }
         }
+        self.merge_parallel_to_relation();
+
         from
     }
 

src/tests/sqlsmith/src/relation.rs

@@ -33,13 +33,13 @@ fn create_join_on_clause(left: String, right: String) -> Expr {
 impl<'a, R: Rng> SqlGenerator<'a, R> {
     /// A relation specified in the FROM clause.
     pub(crate) fn gen_from_relation(&mut self) -> TableWithJoins {
-        match self.rng.gen_range(0..=9) {
-            0..=7 => self.gen_simple_table(),
-            8..=8 => self.gen_time_window_func(),
+        match self.rng.gen_range(0..=50) {
+            0..=39 => self.gen_simple_table(),
+            40..=44 => self.gen_time_window_func(),
             // TODO: Enable after resolving: <https://github.com/singularity-data/risingwave/issues/2771>.
-            9..=9 => self.gen_equijoin_clause(),
-            // TODO: Currently `gen_subquery` will cause panic due to some wrong assertions.
-            10..=10 => self.gen_subquery(),
+            45..=49 => self.gen_equijoin_clause(),
+            // TODO: Enable after resolving:
+            50..=50 => self.gen_table_subquery(),
             _ => unreachable!(),
         }
     }
@@ -120,8 +120,8 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
         }
     }
 
-    fn gen_subquery(&mut self) -> TableWithJoins {
-        let (subquery, columns) = self.gen_query();
+    fn gen_table_subquery(&mut self) -> TableWithJoins {
+        let (subquery, columns) = self.gen_local_query();
         let alias = self.gen_table_name_with_prefix("sq");
         let table = Table {
             name: alias.clone(),
@@ -141,4 +141,8 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
         self.add_relation_to_context(table);
         relation
     }
+
+    // fn gen_row_subquery(){
+
+    // }
 }

src/tests/sqlsmith/src/utils.rs

@@ -21,19 +21,26 @@ use risingwave_sqlparser::ast::{
 
 use crate::{Column, Expr, Ident, ObjectName, SqlGenerator, Table};
 
-type Context = (Vec<Column>, Vec<Table>);
+type Context = (Vec<Column>, Vec<Table>, Vec<Table>);
 
 /// Context utils
 impl<'a, R: Rng> SqlGenerator<'a, R> {
     pub(crate) fn new_local_context(&mut self) -> Context {
         let current_bound_relations = mem::take(&mut self.bound_relations);
         let current_bound_columns = mem::take(&mut self.bound_columns);
-        (current_bound_columns, current_bound_relations)
+        let current_parallel_relations = mem::take(&mut self.parallel_relations);
+        self.parallel_relations.clear();
+        (
+            current_bound_columns,
+            current_bound_relations,
+            current_parallel_relations,
+        )
     }
 
-    pub(crate) fn restore_context(&mut self, (old_cols, old_rels): Context) {
+    pub(crate) fn restore_context(&mut self, (old_cols, old_rels, old_parallel): Context) {
         self.bound_relations = old_rels;
         self.bound_columns = old_cols;
+        self.parallel_relations = old_parallel;
     }
 }