Scala中不使用collect对DataFrame列每个值应用返回DataFrame的函数

不需要调用collect拉取数据到Driver端，通过分布式行转换+展开即可实现需求，根据你是否可修改原有genDataFrame函数，可选两种方案：

方案1：内置函数实现（性能最优，适合逻辑可改写场景） #

你示例中的genDataFrame逻辑是固定的：每个输入日期返回2行固定规则的结果，这种场景不需要逐行创建DataFrame，直接用Spark内置的数组生成+展开函数即可完成分布式计算，性能最高。

import org.apache.spark.sql.functions._
import java.time.LocalDate

val result = df
  // 为每一行生成包含2个结构体的数组，对应genDataFrame返回的2行结果
  .withColumn("res_rows", array(
    struct(col("date").as("col_A"), date_sub(col("date"), 7).as("col_B")),
    struct(date_add(col("date"), 3).as("col_A"), date_add(col("date"), 7).as("col_B"))
  ))
  // 展开数组，每个结构体转为一行
  .withColumn("res", explode(col("res_rows")))
  // 提取结构体字段为最终列
  .select("res.*")

执行后输出的结果和你给出的预期完全一致，仅分布式计算可能导致行顺序变化，如需固定顺序可增加排序逻辑。

方案2：复用现有函数（通用场景，适合不可修改genDataFrame的情况） #

如果你需要保留原有genDataFrame的逻辑不做改动，不要直接在算子中调用返回DataFrame的方法（Executor端无SparkSession上下文，会报运行时错误），可以先抽离核心的行生成逻辑，再通过flatMap分布式执行：

import org.apache.spark.sql.{DataFrame, Encoders}
import java.time.LocalDate

// 抽离核心的行生成逻辑，返回普通Seq结构，可在Executor端正常执行
def genRows(refDate: LocalDate): Seq[(LocalDate, LocalDate)] = {
  Seq(
    (refDate, refDate.minusDays(7)),
    (refDate.plusDays(3), refDate.plusDays(7))
  )
}

// 原有genDataFrame逻辑保留，不影响历史调用
def genDataFrame(refDate: LocalDate): DataFrame = {
  genRows(refDate).toDF("col_A", "col_B")
}

// 分布式处理，全程不调用collect
implicit val dateTupleEncoder = Encoders.tuple(Encoders.localDate(), Encoders.localDate())
val result = df
  .as[LocalDate]
  .flatMap(refDate => genRows(refDate))
  .toDF("col_A", "col_B")

该方案所有计算都在Executor端分布式完成，不会将全量数据拉取到Driver节点，符合要求。

注意事项 #

不要采用Driver端循环遍历日期值、逐次调用genDataFrame再union的实现：这种方式必须先调用collect拿到所有日期值，还会产生大量小DataFrame做union，性能极差，数据量大时会直接导致Driver内存溢出。

内容的提问来源于stack exchange，提问作者Danusio Gadelha Filho