Polars实现COUNT(DISTINCT)及LazyFrame唯一键校验优化方案问询

问题背景 #

在数据处理流程中，我需要在每步操作后校验数据是否保留预期的唯一键（比如列组合(a, b)），常规方式是判断该组合的去重后数量是否等于总行数。

之前在Polars里我用这种方式实现COUNT(DISTINCT ...)：

(
  df
  .select('a', 'b')
  .unique()
  .height
)

但这个方法有两个问题：一是height无法作用于LazyFrame，二是需要物化全部去重后的数据，效率偏低。参考R语言data.table的uniqueN实现：

mtc_dt <- data.table::as.data.table(mtcars)
stopifnot(data.table::uniqueN(mtc_dt[, .(mpg, disp)]) == nrow(mtc_dt))

想知道Polars里有没有更优的方案？

最优实现方案 #

Polars提供了n_unique()方法，完美对应data.table的uniqueN，支持多列组合计算，且对DataFrame和LazyFrame都适用，不需要物化全部去重数据，执行效率更高。

1. 针对Eager DataFrame

直接通过n_unique()计算组合的唯一数量，再和总行数对比：

# 直接返回布尔值，判断(a,b)是否唯一
is_unique = df.select(pl.col(["a", "b"]).n_unique() == pl.len()).item()

# 如果需要查看具体数值（总行数、唯一组合数）
result = df.select(
    pl.len().alias("total_rows"),
    pl.col(["a", "b"]).n_unique().alias("unique_pairs"),
    (pl.len() == pl.col(["a", "b"]).n_unique()).alias("is_unique")
)

2. 针对LazyFrame

不需要提前物化数据，直接在懒执行计划中加入统计计算，最后再collect()获取结果：

df_lazy = pl.scan_parquet("your_data.parquet")  # 示例LazyFrame

is_unique = (
    df_lazy
    .select(pl.col(["a", "b"]).n_unique() == pl.len())
    .collect()
    .item()
)

# 查看详细统计的版本
result = (
    df_lazy
    .select(
        pl.len().alias("total_rows"),
        pl.col(["a", "b"]).n_unique().alias("unique_pairs")
    )
    .with_columns((pl.col("total_rows") == pl.col("unique_pairs")).alias("is_unique"))
    .collect()
)

为什么这个方案更好 #

• 避免物化大量数据：n_unique()直接计算唯一值数量，不需要生成完整的去重数据集，内存占用和执行速度都远优于unique().height。
• 兼容懒执行：完全支持LazyFrame，可以融入整个懒执行流水线，Polars会自动优化执行计划，只计算必要的统计量。
• 语法简洁：和data.table的uniqueN逻辑对齐，可读性强。

致所有开发者：感谢Polars这款优秀工具，它大幅提升了我的工作效率！

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