基于Mongo-Spark-Connector 2.4.1实现Spark自定义写入MongoDB的方案：满足更新插入规则及只读字段保留需求

我之前在项目里刚好遇到过一模一样的需求，用Mongo-Spark Connector 2.4.1实现这种自定义的Upsert逻辑，确实得绕开replaceDocument的固有局限，下面是我验证过的可行方案：

先明确你遇到的核心矛盾：

• replaceDocument=false：只能更新已存在的字段，无法新增字段
• replaceDocument=true：会覆盖整个文档，导致未出现在新数据中的旧字段丢失
• 还要保护first_load_date只读，自动维护updated_on时间戳

利用MongoDB的原生updateOne命令配合操作符来实现精准控制：

• $set：更新已存在的字段，同时新增不存在的字段（完美解决replaceDocument的两个问题）
• $setOnInsert：仅在插入新文档时设置first_load_date，保证该字段只读
• $currentDate：自动将updated_on设为当前时间（更新/插入操作都会触发）
• upsert: true：实现“存在则更新，不存在则插入”的逻辑

1. 自定义字段排除函数 #

首先需要一个工具函数，用来把不需要更新的字段（比如first_load_date和唯一匹配键）从$set操作中排除：

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

def excludeCol(structCol: Column, colsToExclude: String*): Column = {
  val fields = structCol.schema.fields.map(_.name).filterNot(colsToExclude.contains)
  struct(fields.map(name => structCol.getField(name)): _*)
}

2. 构造更新命令DataFrame #

把你的业务DataFrame转换成MongoDB可识别的updateOne命令结构：

// 假设你的唯一匹配字段是`user_id`，替换成你实际的主键/唯一索引字段
val updateDf = yourOriginalDf
  // 1. 构造唯一匹配条件（必须唯一，避免批量更新错误）
  .withColumn("filter", struct(col("user_id").alias("user_id")))
  // 2. 提取所有业务字段，排除不需要更新的字段
  .withColumn("setFields", struct("*"))
  .withColumn("setFields", excludeCol(col("setFields"), "user_id", "first_load_date"))
  // 3. 组装更新逻辑
  .withColumn("update", struct(
    col("setFields").alias("$set"), // 更新已有字段 + 新增字段
    current_timestamp().alias("$setOnInsert.first_load_date"), // 仅插入新文档时设置
    lit(true).alias("$currentDate.updated_on") // 自动更新时间戳
  ))
  // 4. 开启Upsert模式
  .withColumn("upsert", lit(true))
  // 5. 包装成MongoDB可执行的命令格式
  .select(struct(
    struct(col("filter"), col("update"), col("upsert")).alias("updateOne")
  ).alias("command"))

3. 执行写入 #

把命令DataFrame写入MongoDB：

updateDf.write
  .format("mongo")
  .mode("append")
  .option("database", "db1")
  .option("collection", "my_collection")
  .option("replaceDocument", "false") // 设置为false避免意外覆盖文档
  .save()

如果用PySpark，逻辑完全一致，仅语法调整：

from pyspark.sql import functions as F

def exclude_col(struct_col, cols_to_exclude):
    fields = [field.name for field in struct_col.schema.fields if field.name not in cols_to_exclude]
    return F.struct(*[struct_col[field] for field in fields])

# 构造更新命令DataFrame
update_df = your_original_df \
    .withColumn("filter", F.struct(F.col("user_id").alias("user_id"))) \
    .withColumn("set_fields", F.struct("*")) \
    .withColumn("set_fields", exclude_col(F.col("set_fields"), ["user_id", "first_load_date"])) \
    .withColumn("update", F.struct(
        F.col("set_fields").alias("$set"),
        F.current_timestamp().alias("$setOnInsert.first_load_date"),
        F.lit(True).alias("$currentDate.updated_on")
    )) \
    .withColumn("upsert", F.lit(True)) \
    .select(F.struct(F.struct(F.col("filter"), F.col("update"), F.col("upsert")).alias("updateOne")).alias("command"))

# 写入MongoDB
update_df.write \
    .format("mongo") \
    .mode("append") \
    .option("database", "db1") \
    .option("collection", "my_collection") \
    .option("replaceDocument", "false") \
    .save()

• 唯一匹配条件：一定要确保filter里的字段是唯一的（比如主键_id或带唯一索引的字段），否则会一次更新多个文档，导致数据错误。
• 字段排除：务必把first_load_date从$set中排除，否则会覆盖掉原来的只读值。
• 性能：这种方式是批量执行MongoDB的updateOne命令，性能和直接写入相当，不会有明显损耗。

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