risingwave_sqlsmith/sql_gen/
expr.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use itertools::Itertools;
16use rand::Rng;
17use rand::seq::IndexedRandom;
18use risingwave_common::types::{DataType, DataTypeName, StructType};
19use risingwave_expr::sig::FUNCTION_REGISTRY;
20use risingwave_frontend::expr::cast_sigs;
21use risingwave_sqlparser::ast::{Expr, OrderByExpr, Value};
22
23use crate::sql_gen::types::data_type_to_ast_data_type;
24use crate::sql_gen::{SqlGenerator, SqlGeneratorContext};
25
26static STRUCT_FIELD_NAMES: [&str; 26] = [
27    "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s",
28    "t", "u", "v", "w", "x", "y", "z",
29];
30
31impl<R: Rng> SqlGenerator<'_, R> {
32    /// In generating expression, there are two execution modes:
33    /// 1) Can have Aggregate expressions (`can_agg` = true)
34    ///    We can have aggregate of all bound columns (those present in GROUP BY and otherwise).
35    ///    Not all GROUP BY columns need to be aggregated.
36    /// 2) Can't have Aggregate expressions (`can_agg` = false)
37    ///    Only columns present in GROUP BY can be selected.
38    ///
39    /// `inside_agg` indicates if we are calling `gen_expr` inside an aggregate.
40    pub(crate) fn gen_expr(&mut self, typ: &DataType, context: SqlGeneratorContext) -> Expr {
41        if !self.can_recurse() {
42            // Stop recursion with a simple scalar or column.
43            // Weight it more towards columns, scalar has much higher chance of being generated,
44            // since it is usually used as fail-safe expression.
45            return match self.rng.random_bool(0.1) {
46                true => self.gen_simple_scalar(typ),
47                false => self.gen_col(typ, context),
48            };
49        }
50
51        if *typ == DataType::Boolean && self.rng.random_bool(0.05) {
52            return match self.rng.random_bool(0.5) {
53                true => {
54                    let (ty, expr) = self.gen_arbitrary_expr(context);
55                    let n = self.rng.random_range(1..=10);
56                    Expr::InList {
57                        expr: Box::new(Expr::Nested(Box::new(expr))),
58                        list: self.gen_n_exprs_with_type(n, &ty, context),
59                        negated: self.flip_coin(),
60                    }
61                }
62                false => {
63                    // TODO: InSubquery expression may not be always bound in all context.
64                    // Parts labelled workaround can be removed or
65                    // generalized if it is bound in all contexts.
66                    // https://github.com/risingwavelabs/risingwave/issues/1343
67                    let old_ctxt = self.new_local_context(); // WORKAROUND
68                    let (query, column) = self.gen_single_item_query();
69                    let ty = column.data_type;
70                    let expr = self.gen_simple_scalar(&ty); // WORKAROUND
71                    let in_subquery_expr = Expr::InSubquery {
72                        expr: Box::new(Expr::Nested(Box::new(expr))),
73                        subquery: Box::new(query),
74                        negated: self.flip_coin(),
75                    };
76                    self.restore_context(old_ctxt); // WORKAROUND
77                    in_subquery_expr
78                }
79            };
80        }
81
82        // NOTE:
83        // We generate AST first, then use its `Display` trait
84        // to generate an sql string.
85        // That may erase nesting context.
86        // For instance `IN(a, b)` is `a IN b`.
87        // this can lead to ambiguity, if `a` is an
88        // INFIX/POSTFIX compound expression too:
89        // - `a1 IN a2 IN b`
90        // - `a1 >= a2 IN b`
91        // ...
92        // We just nest compound expressions to avoid this.
93        let range = if context.can_gen_agg() { 100 } else { 50 };
94        match self.rng.random_range(0..=range) {
95            0..=35 => Expr::Nested(Box::new(self.gen_func(typ, context))),
96            36..=40 => self.gen_exists(typ, context),
97            41..=50 => self.gen_explicit_cast(typ, context),
98            51..=100 => self.gen_agg(typ),
99            _ => unreachable!(),
100        }
101    }
102
103    fn gen_data_type(&mut self) -> DataType {
104        // Depth of struct/list nesting
105        let depth = self.rng.random_range(0..=1);
106        self.gen_data_type_inner(depth)
107    }
108
109    fn gen_data_type_inner(&mut self, depth: usize) -> DataType {
110        match self.rng.random_bool(0.8) {
111            true if !self.bound_columns.is_empty() => self
112                .bound_columns
113                .choose(&mut self.rng)
114                .unwrap()
115                .data_type
116                .clone(),
117            _ => {
118                use DataType as S;
119                use DataTypeName as T;
120                let mut candidate_ret_types = vec![
121                    T::Boolean,
122                    T::Int16,
123                    T::Int32,
124                    T::Int64,
125                    T::Decimal,
126                    T::Float32,
127                    T::Float64,
128                    T::Varchar,
129                    T::Date,
130                    T::Timestamp,
131                    // ENABLE: https://github.com/risingwavelabs/risingwave/issues/5826
132                    // T::Timestamptz,
133                    T::Time,
134                    T::Interval,
135                ];
136                if depth > 0 {
137                    candidate_ret_types.push(T::Struct);
138                    candidate_ret_types.push(T::List);
139                }
140                let typ_name = candidate_ret_types.choose(&mut self.rng).unwrap();
141                match typ_name {
142                    T::Boolean => S::Boolean,
143                    T::Int16 => S::Int16,
144                    T::Int32 => S::Int32,
145                    T::Int64 => S::Int64,
146                    T::Decimal => S::Decimal,
147                    T::Float32 => S::Float32,
148                    T::Float64 => S::Float64,
149                    T::Varchar => S::Varchar,
150                    T::Date => S::Date,
151                    T::Timestamp => S::Timestamp,
152                    T::Timestamptz => S::Timestamptz,
153                    T::Time => S::Time,
154                    T::Interval => S::Interval,
155                    T::Struct => self.gen_struct_data_type(depth - 1),
156                    T::List => self.gen_list_data_type(depth - 1),
157                    _ => unreachable!(),
158                }
159            }
160        }
161    }
162
163    pub(crate) fn gen_list_data_type(&mut self, depth: usize) -> DataType {
164        DataType::list(self.gen_data_type_inner(depth))
165    }
166
167    fn gen_struct_data_type(&mut self, depth: usize) -> DataType {
168        let num_fields = self.rng.random_range(1..4);
169        DataType::Struct(StructType::new(
170            STRUCT_FIELD_NAMES[0..num_fields]
171                .iter()
172                .map(|s| (s.to_string(), self.gen_data_type_inner(depth))),
173        ))
174    }
175
176    /// Generates an arbitrary expression, but biased towards datatypes present in bound columns.
177    pub(crate) fn gen_arbitrary_expr(&mut self, context: SqlGeneratorContext) -> (DataType, Expr) {
178        let ret_type = self.gen_data_type();
179        let expr = self.gen_expr(&ret_type, context);
180        (ret_type, expr)
181    }
182
183    fn gen_col(&mut self, typ: &DataType, context: SqlGeneratorContext) -> Expr {
184        let columns = if context.is_inside_agg() {
185            if self.bound_relations.is_empty() {
186                return self.gen_simple_scalar(typ);
187            }
188            self.bound_relations
189                .choose(self.rng)
190                .unwrap()
191                .get_qualified_columns()
192        } else {
193            if self.bound_columns.is_empty() {
194                return self.gen_simple_scalar(typ);
195            }
196            self.bound_columns.clone()
197        };
198
199        let matched_cols = columns
200            .iter()
201            .filter(|col| col.data_type == *typ)
202            .collect::<Vec<_>>();
203        if matched_cols.is_empty() {
204            self.gen_simple_scalar(typ)
205        } else {
206            let col_def = matched_cols.choose(&mut self.rng).unwrap();
207            col_def.name_expr()
208        }
209    }
210
211    /// Generates `n` expressions of type `ret`.
212    pub(crate) fn gen_n_exprs_with_type(
213        &mut self,
214        n: usize,
215        ret: &DataType,
216        context: SqlGeneratorContext,
217    ) -> Vec<Expr> {
218        (0..n).map(|_| self.gen_expr(ret, context)).collect()
219    }
220
221    fn gen_exists(&mut self, ret: &DataType, context: SqlGeneratorContext) -> Expr {
222        if *ret != DataType::Boolean || context.can_gen_agg() {
223            return self.gen_simple_scalar(ret);
224        };
225        // Generating correlated subquery tends to create queries which cannot be unnested.
226        // we still want to test it, but reduce the chance it occurs.
227        let (subquery, _) = match self.rng.random_bool(0.05) {
228            true => self.gen_correlated_query(),
229            false => self.gen_local_query(),
230        };
231        Expr::Exists(Box::new(subquery))
232    }
233
234    /// Generate ORDER BY expressions by choosing from available bound columns.
235    pub(crate) fn gen_order_by(&mut self) -> Vec<OrderByExpr> {
236        if self.bound_columns.is_empty() {
237            return vec![];
238        }
239        let mut order_by = vec![];
240        while self.flip_coin() {
241            let column = self.bound_columns.choose(&mut self.rng).unwrap();
242            order_by.push(OrderByExpr {
243                expr: column.name_expr(),
244                asc: if self.rng.random_bool(0.3) {
245                    None
246                } else {
247                    Some(self.rng.random_bool(0.5))
248                },
249                nulls_first: if self.rng.random_bool(0.3) {
250                    None
251                } else {
252                    Some(self.rng.random_bool(0.5))
253                },
254            })
255        }
256        order_by
257    }
258
259    /// Generate ORDER BY expressions by choosing from given expressions.
260    pub(crate) fn gen_order_by_within(&mut self, exprs: &[Expr]) -> Vec<OrderByExpr> {
261        let exprs = exprs
262            .iter()
263            .filter(|e| matches!(e, Expr::Identifier(_) | Expr::Value(_)))
264            .cloned()
265            .collect::<Vec<_>>();
266        if exprs.is_empty() {
267            return vec![];
268        }
269        let mut order_by = vec![];
270        while self.flip_coin() {
271            let expr = exprs.choose(&mut self.rng).unwrap();
272            order_by.push(OrderByExpr {
273                expr: expr.clone(),
274                asc: if self.rng.random_bool(0.3) {
275                    None
276                } else {
277                    Some(self.rng.random_bool(0.5))
278                },
279                nulls_first: if self.rng.random_bool(0.3) {
280                    None
281                } else {
282                    Some(self.rng.random_bool(0.5))
283                },
284            })
285        }
286        order_by
287    }
288}
289
290pub(crate) fn typed_null(ty: &DataType) -> Expr {
291    Expr::Cast {
292        expr: Box::new(sql_null()),
293        data_type: data_type_to_ast_data_type(ty),
294    }
295}
296
297/// Generates a `NULL` value.
298pub(crate) fn sql_null() -> Expr {
299    Expr::Value(Value::Null)
300}
301
302// TODO(kwannoel):
303// Add variadic function signatures. Can add these functions
304// to a FUNC_TABLE too.
305pub fn print_function_table() -> String {
306    let func_str = FUNCTION_REGISTRY
307        .iter_scalars()
308        .map(|sign| {
309            format!(
310                "{}({}) -> {}",
311                sign.name,
312                sign.inputs_type.iter().format(", "),
313                sign.ret_type,
314            )
315        })
316        .join("\n");
317
318    let agg_func_str = FUNCTION_REGISTRY
319        .iter_aggregates()
320        .map(|sign| {
321            format!(
322                "{}({}) -> {}",
323                sign.name,
324                sign.inputs_type.iter().format(", "),
325                sign.ret_type,
326            )
327        })
328        .join("\n");
329
330    let cast_str = cast_sigs()
331        .map(|sig| {
332            format!(
333                "{:?} CAST {:?} -> {:?}",
334                sig.context, sig.from_type, sig.to_type,
335            )
336        })
337        .sorted()
338        .join("\n");
339
340    format!(
341        "
342==== FUNCTION SIGNATURES
343{}
344
345==== AGGREGATE FUNCTION SIGNATURES
346{}
347
348==== CAST SIGNATURES
349{}
350",
351        func_str, agg_func_str, cast_str
352    )
353}
risingwave_sqlsmith/sql_gen/expr.rs

risingwave_sqlsmith/sql_gen/
expr.rs
