Skip to content

Commit

Permalink
feat: PERCENTILE_CONT function support
Browse files Browse the repository at this point in the history
  • Loading branch information
MazterQyou committed Aug 14, 2024
1 parent 6a54d27 commit f1a97af
Show file tree
Hide file tree
Showing 3 changed files with 60 additions and 0 deletions.
22 changes: 22 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 `PERCENTILE_CONT` function `SELECT PERCENTILE_CONT(...) WITHIN GROUP (ORDER BY ...)`
PercentileCont(PercentileCont),
/// 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::PercentileCont(percentilecont) => write!(f, "{}", percentilecont),
Expr::GroupingSets(sets) => {
write!(f, "GROUPING SETS (")?;
let mut sep = "";
Expand Down Expand Up @@ -2523,6 +2526,25 @@ impl fmt::Display for ArrayAgg {
}
}

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

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

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum ObjectType {
Expand Down
14 changes: 14 additions & 0 deletions src/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -464,6 +464,7 @@ impl<'a> Parser<'a> {
Keyword::TRIM => self.parse_trim_expr(),
Keyword::INTERVAL => self.parse_literal_interval(),
Keyword::LISTAGG => self.parse_listagg_expr(),
Keyword::PERCENTILE_CONT => self.parse_percentile_cont_expr(),
// Treat ARRAY[1,2,3] as an array [1,2,3], otherwise try as function call
Keyword::ARRAY if self.peek_token() == Token::LBracket => {
self.expect_token(&Token::LBracket)?;
Expand Down Expand Up @@ -1056,6 +1057,19 @@ impl<'a> Parser<'a> {
}))
}

/// Parse a SQL PERCENTILE_CONT expression, e.g. `PERCENTILE_CONT(...) WITHIN GROUP (ORDER BY ...)`.
pub fn parse_percentile_cont_expr(&mut self) -> Result<Expr, ParserError> {
self.expect_token(&Token::LParen)?;
let expr = Box::new(self.parse_expr()?);
self.expect_token(&Token::RParen)?;
self.expect_keywords(&[Keyword::WITHIN, Keyword::GROUP])?;
self.expect_token(&Token::LParen)?;
self.expect_keywords(&[Keyword::ORDER, Keyword::BY])?;
let within_group = Box::new(self.parse_order_by_expr()?);
self.expect_token(&Token::RParen)?;
Ok(Expr::PercentileCont(PercentileCont { expr, within_group }))
}

// This function parses date/time fields for both the EXTRACT function-like
// operator and interval qualifiers. EXTRACT supports a wider set of
// date/time fields than interval qualifiers, so this function may need to
Expand Down
24 changes: 24 additions & 0 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1736,6 +1736,30 @@ fn parse_array_agg_func() {
}
}

#[test]
fn parse_percentile_cont() {
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".to_string();
let expr = Box::new(Expr::Value(Value::Number(value, false)));
let within_group = Box::new(OrderByExpr {
expr: Expr::Identifier(Ident {
value: "name".to_string(),
quote_style: None,
}),
asc: Some(true),
nulls_first: Some(true),
});
assert_eq!(
&Expr::PercentileCont(PercentileCont { expr, 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 f1a97af

Please sign in to comment.