PySpark写入Delta Lake时DataFrame意外副作用的解决方案咨询

背景 #

我正在为PySpark流水线实现状态存储，选用Delta Lake保存和恢复状态——处理的都是极小的DataFrame，这类DataFrame可能在当前Spark会话内及跨批次被多次更新和读取。

初期可用代码 #

加载状态 #

def load_state(self, state_id: str, id: str):
    path = os.path.join(self.state_folder, state_id, msn)
    if DeltaTable.isDeltaTable(self.spark, path):
        return self.spark.read.format("delta").load(path)
    else:
        return self.spark.createDataFrame([], StructType([]))

保存状态（初始版本） #

def save_state(self, state_id: str, id: str, new_state: DataFrame):
    path = os.path.join(self.state_folder, state_id, id)
    write_options = {
        "overwriteSchema": "true",  # 支持schema演进
        "deltaLog.deleteAllCheckpoints": "true"  # 写入时删除所有检查点，避免生成新的
    }
    isolated_df = new_state.select("*").cache()
    isolated_df.count()
    (isolated_df.write.mode("overwrite")
     .options(**write_options)
     .format("delta").save(path))
    isolated_df.unpersist()

遇到的问题 #

当new_state由load_state生成（用于更新状态）时，执行写入操作后，new_state会出现额外行——原状态的行被复制到DataFrame中。尝试通过select("*").cache()隔离new_state以避免副作用，但未成功：写入后（unpersist前），isolated_df和new_state均被修改。

核心疑问 #

1. 是否存在真正隔离new_state以避免副作用的方法？
2. 如何保留DataFrame的不变副本/状态，并能随时将其持久化/保存为Delta格式（类似窃听机制）？

当前临时解决方案 #

借助额外副作用：写入后读回Delta表并修改new_state，虽能正常工作，但性能和代码实现都不理想。

def save_state(self, state_id: str, msn: str, new_state: DataFrame):
    """
    保存/覆盖指定MSN的状态

    :param state_id: 标识要保存的状态
    :param msn: 标识飞行器
    :param new_state: 要保存的状态，会覆盖现有状态
    :return:
    """
    path = os.path.join(self.state_folder, state_id, msn)
    write_options = {
        "overwriteSchema": "true",  # 支持schema演进
        "deltaLog.deleteAllCheckpoints": "true"  # 写入时删除所有检查点，避免生成新的
    }
    (new_state.write.mode("overwrite")
     .options(**write_options)
     .format("delta").save(path))
    read_back = self.spark.read.format("delta").load(path)
    new_state._jdf = read_back._jdf
    new_state._schema = read_back._schema

失败的测试用例 #

def test_state_manager(self):
    # ... 测试初始化
    
    new_state = self.spark.createDataFrame([(1, 1.0), (2, 2.0), (3, 3.0)], schema)
    state_manager.save_state(state_id, "pc24_666", new_state)

    loaded_state = state_manager.load_state(state_id, "pc24_666")
    assert loaded_state.count() == 3

    new_state1 = self.spark.createDataFrame([(4, 4.0), (5, 5.0), (6, 6.0)], schema)
    updated_state = loaded_state.union(new_state1)
    
    assert updated_state.count() == 6
    state_manager.save_state(state_id, "pc24_666", updated_state) # 此时updated_state被修改

    update_read = state_manager.load_state(state_id, "pc24_666")

    assert updated_state.count() == update_read.count() # update_read正确为6条，但updated_state变成9条——原状态行被复制

解决方案 #

1. 彻底隔离DataFrame的方法 #

Spark DataFrame是不可变的逻辑计划，出现副作用的核心原因是原DataFrame的血缘关联到了Delta表，写入操作修改Delta表后，原DataFrame的后续执行（比如count()）会重新读取最新的Delta数据。要彻底隔离，需打破原血缘：

• 针对极小DataFrame，直接收集到本地再重新创建，性能开销可忽略：

  def get_isolated_df(df: DataFrame) -> DataFrame:
      if df.rdd.isEmpty():
          return df.sparkSession.createDataFrame([], df.schema)
      local_data = df.collect()
      return df.sparkSession.createDataFrame(local_data, df.schema)
  

• 也可以用new_state.toDF(*new_state.columns)生成新的逻辑计划副本，但collect()方式能完全切断与原Delta表的关联。

2. 无副作用的状态保存实现 #

修改save_state方法，仅操作隔离后的副本，绝不修改传入的new_state：

def save_state(self, state_id: str, id: str, new_state: DataFrame):
    path = os.path.join(self.state_folder, state_id, id)
    write_options = {
        "overwriteSchema": "true",
        "deltaLog.deleteAllCheckpoints": "true"
    }
    # 生成完全独立的副本
    isolated_df = get_isolated_df(new_state)
    # 写入Delta表
    (isolated_df.write.mode("overwrite")
     .options(**write_options)
     .format("delta").save(path))
    # 调用端的new_state不受任何影响

3. 测试用例修正 #

无需依赖写入后修改原DataFrame，直接验证读取结果即可，原updated_state的count始终保持6：

def test_state_manager(self):
    # ... 测试初始化
    
    new_state = self.spark.createDataFrame([(1, 1.0), (2, 2.0), (3, 3.0)], schema)
    state_manager.save_state(state_id, "pc24_666", new_state)

    loaded_state = state_manager.load_state(state_id, "pc24_666")
    assert loaded_state.count() == 3

    new_state1 = self.spark.createDataFrame([(4, 4.0), (5, 5.0), (6, 6.0)], schema)
    updated_state = loaded_state.union(new_state1)
    
    assert updated_state.count() == 6
    state_manager.save_state(state_id, "pc24_666", updated_state)

    update_read = state_manager.load_state(state_id, "pc24_666")
    assert update_read.count() == 6
    # 原updated_state不受影响，依然是6条
    assert updated_state.count() == 6

关键原理 #

Spark DataFrame的不可变性是逻辑层面的，当血缘关联到Delta表时，Delta的元数据变化会导致原DataFrame执行时读取最新数据，产生“副作用”。通过生成完全脱离原血缘的副本，就能彻底避免这种问题。

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