PySpark cache/persist首次运行无性能提升？原因及优化方案问询

核心问题 #

使用PySpark优化模型运行，采用sparkdf.groupBy().applyInPandas()和cache方法，但发现cache仅在第二次运行代码时提升性能，首次运行无效果。希望了解原因，以及如何让cache在首次运行就带来性能提升（可接受前置步骤付出缓存成本）。本次性能测试目标为评分并保存评分数据集的耗时。

首次运行耗时 #

Processing partition 1 at 2025-03-06 11:30:20.263913-05:00
Partition 1 processed in 255.919597864151 seconds            ??? 猜测是加载到内存的过程？
Processing partition 2 at 2025-03-06 11:34:36.183714-05:00
Partition 2 processed in 28.266658306121826 seconds
Processing partition 3 at 2025-03-06 11:35:04.450641-05:00
Partition 3 processed in 19.673797845840454 seconds
...
...
Processing partition 12 at 2025-03-06 11:38:41.491604-05:00
Partition 12 processed in 18.98810386657715 seconds
All partitions processed in 520.2162759304047 seconds - at an average of 43.35135632753372 seconds per partition.

第二次运行耗时 #

Processing partition 1 at 2025-03-06 11:45:15.390380-05:00
Partition 1 processed in 2.5390563011169434 seconds
Processing partition 2 at 2025-03-06 11:45:17.929636-05:00
Partition 2 processed in 2.3451147079467773 seconds
Processing partition 3 at 2025-03-06 11:45:20.274882-05:00
Partition 3 processed in 2.273413896560669 seconds
...
...
Processing partition 12 at 2025-03-06 11:45:41.702098-05:00
Partition 12 processed in 2.234912157058716 seconds
All partitions processed in 28.547115802764893 seconds - at an average of 2.3789263168970742 seconds per partition.

代码示例 #

results_df = spark.createDataFrame(data = [], schema = results_schema)
results_df.write.partitionBy("scenario_name").mode("overwrite").option("overwriteSchema", "true").saveAsTable("credit_risk.test_schema.perf_output_table1")
partition_df.cache()

beg_time = time.time()
for n_part in range(1, max_partitions + 1):
    print(f"Processing partition {n_part} at {datetime.datetime.fromtimestamp(time.time(), pytz.timezone('US/Eastern'))}")
    loop_beg_time = time.time()
    partition_data = partition_df.filter((f.col("partition_num") == n_part)).drop('partition_num')
    ecl_result = partition_data.groupBy('scenario_name').applyInPandas(cashflow_loop, results_schema)
    ecl_result.cache()
    
    ecl_result.write.partitionBy("scenario_name").mode("append").saveAsTable("credit_risk.test_schema.perf_output_table1")

    print(f"Partition {n_part} processed in {time.time() - loop_beg_time} seconds")

end_time = time.time()
print(f"All partitions processed in {end_time - beg_time} seconds - at an average of {(end_time - beg_time) / max_partitions} seconds per partition.")

节点性能截图 #

其他思考 #

• 曾猜测是I/O受限，试图将数据全部存入内存提升性能，但查看性能指标后并非如此。在1节点和2节点集群上运行，耗时无明显变化，排除核心/并行化问题。
• 未直接对partition_num使用applyInPandas的原因：DEV环境小型节点上遇到多种异常（Python worker崩溃、OOM错误等），内部测试环境无此类问题，推测是进程内存占用与其他作业冲突导致任务失败。
  • 测试环境任务可完成，但耗时波动极大，在120秒到600秒以上。
  • 希望了解缓存数据集的具体工作机制，以及checkpoint是否适用。

问题解答 #

一、首次运行缓存无效果的原因 #

PySpark的cache()是惰性执行的：调用cache()仅会标记数据集需要被缓存，但并不会立即触发数据加载和缓存操作。只有当数据集被触发行动算子（如count()、write()、show()等）时，才会实际执行计算并将数据写入内存/磁盘缓存。

你的首次运行流程中：

1. 调用partition_df.cache()后，并未立即触发缓存，而是进入循环逐个过滤分区
2. 第一次过滤partition_df并执行applyInPandas+write时，才会实际计算partition_df的全量数据（或对应分区数据），同时完成缓存写入
3. 后续分区处理时，虽然可以复用部分缓存，但首次处理第一个分区时，缓存还未完成构建，因此耗时极长；而第二次运行时，partition_df的缓存已经存在，所有分区处理都直接读取缓存，性能大幅提升。

另外，你在循环内对ecl_result调用cache()但随后立即执行write()，这种缓存几乎无意义：因为write()是行动算子，执行后ecl_result的生命周期就结束了，后续不会再复用这个数据集，缓存资源白白浪费。

二、让缓存在首次运行就生效的方案 #

要让缓存在首次运行的业务计算前就完成构建，需要显式触发缓存的物化操作，具体步骤如下：

1. 提前物化partition_df的缓存

在计时开始前，调用行动算子触发partition_df的缓存写入，比如：

partition_df.cache()
# 触发缓存物化，count()是轻量级行动算子，不会额外产生大量计算开销
partition_df.count()
# 或者如果需要验证缓存状态，可以查看存储级别
print(f"缓存状态: {partition_df.storageLevel}")

这样在进入循环处理之前，partition_df已经完全加载到缓存中，首次运行的每个分区过滤都会直接读取缓存数据。

2. 优化循环内的缓存逻辑

移除循环内的ecl_result.cache()，因为该数据集仅用于一次写入，缓存没有复用价值，反而会占用内存资源影响其他操作。

3. 调整分区处理方式（可选）

如果集群资源允许，可以考虑将循环改为并行处理，但需要注意避免DEV环境的OOM问题。比如可以按partition_num进行分组后批量处理，或者使用foreachPartition替代循环，但需要结合你的业务逻辑调整。

三、缓存机制与Checkpoint的区别 #

• 缓存（Cache/Persist）：
  • 存储位置：默认内存，可配置为内存+磁盘、仅磁盘等（通过persist(StorageLevel)指定）
  • 生命周期：依赖Spark应用的生命周期，应用结束后缓存自动清除
  • 用途：适用于需要重复复用的中间数据集，加速多次计算
• Checkpoint：
  • 存储位置：写入到指定的持久化存储路径
  • 生命周期：独立于Spark应用，应用结束后数据依然存在
  • 用途：用于切断RDD/DataFrame的血缘关系，避免因血缘过长导致的调度开销，同时实现数据的持久化备份

针对你的场景：

• 如果只是为了加速同一次应用内的多次数据读取，缓存足够，不需要Checkpoint
• 如果需要在多个应用间复用数据，或者数据集的血缘关系非常复杂导致调度缓慢，可以考虑使用Checkpoint替代缓存，但需要额外的存储开销和写入时间

四、额外性能优化建议 #

1. 预分区调整：如果partition_num是自定义的分区字段，可以提前对partition_df按partition_num进行分区（repartition("partition_num")），这样过滤时可以直接定位到对应的分区，避免全表扫描
2. 内存配置优化：确保Spark的executor内存足够容纳partition_df的缓存数据，避免缓存溢出到磁盘导致性能下降（可通过spark.executor.memory参数调整）
3. Python Worker资源隔离：针对DEV环境的OOM问题，可以调整spark.python.worker.memory参数限制Python worker的内存占用，避免与其他作业冲突

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