基于Scala实现Spark DataFrame的记录删除与插入操作

示例数据 #

// original_df 列数可变，示例结构
val original_df = spark.createDataFrame(Seq(
  ("s1", "o1", "o2", "o7"),
  ("s1", "o2", "o2", "o7"),
  ("s2", "o3", null, "o4"),
  ("s3", null, "o5", "o6"),
  ("s4", null, null, "o6")
)).toDF("s", "p1", "p2", "p3")

val delete_df = spark.createDataFrame(Seq(
  ("s1", "p3", "o7"),
  ("s2", "p1", "o3"),
  ("s4", "p3", "o6")
)).toDF("s", "p", "o")

val insert_df = spark.createDataFrame(Seq(
  ("s1", "p3", "o8"),
  ("s5", "p2", "o9")
)).toDF("s", "p", "o")

实现步骤 #

步骤1：处理删除规则

核心逻辑：匹配delete_df的(s,p,o)记录，将original_df对应行的指定P列置为null。因为P列是动态可变的，所以需要动态生成处理逻辑：

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

// 提取所有P列（排除s列）
val pColumns = original_df.columns.filter(_ != "s")

// 将delete_df转换为广播变量，避免大表关联开销
val deleteMap = delete_df
  .groupBy("s")
  .agg(collect_set(struct($"p", $"o")).as("delete_pairs"))
  .collectAsMap()
val broadcastDeleteMap = spark.sparkContext.broadcast(deleteMap)

// 动态生成每个P列的删除判断逻辑
val afterDeleteCols = pColumns.map(colName => {
  when(
    broadcastDeleteMap.value.get(col("s")).exists(_.contains((colName, col(colName).cast(StringType)))),
    lit(null)
  ).otherwise(col(colName)).alias(colName)
})

// 生成删除处理后的DataFrame
val afterDeleteDf = original_df.select(col("s") +: afterDeleteCols: _*)

步骤2：处理插入规则

核心逻辑：将insert_df转成宽表，与删除后的DataFrame做全外连接，用插入值覆盖原有值，同时新增不存在的s行：

// 将insert_df转成宽表，确保列顺序与original_df一致
val insertWideDf = insert_df
  .groupBy("s")
  .pivot("p", pColumns)
  .agg(first($"o"))

// 合并数据：优先用insert的值，保留original的行，新增insert的行
val afterInsertDf = afterDeleteDf
  .join(insertWideDf, Seq("s"), "full_outer")
  .select(col("s") +: pColumns.map(colName => {
    coalesce(insertWideDf(colName), afterDeleteDf(colName)).alias(colName)
  }): _*)

步骤3：移除全空记录

过滤掉所有P列均为null的行：

// 生成全空判断条件
val allNullCondition = pColumns.map(col(_).isNull).reduce(_ && _)

// 过滤得到最终结果
val final_df = afterInsertDf.filter(not(allNullCondition))

完整代码 #

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

// 初始化示例数据
val original_df = spark.createDataFrame(Seq(
  ("s1", "o1", "o2", "o7"),
  ("s1", "o2", "o2", "o7"),
  ("s2", "o3", null, "o4"),
  ("s3", null, "o5", "o6"),
  ("s4", null, null, "o6")
)).toDF("s", "p1", "p2", "p3")

val delete_df = spark.createDataFrame(Seq(
  ("s1", "p3", "o7"),
  ("s2", "p1", "o3"),
  ("s4", "p3", "o6")
)).toDF("s", "p", "o")

val insert_df = spark.createDataFrame(Seq(
  ("s1", "p3", "o8"),
  ("s5", "p2", "o9")
)).toDF("s", "p", "o")

// 提取P列
val pColumns = original_df.columns.filter(_ != "s")

// 处理删除逻辑
val deleteMap = delete_df
  .groupBy("s")
  .agg(collect_set(struct($"p", $"o")).as("delete_pairs"))
  .collectAsMap()
val broadcastDeleteMap = spark.sparkContext.broadcast(deleteMap)

val afterDeleteCols = pColumns.map(colName => {
  when(
    broadcastDeleteMap.value.get(col("s")).exists(_.contains((colName, col(colName).cast(StringType)))),
    lit(null)
  ).otherwise(col(colName)).alias(colName)
})
val afterDeleteDf = original_df.select(col("s") +: afterDeleteCols: _*)

// 处理插入逻辑
val insertWideDf = insert_df
  .groupBy("s")
  .pivot("p", pColumns)
  .agg(first($"o"))

val afterInsertDf = afterDeleteDf
  .join(insertWideDf, Seq("s"), "full_outer")
  .select(col("s") +: pColumns.map(colName => {
    coalesce(insertWideDf(colName), afterDeleteDf(colName)).alias(colName)
  }): _*)

// 移除全空记录
val allNullCondition = pColumns.map(col(_).isNull).reduce(_ && _)
val final_df = afterInsertDf.filter(not(allNullCondition))

// 查看结果
final_df.show()

效率说明 #

• 用广播变量存储删除规则，避免大表shuffle关联，提升处理速度
• pivot时指定pColumns参数，避免自动排序带来的额外开销
• 全程使用Spark内置函数，保证分布式计算的高效性

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