如何加快 PySpark 单元测试？

2026-11-04数据小陈

结论

两行数据通常不足以解释每个测试要跑 3–5 分钟。先查清耗时发生在 Spark 启动、测试操作，还是 pytest fixture 的初始化和清理阶段。scope="session" 会让 fixture 在每个 pytest 进程中复用。如果每个测试都启动了 Spark，检查测试是否正确使用了 fixture，以及 CI 是否通过多个独立进程运行测试。

先拆分计时

用 pytest 找出最慢的测试，再分别记录 fixture 创建和测试中 Spark 操作的耗时：

pytest --durations=20

例如，在 fixture 中记录 SparkSession.builder...getOrCreate() 的耗时；在测试里分别记录 DataFrame 创建、转换和触发 action 的耗时。Spark 转换通常是惰性执行的，count()、collect()、show() 这类 action 才会触发计算。同一个测试反复调用 action，可能会重复执行工作。

确认 fixture 只创建一次

确认测试函数通过参数使用 spark fixture，并且 fixture 的 yield 和清理逻辑都缩进在函数体内。测试结束后再调用 spark.stop()：

@pytest.fixture(scope="session")
def spark():
    session = (
        SparkSession.builder
        .master("local[2]")
        .appName("pytest-pyspark")
        .config("spark.sql.shuffle.partitions", "1")
        .config("spark.default.parallelism", "1")
        .config("spark.ui.enabled", "false")
        .getOrCreate()
    )
    yield session
    session.stop()

local[2] 可以作为起点，但不一定适合所有环境。local[*] 会使用可用核心。在资源有限或 CI 并发较高时，可以比较 local[1]、local[2] 和当前配置的实际耗时。

pytest-xdist 等并行执行方式会启动多个 worker 进程。每个进程都有自己的 session-scoped fixture，可能各自启动 Spark，也可能争用 CPU 和内存。可以比较单进程与并行运行时的耗时和资源占用，再确定 CI 中的 worker 数量。

减少测试里的 Spark 工作

将不依赖 Spark 的逻辑拆成普通 Python 函数，用常规单元测试覆盖。只在需要验证 Spark 行为时启动 Spark。
创建 DataFrame 时提供明确 schema，避免依赖类型推断：
```
df = spark.createDataFrame(rows, schema)
```
检查测试中是否有多余的 count()、collect()、show() 或重复 action。尽量一次完成断言，避免为了查看结果多次触发计算。
对小型测试数据，尽量减少不必要的 shuffle、repartition 和复杂查询计划。spark.sql.shuffle.partitions=1 可以减少 SQL shuffle 分区，实际效果取决于测试执行的操作。
spark.sql.execution.arrow.pyspark.enabled=true 主要影响使用 Arrow 的相关路径。测试没有走这些路径时，不要期待它能普遍加速 Spark SQL 测试。

如果仍然很慢

如果计时发现单次 getOrCreate() 或某个 action 本身耗时很长，记录 Spark 和 Python 版本、运行命令、CI worker 数量，以及各阶段的具体耗时。再检查本机和 CI 的执行方式是否一致，以及是否存在资源争用或外部连接配置。现有信息不足以判断具体瓶颈，数据只有两行并不能确定问题所在。