Polars中按组统计子组唯一元素及相关技术疑问

原始DataFrame #

df = (polars
  .DataFrame(
    dict(
      j=[1,1,1,1,2,2,3,3,3,3,3,3,3],
      k=[1,1,2,2,3,3,4,4,5,5,6,6,6],
      l=[1,2,1,2,2,2,3,4,3,3,3,4,3],
      u=[1,1,1,1,2,2,3,3,3,3,3,3,3],
      )
    )
  )

对应数据结构：

j    k    l    u
 i64  i64  i64  i64
 1    1    1    1
 1    1    2    1
 1    2    1    1
 1    2    2    1
 2    3    2    2
 2    3    2    2
 3    4    3    3
 3    4    4    3
 3    5    3    3
 3    5    3    3
 3    6    3    3
 3    6    4    3
 3    6    3    3
shape: (13, 4)

问题1：agg聚合得到的struct列表未按原数据顺序排列，如何修正？ #

原因 #

Polars聚合操作默认会对结果做哈希排序优化，导致原数据的出现顺序丢失，尤其是去重类聚合会优先保证性能而非顺序。

解决方案 #

在聚合时使用maintain_order=True参数，强制保留数据首次出现的顺序：

(df
 .group_by('j', 'k')
 .agg(pl.struct(['l','u']).unique(maintain_order=True))
)

问题2：list[struct[2]]类型无法执行unique操作？ #

原因 #

早期Polars版本对嵌套复合类型（如包含struct的list）的哈希计算支持不完善，struct作为非原子类型，默认无法直接被判定为可哈希对象，导致unique()操作失败；另外struct的字段顺序差异也会被判定为不同值。

解决方案 #

1. 升级Polars到最新版本，新版本已优化嵌套类型的unique支持；
2. 手动将struct转换为可哈希的tuple类型后再执行unique：

(df
 .with_columns(pl.struct(['l','u']).map_elements(lambda x: tuple(x.items()), return_dtype=pl.Tuple))
 .group_by('j')
 .agg(pl.col('tuple').list.unique())
)

问题3：按j列分组，获取每个组内子组k对应的剩余列唯一值集合列表 #

实现代码 #

先按j+k分组获取每个子组的唯一struct列表，再按j分组收集结果，并合并重复的子列表：

result = (df
 .group_by('j', 'k', maintain_order=True)
 .agg(pl.struct(['l','u']).unique(maintain_order=True))
 .group_by('j', maintain_order=True)
 .agg(pl.col('struct').unique(maintain_order=True).alias('i'))
)
print(result)

输出结果 #

j    i
 i64  list[list[struct[2]]]
 1    [[{1,1}, {2,1}]]
 2    [[{2,2}]]
 3    [[{3,3}, {4,3}], [{3,3}]]
shape: (3, 2)

问题4：列出同一j组内剩余列值集合不同的子组k相关数据 #

实现代码 #

先获取每个(j,k)子组的唯一值集合，再筛选出同一j下集合存在差异的子组：

# 1. 生成每个(j,k)对应的唯一值集合
grouped = (df
 .group_by('j', 'k', maintain_order=True)
 .agg(pl.struct(['l','u']).unique().alias('unique_set'))
)

# 2. 筛选出存在不同集合的j组
j_with_diff = (grouped
 .group_by('j')
 .agg(pl.col('unique_set').unique().count().alias('set_count'))
 .filter(pl.col('set_count') > 1)
 .select('j')
)

# 3. 关联得到目标数据
result = grouped.join(j_with_diff, on='j', how='inner')
print(result)

输出结果 #

j    k    unique_set
 i64  i64  list[struct[2]]
 3    4    [{3,3}, {4,3}]
 3    5    [{3,3}]
 3    6    [{3,3}, {4,3}]
shape: (3, 3)

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