如何用PySpark窗口函数为同组记录分配相同UUID

你的核心问题是原方案把集群ID拉到Driver端循环处理，导致无法扩展到百万级数据。完全可以通过分布式分组+关联或者窗口函数实现，全程不需要循环，性能可以线性扩展。

方法一：分组生成UUID后关联（推荐，性能最优） #

这种方式通过Spark分布式分组生成每个唯一组合的UUID，再关联回原表，避免Driver端处理大量数据：

from pyspark.sql import functions as f
import uuid

# 过滤空值（保留原逻辑）
df = df.filter("dob IS NOT NULL AND patient_ssn IS NOT NULL AND dodid IS NOT NULL")

# 为每个(dob, patient_ssn, dodid)组合生成唯一UUID
group_uuid_df = df.groupBy("dob", "patient_ssn", "dodid") \
                  .agg(f.lit(str(uuid.uuid4())).alias("uuid"))

# 关联回原DataFrame，为同组记录分配相同UUID
df = df.join(group_uuid_df, on=["dob", "patient_ssn", "dodid"], how="left")

方法二：窗口函数实现 #

通过partitionBy窗口分组，为每组生成一次UUID并复用：

from pyspark.sql import functions as f
from pyspark.sql.window import Window
import uuid

# 过滤空值
df = df.filter("dob IS NOT NULL AND patient_ssn IS NOT NULL AND dodid IS NOT NULL")

# 按目标字段分组的窗口（无需排序，避免全局shuffle）
group_window = Window.partitionBy("dob", "patient_ssn", "dodid")

# 为组内所有记录分配同一个UUID
df = df.withColumn("uuid", f.first(f.lit(str(uuid.uuid4()))).over(group_window))

原方案性能瓶颈分析 #

1. 全局排序开销：原代码用Window.orderBy做dense_rank，会触发全局shuffle排序，百万级数据下耗时极长。
2. Driver端压力：collect()把所有集群ID拉到Driver内存，分组数过多时会直接OOM。
3. 循环低效：循环过滤每个分组生成DataFrame再unionAll，会触发多次查询和shuffle，性能呈指数级下降。

注意事项 #

• 推荐用uuid.uuid4()而非uuid.uuid1()，前者是随机UUID，后者依赖MAC地址和时间，分布式场景下重复风险更低。
• 如果Spark版本≥3.1，可以用内置函数f.expr("uuid()")生成UUID，但需要配合分组聚合确保同组唯一：

  group_uuid_df = df.groupBy("dob", "patient_ssn", "dodid") \
                    .agg(f.expr("uuid()").alias("uuid"))
  

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