Skip to content

Commit

Permalink
feat: expose between (apache#868)
Browse files Browse the repository at this point in the history
closes apache#809
  • Loading branch information
mesejo authored Sep 11, 2024
1 parent 89b77ab commit 02d4453
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 0 deletions.
16 changes: 16 additions & 0 deletions python/datafusion/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -393,6 +393,22 @@ def cast(

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

def between(self, low: Any, high: Any, negated: bool = False) -> Expr:
"""Returns ``True`` if this expression is between a given range.
Args:
low: lower bound of the range (inclusive).
high: higher bound of the range (inclusive).
negated: negates whether the expression is between a given range
"""
if not isinstance(low, Expr):
low = Expr.literal(low)

if not isinstance(high, Expr):
high = Expr.literal(high)

return Expr(self.expr.between(low.expr, high.expr, negated=negated))

def rex_type(self) -> RexType:
"""Return the Rex Type of this expression.
Expand Down
31 changes: 31 additions & 0 deletions python/datafusion/tests/test_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -1024,3 +1024,34 @@ def test_cast(df, python_datatype, name: str, expected):
result = df.collect()
result = result[0]
assert result.column(0) == result.column(1)


@pytest.mark.parametrize(
"negated, low, high, expected",
[
pytest.param(False, 3, 5, {"filtered": [4, 5]}),
pytest.param(False, 4, 5, {"filtered": [4, 5]}),
pytest.param(True, 3, 5, {"filtered": [6]}),
pytest.param(True, 4, 6, []),
],
)
def test_between(df, negated, low, high, expected):
df = df.filter(column("b").between(low, high, negated=negated)).select(
column("b").alias("filtered")
)

actual = df.collect()

if expected:
actual = actual[0].to_pydict()
assert actual == expected
else:
assert len(actual) == 0 # the rows are empty


def test_between_default(df):
df = df.filter(column("b").between(3, 5)).select(column("b").alias("filtered"))
expected = {"filtered": [4, 5]}

actual = df.collect()[0].to_pydict()
assert actual == expected
11 changes: 11 additions & 0 deletions src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -293,6 +293,17 @@ impl PyExpr {
expr.into()
}

#[pyo3(signature = (low, high, negated=false))]
pub fn between(&self, low: PyExpr, high: PyExpr, negated: bool) -> PyExpr {
let expr = Expr::Between(Between::new(
Box::new(self.expr.clone()),
negated,
Box::new(low.into()),
Box::new(high.into()),
));
expr.into()
}

/// A Rex (Row Expression) specifies a single row of data. That specification
/// could include user defined functions or types. RexType identifies the row
/// as one of the possible valid `RexTypes`.
Expand Down

0 comments on commit 02d4453

Please sign in to comment.