Spark并行化map函数遇序列化错误及DynamoDB并发控制问题

一、序列化错误原因及修复 #

你遇到的PicklingError是因为直接将DynamoDB客户端/表对象（ddb_table）传递给了UDF或RDD的map函数。这类AWS客户端对象包含线程锁等无法序列化的资源，Spark在将任务分发到Worker节点时，需要序列化所有传递给任务的对象，因此触发错误。

修复方案：Worker端延迟初始化客户端 #

不要在Driver端创建DynamoDB客户端，而是在UDF或map函数内部（Worker进程中）初始化客户端，避免序列化问题。

示例1：UDF写法

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType
import boto3

def func1(id_val):
    # 在Worker进程内初始化DynamoDB客户端
    ddb = boto3.resource('dynamodb')
    ddb_table = ddb.Table('你的表名')
    
    # 执行DynamoDB查询逻辑
    response = ddb_table.get_item(Key={'id': id_val})
    if 'Item' in response:
        # 调用其他AWS服务（同样在内部初始化对应客户端）
        sns = boto3.client('sns')
        sns.publish(TopicArn='你的SNS主题ARN', Message=f"ID {id_val} 存在")
    return "处理完成"

udf1 = udf(func1, StringType())
df_test = df.withColumn("response", udf1(df["id"]))
df_test.show()

示例2：RDD写法

import boto3

def func1(row):
    id_val = row['id']
    # Worker端初始化客户端
    ddb = boto3.resource('dynamodb')
    ddb_table = ddb.Table('你的表名')
    
    response = ddb_table.get_item(Key={'id': id_val})
    if 'Item' in response:
        # 调用其他AWS服务
        pass
    return (id_val, "处理完成")

df_rdd = df.repartition(20)
final_rdd = df_rdd.rdd.map(func1)
# 触发执行
final_rdd.collect()

二、DynamoDB并发控制实现 #

Spark的RDD/UDF是分布式并行执行的，默认并发度由分区数和集群资源决定。要实现每秒50次调用的限流，需要在Worker进程内做单进程限流，再根据Worker数量调整单进程阈值（比如集群有5个Worker进程，每个进程限流10次/秒，总并发就是50次/秒）。

方案1：使用限流库（推荐） #

可以用ratelimit库实现单进程限流，需要在Glue作业中安装该库（通过--additional-python-modules ratelimit参数配置）。

from ratelimit import limits, sleep_and_retry
import boto3

# 单进程每秒允许10次调用（根据Worker数量调整）
CALLS_PER_SECOND = 10
PERIOD = 1

@sleep_and_retry
@limits(calls=CALLS_PER_SECOND, period=PERIOD)
def query_ddb(id_val):
    ddb = boto3.resource('dynamodb')
    ddb_table = ddb.Table('你的表名')
    return ddb_table.get_item(Key={'id': id_val})

def func1(row):
    id_val = row['id']
    response = query_ddb(id_val)
    if 'Item' in response:
        # 调用其他AWS服务
        pass
    return (id_val, "处理完成")

df_rdd = df.repartition(20)
final_rdd = df_rdd.rdd.map(func1)
final_rdd.collect()

方案2：批量查询优化（更高效） #

优先使用DynamoDB的batch_get_item批量查询，减少API调用次数，从根源降低并发压力。

from pyspark.sql.functions import collect_list, explode, udf
from pyspark.sql.types import ArrayType, StringType
import boto3

def batch_query_ddb(ids):
    ddb = boto3.resource('dynamodb')
    ddb_table = ddb.Table('你的表名')
    # 批量查询最多支持100个Key，可自行拆分
    response = ddb_table.batch_get_item(RequestItems={
        '你的表名': {'Keys': [{'id': val} for val in ids[:100]]}
    })
    # 提取存在的ID
    existing_ids = [item['id'] for item in response['Responses']['你的表名']]
    return existing_ids

# 将ID分组，批量查询
grouped_df = df.groupBy().agg(collect_list('id').alias('id_list'))
batch_udf = udf(batch_query_ddb, ArrayType(StringType()))
result_df = grouped_df.withColumn('existing_ids', batch_udf('id_list'))

# 展开结果并调用其他服务
final_df = result_df.select(explode('existing_ids').alias('existing_id'))

def call_other_service(id_val):
    sns = boto3.client('sns')
    sns.publish(TopicArn='你的SNS主题ARN', Message=f"ID {id_val} 存在")
    return "服务调用完成"

service_udf = udf(call_other_service, StringType())
final_df = final_df.withColumn('service_response', service_udf('existing_id'))

三、Glue作业配置注意事项 #

• 如果使用第三方限流库，需要在Glue作业的作业参数中添加--additional-python-modules ratelimit，让Glue自动安装依赖。
• 调整分区数时，避免分区过多导致并发过高，分区数建议与Worker核心数匹配。

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