Support for Snowflake ASOF joins (apache#1288)

Author: jmhain

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 <https://docs.snowflake.com/en/sql-reference/constructs/asof-join>.
+    AsOf {
+        match_condition: Expr,
+        constraint: JoinConstraint,
+    },
 }
 
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]

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,

src/parser/mod.rs

@@ -3196,6 +3196,16 @@ impl<'a> Parser<'a> {
         Ok(values)
     }
 
+    pub fn parse_parenthesized<T, F>(&mut self, mut f: F) -> Result<T, ParserError>
+    where
+        F: FnMut(&mut Parser<'a>) -> Result<T, ParserError>,
+    {
+        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<T, F>(&mut self, f: F) -> Result<Vec<T>, 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)?;
 

src/test_utils.rs

@@ -312,6 +312,20 @@ pub fn table(name: impl Into<String>) -> TableFactor {
     }
 }
 
+pub fn table_with_alias(name: impl Into<String>, alias: impl Into<String>) -> 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,

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",
+    ));
+}