PySpark：如何让生成的加密密钥列值保持不变？

问题根源 #

你的代码中，UDF属于懒求值逻辑，Spark每次对DataFrame执行动作（如show()、count()）时，都会重新调用generate_key()生成新密钥。用lit(get_key())也无法解决，因为get_key()作为UDF，本质还是在计算阶段才会执行，而非提前生成固定值。

解决方案 #

分两种场景给出实现：

场景1：所有行使用同一个固定密钥

提前在Driver端生成密钥，再通过lit()将固定值添加为列，确保密钥不会随查询重复生成：

from cryptography.fernet import Fernet
from pyspark.sql.functions import lit

def generate_key():
    return Fernet.generate_key().decode('utf-8')

# 提前生成固定密钥（仅执行一次）
fixed_enc_key = generate_key()

def key_table(df):
    return df.select("id").withColumn('encKey', lit(fixed_enc_key))

场景2：每行使用唯一且固定的密钥

如果需要为每行生成唯一密钥且后续查询不再变更，需提前生成所有密钥并与原DataFrame关联，同时持久化结果避免重复计算：

from cryptography.fernet import Fernet
from pyspark.sql.functions import monotonically_increasing_id

def generate_key_list(num_rows):
    # 生成对应行数的唯一密钥列表
    return [Fernet.generate_key().decode('utf-8') for _ in range(num_rows)]

def key_table(df):
    # 1. 获取原DataFrame行数，生成对应数量的密钥
    row_count = df.count()
    keys_df = spark.createDataFrame(
        zip(range(row_count), generate_key_list(row_count)),
        schema=["row_id", "encKey"]
    )
    # 2. 给原DataFrame添加唯一行号，用于关联密钥
    df_with_rowid = df.select("id").withColumn("row_id", monotonically_increasing_id())
    # 3. 关联密钥并清理临时列
    result_df = df_with_rowid.join(keys_df, on="row_id", how="inner").drop("row_id")
    # 4. 持久化结果，避免后续查询重新生成密钥
    result_df.cache()
    return result_df

关键说明 #

• 避免在UDF中生成密钥：UDF的执行时机由Spark的作业调度决定，每次触发计算都会重新执行。
• 持久化DataFrame：如果是场景2，使用cache()或persist()将结果存储在内存/磁盘，后续查询直接读取持久化数据，不再重新生成密钥。

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