def test_compare_frame_equal_nans() -> None:
# NaN values do not _compare_ equal, but should _assert_ as equal here
nan = float("NaN")
df1 = pl.DataFrame(
data={
"x": [1.0, nan],
"y": [nan, 2.0]
},
columns=[("x", pl.Float32), ("y", pl.Float64)],
)
assert_frame_equal(df1, df1, check_exact=True)
df2 = pl.DataFrame(
data={
"x": [1.0, nan],
"y": [None, 2.0]
},
columns=[("x", pl.Float32), ("y", pl.Float64)],
)
with pytest.raises(AssertionError):
assert_frame_equal(df1, df2, check_exact=True)
def test_slice() -> None:
ldf = pl.DataFrame({"a": [1, 2, 3, 4], "b": ["a", "b", "c", "d"]}).lazy()
expected = pl.DataFrame({"a": [3, 4], "b": ["c", "d"]}).lazy()
for slice_params in (
[2, 10], # slice > len(df)
[2, 4], # slice == len(df)
[2], # optional len
):
assert_frame_equal(ldf.slice(*slice_params), expected)
for py_slice in (
slice(1, 2),
slice(0, 3, 2),
slice(-3, None),
slice(None, 2, 2),
slice(3, None, -1),
slice(1, None, -2),
):
# confirm frame slice matches python slice
assert ldf[py_slice].collect().rows() == ldf.collect().rows()[py_slice]
assert_frame_equal(ldf[::-1], ldf.reverse())
assert_frame_equal(ldf[::-2], ldf.reverse().take_every(2))
def test_repr(df: pl.DataFrame) -> None:
assert isinstance(repr(df), str)
assert_frame_equal(df, df, check_exact=True)
def test_take_every() -> None:
df = pl.DataFrame({"a": [1, 2, 3, 4], "b": ["w", "x", "y", "z"]}).lazy()
expected_df = pl.DataFrame({"a": [1, 3], "b": ["w", "y"]})
assert_frame_equal(expected_df, df.take_every(2).collect())
def test_assert_frame_equal_column_mismatch_order() -> None:
df1 = pl.DataFrame({"b": [3, 4], "a": [1, 2]})
df2 = pl.DataFrame({"a": [1, 2], "b": [3, 4]})
with pytest.raises(AssertionError):
assert_frame_equal(df1, df2)
assert_frame_equal(df1, df2, check_column_names=False)
def test_assert_frame_equal_column_mismatch2() -> None:
df1 = pl.DataFrame({"a": [1, 2]})
df2 = pl.DataFrame({"a": [1, 2], "b": [3, 4]})
with pytest.raises(AssertionError):
assert_frame_equal(df1, df2)
def test_assert_frame_equal_length_mismatch() -> None:
df1 = pl.DataFrame({"a": [1, 2]})
df2 = pl.DataFrame({"a": [1, 2, 3]})
with pytest.raises(AssertionError):
assert_frame_equal(df1, df2)
def test_assert_frame_equal_types() -> None:
df1 = pl.DataFrame({"a": [1, 2]})
srs1 = pl.Series(values=[1, 2], name="a")
with pytest.raises(AssertionError):
assert_frame_equal(df1, srs1) # type: ignore[arg-type]
def test_assert_frame_equal_pass() -> None:
df1 = pl.DataFrame({"a": [1, 2]})
df2 = pl.DataFrame({"a": [1, 2]})
assert_frame_equal(df1, df2)