按组内的唯一计数进行过滤

Question

我正在尝试按组内唯一项目的数量进行过滤。

例如，假设我有以下数据集：

df = pl.DataFrame({
    'data': [1,1,1,1,1,2,2],
    'group': [1,1,1,2,2,1,1],
    'id': [1,2,3,4,5,1,2],
    'x': [1,1,2,1,2,1,1]
})

其中 (data, group) 是项目组的复合键，每个项目都有一个“id”和一个“x”值。我想过滤数据集以仅保留至少具有两个不同“x”值（1 和 2）的组。

我尝试了以下操作，但收到错误消息：

df.filter(pl.col('x').unique_counts().over(['data', 'group']) >= 2)

ComputeError                              Traceback (most recent call last)
<ipython-input-24-e6f9db2d2a13> in <cell line: 1>()
----> 1 df.filter(pl.col('x').unique_counts().over(['data', 'group']) >= 2)

1 frames
/usr/local/lib/python3.10/dist-packages/polars/lazyframe/frame.py in collect(self, type_coercion, predicate_pushdown, projection_pushdown, simplify_expression, slice_pushdown, comm_subplan_elim, comm_subexpr_elim, no_optimization, streaming, _eager)
   1704             _eager,
   1705         )
-> 1706         return wrap_df(ldf.collect())
   1707 
   1708     @overload

ComputeError: the length of the window expression did not match that of the group

有人可以帮助我理解我做错了什么，或者如何实现这个目标吗？

Answer 1

.n_unique()

将为您提供唯一值的数量。

df.filter(pl.col('x').n_unique().over('data', 'group') >= 2)

shape: (5, 4)
┌──────┬───────┬─────┬─────┐
│ data ┆ group ┆ id  ┆ x   │
│ ---  ┆ ---   ┆ --- ┆ --- │
│ i64  ┆ i64   ┆ i64 ┆ i64 │
╞══════╪═══════╪═════╪═════╡
│ 1    ┆ 1     ┆ 1   ┆ 1   │
│ 1    ┆ 1     ┆ 2   ┆ 1   │
│ 1    ┆ 1     ┆ 3   ┆ 2   │
│ 1    ┆ 2     ┆ 4   ┆ 1   │
│ 1    ┆ 2     ┆ 5   ┆ 2   │
└──────┴───────┴─────┴─────┘

.unique_counts()

的问题是它返回 group_by 上下文中的值列表。

df.group_by("data", "group").agg(pl.col("x").unique_counts())

shape: (3, 3)
┌──────┬───────┬───────────┐
│ data ┆ group ┆ x         │
│ ---  ┆ ---   ┆ ---       │
│ i64  ┆ i64   ┆ list[u32] │
╞══════╪═══════╪═══════════╡
│ 1    ┆ 1     ┆ [2, 1]    │
│ 2    ┆ 1     ┆ [2]       │
│ 1    ┆ 2     ┆ [1, 1]    │
└──────┴───────┴───────────┘