Pyspark将列转为JSON传入API时col should be Column报错解决

现有PySpark作业处理流程如下：

• 首先读取CSV数据源，对DataFrame按4K batch粒度拆分记录，该部分拆分逻辑运行正常
• 后续需要以JSON格式发起批量API请求：将DataFrame指定列构建为struct类型后转为JSON，传入call_to_cust_bulk_api方法，方法内部组装请求体custRequestBody后向目标API发送请求
• 目标是将API返回的响应作为新列写入DataFrame，最终持久化输出结果

原始实现代码如下：

def call_to_cust_bulk_api(url, payload):
    print("Calling Bulk API")
    try:
        print(payload)
        token = get_token(tokenUrl, tokenBody)
        custRequestBody = {
            "Token": token,
            "CustomerName": "",
            "Object": "",
            "Data": payload
        }
        headers = {'content-type': 'application/json'}
        print(":::::::jsn load::::")
        # 原代码此处传参错误：直接传字典给data参数不会自动序列化为JSON
        response = requests.post(url, data=custRequestBody, headers=headers)
        print(":::Response::::", response)
        data = response.json()
        return data
    except Exception as e:
        print('ExceptionInPushing ' + str(e))
          
df = spark.read.option("header", "true").csv(
    ".csv",
    sep="~")

df = df.withColumn("uniqueID", lit("1"))

df = df.withColumn("row_num", row_number().over(
    Window.partitionBy(col("uniqueID")).orderBy(col("uniqueID"))
))  

total_count = df.count()  

i = 1
while i < total_count:
    rangeNum = i + 3999
    print("Range Num:::")
    print(rangeNum)
    df1 = df.filter((col("row_num") >= i) & (col("row_num") <= rangeNum))
    finalDF = df1.drop("row_num", "edl_timestamp", "uniqueID")
    colsListToBePassed = finalDF.columns
    print("finalDF count:::", finalDF.count())
    # 原代码此处触发类型错误
    finalDF = finalDF.repartition(finalDF.rdd.getNumPartitions()).withColumn("status_for_batch",
                                                                             call_to_cust_bulk_api(policyUrl,
                                                                                                              to_json(
                                                                                                                  struct(*colsListToBePassed))))

代码运行后抛出如下类型错误：

Traceback (most recent call last):
  File "/home/lumiq/IdeaProjects/pyspark_python_test/com/apitest3.py", line 116, in <module>
    finalDF = finalDF.repartition(finalDF.rdd.getNumPartitions()).withColumn("status_for_batch",
  File "/home/lumiq/Downloads/spark-3.2.1-bin-hadoop3.2/python/pyspark/sql/dataframe.py", line 2477, in withColumn
    raise TypeError("col should be Column")
TypeError: col should be Column

1. 直接触发报错的原因：withColumnAPI要求第二个传入参数必须是Spark Column类型对象。原代码直接在Driver端调用普通Python函数call_to_cust_bulk_api，传入的to_json(struct(*colsListToBePassed))是Spark的Column表达式（不是实际计算后的JSON值），函数最终返回的是普通Python对象（字典/None），不符合withColumn的参数类型要求，直接触发类型校验报错。
2. 执行逻辑错误：Spark转换算子（比如withColumn、select）是惰性执行的，直接把Column表达式传给普通Python函数，不会触发DataFrame的实际计算，根本拿不到数据集里的真实数据，就算绕过类型检查也无法正确生成API请求需要的payload。
3. API请求参数错误：原代码调用requests.post时传入data=custRequestBody，requests库不会自动把Python字典序列化为JSON格式，请求会以form表单格式发送，不符合接口要求的application/json格式。

由于调用的是批量API（每4K条记录只需要发起1次请求，不是逐行调用），不需要用UDF逐行执行，直接在Driver端按批次收集数据、调用API、将结果关联回对应批次的DataFrame即可，逻辑更简单性能更好。
修正后的代码如下：

import json
from functools import reduce
from pyspark.sql import functions as F
from pyspark.sql.window import Window
import requests

def call_to_cust_bulk_api(url, payload):
    print("Calling Bulk API")
    try:
        print(payload)
        token = get_token(tokenUrl, tokenBody)
        custRequestBody = {
            "Token": token,
            "CustomerName": "",
            "Object": "",
            "Data": payload
        }
        headers = {'content-type': 'application/json'}
        print(":::::::jsn load::::")
        # 修正1：传json参数自动序列化请求体为JSON格式
        response = requests.post(url, json=custRequestBody, headers=headers)
        print(":::Response::::", response)
        response.raise_for_status() # 主动抛出HTTP错误异常
        data = response.json()
        # 把返回结果转成JSON字符串，方便存入DF列
        return json.dumps(data)
    except Exception as e:
        print('ExceptionInPushing ' + str(e))
        return json.dumps({"error": str(e)})
          
df = spark.read.option("header", "true").csv(
    ".csv",
    sep="~")

df = df.withColumn("uniqueID", F.lit("1"))

df = df.withColumn("row_num", F.row_number().over(
    Window.partitionBy(F.col("uniqueID")).orderBy(F.col("uniqueID"))
))  

total_count = df.count()  
processed_dfs = [] # 存每个批次处理后的DF
i = 1
while i < total_count:
    rangeNum = i + 3999
    print("Range Num:::", rangeNum)
    df1 = df.filter((F.col("row_num") >= i) & (F.col("row_num") <= rangeNum))
    finalDF = df1.drop("row_num", "edl_timestamp", "uniqueID")
    colsListToBePassed = finalDF.columns
    print("finalDF count:::", finalDF.count())
    
    # 修正2：收集当前批次的所有记录，转为JSON格式列表作为payload
    batch_payload = [row["json_data"] for row in finalDF.select(
        F.to_json(F.struct(*colsListToBePassed)).alias("json_data")
    ).collect()]
    
    # 调用批量API拿到返回结果
    batch_resp = call_to_cust_bulk_api(policyUrl, batch_payload)
    
    # 修正3：把API返回结果作为常量列加到当前批次DF，存入待合并列表
    batch_df = finalDF.withColumn("status_for_batch", F.lit(batch_resp))
    processed_dfs.append(batch_df)
    
    i = rangeNum + 1

# 修正4：合并所有批次的DF，后续即可做持久化输出
final_result_df = reduce(lambda df1, df2: df1.unionByName(df2), processed_dfs)
# 后续执行final_result_df.write.xxx 即可持久化

注意：如果单批次4K条记录的数据量很大，collect到Driver端可能造成Driver内存溢出，可以根据实际集群资源调整批次大小，或者改用mapPartitions按分区在Executor端调用API，但要注意做好异常重试、请求频率控制，避免把API服务打挂。

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