如何通过单次API调用将JSON双键值对存入PySpark DataFrame（基于Requests）

嘿，我明白你现在的困扰——两次调用API太冗余了对吧？其实咱们可以用**PySpark的自定义函数（UDF）**来搞定，一次API调用就能把两个键值对同时塞进DataFrame里，效率高多了！

核心思路是：把API调用逻辑封装成一个UDF，让它接收每行的businessId，一次性获取完整的JSON响应，然后返回一个包含两个字段的结构化结果，最后把这个结构拆成DataFrame的两列就行。

具体步骤和代码示例 #

1. 导入必要的模块

首先得把需要的库都导入进来：

import requests
from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StructType, StructField, StringType

2. 初始化SparkSession（如果还没创建的话）

spark = SparkSession.builder.appName("APItoDataFrame").getOrCreate()

3. 创建你的原始DataFrame

这里模拟一下你只有businessId列的DataFrame：

data = [("dksldfaw2",), ("kkldsdok3",), ("djdfkdfk3",), ("23lksdlk8",)]
df = spark.createDataFrame(data, ["businessId"])

4. 编写API调用函数

这个函数会接收businessId，调用API并返回两个键值对的结果，记得加上异常处理，避免单个请求失败导致整个任务挂掉：

def fetch_api_data(business_id):
    try:
        # 替换成你的实际API地址，这里假设是GET请求，把businessId作为参数传递
        api_url = f"https://your-api-endpoint.com?businessId={business_id}"
        # 如果需要认证或者自定义头，在这里添加headers参数
        response = requests.get(api_url)
        response.raise_for_status()  # 检查请求是否成功（比如4xx/5xx错误）
        json_response = response.json()
        
        # 假设API返回的两个键是"keyA"和"keyB"，替换成你实际的键名
        return (json_response.get("keyA"), json_response.get("keyB"))
    except Exception as e:
        # 打印错误日志，同时返回默认值（比如None）保证任务继续执行
        print(f"Failed to fetch data for businessId {business_id}: {str(e)}")
        return (None, None)

5. 定义UDF的返回结构

PySpark需要知道UDF返回的结构化数据类型，所以我们先定义一个StructType：

api_response_schema = StructType([
    StructField("column_a", StringType(), nullable=True),  # 对应API的keyA，自定义列名
    StructField("column_b", StringType(), nullable=True)   # 对应API的keyB，自定义列名
])

6. 注册UDF并应用到DataFrame

把刚才的函数注册成UDF，然后应用到原始DataFrame上，最后拆分结构化字段：

# 注册UDF
api_fetch_udf = udf(fetch_api_data, api_response_schema)

# 应用UDF并拆分列
result_df = df.withColumn("api_result", api_fetch_udf(col("businessId"))) \
              .select(
                  "businessId",
                  col("api_result.column_a").alias("your_first_column"),
                  col("api_result.column_b").alias("your_second_column")
              )

# 查看最终结果
result_df.show(truncate=False)

一些实用的注意事项 #

• 性能优化：如果你的数据量很大，行级UDF的API调用可能会比较慢。如果你的API支持批量请求（比如一次传多个businessId），建议改成批量调用，能大幅减少请求次数。
• 重试机制：可以给requests添加重试逻辑（比如用tenacity库），处理临时的网络波动或者API限流问题。
• 认证与安全：如果API需要认证（比如Token），不要把敏感信息硬编码，建议通过环境变量或者Spark配置传递。
• 数据类型适配：如果API返回的是数字、日期等类型，记得把StringType改成对应的类型（比如IntegerType、DateType）。

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