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

Question

I'm trying to use aws encryption sdk to encrypt pyspark columns but running into this error. Here's the code

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") 

Is there a way to resolve this ?

Answer 1

Had a similar issue with boto3, solved by including the boto3 client within the UDF.

# Define a Python function
def my_lambda_function(row_struct):
    # Call the AWS Lambda function using the Boto3 client
    lambda_client = boto3.client('lambda')
    response = lambda_client.invoke(
        FunctionName='my-lambda',
        Payload={
            'arg1': row_struct,
        }
    )
    # Extract the result from the Lambda response
    result = response['Payload'].read().decode('utf-8')
    return result
    
 my_udf = udf(my_lambda_function)