如何在PySpark批处理间读取/恢复checkpointed DataFrame

首先要明确：DataFrame.checkpoint不是用来跨批次持久化状态的工具。它的核心作用是截断DataFrame的执行血统（lineage），避免复杂作业的执行计划过于冗长，本质是Spark的内部优化手段。默认情况下checkpoint的数据是临时的，作业结束后会被自动清理；即使指定了checkpoint目录，它也是为流处理容错、RDD故障恢复设计的，并不支持手动读取复用。

正确实现跨批次状态持久化与读取 #

要在后续批次复用之前的计算结果，你需要将数据持久化到可读取的存储介质（比如Parquet、CSV），而不是依赖checkpoint。修改你的测试代码如下：

import pytest
from pyspark.sql import functions as f
import os

class TestCheckpoint:

    @pytest.fixture(autouse=True)
    def init_test(self, spark_unit_test_fixture, data_dir, tmp_path):
        self.spark = spark_unit_test_fixture
        self.dir = data_dir("")
        self.state_dir = tmp_path
        # 为求和状态单独创建存储路径
        self.sum_state_path = os.path.join(self.state_dir, "sum_state")

    def test_first(self):
        df = (self.spark.read.format("csv")
              .option("pathGlobFilter", "numbers.csv")
              .load(self.dir))

        sum_df = df.agg(f.sum("_c1").alias("sum"))
        # 将求和结果写入Parquet（列式存储，读写高效）
        sum_df.write.mode("overwrite").parquet(self.sum_state_path)
        assert sum_df.first()["sum"] == 3  # 验证初始求和

    def test_second(self):
        # 读取新批次数据
        df = (self.spark.read.format("csv")
              .option("pathGlobFilter", "numbers2.csv")
              .load(self.dir))

        # 读取之前保存的求和状态
        prev_sum_df = self.spark.read.parquet(self.sum_state_path)
        prev_sum = prev_sum_df.first()["sum"]

        # 计算新数据的求和并累加
        new_sum = df.agg(f.sum("_c1").alias("new_sum")).first()["new_sum"]
        total_sum = prev_sum + new_sum

        # 验证累加结果（示例逻辑）
        assert total_sum > prev_sum

多状态（求和、平均值）的存储处理 #

如果需要同时存储多个状态（比如求和、平均值），只需为每个状态分配独立的存储路径即可：

# 在fixture中定义多状态路径
self.sum_state_path = os.path.join(self.state_dir, "sum_state")
self.avg_state_path = os.path.join(self.state_dir, "avg_state")

# 在test_first中分别写入
sum_df.write.mode("overwrite").parquet(self.sum_state_path)
avg_df.write.mode("overwrite").parquet(self.avg_state_path)

# 在test_second中分别读取
prev_sum_df = self.spark.read.parquet(self.sum_state_path)
prev_avg_df = self.spark.read.parquet(self.avg_state_path)

跨批次状态存储的更优方案 #

根据你的实际场景复杂度，可选择以下方案：

• 列式存储（Parquet/ORC）：适合批量数据场景，读写效率高、压缩比大，是Spark的原生推荐格式。
• Delta Lake：支持ACID事务、版本控制和数据回溯，适合需要保证数据一致性、多批次并发读写的复杂场景。
• 结构化流状态管理：如果你的业务是流处理场景，Spark Structured Streaming内置了状态管理（支持update/complete模式），自动处理状态的持久化与恢复。
• 数据库存储：对于需要低延迟读写、频繁更新的状态，可选择HBase、Cassandra等NoSQL数据库，或MySQL等关系型数据库。

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