如何高效统计Dask DataFrame各列的唯一值数量？

性能问题原因 #

原有实现每次循环都调用一次compute()，相当于对全表执行120次扫描，每次计算都需要重新加载、处理数据，同时集群任务调度开销也被放大了120倍，因此耗时极高。

最优解决方案 #

使用agg聚合方法一次性对所有列提交nunique计算任务，仅触发一次全局计算即可。Dask会自动优化任务流，仅需要1次全表扫描就能得到所有列的唯一值统计结果，耗时会大幅降低。
示例代码如下：

from dask import dataframe as dd
dask_df = dd.read_csv("train.csv")

# 一次性聚合所有列的唯一值数量，仅调用一次compute
nunique_result = dask_df.agg(func="nunique").compute()

# 输出结果，为pandas Series类型，索引为列名，值为对应列的唯一值数量
print(nunique_result)

可选优化建议 #

• 如果只需要统计特定类型列的唯一值，可先通过select_dtypes筛选列再聚合，比如仅统计数值列：

  # 仅统计数值类型列的唯一值数量
  num_nunique = dask_df.select_dtypes(include=['int64', 'float64']).agg('nunique').compute()
  

• 若DataFrame分区数过多导致调度开销大，可以在读取CSV时通过blocksize参数调整分区大小，比如设置单分区大小为64MB：

  dask_df = dd.read_csv("train.csv", blocksize="64MB")
  

• 如不需要完全精确的统计结果，可使用approx_count_distinct近似聚合函数，计算速度会更快。

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