Skip to content

Commit

Permalink
test: Add tests for floor/ceil on integers (#20479)
Browse files Browse the repository at this point in the history
  • Loading branch information
lukemanley authored Dec 27, 2024
1 parent 6a7adce commit 208d530
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 4 deletions.
18 changes: 16 additions & 2 deletions py-polars/tests/unit/dataframe/test_df.py
Original file line number Diff line number Diff line change
Expand Up @@ -2790,10 +2790,24 @@ def test_deadlocks_3409() -> None:
) == {"col1": [0, 0, 0]}


def test_ceil() -> None:
df = pl.DataFrame({"a": [1.8, 1.2, 3.0]})
result = df.select(pl.col("a").ceil())
assert_frame_equal(result, pl.DataFrame({"a": [2.0, 2.0, 3.0]}))

df = pl.DataFrame({"a": [1, 2, 3]})
result = df.select(pl.col("a").ceil())
assert_frame_equal(df, result)


def test_floor() -> None:
df = pl.DataFrame({"a": [1.8, 1.2, 3.0]})
col_a_floor = df.select(pl.col("a").floor())["a"]
assert_series_equal(col_a_floor, pl.Series("a", [1, 1, 3]).cast(pl.Float64))
result = df.select(pl.col("a").floor())
assert_frame_equal(result, pl.DataFrame({"a": [1.0, 1.0, 3.0]}))

df = pl.DataFrame({"a": [1, 2, 3]})
result = df.select(pl.col("a").floor())
assert_frame_equal(df, result)


def test_floor_divide() -> None:
Expand Down
19 changes: 17 additions & 2 deletions py-polars/tests/unit/lazyframe/test_lazyframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -504,9 +504,24 @@ def test_cum_agg() -> None:
)


def test_ceil() -> None:
ldf = pl.LazyFrame({"a": [1.8, 1.2, 3.0]})
result = ldf.select(pl.col("a").ceil()).collect()
assert_frame_equal(result, pl.DataFrame({"a": [2.0, 2.0, 3.0]}))

ldf = pl.LazyFrame({"a": [1, 2, 3]})
result = ldf.select(pl.col("a").ceil()).collect()
assert_frame_equal(ldf.collect(), result)


def test_floor() -> None:
ldf = pl.LazyFrame({"a": [1.8, 1.2, 3.0]}).select(pl.col("a").floor())
assert_series_equal(ldf.collect()["a"], pl.Series("a", [1, 1, 3]).cast(pl.Float64))
ldf = pl.LazyFrame({"a": [1.8, 1.2, 3.0]})
result = ldf.select(pl.col("a").floor()).collect()
assert_frame_equal(result, pl.DataFrame({"a": [1.0, 1.0, 3.0]}))

ldf = pl.LazyFrame({"a": [1, 2, 3]})
result = ldf.select(pl.col("a").floor()).collect()
assert_frame_equal(ldf.collect(), result)


@pytest.mark.parametrize(
Expand Down

0 comments on commit 208d530

Please sign in to comment.