如何在polars的group_by中将数据框行分组为列表

Question

import polars as pl

df = pl.DataFrame(
    {
        'Letter': ['A', 'A', 'B', 'B', 'B', 'C', 'C', 'D','D','E'],
        'Value': [1, 2, 3, 4, 5, 6, 7, 8, 9,10]
    }
)

我想把Letter分组，并把它们对应的Value收集到一个列表里。

相关的Pandas问题：如何在Pandas的groupby中将数据框的行分组到列表中

我知道这里的pandas代码是行不通的：

df.group_by('a')['b'].apply(list)

TypeError: 'GroupBy' object is not subscriptable

输出结果将是：

| A      ┆ [1, 2]    │
│ B      ┆ [3, 4, 5] │
│ C      ┆ [6, 7]    │
│ D      ┆ [8, 9]    │
│ E      ┆ [10]      |

Answer 1

很简单：如果你在使用 pl.DataFrame.group_by().agg() 的时候没有指定聚合函数，那么同一组里的所有值会被放到一个列表里。

df.group_by("Letter").agg("Value")

shape: (5, 2)
┌────────┬───────────┐
│ Letter ┆ Value     │
│ ---    ┆ ---       │
│ str    ┆ list[i64] │
╞════════╪═══════════╡
│ A      ┆ [1, 2]    │
│ B      ┆ [3, 4, 5] │
│ C      ┆ [6, 7]    │
│ D      ┆ [8, 9]    │
│ E      ┆ [10]      │
└────────┴───────────┘

Answer 2

你可以这样做。maintain_order=True 是必须的，如果你想让分组的顺序和输入的数据保持一致。

import polars as pl

df = pl.DataFrame(
    {
        'Letter': ['A', 'A', 'B', 'B', 'B', 'C', 'C', 'D','D','E'],
        'Value': [1, 2, 3, 4, 5, 6, 7, 8, 9,10]
    }
)
g = df.group_by('Letter', maintain_order=True).agg(pl.col('Value'))
print(g)

这将会打印出

┌────────┬───────────┐
│ Letter ┆ Value     │
│ ---    ┆ ---       │
│ str    ┆ list[i64] │
╞════════╪═══════════╡
│ A      ┆ [1, 2]    │
│ B      ┆ [3, 4, 5] │
│ C      ┆ [6, 7]    │
│ D      ┆ [8, 9]    │
│ E      ┆ [10]      │
└────────┴───────────┘