使用AWS Encryption SDK加密PySpark列时遇PicklingError问题求助

问题场景 #

使用AWS Encryption SDK对PySpark列加密时，触发如下序列化错误：

PicklingError: Could not serialize object: TypeError: can't pickle SSLContext objects

用户的原始代码如下：

from pyspark.sql import functions as F
from pyspark.sql.functions import udf, col, lit
from pyspark.context import SparkContext
import aws_encryption_sdk
from aws_encryption_sdk import CommitmentPolicy

client = aws_encryption_sdk.EncryptionSDKClient(
    commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
)

kms_kwargs = dict(key_ids=[key_arn])
global master_key_provider
master_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(**kms_kwargs)
df = spark.read.csv('test.csv', inferSchema = True, header = True)

def encrypt_string(text):
    encrypted_text, encryptor_header = client.encrypt(
        source=text, key_provider=master_key_provider
    )
    return encrypted_text

udf_encrypt = udf(lambda text: encrypt_string(text))
    
def spark_encrypt(df, colmn):
    return df.withColumn("segment_encrypt", udf_encrypt(col(colmn)))

df_out = spark_encrypt(df, "segment") 

解决办法 #

1. 核心原因 #

EncryptionSDKClient和StrictAwsKmsMasterKeyProvider内部包含SSLContext这类无法被Python序列化的对象。PySpark执行UDF时，会尝试将UDF依赖的全局对象序列化后分发到各个Executor节点，直接引用这些对象就会触发序列化失败。

2. 基础修复方案：Executor端延迟初始化加密组件 #

将加密客户端和密钥提供者的初始化逻辑放到UDF内部，让每个Executor节点独立创建这些对象，避免序列化传递：

from pyspark.sql import functions as F
from pyspark.sql.functions import udf, col
import aws_encryption_sdk
from aws_encryption_sdk import CommitmentPolicy

df = spark.read.csv('test.csv', inferSchema=True, header=True)

def encrypt_string(text, key_arn):
    # 在Executor节点上初始化客户端和密钥提供者
    client = aws_encryption_sdk.EncryptionSDKClient(
        commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
    )
    master_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(
        key_ids=[key_arn]
    )
    encrypted_text, _ = client.encrypt(source=text, key_provider=master_key_provider)
    return encrypted_text

# 将key_arn作为参数传入UDF
udf_encrypt = udf(lambda text: encrypt_string(text, key_arn))

def spark_encrypt(df, colmn):
    return df.withColumn("segment_encrypt", udf_encrypt(col(colmn)))

df_out = spark_encrypt(df, "segment")

3. 性能优化：用广播变量传递配置 #

如果担心每个UDF调用初始化客户端影响性能，可以将key_arn这类静态配置通过广播变量传递，避免重复传递参数：

from pyspark.sql import functions as F
from pyspark.sql.functions import udf, col
import aws_encryption_sdk
from aws_encryption_sdk import CommitmentPolicy

df = spark.read.csv('test.csv', inferSchema=True, header=True)

# 广播KMS密钥ARN到所有Executor
broadcast_key_arn = spark.sparkContext.broadcast(key_arn)

def encrypt_string(text):
    key_arn = broadcast_key_arn.value
    client = aws_encryption_sdk.EncryptionSDKClient(
        commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
    )
    master_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(
        key_ids=[key_arn]
    )
    encrypted_text, _ = client.encrypt(source=text, key_provider=master_key_provider)
    return encrypted_text

udf_encrypt = udf(encrypt_string)

def spark_encrypt(df, colmn):
    return df.withColumn("segment_encrypt", udf_encrypt(col(colmn)))

df_out = spark_encrypt(df, "segment")

4. 权限验证 #

确保PySpark的Executor节点拥有指定KMS密钥的kms:Encrypt权限，避免初始化密钥提供者时出现权限错误。

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