def test_selection():
df = DataFrame({
"a": [1, 2, 3],
"b": [1.0, 2.0, 3.0],
"c": ["a", "b", "c"]
})
assert df["a"].dtype == Int64
assert df["b"].dtype == Float64
assert df["c"].dtype == Utf8
assert df[["a", "b"]].columns == ["a", "b"]
assert df[[True, False, True]].height == 2
assert df[[True, False, True], "b"].shape == (2, 1)
assert df[[True, False, False], ["a", "b"]].shape == (1, 2)
assert df[[0, 1], "b"].shape == (2, 1)
assert df[[2], ["a", "b"]].shape == (1, 2)
assert df.select_at_idx(0).name == "a"
assert (df.a == df["a"]).sum() == 3
assert (df.c == df["a"]).sum() == 0
assert df[:, "a":"b"].shape == (3, 2)
assert df[:, "a":"c"].columns == ["a", "b", "c"]
expect = pl.DataFrame({"c": ["b"]})
assert df[1, [2]].frame_equal(expect)
expect = pl.DataFrame({"b": [1.0, 3.0]})
assert df[[0, 2], [1]].frame_equal(expect)
def test_null_count(df: pl.DataFrame) -> None:
# note: the zero-row and zero-col cases are always passed as explicit examples
null_count, ncols = df.null_count(), len(df.columns)
if ncols == 0:
assert null_count.shape == (0, 0)
else:
assert null_count.shape == (1, ncols)
for idx, count in enumerate(null_count.rows()[0]):
assert count == sum(v is None
for v in df.select_at_idx(idx).to_list())
print(null_count.rows())
def test_selection():
df = DataFrame({"a": [1, 2, 3], "b": [1.0, 2.0, 3.0], "c": ["a", "b", "c"]})
# column selection by string(s) in first dimension
assert df["a"] == [1, 2, 3]
assert df["b"] == [1.0, 2.0, 3.0]
assert df["c"] == ["a", "b", "c"]
# row selection by integers(s) in first dimension
assert df[0].frame_equal(pl.DataFrame({"a": [1], "b": [1.0], "c": ["a"]}))
assert df[-1].frame_equal(pl.DataFrame({"a": [3], "b": [3.0], "c": ["c"]}))
# row, column selection when using two dimensions
assert df[:, 0] == [1, 2, 3]
assert df[:, 1] == [1.0, 2.0, 3.0]
assert df[:2, 2] == ["a", "b"]
assert df[[1, 2]].frame_equal(
pl.DataFrame({"a": [2, 3], "b": [2.0, 3.0], "c": ["b", "c"]})
)
assert df[[-1, -2]].frame_equal(
pl.DataFrame({"a": [3, 2], "b": [3.0, 2.0], "c": ["c", "b"]})
)
assert df[[True, False, True]].frame_equal(
pl.DataFrame({"a": [1, 3], "b": [1.0, 3.0], "c": ["a", "c"]})
)
assert df[["a", "b"]].columns == ["a", "b"]
assert df[[1, 2], [1, 2]].frame_equal(
pl.DataFrame({"b": [2.0, 3.0], "c": ["b", "c"]})
)
assert df[1, 2] == "b"
assert df[1, 1] == 2.0
assert df[2, 0] == 3
assert df[[True, False, True], "b"].shape == (2, 1)
assert df[[True, False, False], ["a", "b"]].shape == (1, 2)
assert df[[0, 1], "b"].shape == (2, 1)
assert df[[2], ["a", "b"]].shape == (1, 2)
assert df.select_at_idx(0).name == "a"
assert (df.a == df["a"]).sum() == 3
assert (df.c == df["a"]).sum() == 0
assert df[:, "a":"b"].shape == (3, 2)
assert df[:, "a":"c"].columns == ["a", "b", "c"]
expect = pl.DataFrame({"c": ["b"]})
assert df[1, [2]].frame_equal(expect)
expect = pl.DataFrame({"b": [1.0, 3.0]})
assert df[[0, 2], [1]].frame_equal(expect)
assert df[0, "c"] == "a"
assert df[1, "c"] == "b"
assert df[2, "c"] == "c"
assert df[0, "a"] == 1
