Skip to content

Commit

Permalink
feat : added arrow is_between and is_between_test.py (#609)
Browse files Browse the repository at this point in the history
  • Loading branch information
anopsy authored Jul 24, 2024
1 parent 1afdd76 commit c52f585
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 0 deletions.
23 changes: 23 additions & 0 deletions narwhals/_arrow/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,29 @@ def all(self) -> bool:
pc = get_pyarrow_compute()
return pc.all(self._native_series) # type: ignore[no-any-return]

def is_between(self, lower_bound: Any, upper_bound: Any, closed: str = "both") -> Any:
pc = get_pyarrow_compute()
ser = self._native_series
if closed == "left":
ge = pc.greater_equal(ser, lower_bound)
lt = pc.less(ser, upper_bound)
res = pc.and_kleene(ge, lt)
elif closed == "right":
gt = pc.greater(ser, lower_bound)
le = pc.less_equal(ser, upper_bound)
res = pc.and_kleene(gt, le)
elif closed == "none":
gt = pc.greater(ser, lower_bound)
lt = pc.less(ser, upper_bound)
res = pc.and_kleene(gt, lt)
elif closed == "both":
ge = pc.greater_equal(ser, lower_bound)
le = pc.less_equal(ser, upper_bound)
res = pc.and_kleene(ge, le)
else: # pragma: no cover
raise AssertionError
return self._from_native_series(res)

def is_empty(self) -> bool:
return len(self) == 0

Expand Down
25 changes: 25 additions & 0 deletions tests/series_only/is_between_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
from __future__ import annotations

from typing import Any

import pytest

import narwhals.stable.v1 as nw
from tests.utils import compare_dicts

data = [1, 4, 2, 5]


@pytest.mark.parametrize(
("closed", "expected"),
[
("left", [True, True, True, False]),
("right", [False, True, True, True]),
("both", [True, True, True, True]),
("none", [False, True, True, False]),
],
)
def test_is_between(constructor_series: Any, closed: str, expected: list[bool]) -> None:
ser = nw.from_native(constructor_series(data), series_only=True)
result = ser.is_between(1, 5, closed=closed)
compare_dicts({"a": result}, {"a": expected})

0 comments on commit c52f585

Please sign in to comment.