Skip to content

Commit

Permalink
[FEAT] Filter predicates in SQL join (#3371)
Browse files Browse the repository at this point in the history
Adds support for things like:
```sql
SELECT * FROM a JOIN b ON a.x = b.x AND a.y > 0
```

Enables TPC-H q13
  • Loading branch information
kevinzwang authored Nov 21, 2024
1 parent f6eb993 commit 2c0f3cd
Show file tree
Hide file tree
Showing 2 changed files with 193 additions and 125 deletions.
30 changes: 30 additions & 0 deletions src/daft-sql/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ mod tests {
Schema::new(vec![
Field::new("text", DataType::Utf8),
Field::new("id", DataType::Int32),
Field::new("val", DataType::Int32),
])
.unwrap(),
);
Expand Down Expand Up @@ -138,6 +139,7 @@ mod tests {
#[case::slice("select list_utf8[0:2] from tbl1")]
#[case::join("select * from tbl2 join tbl3 on tbl2.id = tbl3.id")]
#[case::null_safe_join("select * from tbl2 left join tbl3 on tbl2.id <=> tbl3.id")]
#[case::join_with_filter("select * from tbl2 join tbl3 on tbl2.id = tbl3.id and tbl2.val > 0")]
#[case::from("select tbl2.text from tbl2")]
#[case::using("select tbl2.text from tbl2 join tbl3 using (id)")]
#[case(
Expand Down Expand Up @@ -301,6 +303,34 @@ mod tests {
Ok(())
}

#[rstest]
fn test_join_with_filter(
mut planner: SQLPlanner,
tbl_2: LogicalPlanRef,
tbl_3: LogicalPlanRef,
) -> SQLPlannerResult<()> {
let sql = "select * from tbl2 join tbl3 on tbl2.id = tbl3.id and tbl2.val > 0";
let plan = planner.plan_sql(&sql)?;

let expected = LogicalPlanBuilder::new(tbl_2, None)
.filter(col("val").gt(lit(0 as i64)))?
.join_with_null_safe_equal(
tbl_3,
vec![col("id")],
vec![col("id")],
Some(vec![false]),
JoinType::Inner,
None,
None,
Some("tbl3."),
true,
)?
.select(vec![col("*")])?
.build();
assert_eq!(plan, expected);
Ok(())
}

#[rstest]
#[case::abs("select abs(i32) as abs from tbl1")]
#[case::ceil("select ceil(i32) as ceil from tbl1")]
Expand Down
Loading

0 comments on commit 2c0f3cd

Please sign in to comment.