Spark基于另一DataFrame字段更新列的Scala及Spark SQL实现方法

前置说明 #

Spark本身不支持传统关系型数据库的UPDATE FROM行级更新语法（普通表/临时视图场景），常规实现思路是通过两表关联后重新构造目标字段值，最终覆盖写入目标表实现等效逻辑；如果使用Delta Lake表则可以直接用MERGE语法实现更高效的行级更新。

1. Spark SQL实现 #

普通表场景（非Delta） #

INSERT OVERWRITE TABLE tbl1
SELECT
  a.f,
  -- 匹配到b表的行更新value为1，未匹配保留原值
  CASE WHEN b.f IS NOT NULL THEN '1' ELSE a.value END AS value,
  -- 匹配到b表的行更新name为b表的old-name，未匹配保留原值
  CASE WHEN b.f IS NOT NULL THEN b.`old-name` ELSE a.name END AS name,
  -- 此处按原tbl1的字段顺序，补全剩下所有未修改的字段，避免字段丢失
  a.col1,
  a.col2,
  a.other_cols
FROM tbl1 a
LEFT JOIN tbl2 b ON a.f = b.f;

Delta Lake表场景 #

可以直接使用MERGE语法，和原SQL逻辑完全对齐：

MERGE INTO tbl1 a
USING tbl2 b
ON a.f = b.f
WHEN MATCHED THEN UPDATE SET
  a.value = '1',
  a.name = b.`old-name`;

2. Spark Scala代码实现 #

普通DataFrame场景 #

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

// df1对应目标表tbl1，df2对应关联表tbl2
val df1 = spark.table("tbl1")
val df2 = spark.table("tbl2")

// 先对df2按关联键f去重，避免关联后数据膨胀
val distinctDf2 = df2.dropDuplicates("f")

val updatedDf = df1.alias("a")
  .join(distinctDf2.alias("b"), col("a.f") === col("b.f"), "left")
  .select(
    col("a.f"),
    when(col("b.f").isNotNull, lit("1")).otherwise(col("a.value")).alias("value"),
    when(col("b.f").isNotNull, col("b.`old-name`")).otherwise(col("a.name")).alias("name"),
    // 补全tbl1其余未修改的字段，也可以通过遍历df1.columns批量生成
    col("a.col1"),
    col("a.col2"),
    col("a.other_cols")
  )

// 覆盖写入原表完成更新
updatedDf.write.mode("overwrite").saveAsTable("tbl1")

Delta Lake表场景 #

import io.delta.tables._

val deltaTable = DeltaTable.forName(spark, "tbl1")
val df2 = spark.table("tbl2")

deltaTable.as("a")
  .merge(
    df2.as("b"),
    "a.f = b.f"
  )
  .whenMatched()
  .updateExpr(Map(
    "value" -> "'1'",
    "name" -> "b.`old-name`"
  ))
  .execute()

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