PySpark UDF按概率生成随机字符串返回常量的问题排查与修复

问题根源 #

你的代码中，random.choices 仅在调用 sim_strings 时执行一次，生成的str_sampled是固定值。后续UDF只是反复返回这个固定值，导致所有行的结果完全相同。

修复方案1：修正UDF实现 #

将随机采样逻辑移到内部函数中，确保每行数据处理时都重新采样：

from pyspark.sql import functions as F

def sim_strings(lst_choices, lst_probs):
    import random
    
    def f(x):
        # 每行调用时重新执行采样逻辑
        return random.choices(lst_choices, weights=lst_probs)[0]
    
    return F.udf(f)

lst_choices_ = ['A', 'B', 'C']
lst_probs_ = [0.5, 0.45, 0.05]

df.withColumn('newcol', sim_strings(lst_choices_, lst_probs_)(F.col('existingcol'))) \
  .select('newcol') \
  .show(100)

注意：使用Python标准库的random模块在Spark UDF中可能存在执行器间随机状态不一致的问题，且UDF性能不如Spark原生函数。

修复方案2：使用Spark原生函数（推荐） #

利用Spark内置的rand()函数生成随机值，结合累积概率映射到目标字符串，这种方式更高效且符合分布式计算特性：

from pyspark.sql import functions as F

lst_choices_ = ['A', 'B', 'C']
lst_probs_ = [0.5, 0.45, 0.05]

# 计算累积概率
cumulative_probs = []
current_total = 0.0
for prob in lst_probs_:
    current_total += prob
    cumulative_probs.append(current_total)

# 生成每行的随机值（仅计算一次，避免重复生成不一致的随机数）
df = df.withColumn('_rand_val', F.rand())

# 构建条件映射链
result_col = F.when(F.col('_rand_val') < cumulative_probs[0], lst_choices_[0])
for idx in range(1, len(lst_choices_)):
    result_col = result_col.when(
        (F.col('_rand_val') >= cumulative_probs[idx-1]) & (F.col('_rand_val') < cumulative_probs[idx]),
        lst_choices_[idx]
    )
result_col = result_col.otherwise(lst_choices_[-1])

# 添加结果列并删除临时随机值列
df = df.withColumn('newcol', result_col).drop('_rand_val')
df.select('newcol').show(100)

这种方法避免了UDF的序列化开销，且随机值生成逻辑完全由Spark分布式处理，结果更稳定。

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