从 Polars 中的每个组中获取元素

Question

如何在 Polars DataFrame 的每组中按索引获取元素？例如，如果我想获取每个组的第一个和第三个元素，我可能会尝试这样的操作：

import polars as pl

df = pl.DataFrame(dict(x=[1,0,1,0,1,0], y=[1,2,3,4,5,6]))

df.groupby('x').take([0,2])
# AttributeError: 'GroupBy' object has no attribute 'take'

但这显然行不通。

Answer 1

df.groupby("x").agg(pl.all().take([0, 2]))

shape: (2, 2)
┌─────┬────────────┐
│ x   ┆ y          │
│ --- ┆ ---        │
│ i64 ┆ list [i64] │
╞═════╪════════════╡
│ 1   ┆ [1, 5]     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 0   ┆ [2, 6]     │
└─────┴────────────┘

如果需要，您可以使用 explode 来展平列表列。

df.groupby("x").agg(pl.all().take([0, 2])).explode('y')

shape: (4, 2)
┌─────┬─────┐
│ x   ┆ y   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 1   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 1   ┆ 5   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 0   ┆ 2   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 0   ┆ 6   │
└─────┴─────┘

take 表达式的文档具有类似的 example。

从 Polars 中的每个组中获取元素

Take elements from each group in Polars

python-polars