如何用PySpark/Spark SQL将DataFrame列转为HyperLogLog Sketch加速近似去重计数？

一、核心思路 #

预计算每个大型DataFrame的HyperLogLog Sketch（一种紧凑的分布式存储结构），后续通过合并Sketch即可快速得到并集基数或重叠度，避免重复扫描全量原始数据。Spark原生的approx_count_distinct已基于HyperLogLog实现，但显式生成并复用Sketch能进一步降低计算开销，尤其适合多次合并分析的场景。

二、实现方案 #

1. 利用Spark 3.2+原生sketch函数（推荐） #

Spark 3.2及以上版本在内置函数中提供了sketch、sketch_merge、sketch_estimate一套API，可直接生成、合并Sketch并估算基数，无需额外依赖：

步骤1：生成各DataFrame的Sketch

from pyspark.sql import functions as F

# 生成DataFrame A的HyperLogLog Sketch
df_a_sketch = df_a.groupBy().agg(F.sketch("user_id").alias("a_sketch"))

# 生成DataFrame B的HyperLogLog Sketch
df_b_sketch = df_b.groupBy().agg(F.sketch("user_id").alias("b_sketch"))

步骤2：合并Sketch计算并集基数

# 合并两个Sketch得到并集的Sketch
combined_sketch_df = df_a_sketch.crossJoin(df_b_sketch).agg(
    F.sketch_merge("a_sketch", "b_sketch").alias("union_sketch")
)

# 估算并集的distinct count
union_distinct_count = combined_sketch_df.select(
    F.sketch_estimate("union_sketch").alias("union_count")
).collect()[0][0]

步骤3：计算重叠度（交集基数）

利用容斥原理推导交集基数，再计算重叠占比：

# 估算DataFrame A的distinct count
a_distinct_count = df_a_sketch.select(
    F.sketch_estimate("a_sketch").alias("a_count")
).collect()[0][0]

# 估算DataFrame B的distinct count
b_distinct_count = df_b_sketch.select(
    F.sketch_estimate("b_sketch").alias("b_count")
).collect()[0][0]

# 计算交集基数
intersection_count = a_distinct_count + b_distinct_count - union_distinct_count

# 计算重叠度（交集占较小集合的比例）
overlap_ratio = intersection_count / min(a_distinct_count, b_distinct_count)

2. 自定义UDAF兼容低版本Spark（Spark < 3.2） #

如果你的Spark版本低于3.2，可借助Python的hyperloglog库实现自定义UDAF，需确保集群所有节点安装该依赖：

步骤1：集群安装依赖

在所有集群节点执行：

pip install hyperloglog

步骤2：实现自定义UDF

from pyspark.sql.functions import udf
from pyspark.sql.types import BinaryType, LongType
import hyperloglog
import pickle

# 生成单分区HyperLogLog Sketch的UDF
def create_hll_partition_sketch(ids):
    hll = hyperloglog.HyperLogLog(0.05)  # 误差率5%，与Spark默认一致
    for uid in ids:
        hll.add(str(uid))
    return pickle.dumps(hll)

create_hll_part_udf = udf(create_hll_partition_sketch, BinaryType())

# 合并多个Sketch的UDF
def merge_hll_sketches(sketch_list):
    merged_hll = hyperloglog.HyperLogLog(0.05)
    for sketch_bytes in sketch_list:
        hll = pickle.loads(sketch_bytes)
        merged_hll.update(hll)
    return pickle.dumps(merged_hll)

merge_hll_udf = udf(merge_hll_sketches, BinaryType())

# 估算Sketch基数的UDF
def estimate_hll_cardinality(sketch_bytes):
    hll = pickle.loads(sketch_bytes)
    return int(hll.cardinality())

estimate_hll_udf = udf(estimate_hll_cardinality, LongType())

步骤3：处理数据生成并合并Sketch

# 生成DataFrame A的全局Sketch（先按分区聚合，再合并分区Sketch）
df_a_part_sketches = df_a.rdd.map(lambda row: row.user_id)\
    .mapPartitions(lambda iter: [create_hll_partition_sketch(iter)])\
    .toDF("part_sketch")
df_a_sketch = df_a_part_sketches.groupBy().agg(
    merge_hll_udf(F.collect_list("part_sketch")).alias("a_sketch")
)

# 同理生成DataFrame B的全局Sketch
df_b_part_sketches = df_b.rdd.map(lambda row: row.user_id)\
    .mapPartitions(lambda iter: [create_hll_partition_sketch(iter)])\
    .toDF("part_sketch")
df_b_sketch = df_b_part_sketches.groupBy().agg(
    merge_hll_udf(F.collect_list("part_sketch")).alias("b_sketch")
)

# 后续计算并集基数、交集基数、重叠度的逻辑与方案1一致

三、性能优化建议 #

• 持久化Sketch：将生成的Sketch DataFrame写入Parquet等列式存储格式，后续分析直接加载Sketch，无需重复处理原始数据：

  df_a_sketch.write.mode("overwrite").parquet("/path/to/persisted/a_sketch.parquet")
  

• 调整误差率：若业务允许更大误差（如10%），可降低HyperLogLog的精度参数，进一步缩小Sketch体积，提升合并速度。
• 优化分区数：确保原始DataFrame有合理的分区数，让Sketch生成过程充分利用集群并行资源。

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