Skip to content

Commit 80d9e78

Browse files
TCeasonTCeason
authored
chore(sqlsmith): support select stmt having && modify some err msg (#12959)
* chore(sqlsmith): support select stmt having && modify some err msg * fix conversation * sqlsmith support generate window_list * fix ci err
1 parent 510a6da commit 80d9e78

File tree

6 files changed

+126
-41
lines changed

6 files changed

+126
-41
lines changed

โ€Žsrc/query/expression/src/type_check.rs

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -282,12 +282,7 @@ pub fn check_number<Index: ColumnIndex, T: Number>(
282282
ErrorCode::InvalidArgument(format!("Expect {}, but got {}", T::data_type(), origin_ty))
283283
.set_span(span)
284284
}),
285-
_ => Err(ErrorCode::InvalidArgument(format!(
286-
"Expect {}, but got {}",
287-
T::data_type(),
288-
origin_ty
289-
))
290-
.set_span(span)),
285+
_ => Err(ErrorCode::InvalidArgument("Need constant number.").set_span(span)),
291286
}
292287
}
293288

โ€Žsrc/query/sql/src/planner/semantic/type_check.rs

Lines changed: 67 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ use common_expression::types::NumberDataType;
5454
use common_expression::types::NumberScalar;
5555
use common_expression::ColumnIndex;
5656
use common_expression::ConstantFolder;
57+
use common_expression::Expr as EExpr;
5758
use common_expression::FunctionContext;
5859
use common_expression::FunctionKind;
5960
use common_expression::RawExpr;
@@ -1410,19 +1411,28 @@ impl<'a> TypeChecker<'a> {
14101411
arg_types: &[DataType],
14111412
) -> Result<WindowFuncType> {
14121413
if args.is_empty() || args.len() > 3 {
1413-
return Err(ErrorCode::InvalidArgument(
1414-
"Argument number is invalid".to_string(),
1415-
));
1414+
return Err(ErrorCode::InvalidArgument(format!(
1415+
"Function {:?} only support 1 to 3 arguments",
1416+
func_name
1417+
)));
14161418
}
14171419

14181420
let offset = if args.len() >= 2 {
14191421
let off = args[1].as_expr()?;
1420-
Some(check_number::<_, i64>(
1421-
off.span(),
1422-
&self.func_ctx,
1423-
&off,
1424-
&BUILTIN_FUNCTIONS,
1425-
)?)
1422+
match off {
1423+
EExpr::Constant { .. } => Some(check_number::<_, i64>(
1424+
off.span(),
1425+
&self.func_ctx,
1426+
&off,
1427+
&BUILTIN_FUNCTIONS,
1428+
)?),
1429+
_ => {
1430+
return Err(ErrorCode::InvalidArgument(format!(
1431+
"The second argument to the function {:?} must be a constant",
1432+
func_name
1433+
)));
1434+
}
1435+
}
14261436
} else {
14271437
None
14281438
};
@@ -1473,9 +1483,10 @@ impl<'a> TypeChecker<'a> {
14731483
Ok(match func_name {
14741484
"first_value" | "first" => {
14751485
if args.len() != 1 {
1476-
return Err(ErrorCode::InvalidArgument(
1477-
"Argument number is invalid".to_string(),
1478-
));
1486+
return Err(ErrorCode::InvalidArgument(format!(
1487+
"The function {:?} must take one argument",
1488+
func_name
1489+
)));
14791490
}
14801491
let return_type = arg_types[0].wrap_nullable();
14811492
WindowFuncType::NthValue(NthValueFunction {
@@ -1486,9 +1497,10 @@ impl<'a> TypeChecker<'a> {
14861497
}
14871498
"last_value" | "last" => {
14881499
if args.len() != 1 {
1489-
return Err(ErrorCode::InvalidArgument(
1490-
"Argument number is invalid".to_string(),
1491-
));
1500+
return Err(ErrorCode::InvalidArgument(format!(
1501+
"The function {:?} must take one argument",
1502+
func_name
1503+
)));
14921504
}
14931505
let return_type = arg_types[0].wrap_nullable();
14941506
WindowFuncType::NthValue(NthValueFunction {
@@ -1501,17 +1513,24 @@ impl<'a> TypeChecker<'a> {
15011513
// nth_value
15021514
if args.len() != 2 {
15031515
return Err(ErrorCode::InvalidArgument(
1504-
"Argument number is invalid".to_string(),
1516+
"The function nth_value must take two arguments".to_string(),
15051517
));
15061518
}
15071519
let return_type = arg_types[0].wrap_nullable();
15081520
let n_expr = args[1].as_expr()?;
1509-
let n = check_number::<_, u64>(
1510-
n_expr.span(),
1511-
&self.func_ctx,
1512-
&n_expr,
1513-
&BUILTIN_FUNCTIONS,
1514-
)?;
1521+
let n = match n_expr {
1522+
EExpr::Constant { .. } => check_number::<_, u64>(
1523+
n_expr.span(),
1524+
&self.func_ctx,
1525+
&n_expr,
1526+
&BUILTIN_FUNCTIONS,
1527+
)?,
1528+
_ => {
1529+
return Err(ErrorCode::InvalidArgument(
1530+
"The count of `nth_value` must be constant positive integer",
1531+
));
1532+
}
1533+
};
15151534
if n == 0 {
15161535
return Err(ErrorCode::InvalidArgument(
15171536
"nth_value should count from 1".to_string(),
@@ -1534,12 +1553,21 @@ impl<'a> TypeChecker<'a> {
15341553
) -> Result<WindowFuncType> {
15351554
if args.len() != 1 {
15361555
return Err(ErrorCode::InvalidArgument(
1537-
"Argument number is invalid".to_string(),
1556+
"Function ntile can only take one argument".to_string(),
15381557
));
15391558
}
15401559
let n_expr = args[0].as_expr()?;
15411560
let return_type = DataType::Number(NumberDataType::UInt64);
1542-
let n = check_number::<_, u64>(n_expr.span(), &self.func_ctx, &n_expr, &BUILTIN_FUNCTIONS)?;
1561+
let n = match n_expr {
1562+
EExpr::Constant { .. } => {
1563+
check_number::<_, u64>(n_expr.span(), &self.func_ctx, &n_expr, &BUILTIN_FUNCTIONS)?
1564+
}
1565+
_ => {
1566+
return Err(ErrorCode::InvalidArgument(
1567+
"The argument of `ntile` must be constant".to_string(),
1568+
));
1569+
}
1570+
};
15431571
if n == 0 {
15441572
return Err(ErrorCode::InvalidArgument(
15451573
"ntile buckets must be greater than 0".to_string(),
@@ -1981,7 +2009,7 @@ impl<'a> TypeChecker<'a> {
19812009
)
19822010
.await
19832011
}
1984-
_ => Err(ErrorCode::SemanticError("Only these interval types are currently supported: [year, month, day, hour, minute, second]".to_string()).set_span(span)),
2012+
_ => Err(ErrorCode::SemanticError("Only these interval types are currently supported: [year, quarter, month, day, hour, minute, second]".to_string()).set_span(span)),
19852013
}
19862014
}
19872015

@@ -2247,7 +2275,20 @@ impl<'a> TypeChecker<'a> {
22472275
let box (scalar, _) = self.resolve(args[0]).await?;
22482276

22492277
let expr = scalar.as_expr()?;
2250-
check_number::<_, i64>(span, &self.func_ctx, &expr, &BUILTIN_FUNCTIONS)?
2278+
match expr {
2279+
EExpr::Constant { .. } => check_number::<_, i64>(
2280+
span,
2281+
&self.func_ctx,
2282+
&expr,
2283+
&BUILTIN_FUNCTIONS,
2284+
)?,
2285+
_ => {
2286+
return Some(Err(ErrorCode::BadArguments(
2287+
"last_query_id argument only support constant",
2288+
)
2289+
.set_span(span)));
2290+
}
2291+
}
22512292
}
22522293
};
22532294

โ€Žsrc/tests/sqlsmith/src/sql_gen/func.rs

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use common_ast::ast::Window;
2121
use common_ast::ast::WindowFrame;
2222
use common_ast::ast::WindowFrameBound;
2323
use common_ast::ast::WindowFrameUnits;
24+
use common_ast::ast::WindowRef;
2425
use common_ast::ast::WindowSpec;
2526
use common_expression::types::DataType;
2627
use common_expression::types::DecimalDataType::Decimal128;
@@ -615,6 +616,27 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
615616
}
616617

617618
fn gen_window(&mut self) -> Option<Window> {
619+
if self.rng.gen_bool(0.2) && !self.windows_name.is_empty() {
620+
let len = self.windows_name.len();
621+
let name = if len == 1 {
622+
self.windows_name[0].to_string()
623+
} else {
624+
self.windows_name[self.rng.gen_range(0..=len - 1)].to_string()
625+
};
626+
Some(Window::WindowReference(WindowRef {
627+
window_name: Identifier {
628+
name,
629+
quote: None,
630+
span: None,
631+
},
632+
}))
633+
} else {
634+
let window_spec = self.gen_window_spec();
635+
Some(Window::WindowSpec(window_spec))
636+
}
637+
}
638+
639+
pub(crate) fn gen_window_spec(&mut self) -> WindowSpec {
618640
let ty = self.gen_data_type();
619641
let expr1 = self.gen_scalar_value(&ty);
620642
let expr2 = self.gen_scalar_value(&ty);
@@ -633,7 +655,7 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
633655
nulls_first: Some(true),
634656
},
635657
];
636-
Some(Window::WindowSpec(WindowSpec {
658+
WindowSpec {
637659
existing_window_name: None,
638660
partition_by: vec![expr3, expr4],
639661
order_by,
@@ -646,7 +668,7 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
646668
end_bound: WindowFrameBound::CurrentRow,
647669
})
648670
},
649-
}))
671+
}
650672
}
651673

652674
fn gen_func(

โ€Žsrc/tests/sqlsmith/src/sql_gen/query.rs

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ use common_ast::ast::SelectTarget;
2828
use common_ast::ast::SetExpr;
2929
use common_ast::ast::TableAlias;
3030
use common_ast::ast::TableReference;
31+
use common_ast::ast::WindowDefinition;
3132
use common_ast::ast::With;
3233
use common_ast::ast::CTE;
3334
use common_expression::infer_schema_type;
@@ -289,28 +290,52 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
289290
}
290291

291292
fn gen_select(&mut self) -> SelectStmt {
293+
self.windows_name.clear();
292294
let from = self.gen_from();
293295
let group_by = self.gen_group_by();
294296
self.group_by = group_by.clone();
297+
let window_list = self.gen_window_list();
298+
if let Some(window_list) = window_list {
299+
for window in window_list {
300+
self.windows_name.push(window.name.name)
301+
}
302+
}
295303
let select_list = self.gen_select_list(&group_by);
296304
let selection = self.gen_selection();
297305
SelectStmt {
298306
span: None,
299307
// TODO
300308
hints: None,
301-
// TODO
302-
distinct: false,
309+
distinct: self.rng.gen_bool(0.7),
303310
select_list,
304311
from,
305312
selection,
306313
group_by,
307-
// TODO
308-
having: None,
309-
// TODO
310-
window_list: None,
314+
having: self.gen_selection(),
315+
window_list: self.gen_window_list(),
311316
}
312317
}
313318

319+
fn gen_window_list(&mut self) -> Option<Vec<WindowDefinition>> {
320+
if self.rng.gen_bool(0.1) {
321+
let mut res = vec![];
322+
for _ in 0..self.rng.gen_range(1..3) {
323+
let name: String = (0..4)
324+
.map(|_| self.rng.sample(Alphanumeric) as char)
325+
.collect();
326+
let window_name = format!("w_{}", name);
327+
let spec = self.gen_window_spec();
328+
let window_def = WindowDefinition {
329+
name: Identifier::from_name(window_name),
330+
spec,
331+
};
332+
res.push(window_def);
333+
}
334+
return Some(res);
335+
}
336+
None
337+
}
338+
314339
fn gen_group_by(&mut self) -> Option<GroupBy> {
315340
if self.rng.gen_bool(0.8) {
316341
return None;

โ€Žsrc/tests/sqlsmith/src/sql_gen/sql_generator.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ pub(crate) struct SqlGenerator<'a, R: Rng> {
4848
pub(crate) scalar_func_sigs: Vec<FunctionSignature>,
4949
pub(crate) rng: &'a mut R,
5050
pub(crate) group_by: Option<GroupBy>,
51+
pub(crate) windows_name: Vec<String>,
5152
}
5253

5354
impl<'a, R: Rng> SqlGenerator<'a, R> {
@@ -82,6 +83,7 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
8283
scalar_func_sigs,
8384
rng,
8485
group_by: None,
86+
windows_name: vec![],
8587
}
8688
}
8789
}

โ€Žtests/sqllogictests/suites/query/02_function/02_0000_function_aggregate_mix

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -323,7 +323,7 @@ select group_array_moving_avg(k), group_array_moving_avg(2)(v) from aggr;
323323
----
324324
[0.09090909090909091,0.2727272727272727,0.45454545454545453,0.6363636363636364,0.8181818181818182,1.0,1.1818181818181819,1.3636363636363635,1.5454545454545454,1.7272727272727273,1.9090909090909092] [5.0,10.0,10.0,10.0,15.0,20.0,22.5,27.5,30.0,30.0,30.0]
325325

326-
statement error Expect UInt64, but got String
326+
statement error Need constant number
327327
SELECT group_array_moving_sum('x')(-1130932975.87767);
328328

329329
query TTT

0 commit comments

Comments
ย (0)