使用Pandas Series与Spark UDF批量发HTTP请求的列表聚合问题

问题原因 #

你的代码使用pandas_udf时，以逐行方式调用UDF（传入col("name")），导致insert函数每次仅处理单个名称，自然生成的records列表每次只有一个元素。此外，原代码中不必要的分布式处理逻辑，反而阻碍了批量生成完整的records列表。

解决方案 #

如果数据量不大（如示例中的5条），最直接的方式是在Driver端一次性收集所有名称，生成完整的records列表后再处理请求。以下是修改后的完整代码：

from pyspark.sql import SparkSession
import argparse
from faker import Faker
import pandas as pd
import json
import requests

# Create a Spark session
spark = SparkSession.builder.appName('SparkByExamples.com').getOrCreate()

parser = argparse.ArgumentParser(description="Program Arguments for Insert")
subparsers = parser.add_subparsers(dest='subparser')

parser_1 = subparsers.add_parser('Insert-Records', help='Insert Records')
parser_1.add_argument("--numberOfRecordsToBeInserted", type=int, default=500,
                      help="The number of records to be inserted into the newly created vault")

args = parser.parse_args()

bearerToken = 'abc'
insertURL = '你的目标API地址'  # 替换为实际的API URL

def insert_all_names(names_list):
    # 生成包含所有名称的records列表
    records = [{"fields": {"name": name}} for name in names_list]
    print(records)
    
    try:
        payload = json.dumps({"records": records})
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
            'Authorization': f'Bearer {bearerToken}'
        }
        
        response = requests.request("POST", insertURL, headers=headers, data=payload)
        dataRes = response.json()
        print("请求响应:", dataRes)
        return ["Success"] * len(names_list)
    except Exception as e:
        print("Error Occurred:", e)
        return ["Failed"] * len(names_list)

def main():
    match args.subparser:
        case 'Insert-Records':
            fake = Faker()
            numberOfRecords = args.numberOfRecordsToBeInserted

            columns = ["seqno", "name"]
            data = []
            for i in range(numberOfRecords):
                data.append((i, fake.name()))

            df = spark.createDataFrame(pd.DataFrame(data), schema=columns)
            
            # 一次性收集所有名称到Driver端
            names_list = df.select("name").rdd.flatMap(lambda x: x).collect()
            
            # 生成完整records并发送请求
            results = insert_all_names(names_list)
            
            # 将结果合并到DataFrame中
            result_df = df.withColumn("response", pd.Series(results).astype(str))
            result_df.show(truncate=False)

        case _:
            print("Invalid operation")

if __name__ == "__main__":
    main()

关键修改点 #

• 移除了不必要的pandas_udf逻辑，直接在Driver端收集所有名称
• 新增insert_all_names函数，接收完整的名称列表，一次性生成包含所有对象的records
• 将请求逻辑移到Driver端，确保只发送一次包含所有数据的POST请求

大数据量场景优化（可选） #

如果数据量很大，直接收集到Driver端会占用过多内存，可以改用分区批量处理：每个分区生成一个records列表并发送请求，避免Driver压力。示例代码如下：

# 替换main函数中的收集和处理部分
def process_partition(partition):
    names = [row.name for row in partition]
    records = [{"fields": {"name": name}} for name in names]
    print(f"分区records: {records}")
    # 这里添加当前分区的POST请求逻辑
    return [(name, "Success") for name in names]

# 使用mapPartitions处理每个分区
result_rdd = df.rdd.mapPartitions(process_partition)
result_df = spark.createDataFrame(result_rdd, ["name", "response"])
result_df.show(truncate=False)

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