Skip to content

Commit 30c41b5

Browse files
sundy-lisundy-li
authored
fix(query): support count(t.*) (#18113)
* fix(query): support count(t.*) * fix(query): support count(t.*)
1 parent 7446dbc commit 30c41b5

File tree

11 files changed

+154
-10
lines changed

11 files changed

+154
-10
lines changed

‎src/query/ast/src/ast/expr.rs‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,9 @@ use super::TimeTravelPoint;
3838
use crate::ast::display_decimal_256;
3939
use crate::ast::quote::QuotedString;
4040
use crate::ast::write_comma_separated_list;
41+
use crate::ast::write_dot_separated_list;
4142
use crate::ast::Identifier;
43+
use crate::ast::Indirection;
4244
use crate::ast::Query;
4345
use crate::ast::SetExpr;
4446
use crate::span::merge_span;
@@ -175,6 +177,7 @@ pub enum Expr {
175177
/// `COUNT(*)` expression
176178
CountAll {
177179
span: Span,
180+
qualified: Vec<Indirection>,
178181
window: Option<Window>,
179182
},
180183
/// `(foo, bar)`
@@ -713,8 +716,12 @@ impl Display for Expr {
713716
Expr::Literal { value, .. } => {
714717
write!(f, "{value}")?;
715718
}
716-
Expr::CountAll { window, .. } => {
717-
write!(f, "COUNT(*)")?;
719+
Expr::CountAll {
720+
window, qualified, ..
721+
} => {
722+
write!(f, "COUNT(")?;
723+
write_dot_separated_list(f, qualified)?;
724+
write!(f, ")")?;
718725
if let Some(window) = window {
719726
write!(f, " OVER {window}")?;
720727
}

‎src/query/ast/src/parser/expr.rs‎

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,7 @@ pub enum ExprElement {
241241
},
242242
/// `Count(*)` expression
243243
CountAll {
244+
qualified: QualifiedName,
244245
window: Option<Window>,
245246
},
246247
/// `(foo, bar)`
@@ -569,8 +570,9 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
569570
span: transform_span(elem.span.tokens),
570571
value,
571572
},
572-
ExprElement::CountAll { window } => Expr::CountAll {
573+
ExprElement::CountAll { qualified, window } => Expr::CountAll {
573574
span: transform_span(elem.span.tokens),
575+
qualified,
574576
window,
575577
},
576578
ExprElement::Tuple { exprs } => Expr::Tuple {
@@ -1036,10 +1038,28 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
10361038

10371039
let count_all_with_window = map(
10381040
rule! {
1039-
COUNT ~ "(" ~ "*" ~ ")" ~ ( OVER ~ #window_spec_ident )?
1040-
},
1041-
|(_, _, _, _, window)| ExprElement::CountAll {
1042-
window: window.map(|w| w.1),
1041+
COUNT ~ "(" ~ ( #ident ~ "." ~ ( #ident ~ "." )? )? ~ "*" ~ ")" ~ ( OVER ~ #window_spec_ident )?
1042+
},
1043+
|(_, _, res, star, _, window)| match res {
1044+
Some((fst, _, Some((snd, _)))) => ExprElement::CountAll {
1045+
qualified: vec![
1046+
Indirection::Identifier(fst),
1047+
Indirection::Identifier(snd),
1048+
Indirection::Star(Some(star.span)),
1049+
],
1050+
window: window.map(|w| w.1),
1051+
},
1052+
Some((fst, _, None)) => ExprElement::CountAll {
1053+
qualified: vec![
1054+
Indirection::Identifier(fst),
1055+
Indirection::Star(Some(star.span)),
1056+
],
1057+
window: window.map(|w| w.1),
1058+
},
1059+
None => ExprElement::CountAll {
1060+
qualified: vec![Indirection::Star(Some(star.span))],
1061+
window: window.map(|w| w.1),
1062+
},
10431063
},
10441064
);
10451065

‎src/query/ast/tests/it/parser.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1114,6 +1114,7 @@ fn test_query() {
11141114
r#"select ?"#,
11151115
r#"select * exclude c1, b.* exclude (c2, c3, c4) from customer inner join orders on a = b limit 1"#,
11161116
r#"select columns('abc'), columns(a -> length(a) = 3) from t"#,
1117+
r#"select count(t.*) from t"#,
11171118
r#"select * from customer at(offset => -10 * 30)"#,
11181119
r#"select * from customer changes(information => default) at (stream => s) order by a, b"#,
11191120
r#"select * from customer with consume as s"#,

‎src/query/ast/tests/it/testdata/query.txt‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -447,6 +447,89 @@ Query {
447447
}
448448

449449

450+
---------- Input ----------
451+
select count(t.*) from t
452+
---------- Output ---------
453+
SELECT COUNT(t.*) FROM t
454+
---------- AST ------------
455+
Query {
456+
span: Some(
457+
0..24,
458+
),
459+
with: None,
460+
body: Select(
461+
SelectStmt {
462+
span: Some(
463+
0..24,
464+
),
465+
hints: None,
466+
distinct: false,
467+
top_n: None,
468+
select_list: [
469+
AliasedExpr {
470+
expr: CountAll {
471+
span: Some(
472+
7..17,
473+
),
474+
qualified: [
475+
Identifier(
476+
Identifier {
477+
span: Some(
478+
13..14,
479+
),
480+
name: "t",
481+
quote: None,
482+
ident_type: None,
483+
},
484+
),
485+
Star(
486+
Some(
487+
15..16,
488+
),
489+
),
490+
],
491+
window: None,
492+
},
493+
alias: None,
494+
},
495+
],
496+
from: [
497+
Table {
498+
span: Some(
499+
23..24,
500+
),
501+
catalog: None,
502+
database: None,
503+
table: Identifier {
504+
span: Some(
505+
23..24,
506+
),
507+
name: "t",
508+
quote: None,
509+
ident_type: None,
510+
},
511+
alias: None,
512+
temporal: None,
513+
with_options: None,
514+
pivot: None,
515+
unpivot: None,
516+
sample: None,
517+
},
518+
],
519+
selection: None,
520+
group_by: None,
521+
having: None,
522+
window_list: None,
523+
qualify: None,
524+
},
525+
),
526+
order_by: [],
527+
limit: [],
528+
offset: None,
529+
ignore_result: false,
530+
}
531+
532+
450533
---------- Input ----------
451534
select * from customer at(offset => -10 * 30)
452535
---------- Output ---------
@@ -3679,6 +3762,13 @@ Query {
36793762
span: Some(
36803763
19..27,
36813764
),
3765+
qualified: [
3766+
Star(
3767+
Some(
3768+
25..26,
3769+
),
3770+
),
3771+
],
36823772
window: None,
36833773
},
36843774
alias: Some(

‎src/query/ast/tests/it/testdata/script.txt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -287,6 +287,13 @@ Return {
287287
span: Some(
288288
20..28,
289289
),
290+
qualified: [
291+
Star(
292+
Some(
293+
26..27,
294+
),
295+
),
296+
],
290297
window: None,
291298
},
292299
alias: None,

‎src/query/ast/tests/it/testdata/stmt.txt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6001,6 +6001,13 @@ Query(
60016001
span: Some(
60026002
19..27,
60036003
),
6004+
qualified: [
6005+
Star(
6006+
Some(
6007+
25..26,
6008+
),
6009+
),
6010+
],
60046011
window: None,
60056012
},
60066013
alias: None,

‎src/query/sql/src/planner/semantic/aggregating_index_visitor.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -323,7 +323,7 @@ impl RefreshAggregatingIndexRewriter {
323323
self.has_agg_function = true;
324324
name.name = format!("{}_STATE", name.name);
325325
}
326-
Expr::CountAll { span, window } if window.is_none() => {
326+
Expr::CountAll { span, window, .. } if window.is_none() => {
327327
self.has_agg_function = true;
328328
*expr = Expr::FunctionCall {
329329
span: None,

‎src/query/sql/src/planner/semantic/type_check.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -940,7 +940,7 @@ impl<'a> TypeChecker<'a> {
940940
}
941941
}
942942

943-
Expr::CountAll { span, window } => {
943+
Expr::CountAll { span, window, .. } => {
944944
let (new_agg_func, data_type) =
945945
self.resolve_aggregate_function(*span, "count", expr, false, vec![], &[], &[])?;
946946

‎src/tests/sqlsmith/src/query_fuzzer.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -404,9 +404,14 @@ impl QueryVisitor {
404404
.as_ref()
405405
.map(|(trim_where, expr)| (trim_where.clone(), Box::new(self.fuzz_expr(expr)))),
406406
},
407-
Expr::CountAll { span, window } => Expr::CountAll {
407+
Expr::CountAll {
408+
span,
409+
window,
410+
qualified,
411+
} => Expr::CountAll {
408412
span: *span,
409413
window: window.as_ref().map(|window| self.fuzz_window(window)),
414+
qualified: qualified.clone(),
410415
},
411416
Expr::Tuple { span, exprs } => Expr::Tuple {
412417
span: *span,

‎src/tests/sqlsmith/src/sql_gen/expr.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -526,6 +526,7 @@ impl<R: Rng> SqlGenerator<'_, R> {
526526
}
527527
}
528528
2 => Expr::CountAll {
529+
qualified: vec![],
529530
span: None,
530531
window: None,
531532
},

0 commit comments

Comments
 (0)