如何分批运行PySpark UDF处理DataFrame多行直至最后一行？

问题背景 #

通过PySpark UDF调用Flask API获取字符串结果，为DataFrame新增列，但处理50万+行数据时频繁抛出内存溢出异常：

java.lang.OutOfMemoryError: Java heap space, java.lang.OutOfMemoryError : GC overhead limit exceeded

相关代码 #

API调用UDF代码

def my_method(input: str):
    url = 'URL'
    payload = {}
    headers = {'Authorization': <SOME_TOKEN>}
    # 修正原代码笔误：url. headers → url, headers；payoad → payload
    resp = requests.request('GET', url, headers=headers, data=payload)
    # 修正原代码笔误：response → resp；= → ==
    if resp.status_code == 200:
        return resp.text
    else:
        # 确保返回值为字符串类型，匹配UDF的StringType
        return str(resp.status_code)

# 修正原代码笔误：lambda参数z → x
my_udf = udf(lambda x: my_method(x), StringType())

数据表结构

CREATE TABLE student (id String, name STRING, age INT);

当前运行方式

df = spark.sql('select * from student')
output_df = df.withColumn('api_output', my_udf(col('id')))
output_df.write.format('parquet').mode('append').saveAsTable('database.tablename')

核心提问 #

是否可以分批处理student表的ID列数据，循环运行UDF直至处理完所有行？

解决方案：分批处理完全可行，且是缓解内存压力的有效手段 #

当然可以通过分批拆分数据，降低单批次处理的数据量，从而减少JVM内存占用，避免溢出问题。以下是几种实用的分批实现方式：

1. 基于ID范围分批（适用于有序ID） #

如果id是有序的字符串（如数字字符串），可以按ID区间拆分批次：

# 获取ID的最小、最大值
min_id = spark.sql("SELECT MIN(id) FROM student").collect()[0][0]
max_id = spark.sql("SELECT MAX(id) FROM student").collect()[0][0]

# 每批处理10000条，可根据内存情况调整
batch_size = 10000

current_start = min_id
while int(current_start) <= int(max_id):
    current_end = str(int(current_start) + batch_size - 1)
    # 读取当前批次数据
    batch_df = spark.sql(f"SELECT * FROM student WHERE id >= '{current_start}' AND id <= '{current_end}'")
    # 调用UDF处理
    processed_df = batch_df.withColumn('api_output', my_udf(col('id')))
    # 追加写入结果表
    processed_df.write.format('parquet').mode('append').saveAsTable('database.tablename')
    # 更新下一批起始ID
    current_start = str(int(current_end) + 1)

2. 基于哈希分桶分批（适用于无序ID） #

如果id无顺序规律，可以通过哈希分桶将数据均匀拆分为固定批次：

# 总批次数，可根据内存和集群资源调整（示例为100批）
total_batches = 100

for batch_num in range(total_batches):
    # 通过hash函数将数据均匀分配到不同批次
    batch_df = spark.sql(f"""
        SELECT * FROM student 
        WHERE MOD(hash(id), {total_batches}) = {batch_num}
    """)
    processed_df = batch_df.withColumn('api_output', my_udf(col('id')))
    processed_df.write.format('parquet').mode('append').saveAsTable('database.tablename')

3. 额外优化建议 #

• UDF性能优化：使用pandas_udf替代普通UDF，利用矢量化处理减少Python与JVM的交互开销，提升处理效率
• API调用优化：在my_method中加入请求重试逻辑，避免单次请求失败导致批次中断；若Flask API支持批量查询，改为批量调用API，减少HTTP请求次数
• Spark内存配置：适当调大spark.driver.memory和spark.executor.memory参数，提升JVM可用内存空间

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