Skip to content

Commit

Permalink
docs(python): add example of using "is_between" with string bounds, a…
Browse files Browse the repository at this point in the history
…nd extend test coverage for the same
  • Loading branch information
alexander-beedie committed Feb 2, 2023
1 parent dea9bbd commit 18731be
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 8 deletions.
33 changes: 27 additions & 6 deletions py-polars/polars/internals/expr/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -3390,19 +3390,19 @@ def repeat_by(self, by: Expr | str) -> Expr:

def is_between(
self,
start: Expr | datetime | date | time | int | float,
end: Expr | datetime | date | time | int | float,
start: Expr | datetime | date | time | int | float | str,
end: Expr | datetime | date | time | int | float | str,
closed: ClosedInterval = "both",
) -> Expr:
"""
Check if this expression is between start and end.
Check if this expression is between the given start and end values.
Parameters
----------
start
Lower bound as primitive type or datetime.
Lower bound value (can be an expression or literal).
end
Upper bound as primitive type or datetime.
Upper bound value (can be an expression or literal).
closed : {'both', 'left', 'right', 'none'}
Define which sides of the interval are closed (inclusive).
Expand All @@ -3427,7 +3427,7 @@ def is_between(
│ 5 ┆ false │
└─────┴────────────┘
Use the ``closed`` argument to include or exclude the values at the bounds.
Use the ``closed`` argument to include or exclude the values at the bounds:
>>> df.with_columns(pl.col("num").is_between(2, 4, closed="left"))
shape: (5, 2)
Expand All @@ -3443,9 +3443,30 @@ def is_between(
│ 5 ┆ false │
└─────┴────────────┘
Can also use strings as well as numeric/temporal values (note: ensure that
string literals are wrapped with ``lit`` so as not to conflate them with
column names):
>>> df = pl.DataFrame({"a": ["a", "b", "c", "d", "e"]})
>>> df.with_columns(
... pl.col("a").is_between(pl.lit("a"), pl.lit("c"), closed="both")
... )
shape: (5, 2)
┌─────┬────────────┐
│ a ┆ is_between │
│ --- ┆ --- │
│ str ┆ bool │
╞═════╪════════════╡
│ a ┆ true │
│ b ┆ true │
│ c ┆ true │
│ d ┆ false │
│ e ┆ false │
└─────┴────────────┘
"""
start = expr_to_lit_or_expr(start, str_to_lit=False)
end = expr_to_lit_or_expr(end, str_to_lit=False)

if closed == "none":
return ((self > start) & (self < end)).alias("is_between")
elif closed == "both":
Expand Down
19 changes: 17 additions & 2 deletions py-polars/tests/unit/test_lazy.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

import os
from datetime import date, datetime
from datetime import date, datetime, time
from functools import reduce
from operator import add
from string import ascii_letters
Expand Down Expand Up @@ -1218,12 +1218,14 @@ def test_is_between_data_types() -> None:
{
"flt": [1.4, 1.2, 2.5],
"int": [2, 3, 4],
"str": ["xyz", "str", "abc"],
"date": [date(2020, 1, 1), date(2020, 2, 2), date(2020, 3, 3)],
"datetime": [
datetime(2020, 1, 1, 0, 0, 0),
datetime(2020, 1, 1, 10, 0, 0),
datetime(2020, 1, 1, 12, 0, 0),
],
"tm": [time(10, 30), time(0, 45), time(15, 15)],
}
)

Expand All @@ -1236,7 +1238,6 @@ def test_is_between_data_types() -> None:
df.select(pl.col("int").is_between(1.5, 3))[:, 0],
pl.Series("is_between", [True, True, False]),
)

assert_series_equal(
df.select(pl.col("date").is_between(date(2019, 1, 1), date(2020, 2, 5)))[:, 0],
pl.Series("is_between", [True, True, False]),
Expand All @@ -1249,6 +1250,20 @@ def test_is_between_data_types() -> None:
)[:, 0],
pl.Series("is_between", [False, True, False]),
)
assert_series_equal(
df.select(
pl.col("str").is_between(pl.lit("str"), pl.lit("zzz"), closed="left")
)[:, 0],
pl.Series("is_between", [True, True, False]),
)
assert_series_equal(
df.select(
pl.col("tm")
.is_between(time(0, 45), time(10, 30), closed="right")
.alias("tm_between")
)[:, 0],
pl.Series("tm_between", [True, False, False]),
)


def test_unique() -> None:
Expand Down

0 comments on commit 18731be

Please sign in to comment.