Skip to content

Commit

Permalink
feat(sqlsmith): enable sub-query generation (#4216)
Browse files Browse the repository at this point in the history 
* Solve bug 1

* Resolve error

* Resolved Comments

* Undo enable agg

* remove lateralcontexts

* Fix formatting

* Control recursion

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
  • Loading branch information
@marvenlee2486 @mergify
marvenlee2486 and mergify[bot] authored Aug 1, 2022
1 parent 161cb22 commit 293b177
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 42 deletions.
23 changes: 15 additions & 8 deletions src/tests/sqlsmith/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -128,10 +128,12 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
}
}

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);
fn add_relations_to_context(&mut self, mut tables: Vec<Table>) {
for rel in &tables {
let mut bound_columns = rel.get_qualified_columns();
self.bound_columns.append(&mut bound_columns);
}
self.bound_relations.append(&mut tables);
}

fn gen_stmt(&mut self) -> Statement {
Expand Down Expand Up @@ -204,7 +206,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();
Expand Down Expand Up @@ -350,7 +352,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, tables) = self.gen_from_relation();
self.add_relations_to_context(tables);
vec![rel]
};

if self.is_mview {
Expand All @@ -359,12 +363,15 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
assert!(!self.tables.is_empty());
return from;
}

let mut lateral_contexts = vec![];
for _ in 0..self.tables.len() {
if self.flip_coin() {
from.push(self.gen_from_relation());
let (table_with_join, mut table) = self.gen_from_relation();
from.push(table_with_join);
lateral_contexts.append(&mut table);
}
}
self.add_relations_to_context(lateral_contexts);
from
}

Expand Down
59 changes: 32 additions & 27 deletions src/tests/sqlsmith/src/relation.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,28 +32,31 @@ 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) {
pub(crate) fn gen_from_relation(&mut self) -> (TableWithJoins, Vec<Table>) {
let range = if self.can_recurse() { 10 } else { 9 };
match self.rng.gen_range(0..=range) {
0..=7 => self.gen_simple_table(),
8..=8 => 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(),
10..=10 => self.gen_table_subquery(),
_ => unreachable!(),
}
}

fn gen_simple_table(&mut self) -> TableWithJoins {
let (relation, _) = self.gen_simple_table_factor();
fn gen_simple_table(&mut self) -> (TableWithJoins, Vec<Table>) {
let (relation, _, table) = self.gen_simple_table_factor();

TableWithJoins {
relation,
joins: vec![],
}
(
TableWithJoins {
relation,
joins: vec![],
},
table,
)
}

fn gen_simple_table_factor(&mut self) -> (TableFactor, Vec<Column>) {
fn gen_simple_table_factor(&mut self) -> (TableFactor, Vec<Column>, Vec<Table>) {
let alias = self.gen_table_name_with_prefix("t");
let mut table = self.tables.choose(&mut self.rng).unwrap().clone();
let table_factor = TableFactor::Table {
Expand All @@ -66,11 +69,10 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
};
table.name = alias; // Rename the table.
let columns = table.get_qualified_columns();
self.add_relation_to_context(table);
(table_factor, columns)
(table_factor, columns, vec![table])
}

fn gen_table_factor(&mut self) -> (TableFactor, Vec<Column>) {
fn gen_table_factor(&mut self) -> (TableFactor, Vec<Column>, Vec<Table>) {
let current_context = self.new_local_context();
let factor = self.gen_table_factor_inner();
self.restore_context(current_context);
Expand All @@ -79,14 +81,14 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {

/// Generates a table factor, and provides bound columns.
/// Generated column names should be qualified by table name.
fn gen_table_factor_inner(&mut self) -> (TableFactor, Vec<Column>) {
fn gen_table_factor_inner(&mut self) -> (TableFactor, Vec<Column>, Vec<Table>) {
// TODO: TableFactor::Derived, TableFactor::TableFunction, TableFactor::NestedJoin
self.gen_simple_table_factor()
}

fn gen_equijoin_clause(&mut self) -> TableWithJoins {
let (left_factor, left_columns) = self.gen_table_factor();
let (right_factor, right_columns) = self.gen_table_factor();
fn gen_equijoin_clause(&mut self) -> (TableWithJoins, Vec<Table>) {
let (left_factor, left_columns, mut left_table) = self.gen_table_factor();
let (right_factor, right_columns, mut right_table) = self.gen_table_factor();

let mut available_join_on_columns = vec![];
for left_column in &left_columns {
Expand All @@ -113,15 +115,18 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
relation: right_factor,
join_operator: JoinOperator::Inner(JoinConstraint::On(join_on_expr)),
};

TableWithJoins {
relation: left_factor,
joins: vec![right_factor_with_join],
}
left_table.append(&mut right_table);
(
TableWithJoins {
relation: left_factor,
joins: vec![right_factor_with_join],
},
left_table,
)
}

fn gen_subquery(&mut self) -> TableWithJoins {
let (subquery, columns) = self.gen_query();
fn gen_table_subquery(&mut self) -> (TableWithJoins, Vec<Table>) {
let (subquery, columns) = self.gen_local_query();
let alias = self.gen_table_name_with_prefix("sq");
let table = Table {
name: alias.clone(),
Expand All @@ -138,7 +143,7 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
},
joins: vec![],
};
self.add_relation_to_context(table);
relation

(relation, vec![table])
}
}
12 changes: 5 additions & 7 deletions src/tests/sqlsmith/src/time_window.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ use crate::{Column, Expr, SqlGenerator, Table};

impl<'a, R: Rng> SqlGenerator<'a, R> {
/// Generates time window functions.
pub(crate) fn gen_time_window_func(&mut self) -> TableWithJoins {
pub(crate) fn gen_time_window_func(&mut self) -> (TableWithJoins, Vec<Table>) {
match self.flip_coin() {
true => self.gen_hop(),
false => self.gen_tumble(),
Expand All @@ -34,7 +34,7 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {

/// Generates `TUMBLE`.
/// TUMBLE(data: TABLE, timecol: COLUMN, size: INTERVAL, offset?: INTERVAL)
fn gen_tumble(&mut self) -> TableWithJoins {
fn gen_tumble(&mut self) -> (TableWithJoins, Vec<Table>) {
let tables = find_tables_with_timestamp_cols(self.tables.clone());
let (source_table_name, time_cols, schema) = tables
.choose(&mut self.rng)
Expand All @@ -52,14 +52,13 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
let relation = create_tvf("tumble", alias, args);

let table = Table::new(table_name, schema.clone());
self.add_relation_to_context(table);

relation
(relation, vec![table])
}

/// Generates `HOP`.
/// HOP(data: TABLE, timecol: COLUMN, slide: INTERVAL, size: INTERVAL, offset?: INTERVAL)
fn gen_hop(&mut self) -> TableWithJoins {
fn gen_hop(&mut self) -> (TableWithJoins, Vec<Table>) {
let tables = find_tables_with_timestamp_cols(self.tables.clone());
let (source_table_name, time_cols, schema) = tables
.choose(&mut self.rng)
Expand All @@ -84,9 +83,8 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
let relation = create_tvf("hop", alias, args);

let table = Table::new(table_name, schema.clone());
self.add_relation_to_context(table);

relation
(relation, vec![table])
}
}

Expand Down

0 comments on commit 293b177

Please sign in to comment.