Support parametric arguments to FUNCTION for ClickHouse dialect (#1315)

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
This commit is contained in:
hulk 2024-06-23 19:36:05 +08:00 committed by GitHub
parent 7a9793b72e
commit a685e11993
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 119 additions and 5 deletions

View file

@ -4695,6 +4695,16 @@ impl fmt::Display for CloseCursor {
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct Function {
pub name: ObjectName,
/// The parameters to the function, including any options specified within the
/// delimiting parentheses.
///
/// Example:
/// ```plaintext
/// HISTOGRAM(0.5, 0.6)(x, y)
/// ```
///
/// [ClickHouse](https://clickhouse.com/docs/en/sql-reference/aggregate-functions/parametric-functions)
pub parameters: FunctionArguments,
/// The arguments to the function, including any options specified within the
/// delimiting parentheses.
pub args: FunctionArguments,
@ -4723,7 +4733,7 @@ pub struct Function {
impl fmt::Display for Function {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}{}", self.name, self.args)?;
write!(f, "{}{}{}", self.name, self.parameters, self.args)?;
if !self.within_group.is_empty() {
write!(

View file

@ -533,6 +533,7 @@ where
/// null_treatment: None,
/// filter: None,
/// over: None,
/// parameters: FunctionArguments::None,
/// within_group: vec![],
/// });
/// }

View file

@ -27,6 +27,7 @@ use core::{
use log::debug;
use recursion::RecursionCounter;
use IsLateral::*;
use IsOptional::*;
@ -146,8 +147,6 @@ mod recursion {
pub struct DepthGuard {}
}
use recursion::RecursionCounter;
#[derive(PartialEq, Eq)]
pub enum IsOptional {
Optional,
@ -1002,6 +1001,7 @@ impl<'a> Parser<'a> {
{
Ok(Expr::Function(Function {
name: ObjectName(vec![w.to_ident()]),
parameters: FunctionArguments::None,
args: FunctionArguments::None,
null_treatment: None,
filter: None,
@ -1058,6 +1058,7 @@ impl<'a> Parser<'a> {
self.expect_token(&Token::RParen)?;
Ok(Expr::Function(Function {
name: ObjectName(vec![w.to_ident()]),
parameters: FunctionArguments::None,
args: FunctionArguments::Subquery(query),
filter: None,
null_treatment: None,
@ -1293,6 +1294,7 @@ impl<'a> Parser<'a> {
self.expect_token(&Token::RParen)?;
return Ok(Expr::Function(Function {
name,
parameters: FunctionArguments::None,
args: FunctionArguments::Subquery(subquery),
filter: None,
null_treatment: None,
@ -1301,7 +1303,16 @@ impl<'a> Parser<'a> {
}));
}
let args = self.parse_function_argument_list()?;
let mut args = self.parse_function_argument_list()?;
let mut parameters = FunctionArguments::None;
// ClickHouse aggregations support parametric functions like `HISTOGRAM(0.5, 0.6)(x, y)`
// which (0.5, 0.6) is a parameter to the function.
if dialect_of!(self is ClickHouseDialect | GenericDialect)
&& self.consume_token(&Token::LParen)
{
parameters = FunctionArguments::List(args);
args = self.parse_function_argument_list()?;
}
let within_group = if self.parse_keywords(&[Keyword::WITHIN, Keyword::GROUP]) {
self.expect_token(&Token::LParen)?;
@ -1350,6 +1361,7 @@ impl<'a> Parser<'a> {
Ok(Expr::Function(Function {
name,
parameters,
args: FunctionArguments::List(args),
null_treatment,
filter,
@ -1382,6 +1394,7 @@ impl<'a> Parser<'a> {
};
Ok(Expr::Function(Function {
name,
parameters: FunctionArguments::None,
args,
filter: None,
over: None,
@ -6470,6 +6483,7 @@ impl<'a> Parser<'a> {
} else {
Ok(Statement::Call(Function {
name: object_name,
parameters: FunctionArguments::None,
args: FunctionArguments::None,
over: None,
filter: None,
@ -8092,7 +8106,7 @@ impl<'a> Parser<'a> {
pub fn parse_query_body(&mut self, precedence: u8) -> Result<SetExpr, ParserError> {
// We parse the expression using a Pratt parser, as in `parse_expr()`.
// Start by parsing a restricted SELECT or a `(subquery)`:
let mut expr = if self.parse_keyword(Keyword::SELECT) {
let expr = if self.parse_keyword(Keyword::SELECT) {
SetExpr::Select(self.parse_select().map(Box::new)?)
} else if self.consume_token(&Token::LParen) {
// CTEs are not allowed here, but the parser currently accepts them
@ -8111,6 +8125,17 @@ impl<'a> Parser<'a> {
);
};
self.parse_remaining_set_exprs(expr, precedence)
}
/// Parse any extra set expressions that may be present in a query body
///
/// (this is its own function to reduce required stack size in debug builds)
fn parse_remaining_set_exprs(
&mut self,
mut expr: SetExpr,
precedence: u8,
) -> Result<SetExpr, ParserError> {
loop {
// The query can be optionally followed by a set operator:
let op = self.parse_set_operator(&self.peek_token().token);

View file

@ -336,6 +336,7 @@ pub fn join(relation: TableFactor) -> Join {
pub fn call(function: &str, args: impl IntoIterator<Item = Expr>) -> Expr {
Expr::Function(Function {
name: ObjectName(vec![Ident::new(function)]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: args

View file

@ -183,6 +183,7 @@ fn parse_delimited_identifiers() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::with_quote('"', "myfun")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![],
@ -553,6 +554,55 @@ fn parse_select_star_except() {
clickhouse().verified_stmt("SELECT * EXCEPT (prev_status) FROM anomalies");
}
#[test]
fn parse_select_parametric_function() {
match clickhouse_and_generic().verified_stmt("SELECT HISTOGRAM(0.5, 0.6)(x, y) FROM t") {
Statement::Query(query) => {
let projection: &Vec<SelectItem> = query.body.as_select().unwrap().projection.as_ref();
assert_eq!(projection.len(), 1);
match &projection[0] {
UnnamedExpr(Expr::Function(f)) => {
let args = match &f.args {
FunctionArguments::List(ref args) => args,
_ => unreachable!(),
};
assert_eq!(args.args.len(), 2);
assert_eq!(
args.args[0],
FunctionArg::Unnamed(FunctionArgExpr::Expr(Identifier(Ident::from("x"))))
);
assert_eq!(
args.args[1],
FunctionArg::Unnamed(FunctionArgExpr::Expr(Identifier(Ident::from("y"))))
);
let parameters = match f.parameters {
FunctionArguments::List(ref args) => args,
_ => unreachable!(),
};
assert_eq!(parameters.args.len(), 2);
assert_eq!(
parameters.args[0],
FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(Value::Number(
"0.5".parse().unwrap(),
false
))))
);
assert_eq!(
parameters.args[1],
FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(Value::Number(
"0.6".parse().unwrap(),
false
))))
);
}
_ => unreachable!(),
}
}
_ => unreachable!(),
}
}
#[test]
fn parse_select_star_except_no_parens() {
clickhouse().one_statement_parses_to(

View file

@ -1045,6 +1045,7 @@ fn parse_select_count_wildcard() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("COUNT")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Wildcard)],
@ -1066,6 +1067,7 @@ fn parse_select_count_distinct() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("COUNT")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: Some(DuplicateTreatment::Distinct),
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::UnaryOp {
@ -2151,6 +2153,7 @@ fn parse_select_having() {
Some(Expr::BinaryOp {
left: Box::new(Expr::Function(Function {
name: ObjectName(vec![Ident::new("COUNT")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Wildcard)],
@ -2180,6 +2183,7 @@ fn parse_select_qualify() {
Some(Expr::BinaryOp {
left: Box::new(Expr::Function(Function {
name: ObjectName(vec![Ident::new("ROW_NUMBER")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![],
@ -2523,6 +2527,7 @@ fn parse_listagg() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("LISTAGG")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: Some(DuplicateTreatment::Distinct),
args: vec![
@ -4227,6 +4232,7 @@ fn parse_named_argument_function() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("FUN")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![
@ -4265,6 +4271,7 @@ fn parse_named_argument_function_with_eq_operator() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("FUN")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![
@ -4337,6 +4344,7 @@ fn parse_window_functions() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("row_number")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![],
@ -4465,6 +4473,7 @@ fn test_parse_named_window() {
value: "MIN".to_string(),
quote_style: None,
}]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
@ -4494,6 +4503,7 @@ fn test_parse_named_window() {
value: "MAX".to_string(),
quote_style: None,
}]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
@ -8089,6 +8099,7 @@ fn parse_time_functions() {
let select = verified_only_select(&sql);
let select_localtime_func_call_ast = Function {
name: ObjectName(vec![Ident::new(func_name)]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![],
@ -9017,6 +9028,7 @@ fn parse_call() {
assert_eq!(
verified_stmt("CALL my_procedure('a')"),
Statement::Call(Function {
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(
@ -9418,6 +9430,7 @@ fn test_selective_aggregation() {
vec![
SelectItem::UnnamedExpr(Expr::Function(Function {
name: ObjectName(vec![Ident::new("ARRAY_AGG")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
@ -9435,6 +9448,7 @@ fn test_selective_aggregation() {
SelectItem::ExprWithAlias {
expr: Expr::Function(Function {
name: ObjectName(vec![Ident::new("ARRAY_AGG")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(

View file

@ -488,6 +488,7 @@ fn test_duckdb_named_argument_function_with_assignment_operator() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("FUN")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![

View file

@ -381,6 +381,7 @@ fn parse_delimited_identifiers() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::with_quote('"', "myfun")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![],

View file

@ -354,6 +354,7 @@ fn parse_delimited_identifiers() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::with_quote('"', "myfun")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![],

View file

@ -2369,6 +2369,7 @@ fn parse_array_subquery_expr() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("ARRAY")]),
parameters: FunctionArguments::None,
args: FunctionArguments::Subquery(Box::new(Query {
with: None,
body: Box::new(SetExpr::SetOperation {
@ -2729,6 +2730,7 @@ fn test_composite_value() {
Ident::new("information_schema"),
Ident::new("_pg_expandarray")
]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Array(
@ -2955,6 +2957,7 @@ fn parse_current_functions() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("CURRENT_CATALOG")]),
parameters: FunctionArguments::None,
args: FunctionArguments::None,
null_treatment: None,
filter: None,
@ -2966,6 +2969,7 @@ fn parse_current_functions() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("CURRENT_USER")]),
parameters: FunctionArguments::None,
args: FunctionArguments::None,
null_treatment: None,
filter: None,
@ -2977,6 +2981,7 @@ fn parse_current_functions() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("SESSION_USER")]),
parameters: FunctionArguments::None,
args: FunctionArguments::None,
null_treatment: None,
filter: None,
@ -2988,6 +2993,7 @@ fn parse_current_functions() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("USER")]),
parameters: FunctionArguments::None,
args: FunctionArguments::None,
null_treatment: None,
filter: None,
@ -3438,6 +3444,7 @@ fn parse_delimited_identifiers() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::with_quote('"', "myfun")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![],

View file

@ -136,6 +136,7 @@ fn parse_delimited_identifiers() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::with_quote('"', "myfun")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![],

View file

@ -892,6 +892,7 @@ fn parse_delimited_identifiers() {
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::with_quote('"', "myfun")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![],

View file

@ -335,6 +335,7 @@ fn parse_window_function_with_filter() {
select.projection,
vec![SelectItem::UnnamedExpr(Expr::Function(Function {
name: ObjectName(vec![Ident::new(func_name)]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(