Skip to content

Commit

Permalink
Support for Snowflake ASOF joins (#1288)
Browse files Browse the repository at this point in the history
  • Loading branch information
jmhain committed May 31, 2024
1 parent 375742d commit 80c03f5
Show file tree
Hide file tree
Showing 5 changed files with 123 additions and 3 deletions.
17 changes: 17 additions & 0 deletions src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
),
}
}
}
Expand All @@ -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)]
Expand Down
2 changes: 2 additions & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ define_keywords!(
AS,
ASC,
ASENSITIVE,
ASOF,
ASSERT,
ASYMMETRIC,
AT,
Expand Down Expand Up @@ -418,6 +419,7 @@ define_keywords!(
MATCH,
MATCHED,
MATCHES,
MATCH_CONDITION,
MATCH_RECOGNIZE,
MATERIALIZED,
MAX,
Expand Down
26 changes: 23 additions & 3 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)?;

Expand Down
14 changes: 14 additions & 0 deletions src/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
67 changes: 67 additions & 0 deletions tests/sqlparser_snowflake.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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",
));
}

0 comments on commit 80c03f5

Please sign in to comment.