PySpark函数单独运行正常，封装为UDF后无法正确生成唯一编码是什么原因

问题根因 #

• PySpark的UDF是分布式执行的，你在Driver端定义的used_codes全局字典，会被序列化后发送到每个Executor节点，每个Executor运行UDF时操作的都是自己本地的字典副本，不同Executor之间的字典数据完全独立，也不会同步回Driver端。
• 小数据量本地模式下运行正常是因为所有计算都在同一个进程内执行，全局字典天然共享；但分布式集群下多Executor运行时，就会出现重复编码无法识别、生成的code_id重复的问题。

解决方案1：小数据量场景直接转Pandas处理 #

如果数据量可以完全放到单节点内存，直接把Spark DataFrame转成Pandas DataFrame处理，避免分布式执行的状态共享问题：

import pandas as pd
import random

used_codes = {}
def generate_random_code():
    return str(random.randint(10000,90000))

def get_valid_code(code):
    global used_codes
    if code != "" and code not in used_codes:
        used_codes[code] = 1
        return code
    new_code = generate_random_code()
    while new_code in used_codes:
        new_code = generate_random_code()
    used_codes[new_code] = 2
    return new_code

# 转pandas处理
pandas_df = df.toPandas()
pandas_df["code_id"] = pandas_df["id"].apply(get_valid_code)
# 转回Spark DataFrame
new_df = spark.createDataFrame(pandas_df)
new_df.show()

解决方案2：大数据量分布式场景（无全局状态实现） #

纯用Spark原生算子实现，避免UDF依赖全局状态，保证分布式下的正确性：

from pyspark.sql.window import Window
import random

# 1. 给每行加全局唯一索引，保证顺序稳定
df = df.withColumn("row_idx", F.monotonically_increasing_id())

# 2. 按id分组，标记每个id的出现顺序
w_id = Window.partitionBy("id").orderBy("row_idx")
df = df.withColumn("id_rn", F.row_number().over(w_id))

# 3. 首次出现的id保留原值作为code_id，重复的id暂留空
df = df.withColumn("code_id", F.when(F.col("id_rn") == 1, F.col("id")).otherwise(None))

# 4. 收集所有已存在的code值，用于去重
exist_codes = set([row[0] for row in df.select("code_id").filter(F.col("code_id").isNotNull()).collect()])
null_cnt = df.filter(F.col("code_id").isNull()).count()

# 5. 生成足够多的不重复随机码
random_codes = []
while len(random_codes) < null_cnt:
    new_code = str(random.randint(10000,90000))
    if new_code not in exist_codes and new_code not in random_codes:
        random_codes.append(new_code)

# 6. 给空的code_id赋值随机码
w_order = Window.orderBy("row_idx")
df = df.withColumn("global_rn", F.row_number().over(w_order))
code_map = spark.createDataFrame([(i+1, random_codes[i]) for i in range(len(random_codes))], schema=["rn", "random_code"])
new_df = df.join(code_map, df.global_rn == code_map.rn, how="left")\
           .withColumn("code_id", F.coalesce("code_id", "random_code"))\
           .select("firstname", "id", "code_id")

new_df.show()

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