From 80c03f5c6a42c2414d93dc411c013155dea676d4 Mon Sep 17 00:00:00 2001 From: Joey Hain Date: Fri, 31 May 2024 02:45:07 -0700 Subject: [PATCH] Support for Snowflake ASOF joins (#1288) --- src/ast/query.rs | 17 +++++++++ src/keywords.rs | 2 ++ src/parser/mod.rs | 26 ++++++++++++-- src/test_utils.rs | 14 ++++++++ tests/sqlparser_snowflake.rs | 67 ++++++++++++++++++++++++++++++++++++ 5 files changed, 123 insertions(+), 3 deletions(-) diff --git a/src/ast/query.rs b/src/ast/query.rs index c0fa738cd..fcd5b970d 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -1562,6 +1562,15 @@ impl fmt::Display for Join { ), JoinOperator::CrossApply => write!(f, " CROSS APPLY {}", self.relation), JoinOperator::OuterApply => write!(f, " OUTER APPLY {}", self.relation), + JoinOperator::AsOf { + match_condition, + constraint, + } => write!( + f, + " ASOF JOIN {} MATCH_CONDITION ({match_condition}){}", + self.relation, + suffix(constraint) + ), } } } @@ -1587,6 +1596,14 @@ pub enum JoinOperator { CrossApply, /// OUTER APPLY (non-standard) OuterApply, + /// `ASOF` joins are used for joining tables containing time-series data + /// whose timestamp columns do not match exactly. + /// + /// See . + AsOf { + match_condition: Expr, + constraint: JoinConstraint, + }, } #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] diff --git a/src/keywords.rs b/src/keywords.rs index 06086297c..6c6c642c3 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -91,6 +91,7 @@ define_keywords!( AS, ASC, ASENSITIVE, + ASOF, ASSERT, ASYMMETRIC, AT, @@ -418,6 +419,7 @@ define_keywords!( MATCH, MATCHED, MATCHES, + MATCH_CONDITION, MATCH_RECOGNIZE, MATERIALIZED, MAX, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index fef307106..123af045a 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -3196,6 +3196,16 @@ impl<'a> Parser<'a> { Ok(values) } + pub fn parse_parenthesized(&mut self, mut f: F) -> Result + where + F: FnMut(&mut Parser<'a>) -> Result, + { + self.expect_token(&Token::LParen)?; + let res = f(self)?; + self.expect_token(&Token::RParen)?; + Ok(res) + } + /// Parse a comma-separated list of 0+ items accepted by `F` pub fn parse_comma_separated0(&mut self, f: F) -> Result, ParserError> where @@ -8505,6 +8515,18 @@ impl<'a> Parser<'a> { relation: self.parse_table_factor()?, join_operator: JoinOperator::OuterApply, } + } else if self.parse_keyword(Keyword::ASOF) { + self.expect_keyword(Keyword::JOIN)?; + let relation = self.parse_table_factor()?; + self.expect_keyword(Keyword::MATCH_CONDITION)?; + let match_condition = self.parse_parenthesized(Self::parse_expr)?; + Join { + relation, + join_operator: JoinOperator::AsOf { + match_condition, + constraint: self.parse_join_constraint(false)?, + }, + } } else { let natural = self.parse_keyword(Keyword::NATURAL); let peek_keyword = if let Token::Word(w) = self.peek_token().token { @@ -8951,9 +8973,7 @@ impl<'a> Parser<'a> { }; self.expect_keyword(Keyword::PATTERN)?; - self.expect_token(&Token::LParen)?; - let pattern = self.parse_pattern()?; - self.expect_token(&Token::RParen)?; + let pattern = self.parse_parenthesized(Self::parse_pattern)?; self.expect_keyword(Keyword::DEFINE)?; diff --git a/src/test_utils.rs b/src/test_utils.rs index 464366ae4..9af9c8098 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -312,6 +312,20 @@ pub fn table(name: impl Into) -> TableFactor { } } +pub fn table_with_alias(name: impl Into, alias: impl Into) -> TableFactor { + TableFactor::Table { + name: ObjectName(vec![Ident::new(name)]), + alias: Some(TableAlias { + name: Ident::new(alias), + columns: vec![], + }), + args: None, + with_hints: vec![], + version: None, + partitions: vec![], + } +} + pub fn join(relation: TableFactor) -> Join { Join { relation, diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index c11f60993..7492802c7 100644 --- a/tests/sqlparser_snowflake.rs +++ b/tests/sqlparser_snowflake.rs @@ -1688,3 +1688,70 @@ fn test_pivot() { "ORDER BY region", )); } + +#[test] +fn asof_joins() { + #[rustfmt::skip] + let query = snowflake_and_generic().verified_only_select(concat!( + "SELECT * ", + "FROM trades_unixtime AS tu ", + "ASOF JOIN quotes_unixtime AS qu ", + "MATCH_CONDITION (tu.trade_time >= qu.quote_time)", + )); + + assert_eq!( + query.from[0], + TableWithJoins { + relation: table_with_alias("trades_unixtime", "tu"), + joins: vec![Join { + relation: table_with_alias("quotes_unixtime", "qu"), + join_operator: JoinOperator::AsOf { + match_condition: Expr::BinaryOp { + left: Box::new(Expr::CompoundIdentifier(vec![ + Ident::new("tu"), + Ident::new("trade_time"), + ])), + op: BinaryOperator::GtEq, + right: Box::new(Expr::CompoundIdentifier(vec![ + Ident::new("qu"), + Ident::new("quote_time"), + ])), + }, + constraint: JoinConstraint::None, + }, + }], + } + ); + + #[rustfmt::skip] + snowflake_and_generic().verified_query(concat!( + "SELECT t.stock_symbol, t.trade_time, t.quantity, q.quote_time, q.price ", + "FROM trades AS t ASOF JOIN quotes AS q ", + "MATCH_CONDITION (t.trade_time >= quote_time) ", + "ON t.stock_symbol = q.stock_symbol ", + "ORDER BY t.stock_symbol", + )); + + #[rustfmt::skip] + snowflake_and_generic().verified_query(concat!( + "SELECT t.stock_symbol, c.company_name, t.trade_time, t.quantity, q.quote_time, q.price ", + "FROM trades AS t ASOF JOIN quotes AS q ", + "MATCH_CONDITION (t.trade_time <= quote_time) ", + "USING(stock_symbol) ", + "JOIN companies AS c ON c.stock_symbol = t.stock_symbol ", + "ORDER BY t.stock_symbol", + )); + + #[rustfmt::skip] + snowflake_and_generic().verified_query(concat!( + "SELECT * ", + "FROM snowtime AS s ", + "ASOF JOIN raintime AS r ", + "MATCH_CONDITION (s.observed >= r.observed) ", + "ON s.state = r.state ", + "ASOF JOIN preciptime AS p ", + "MATCH_CONDITION (s.observed >= p.observed) ", + "ON s.state = p.state ", + "ORDER BY s.observed", + )); +}