如何在Scala/Spark DataFrame中按行使用带条件的withColumn

首先，先明确你给出的DataFrame结构（示例数据省略）：

接下来，我给你分几种常见场景，讲解如何用withColumn结合条件逻辑处理每一行：

1. 基础单条件判断 #

如果只是简单的二选一逻辑，比如根据FFAction的值生成一个标记列：

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

// 当FFAction为"INSERT"时标记为1，否则为0
val dfWithFlag = df.withColumn("ActionFlag", 
  when(col("FFAction") === "INSERT", 1)
    .otherwise(0)
)

2. 多条件分支判断 #

如果有多个不同的条件需要依次匹配，比如给不同的FFAction类型添加描述：

val dfWithActionDesc = df.withColumn("ActionDescription",
  when(col("FFAction") === "INSERT", "新增数据记录")
    .when(col("FFAction") === "UPDATE", "更新已有记录")
    .when(col("FFAction") === "DELETE", "删除指定记录")
    .otherwise("未知操作类型")
)

3. 基于多列的复杂条件 #

如果你的条件需要结合多个列的值，比如判断当前行的时间区间是否有效：

// 判断effectiveFrom不晚于当前时间，且effectiveTo不早于当前时间，标记为"有效"
val dfWithEffectiveStatus = df.withColumn("IsEffective",
  when(
    col("IdentifierValue_effectiveFrom") <= current_timestamp() && 
    col("IdentifierValue_effectiveTo") >= current_timestamp(),
    "有效"
  ).otherwise("无效")
)

4. 自定义UDF处理复杂逻辑 #

如果你的条件逻辑太复杂，没法用Spark内置函数组合实现，可以自定义UDF来处理：

// 定义一个自定义UDF，接收FFAction和effectiveFrom两个参数，返回自定义标记
val customActionUdf = udf((action: String, fromTime: java.sql.Timestamp) => {
  val threshold = java.sql.Timestamp.valueOf("2023-01-01 00:00:00")
  action match {
    case "INSERT" if fromTime.after(threshold) => "2023年后新增"
    case "UPDATE" => "数据更新操作"
    case "DELETE" => "数据删除操作"
    case _ => "其他操作"
  }
})

// 调用UDF生成新列
val dfWithCustomCol = df.withColumn("CustomActionTag", customActionUdf(col("FFAction"), col("IdentifierValue_effectiveFrom")))

需要注意的是，使用UDF时要保证参数类型和DataFrame中列的类型严格匹配，避免出现类型转换错误。

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