Spark中如何在UDF抛出异常终止任务并获取列值调用API？

方案1：UDF标记错误+后续校验终止 #

UDF内不直接抛异常，而是返回包含API结果和状态码的结构体，后续通过Spark内置逻辑检查错误状态，触发任务终止。

操作步骤： #

1. 定义返回结构体的UDF，调用API后记录状态码与结果
2. 通过withColumn调用UDF，拆分出结果和状态字段
3. 检查是否存在异常状态码，若存在则主动抛出异常终止任务

示例代码（Python）：

from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType, IntegerType
import requests

# 定义UDF返回的Schema
api_result_schema = StructType([
    StructField("result", StringType(), nullable=True),
    StructField("status_code", IntegerType(), nullable=False)
])

def call_api_udf(input_val):
    try:
        response = requests.get(f"https://your-api-url.com/{input_val}")
        if response.status_code >= 400:
            return (None, response.status_code)
        return (response.text, response.status_code)
    except Exception as e:
        return (None, 500)

# 注册UDF
call_api = F.udf(call_api_udf, api_result_schema)

# 处理数据
df = df.withColumn("api_response", call_api(F.col("input_column")))
df = df.withColumn("api_result", F.col("api_response.result")) \
       .withColumn("api_status", F.col("api_response.status_code"))

# 检查异常状态，触发任务终止
has_error = df.agg(F.max(F.when(F.col("api_status") >= 400, 1).otherwise(0))).collect()[0][0]
if has_error == 1:
    raise Exception("API调用返回异常状态码，任务终止")

方案2：用mapPartitions直接处理并抛异常 #

mapPartitions支持在分区级别处理数据，既能直接获取列值调用API，遇到异常时直接抛出即可终止整个Spark任务。

示例代码（Python）：

from pyspark.sql import Row
import requests

def process_partition(partition):
    for row in partition:
        input_val = row.input_column
        response = requests.get(f"https://your-api-url.com/{input_val}")
        if response.status_code >= 400:
            raise Exception(f"API调用失败，状态码{response.status_code}，输入值{input_val}")
        # 返回处理后的行
        yield Row(**row.asDict(), api_result=response.text)

# 转换DataFrame
df_processed = df.rdd.mapPartitions(process_partition).toDF()

注意：mapPartitions中抛出的异常会直接终止整个Spark作业，符合终止任务的需求；但要注意控制API调用并发量，避免触发限流。

方案3：结合子查询与raise_error触发终止 #

先通过UDF获取所有请求的状态码，再用子查询判断是否存在异常记录，若存在则调用raise_error终止任务。

示例代码（Python）：

from pyspark.sql import functions as F

# 先调用API获取状态码（UDF仅返回状态码）
df_with_status = df.withColumn("api_status", call_api_status_udf(F.col("input_column")))

# 检查异常并触发终止
df_with_status.createOrReplaceTempView("api_results")
spark.sql("""
    SELECT CASE WHEN EXISTS(SELECT 1 FROM api_results WHERE api_status >= 400) 
                THEN raise_error('API返回异常状态码，任务终止') 
                ELSE 0 END AS check_result
""").collect()

# 继续处理正常数据
final_df = df_with_status.filter(F.col("api_status") < 400)

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