Skip to content

Commit 02d4453

Browse files
authored
feat: expose between (#868)
closes #809
1 parent 89b77ab commit 02d4453

File tree

3 files changed

+58
-0
lines changed

3 files changed

+58
-0
lines changed

python/datafusion/expr.py

+16
Original file line numberDiff line numberDiff line change
@@ -393,6 +393,22 @@ def cast(
393393

394394
return Expr(self.expr.cast(to))
395395

396+
def between(self, low: Any, high: Any, negated: bool = False) -> Expr:
397+
"""Returns ``True`` if this expression is between a given range.
398+
399+
Args:
400+
low: lower bound of the range (inclusive).
401+
high: higher bound of the range (inclusive).
402+
negated: negates whether the expression is between a given range
403+
"""
404+
if not isinstance(low, Expr):
405+
low = Expr.literal(low)
406+
407+
if not isinstance(high, Expr):
408+
high = Expr.literal(high)
409+
410+
return Expr(self.expr.between(low.expr, high.expr, negated=negated))
411+
396412
def rex_type(self) -> RexType:
397413
"""Return the Rex Type of this expression.
398414

python/datafusion/tests/test_functions.py

+31
Original file line numberDiff line numberDiff line change
@@ -1024,3 +1024,34 @@ def test_cast(df, python_datatype, name: str, expected):
10241024
result = df.collect()
10251025
result = result[0]
10261026
assert result.column(0) == result.column(1)
1027+
1028+
1029+
@pytest.mark.parametrize(
1030+
"negated, low, high, expected",
1031+
[
1032+
pytest.param(False, 3, 5, {"filtered": [4, 5]}),
1033+
pytest.param(False, 4, 5, {"filtered": [4, 5]}),
1034+
pytest.param(True, 3, 5, {"filtered": [6]}),
1035+
pytest.param(True, 4, 6, []),
1036+
],
1037+
)
1038+
def test_between(df, negated, low, high, expected):
1039+
df = df.filter(column("b").between(low, high, negated=negated)).select(
1040+
column("b").alias("filtered")
1041+
)
1042+
1043+
actual = df.collect()
1044+
1045+
if expected:
1046+
actual = actual[0].to_pydict()
1047+
assert actual == expected
1048+
else:
1049+
assert len(actual) == 0 # the rows are empty
1050+
1051+
1052+
def test_between_default(df):
1053+
df = df.filter(column("b").between(3, 5)).select(column("b").alias("filtered"))
1054+
expected = {"filtered": [4, 5]}
1055+
1056+
actual = df.collect()[0].to_pydict()
1057+
assert actual == expected

src/expr.rs

+11
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,17 @@ impl PyExpr {
293293
expr.into()
294294
}
295295

296+
#[pyo3(signature = (low, high, negated=false))]
297+
pub fn between(&self, low: PyExpr, high: PyExpr, negated: bool) -> PyExpr {
298+
let expr = Expr::Between(Between::new(
299+
Box::new(self.expr.clone()),
300+
negated,
301+
Box::new(low.into()),
302+
Box::new(high.into()),
303+
));
304+
expr.into()
305+
}
306+
296307
/// A Rex (Row Expression) specifies a single row of data. That specification
297308
/// could include user defined functions or types. RexType identifies the row
298309
/// as one of the possible valid `RexTypes`.

0 commit comments

Comments
 (0)