PySpark中size(collect_set())与approx_count_distinct()的性能对比

需求是统计每个卖家在**D-1时段（不同卖家的D值不同）**内的唯一订单数。由于官方的count_distinct仅支持groupBy操作，无法用于窗口函数，因此有两种实现精确统计的方案，代码示例如下：

from pyspark.sql import Window
from pyspark.sql.functions import col, collect_set, size, approx_count_distinct
from pyspark.sql.types import LongType

my_window = (Window
             .partitionBy("seller_id")
             .orderBy(col("dt_ref").cast(LongType()))
             .rangeBetween(-60*60*24, 0))

df = (df
      .withColumn("total_orders", size(collect_set("order_id")).over(my_window))
      .withColumn("total_orders_2", approx_count_distinct("order_id", rsd=0.0).over(my_window)))

经测试两种方案返回结果一致，两者的性能差异分析如下：

• size(collect_set())：会先在窗口范围内收集所有唯一的order_id存入集合，再计算集合大小。该方式需要将窗口内所有唯一值加载到内存中，当窗口范围大、唯一订单数量多的时候，内存占用会显著升高，极端情况可能触发内存溢出，且集合的构建与去重操作在数据量大时开销较高。
• approx_count_distinct(rsd=0.0)：尽管函数名称包含“approx”，但当设置rsd=0.0时会强制执行精确统计。它基于HyperLogLog算法的精确模式实现，无需存储所有唯一值，而是通过更高效的统计策略计算唯一值数量。在数据量较大、窗口范围较广的场景下，内存占用更低，计算性能更优。

总结：若窗口内唯一订单数较多，优先选择approx_count_distinct("order_id", rsd=0.0)；若窗口范围极小、唯一值数量极少，两者性能差异不大，但前者在扩展性上更具优势。

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