def test_self_join() -> None:
# 2720
df = pl.from_dict(
data={
"employee_id": [100, 101, 102],
"employee_name": ["James", "Alice", "Bob"],
"manager_id": [None, 100, 101],
}
).lazy()
out = (
df.join(ldf=df, left_on="manager_id", right_on="employee_id", how="left")
.select(
exprs=[
pl.col("employee_id"),
pl.col("employee_name"),
pl.col("employee_name_right").alias("manager_name"),
]
)
.fetch()
)
assert set(out.rows()) == {
(100, "James", None),
(101, "Alice", "James"),
(102, "Bob", "Alice"),
}
def test_from_dict_struct() -> None:
data: Dict[str, Union[Dict, Sequence]] = {
"a": {"b": [1, 3], "c": [2, 4]},
"d": [5, 6],
}
df = pl.from_dict(data)
assert df.shape == (2, 2)
assert df["a"][0] == (1, 2)
assert df["a"][1] == (3, 4)
def test_from_dict_struct() -> None:
data: dict[str, dict | Sequence] = {
"a": {
"b": [1, 3],
"c": [2, 4]
},
"d": [5, 6],
}
df = pl.from_dict(data)
assert df.shape == (2, 2)
assert df["a"][0] == {"b": 1, "c": 2}
assert df["a"][1] == {"b": 3, "c": 4}
def test_rank_so_4109() -> None:
df = pl.from_dict(
{
"id": [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4],
"rank": [None, 3, 2, 4, 1, 4, 3, 2, 1, None, 3, 4, 4, 1, None, 3],
}
).sort(by=["id", "rank"])
assert df.groupby("id").agg(pl.col("rank").rank()).to_dict(False) == {
"id": [1, 2, 3, 4],
"rank": [
[1.0, 2.0, 3.0, 4.0],
[1.0, 2.0, 3.0, 4.0],
[1.0, 2.0, 3.0, 4.0],
[1.0, 2.0, 3.0, 4.0],
],
}
def test_from_dict() -> None:
data = {"a": [1, 2], "b": [3, 4]}
df = pl.from_dict(data)
assert df.shape == (2, 2)