python-polars - 在 polars 中计算 pandas n_group 并分配新的 id

Question

我正在寻找在 polars 中执行 pandas 的 df.groupby(["group_a", "group_b"]).ngroup() 并将特定的 ngroup 计数器值分配回的最佳方法相应的组。

df = pl.DataFrame(
     {"group_a": ["a", "aa", "a"], "group_b": ["b", "bb", "b"], "val": [1, 2, 3]}
)
┌─────────┬─────────┬─────┐
│ group_a ┆ group_b ┆ val │
│ ---     ┆ ---     ┆ --- │
│ str     ┆ str     ┆ i64 │
╞═════════╪═════════╪═════╡
│ a       ┆ b       ┆ 1   │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
│ aa      ┆ bb      ┆ 2   │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
│ a       ┆ b       ┆ 3   │
└─────────┴─────────┴─────┘

应该变成

┌─────────┬─────────┬─────┬───────────┐
│ group_a ┆ group_b ┆ val ┆ new_group │
│ ---     ┆ ---     ┆ --- ┆ ---       │
│ str     ┆ str     ┆ i64 ┆ i64       │
╞═════════╪═════════╪═════╪═══════════╡
│ a       ┆ b       ┆ 1   ┆ 0         │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ aa      ┆ bb      ┆ 2   ┆ 1         │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ a       ┆ b       ┆ 3   ┆ 0         │
└─────────┴─────────┴─────┴───────────┘

在 pandas 中，我会使用 ngroups 通过连接来完成此操作，但在 polars 中我不知道。

编辑:我找到了一种可能的解决方法，但我觉得应该有一种更好、更有效的方法，因为它也会因为我的数百万行的实际数据集而被杀死。

group_ids = ["group_a", "group_b"]
df = df.join(
     pl.concat([
        df.unique(subset=group_ids),
        (pl.arange(0, len(df.unique(subset=group_ids)), eager=True, dtype=pl.Int64)
        .alias("new_id")
        .to_frame())], how="horizontal")
        .select(group_ids + ["new_id"]), left_on=group_ids, right_on=group_ids
)

Answer 1

更新:.rank(method="dense") 似乎是生成顺序 ID 的内置方法。

>>> df = pl.DataFrame({
...    "group_a": ["a", "aa", "a", "aaa"], 
...    "group_b": ["b", "bb", "b", "bbb"], 
...    "val": [1, 2, 3, 4] 
... })
...
... group_ids = ["group_a", "group_b"]
...
... (
...    df
...    .with_row_count(name="new_id")
...    .with_column(
...       pl.col("new_id")
...         .first()
...         .over(group_ids)
...         .rank(method="dense") - 1
...    )
... )
shape: (4, 4)
┌────────┬─────────┬─────────┬─────┐
│ new_id | group_a | group_b | val │
│ ---    | ---     | ---     | --- │
│ u32    | str     | str     | i64 │
╞════════╪═════════╪═════════╪═════╡
│ 0      | a       | b       | 1   │
├────────┼─────────┼─────────┼─────┤
│ 1      | aa      | bb      | 2   │
├────────┼─────────┼─────────┼─────┤
│ 0      | a       | b       | 3   │
├────────┼─────────┼─────────┼─────┤
│ 2      | aaa     | bbb     | 4   │
└─//─────┴─//──────┴─//──────┴─//──┘

---

解释:

我们首先生成非连续的组号

>>> (
...    df
...    .with_row_count(name="new_id")
...    .with_column(pl.col("new_id").first().over(group_ids))
... )
shape: (4, 4)
┌────────┬─────────┬─────────┬─────┐
│ new_id | group_a | group_b | val │
│ ---    | ---     | ---     | --- │
│ u32    | str     | str     | i64 │
╞════════╪═════════╪═════════╪═════╡
│ 0      | a       | b       | 1   │
├────────┼─────────┼─────────┼─────┤
│ 1      | aa      | bb      | 2   │
├────────┼─────────┼─────────┼─────┤
│ 0      | a       | b       | 3   │
├────────┼─────────┼─────────┼─────┤
│ 3      | aaa     | bbb     | 4   │
└─//─────┴─//──────┴─//──────┴─//──┘

.rank(method="dense") 生成从1

开始的序号