Skip to content

Commit

Permalink
feat: WITHIN GROUP expression support
Browse files Browse the repository at this point in the history
  • Loading branch information
MazterQyou committed Sep 4, 2024
1 parent 6a54d27 commit 4670853
Show file tree
Hide file tree
Showing 3 changed files with 79 additions and 2 deletions.
23 changes: 23 additions & 0 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -348,6 +348,8 @@ pub enum Expr {
ListAgg(ListAgg),
/// The `ARRAY_AGG` function `SELECT ARRAY_AGG(... ORDER BY ...)`
ArrayAgg(ArrayAgg),
/// The `WITHIN GROUP` expr `... WITHIN GROUP (ORDER BY ...)`
WithinGroup(WithinGroup),
/// The `GROUPING SETS` expr.
GroupingSets(Vec<Vec<Expr>>),
/// The `CUBE` expr.
Expand Down Expand Up @@ -549,6 +551,7 @@ impl fmt::Display for Expr {
Expr::ArraySubquery(s) => write!(f, "ARRAY({})", s),
Expr::ListAgg(listagg) => write!(f, "{}", listagg),
Expr::ArrayAgg(arrayagg) => write!(f, "{}", arrayagg),
Expr::WithinGroup(withingroup) => write!(f, "{}", withingroup),
Expr::GroupingSets(sets) => {
write!(f, "GROUPING SETS (")?;
let mut sep = "";
Expand Down Expand Up @@ -2523,6 +2526,26 @@ impl fmt::Display for ArrayAgg {
}
}

/// A `WITHIN GROUP` invocation `<expr> WITHIN GROUP (ORDER BY <sort_expr> )`
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct WithinGroup {
pub expr: Box<Expr>,
pub order_by: Vec<OrderByExpr>,
}

impl fmt::Display for WithinGroup {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"{} WITHIN GROUP (ORDER BY {})",
self.expr,
display_comma_separated(&self.order_by),
)?;
Ok(())
}
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum ObjectType {
Expand Down
23 changes: 21 additions & 2 deletions src/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -627,14 +627,33 @@ impl<'a> Parser<'a> {
None
};

Ok(Expr::Function(Function {
let within_group = if self.parse_keywords(&[Keyword::WITHIN, Keyword::GROUP]) {
self.expect_token(&Token::LParen)?;
self.expect_keywords(&[Keyword::ORDER, Keyword::BY])?;
let order_by_expr = self.parse_comma_separated(Parser::parse_order_by_expr)?;
self.expect_token(&Token::RParen)?;
Some(order_by_expr)
} else {
None
};

let function = Expr::Function(Function {
name,
args,
over,
distinct,
special: false,
approximate: false,
}))
});

Ok(if let Some(within_group) = within_group {
Expr::WithinGroup(WithinGroup {
expr: Box::new(function),
order_by: within_group,
})
} else {
function
})
}

pub fn parse_time_functions(&mut self, name: ObjectName) -> Result<Expr, ParserError> {
Expand Down
35 changes: 35 additions & 0 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1736,6 +1736,41 @@ fn parse_array_agg_func() {
}
}

#[test]
fn parse_within_group() {
let sql = "SELECT PERCENTILE_CONT(0.0) WITHIN GROUP (ORDER BY name ASC NULLS FIRST)";
let select = verified_only_select(sql);

#[cfg(feature = "bigdecimal")]
let value = bigdecimal::BigDecimal::from(0);
#[cfg(not(feature = "bigdecimal"))]
let value = "0.0".to_string();
let expr = Expr::Value(Value::Number(value, false));
let function = Expr::Function(Function {
name: ObjectName(vec![Ident::new("PERCENTILE_CONT")]),
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(expr))],
over: None,
distinct: false,
special: false,
approximate: false,
});
let within_group = vec![OrderByExpr {
expr: Expr::Identifier(Ident {
value: "name".to_string(),
quote_style: None,
}),
asc: Some(true),
nulls_first: Some(true),
}];
assert_eq!(
&Expr::WithinGroup(WithinGroup {
expr: Box::new(function),
order_by: within_group
}),
expr_from_projection(only(&select.projection))
);
}

#[test]
fn parse_create_table() {
let sql = "CREATE TABLE uk_cities (\
Expand Down

0 comments on commit 4670853

Please sign in to comment.