Skip to content

Commit

Permalink
[FEAT]: sql between (#3062)
Browse files Browse the repository at this point in the history
closes #3057
  • Loading branch information
universalmind303 authored Oct 16, 2024
1 parent 629674c commit 271ec7c
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 1 deletion.
17 changes: 16 additions & 1 deletion src/daft-sql/src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -597,7 +597,22 @@ impl SQLPlanner {
unsupported_sql_err!("IN subquery")
}
SQLExpr::InUnnest { .. } => unsupported_sql_err!("IN UNNEST"),
SQLExpr::Between { .. } => unsupported_sql_err!("BETWEEN"),
SQLExpr::Between {
expr,
negated,
low,
high,
} => {
let expr = self.plan_expr(expr)?;
let low = self.plan_expr(low)?;
let high = self.plan_expr(high)?;
let expr = expr.between(low, high);
if *negated {
Ok(expr.not())
} else {
Ok(expr)
}
}
SQLExpr::Like {
negated,
expr,
Expand Down
13 changes: 13 additions & 0 deletions tests/sql/test_exprs.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,3 +67,16 @@ def test_hash_exprs():

with pytest.raises(Exception, match="num_hashes is required"):
daft.sql("SELECT minhash(a) as hash_a FROM df").collect()


def test_between():
df = daft.from_pydict(
{
"integers": [0, 1, 2, 3, 5],
}
)

actual = daft.sql("SELECT * FROM df where integers between 1 and 4").collect().to_pydict()

expected = df.filter(col("integers").between(1, 4)).collect().to_pydict()
assert actual == expected

0 comments on commit 271ec7c

Please sign in to comment.