Python共享内存与多进程队列解决PySpark UDF重复生成编码问题

• 原实现依赖全局变量used_codes、new_codes存储状态，在PySpark分布式运行场景下，每个worker进程会独立复制一份全局变量，跨进程之间没有状态同步机制，多个进程会同时分配到相同的新编码
• 随机编码生成范围仅为1000-9000，本身碰撞概率较高，进一步放大了重复问题
• UDF内部修改全局变量的操作只会作用于当前worker进程的本地变量，不会同步到driver或其他worker，已使用编码的状态完全无法全局生效

直接用Spark原生算子实现编码映射，完全规避分布式状态同步的复杂度，性能更高且100%不会出现编码重复问题：

import pyspark.sql.functions as F
import pyspark.sql.types as T
from pyspark.sql.window import Window
import random

# 假设这是已经被占用的编码集合，可根据实际业务场景加载
used_codes = [1001, 1002, 1003]
used_codes_bc = spark.sparkContext.broadcast(used_codes)

def generate_unique_codes(required_count, used_codes):
    """生成指定数量的未被占用的唯一编码"""
    unique_codes = set()
    while len(unique_codes) < required_count:
        code = random.randint(1000, 9000)
        if code not in used_codes:
            unique_codes.add(code)
    return list(unique_codes)

# 步骤1：加载待处理数据集，假设数据集中编码字段名为original_code
df = spark.read.parquet("你的数据集路径")

# 步骤2：筛选出需要分配新编码的记录（原有编码已经被占用）
need_reassign_df = df.filter(F.col("original_code").isin(used_codes_bc.value))
reassign_count = need_reassign_df.select("original_code").distinct().count()

# 步骤3：生成足够的唯一新编码，和需要重分配的编码绑定
new_codes_list = generate_unique_codes(reassign_count, used_codes_bc.value)
# 给需要重分配的旧编码按顺序绑定新编码
reassign_map_df = need_reassign_df.select("original_code").distinct()\
    .withColumn("rn", F.row_number().over(Window.orderBy(F.lit(1))))\
    .withColumn("new_code", F.element_at(F.lit(new_codes_list), F.col("rn")).cast(T.StringType()))\
    .select("original_code", "new_code")

# 步骤4：生成最终编码
result_df = df.join(reassign_map_df, on="original_code", how="left")\
    .withColumn("final_code", F.coalesce(F.col("new_code"), F.col("original_code").cast(T.StringType())))\
    .drop("new_code")

如果是流式处理必须支持多进程同时分配编码，可引入外部KV存储（如Redis）维护全局已用编码集合：

• 每次分配编码前先查询Redis判断编码是否已被占用
• 用Redis的原子操作SETNX保证新编码分配的唯一性，避免并发冲突

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