Polars按组内唯一值计数筛选报错，求正确实现方法

错误原因分析 #

你用unique_counts()的问题在于：这个方法返回的是组内每个唯一x值的计数数组，而非组内唯一值的总数量。比如某个组的x包含1和2，unique_counts()会返回类似[3,2]的数组，当你用它和单个数值2比较时，会出现形状不匹配的错误——因为窗口函数返回的数组长度和组的行数不一致，无法和原数据逐行对应。

正确解法 #

方法1：窗口函数直接筛选

把unique_counts()替换为n_unique()，这个方法会直接返回组内x的唯一值总数（单个数值），能和原数据逐行匹配：

import polars as pl

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]
})

# 筛选组内x至少有2个不同值的所有行
result = df.filter(pl.col('x').n_unique().over('data', 'group') >= 2)
print(result)

运行后会保留(data=1, group=1)和(data=1, group=2)这两个组的所有行（这两个组的x包含1和2），而(data=2, group=1)因x全为1会被过滤。

方法2：先分组筛选再关联（适合复杂场景）

如果需要先明确哪些分组符合条件，再关联回原数据，可采用这种更直观的写法：

# 第一步：找出符合条件的分组键
valid_groups = df.group_by('data', 'group')\
                 .agg(pl.col('x').n_unique().alias('x_unique_count'))\
                 .filter(pl.col('x_unique_count') >= 2)

# 第二步：关联回原数据，保留符合条件的组的行
result = df.join(valid_groups, on=['data', 'group'], how='inner')
print(result)

内容的提问来源于stack exchange，提问作者Brian