Spark DataFrame中如何获取列值并调用外部API？

我当前正在通过一个包含外部API请求的函数为DataFrame添加新列。发起该请求需要用到每行的"x0"和"x1"列，但我需要的是这些列的具体值。请问如何在函数内部获取列的值？或者说，如何在调用函数时传入列值？

代码示例：

df.withColumn("result", getDataFromApiRequest(col("x0"), col("x1")))

def getDataFromApiRequest(x0: Column, x1: Column): Column = {
  val x0Value = ???
  val x1Value = ???
   ...
}

我曾尝试用UDF函数解决，但该函数需要返回Map[String, Any]或自定义类型，而这在Spark 2.x中均无法实现，因此我尝试用Scala函数解决此问题。

核心问题说明 #

你直接通过Column参数无法获取实际值，因为Column是Spark的逻辑执行计划表达式，并非运行时的具体数据，只有在Spark执行任务阶段，数据才会被实例化。下面提供两种适配Spark 2.x的可行方案：

方案1：适配Spark 2.x的UDF实现（支持自定义返回结构） #

Spark 2.x确实不直接支持返回Map[String, Any]或未注册的自定义类型，但可以通过将结果转换为Spark内置的Row类型，并提前定义对应的StructType来解决：

1. 先定义API返回结果对应的StructType：

import org.apache.spark.sql.types._
import org.apache.spark.sql.functions._
import org.apache.spark.sql.Row

// 假设API返回的结构包含code和msg两个字段，可根据实际情况调整
val resultSchema = StructType(Seq(
  StructField("code", IntegerType, nullable = false),
  StructField("msg", StringType, nullable = true)
))

2. 编写UDF，将API请求结果封装为Row：

def callApi(x0: String, x1: String): Row = {
  // 替换为实际的API请求逻辑
  val response = yourApiClient.request(x0, x1)
  Row(response.code, response.msg)
}

// 注册UDF并指定返回的Schema
val apiUdf = udf(callApi _, resultSchema)

3. 在DataFrame中使用该UDF：

val dfWithResult = df.withColumn("result", apiUdf(col("x0"), col("x1")))

如果需要展开结果中的字段，可使用getField提取：

dfWithResult.select("x0", "x1", "result.code", "result.msg")

方案2：使用mapPartitions批量调用API（更高效） #

如果API支持批量请求，推荐用mapPartitions按分区处理数据，减少API调用次数，同时可以在分区级别初始化客户端，避免重复创建资源：

import org.apache.spark.sql.Row

val dfWithResult = df.rdd.mapPartitions { iter =>
  // 每个分区初始化一次API客户端
  val apiClient = new YourApiClient()
  
  iter.map { row =>
    val x0 = row.getAs[String]("x0")
    val x1 = row.getAs[String]("x1")
    // 调用API获取结果
    val apiResult = apiClient.request(x0, x1)
    // 将原数据与API结果组合成新Row
    Row.fromSeq(row.toSeq ++ Seq(apiResult.code, apiResult.msg))
  }
}.toDF(df.columns ++ Seq("result_code", "result_msg"): _*)

若API不支持批量，仍可在map中逐行调用，但mapPartitions的资源复用优势依然存在。

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