如何在Polars中实现窗口函数式的多列唯一值对计数？

需求说明 #

需要实现窗口函数，统计指定分组（context列）内column1与column2的唯一值对数量，功能等价于多列组合的n_unique，但支持窗口上下文计算。

示例数据：

import polars as pl

df = pl.DataFrame({
    'context': [1, 1, 1,  2, 2, 2,  3, 3, 3],
    'column1': [1, 1, 0,  1, 0, 0,  1, 0, 1],
    'column2': [1, 0, 0,  0, 1, 1,  1, 0, 1]
})

问题重现 #

直接调用pl.n_unique('column1', 'column2').over('context')会触发ComputeError，因为当前n_unique不支持同时传入多列并结合窗口over使用：

# 报错代码
(
    df
    .with_columns(
        pl.n_unique('column1', 'column2').over('context').alias('n_unique')
    )
)

现有可行方案 #

通过concat_list将两列合并为列表列，再统计列表的唯一值数量，该方法可行但并非最优：

(   
    df
    .with_columns(
        pl.concat_list('column1', 'column2').alias('pair')
    )
    .with_columns(
        pl.n_unique('pair').over('context')
    )
)

更优实现方案 #

以下几种方法在性能或可读性上更具优势：

方案1：用struct替代list #

结构体在Polars中处理效率高于列表，且语义更贴合"键值对"的概念：

(
    df
    .with_columns(
        pl.struct('column1', 'column2').alias('pair')
    )
    .with_columns(
        pl.n_unique('pair').over('context').alias('n_unique')
    )
)

方案2：先聚合再关联（适合大数据量场景） #

先对分组计算唯一值对数量，再通过join将结果映射回原表，避免窗口函数的重复计算开销：

agg_df = df.group_by('context').agg(
    pl.n_unique(pl.struct('column1', 'column2')).alias('n_unique')
)

df.join(agg_df, on='context', how='left')

方案3：用hash组合列后统计唯一值 #

通过pl.hash将两列组合成哈希值，再统计哈希值的唯一数量，性能通常最优：

(
    df
    .with_columns(
        pl.hash('column1', 'column2').alias('pair_hash')
    )
    .with_columns(
        pl.n_unique('pair_hash').over('context').alias('n_unique')
    )
)

方案对比 #

• struct方案：可读性强，性能优于列表方案，适合大多数常规场景
• 聚合后关联：大数据量下性能最优，减少窗口函数的内存开销
• hash方案：性能最高，但哈希存在极低的碰撞风险，需根据业务场景评估使用

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