Skip to content

Commit cf2c20b

Browse files
jamiibenesch
authored andcommitted
Support COUNT(DISTINCT x) and similar
1 parent 4f944dd commit cf2c20b

File tree

3 files changed

+60
-4
lines changed

3 files changed

+60
-4
lines changed

src/sqlast/mod.rs

+17-2
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,9 @@ pub enum ASTNode {
112112
name: SQLObjectName,
113113
args: Vec<ASTNode>,
114114
over: Option<SQLWindowSpec>,
115+
// aggregate functions may specify eg `COUNT(DISTINCT x)`
116+
all: bool,
117+
distinct: bool,
115118
},
116119
/// CASE [<operand>] WHEN <condition> THEN <result> ... [ELSE <result>] END
117120
/// Note we only recognize a complete single expression as <condition>, not
@@ -190,8 +193,20 @@ impl ToString for ASTNode {
190193
format!("{} {}", operator.to_string(), expr.as_ref().to_string())
191194
}
192195
ASTNode::SQLValue(v) => v.to_string(),
193-
ASTNode::SQLFunction { name, args, over } => {
194-
let mut s = format!("{}({})", name.to_string(), comma_separated_string(args));
196+
ASTNode::SQLFunction {
197+
name,
198+
args,
199+
over,
200+
all,
201+
distinct,
202+
} => {
203+
let mut s = format!(
204+
"{}({}{}{})",
205+
name.to_string(),
206+
if *all { "ALL " } else { "" },
207+
if *distinct { "DISTINCT " } else { "" },
208+
comma_separated_string(args)
209+
);
195210
if let Some(o) = over {
196211
s += &format!(" OVER ({})", o.to_string())
197212
}

src/sqlparser.rs

+15-1
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,14 @@ impl Parser {
253253

254254
pub fn parse_function(&mut self, name: SQLObjectName) -> Result<ASTNode, ParserError> {
255255
self.expect_token(&Token::LParen)?;
256+
let all = self.parse_keyword("ALL");
257+
let distinct = self.parse_keyword("DISTINCT");
258+
if all && distinct {
259+
return parser_err!(format!(
260+
"Cannot specify both ALL and DISTINCT in function: {:?}",
261+
name
262+
));
263+
}
256264
let args = self.parse_optional_args()?;
257265
let over = if self.parse_keyword("OVER") {
258266
// TBD: support window names (`OVER mywin`) in place of inline specification
@@ -279,7 +287,13 @@ impl Parser {
279287
None
280288
};
281289

282-
Ok(ASTNode::SQLFunction { name, args, over })
290+
Ok(ASTNode::SQLFunction {
291+
name,
292+
args,
293+
over,
294+
all,
295+
distinct,
296+
})
283297
}
284298

285299
pub fn parse_window_frame(&mut self) -> Result<Option<SQLWindowFrame>, ParserError> {

tests/sqlparser_common.rs

+28-1
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,27 @@ fn parse_select_count_wildcard() {
197197
name: SQLObjectName(vec!["COUNT".to_string()]),
198198
args: vec![ASTNode::SQLWildcard],
199199
over: None,
200+
distinct: false,
201+
all: false,
202+
},
203+
expr_from_projection(only(&select.projection))
204+
);
205+
}
206+
207+
#[test]
208+
fn parse_select_count_distinct() {
209+
let sql = "SELECT COUNT(DISTINCT + x) FROM customer";
210+
let select = verified_only_select(sql);
211+
assert_eq!(
212+
&ASTNode::SQLFunction {
213+
name: SQLObjectName(vec!["COUNT".to_string()]),
214+
args: vec![ASTNode::SQLUnary {
215+
operator: SQLOperator::Plus,
216+
expr: Box::new(ASTNode::SQLIdentifier("x".to_string()))
217+
}],
218+
over: None,
219+
distinct: true,
220+
all: false,
200221
},
201222
expr_from_projection(only(&select.projection))
202223
);
@@ -662,6 +683,8 @@ fn parse_scalar_function_in_projection() {
662683
name: SQLObjectName(vec!["sqrt".to_string()]),
663684
args: vec![ASTNode::SQLIdentifier("id".to_string())],
664685
over: None,
686+
all: false,
687+
distinct: false,
665688
},
666689
expr_from_projection(only(&select.projection))
667690
);
@@ -690,7 +713,9 @@ fn parse_window_functions() {
690713
asc: Some(false)
691714
}],
692715
window_frame: None,
693-
})
716+
}),
717+
all: false,
718+
distinct: false,
694719
},
695720
expr_from_projection(&select.projection[0])
696721
);
@@ -762,6 +787,8 @@ fn parse_delimited_identifiers() {
762787
name: SQLObjectName(vec![r#""myfun""#.to_string()]),
763788
args: vec![],
764789
over: None,
790+
all: false,
791+
distinct: false,
765792
},
766793
expr_from_projection(&select.projection[1]),
767794
);

0 commit comments

Comments
 (0)