如何加速PySpark单元测试？——Databricks环境下小数据集测试性能优化诉求

我之前也碰到过一模一样的问题——明明测试数据只有几行，但每个PySpark测试都要跑好几分钟，简直让人崩溃！后来摸索出几个关键优化点，把测试时间直接砍到了几十秒，分享给你试试：

你的fixture已经做了不少优化，但还能再精简：

• 把master("local[*]")改成master("local[1]")：local[*]会占用机器所有CPU核心，测试根本不需要多线程，反而会增加JVM初始化的线程开销
• 限制JVM内存：给driver和executor分配最小够用的内存（比如512M），减少内存分配和GC的耗时
• 加JVM串行GC参数：测试场景下不需要并发GC，串行GC启动更快、内存占用更低

优化后的Spark fixture示例：

@pytest.fixture(scope="session") 
def spark(): 
    """Ultra-lightweight session-scoped SparkSession optimized for tests."""
    import logging
    # 提前屏蔽py4j的冗余日志
    logger = logging.getLogger("py4j")
    logger.setLevel(logging.ERROR)
    
    spark = ( 
        SparkSession.builder.master("local[1]") 
        .appName("pytest-pyspark-fast") 
        .config("spark.sql.shuffle.partitions", "1") 
        .config("spark.default.parallelism", "1") 
        .config("spark.driver.extraJavaOptions", "-Djava.net.preferIPv4Stack=true -XX:+UseSerialGC -XX:InitialHeapSize=256m -XX:MaxHeapSize=512m") 
        .config("spark.ui.enabled", "false") 
        .config("spark.ui.showConsoleProgress", "false") 
        .config("spark.sql.ui.retainedExecutions", "0") 
        .config("spark.sql.catalogImplementation", "in-memory") 
        .config("spark.driver.host", "127.0.0.1") 
        .config("spark.driver.bindAddress", "127.0.0.1") 
        .config("spark.sql.adaptive.enabled", "false") 
        .config("spark.sql.execution.arrow.pyspark.enabled", "true") 
        .config("spark.driver.memory", "512m") 
        .config("spark.executor.memory", "512m") 
        .getOrCreate() 
    ) 
    # 把Spark日志级别降到ERROR，彻底屏蔽调试日志
    spark.sparkContext.setLogLevel("ERROR")
    yield spark 
    spark.stop()

如果多个测试都用到同一个小DataFrame，把它做成session级的fixture，只创建一次：

@pytest.fixture(scope="session")
def test_df(spark):
    """Reusable test DataFrame shared across all tests."""
    data = [("a", 1, 2.0, True, "2023-01-01"), ("b", 2, 3.0, False, "2023-01-02")]
    schema = ["col1", "col2", "col3", "col4", "col5"]
    df = spark.createDataFrame(data, schema=schema)
    # 如果需要临时表，在这里一次性创建
    df.createOrReplaceTempView("test_table")
    yield df

这样每个测试直接用这个test_df，不用重复执行createDataFrame，能省不少初始化时间。

别自己写collect()然后逐行比较数据，用专门的PySpark测试库（比如chispa）来做断言——它会用Spark原生操作比较DataFrame，不需要把数据拉到driver，速度快很多：

from chispa import assert_df_equals

def test_my_spark_logic(spark, test_df):
    # 调用你的业务逻辑
    result_df = my_spark_transformation(test_df)
    # 构造预期结果
    expected_data = [("a", 1, 2.0, True, "2023-01-01"), ("b", 2, 3.0, False, "2023-01-02")]
    expected_df = spark.createDataFrame(expected_data, schema=test_df.schema)
    # 高效断言
    assert_df_equals(result_df, expected_df)

• 避免对同一个DataFrame多次触发action（比如多次count()、show()），如果需要重复用，先cache()一下（虽然数据小，但积少成多）
• 不要在测试里创建不必要的临时表、UDF或者Spark配置，用完及时清理
• CI/CD环境里，尽量用预构建的Spark镜像，避免每次都下载依赖包

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