Skip to content

Commit

Permalink
fix(python): Preserve Expr name in is_between (#6661)
Browse files Browse the repository at this point in the history
  • Loading branch information
@stinodego
stinodego authored Feb 3, 2023
1 parent eb88988 commit 0982060
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 43 deletions.
20 changes: 12 additions & 8 deletions py-polars/polars/internals/expr/expr.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3414,7 +3414,7 @@ def is_between(
Examples
--------
>>> df = pl.DataFrame({"num": [1, 2, 3, 4, 5]})
>>> df.with_columns(pl.col("num").is_between(2, 4))
>>> df.with_columns(pl.col("num").is_between(2, 4).alias("is_between"))
shape: (5, 2)
┌─────┬────────────┐
│ num ┆ is_between │
Expand All @@ -3430,7 +3430,9 @@ def is_between(
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"))
>>> df.with_columns(
... pl.col("num").is_between(2, 4, closed="left").alias("is_between")
... )
shape: (5, 2)
┌─────┬────────────┐
│ num ┆ is_between │
Expand All @@ -3444,13 +3446,15 @@ def is_between(
│ 5 ┆ false │
└─────┴────────────┘
Can also use strings as well as numeric/temporal values (note: ensure that
You 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")
... pl.col("a")
... .is_between(pl.lit("a"), pl.lit("c"), closed="both")
... .alias("is_between")
... )
shape: (5, 2)
┌─────┬────────────┐
Expand All @@ -3469,13 +3473,13 @@ def is_between(
end = expr_to_lit_or_expr(end, str_to_lit=False)

if closed == "none":
return ((self > start) & (self < end)).alias("is_between")
return (self > start) & (self < end)
elif closed == "both":
return ((self >= start) & (self <= end)).alias("is_between")
return (self >= start) & (self <= end)
elif closed == "right":
return ((self > start) & (self <= end)).alias("is_between")
return (self > start) & (self <= end)
elif closed == "left":
return ((self >= start) & (self < end)).alias("is_between")
return (self >= start) & (self < end)
else:
raise ValueError(
"closed must be one of {'left', 'right', 'both', 'none'},"
Expand Down
50 changes: 17 additions & 33 deletions py-polars/tests/unit/test_lazy.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1178,38 +1178,22 @@ def test_quantile(fruits_cars: pl.DataFrame) -> None:


def test_is_between(fruits_cars: pl.DataFrame) -> None:
result = fruits_cars.select(pl.col("A").is_between(2, 4))["is_between"]
assert_series_equal(
result, pl.Series("is_between", [False, True, True, True, False])
)
result = fruits_cars.select(pl.col("A").is_between(2, 4)).to_series()
assert_series_equal(result, pl.Series("A", [False, True, True, True, False]))

result = fruits_cars.select(pl.col("A").is_between(2, 4, closed="none"))[
"is_between"
]
assert_series_equal(
result, pl.Series("is_between", [False, False, True, False, False])
)
result = fruits_cars.select(pl.col("A").is_between(2, 4, closed="none")).to_series()
assert_series_equal(result, pl.Series("A", [False, False, True, False, False]))

result = fruits_cars.select(pl.col("A").is_between(2, 4, closed="both"))[
"is_between"
]
assert_series_equal(
result, pl.Series("is_between", [False, True, True, True, False])
)
result = fruits_cars.select(pl.col("A").is_between(2, 4, closed="both")).to_series()
assert_series_equal(result, pl.Series("A", [False, True, True, True, False]))

result = fruits_cars.select(pl.col("A").is_between(2, 4, closed="right"))[
"is_between"
]
assert_series_equal(
result, pl.Series("is_between", [False, False, True, True, False])
)
result = fruits_cars.select(
pl.col("A").is_between(2, 4, closed="right")
).to_series()
assert_series_equal(result, pl.Series("A", [False, False, True, True, False]))

result = fruits_cars.select(pl.col("A").is_between(2, 4, closed="left"))[
"is_between"
]
assert_series_equal(
result, pl.Series("is_between", [False, True, True, False, False])
)
result = fruits_cars.select(pl.col("A").is_between(2, 4, closed="left")).to_series()
assert_series_equal(result, pl.Series("A", [False, True, True, False, False]))


def test_is_between_data_types() -> None:
Expand All @@ -1231,29 +1215,29 @@ def test_is_between_data_types() -> None:
# on purpose, for float and int, we pass in a mixture of bound data types
assert_series_equal(
df.select(pl.col("flt").is_between(1, 2.3))[:, 0],
pl.Series("is_between", [True, True, False]),
pl.Series("flt", [True, True, False]),
)
assert_series_equal(
df.select(pl.col("int").is_between(1.5, 3))[:, 0],
pl.Series("is_between", [True, True, False]),
pl.Series("int", [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]),
pl.Series("date", [True, True, False]),
)
assert_series_equal(
df.select(
pl.col("datetime").is_between(
datetime(2020, 1, 1, 5, 0, 0), datetime(2020, 1, 1, 11, 0, 0)
)
)[:, 0],
pl.Series("is_between", [False, True, False]),
pl.Series("datetime", [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]),
pl.Series("str", [True, True, False]),
)
assert_series_equal(
df.select(
Expand Down
4 changes: 2 additions & 2 deletions py-polars/tests/unit/test_series.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2053,8 +2053,8 @@ def test_is_between_datetime() -> None:
expected = pl.Series("a", [False, True])

# only on the expression api
result = s.to_frame().with_columns(pl.col("*").is_between(start, end))["is_between"]
assert_series_equal(result.rename("a"), expected)
result = s.to_frame().with_columns(pl.col("*").is_between(start, end)).to_series()
assert_series_equal(result, expected)


@pytest.mark.parametrize(
Expand Down

0 comments on commit 0982060

Please sign in to comment.