mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-12-23 11:12:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Databricks: support for lambda functions (#1257)

Browse source
This commit is contained in:
Joey Hain 2024-05-06 10:37:48 -07:00 • committed by GitHub
parent a86c58b1c9
commit d9d69a2192
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 210 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

41
src/parser/mod.rs
View file

@ -1018,7 +1018,19 @@ impl<'a> Parser<'a> {
Keyword::CAST => self.parse_cast_expr(CastKind::Cast),
Keyword::TRY_CAST => self.parse_cast_expr(CastKind::TryCast),
Keyword::SAFE_CAST => self.parse_cast_expr(CastKind::SafeCast),
Keyword::EXISTS => self.parse_exists_expr(false),
Keyword::EXISTS
// Support parsing Databricks has a function named `exists`.
if !dialect_of!(self is DatabricksDialect)
|| matches!(
self.peek_nth_token(1).token,
Token::Word(Word {
keyword: Keyword::SELECT | Keyword::WITH,
..
})
) =>
{
self.parse_exists_expr(false)
}
Keyword::EXTRACT => self.parse_extract_expr(),
Keyword::CEIL => self.parse_ceil_floor_expr(true),
Keyword::FLOOR => self.parse_ceil_floor_expr(false),
@ -1036,7 +1048,7 @@ impl<'a> Parser<'a> {
}
Keyword::ARRAY
if self.peek_token() == Token::LParen
&& !dialect_of!(self is ClickHouseDialect) =>
&& !dialect_of!(self is ClickHouseDialect | DatabricksDialect) =>
{
self.expect_token(&Token::LParen)?;
let query = self.parse_boxed_query()?;
@ -1124,6 +1136,13 @@ impl<'a> Parser<'a> {
value: self.parse_introduced_string_value()?,
})
}
Token::Arrow if self.dialect.supports_lambda_functions() => {
self.expect_token(&Token::Arrow)?;
return Ok(Expr::Lambda(LambdaFunction {
params: OneOrManyWithParens::One(w.to_ident()),
body: Box::new(self.parse_expr()?),
}));
}
_ => Ok(Expr::Identifier(w.to_ident())),
},
}, // End of Token::Word
@ -1182,6 +1201,8 @@ impl<'a> Parser<'a> {
if self.parse_keyword(Keyword::SELECT) || self.parse_keyword(Keyword::WITH) {
self.prev_token();
Expr::Subquery(self.parse_boxed_query()?)
} else if let Some(lambda) = self.try_parse_lambda() {
return Ok(lambda);
} else {
let exprs = self.parse_comma_separated(Parser::parse_expr)?;
match exprs.len() {
@ -1231,6 +1252,22 @@ impl<'a> Parser<'a> {
}
}
fn try_parse_lambda(&mut self) -> Option<Expr> {
if !self.dialect.supports_lambda_functions() {
return None;
}
self.maybe_parse(|p| {
let params = p.parse_comma_separated(|p| p.parse_identifier(false))?;
p.expect_token(&Token::RParen)?;
p.expect_token(&Token::Arrow)?;
let expr = p.parse_expr()?;
Ok(Expr::Lambda(LambdaFunction {
params: OneOrManyWithParens::Many(params),
body: Box::new(expr),
}))
})
}
pub fn parse_function(&mut self, name: ObjectName) -> Result<Expr, ParserError> {
self.expect_token(&Token::LParen)?;