PySpark窗口函数致DataFrame写入S3卡顿的问题求助

我来帮你分析下这个问题，你遇到的情况其实在PySpark处理大表+窗口函数时很常见，咱们一步步拆解可能的原因和解决办法：

首先明确核心矛盾：直接写原始DataFrame到S3很快，但做完窗口聚合后写入就卡数小时，你试过的 repartition/coalesce、补缺失值、随机分区都没效果，那咱们从窗口函数的执行逻辑、Spark的计算流程、S3写入配置这几个方向入手。

1. 先排查窗口函数带来的数据倾斜/计算压力 #

你的窗口函数用了partitionBy("state")加rangeBetween(-100,0)，这两个点很容易出问题：

• 数据倾斜：如果某个state对应的行数特别多（比如占总数据的30%以上），那这个分区的任务会比其他任务大得多，直接拖慢整个窗口计算，后续写入也会因为这个分区的数据量过大而卡住。你可以先跑个统计看看state的分布：

  df.groupBy("state").count().orderBy(F.desc("count")).show(20)
  

  如果发现确实有倾斜的state，可以给这些热门键加个随机后缀拆分到多个分区，比如：

  from pyspark.sql.functions import when, rand, concat, lit, floor
  
  # 先统计每个state的行数，关联回原表
  state_count = df.groupBy("state").count()
  df_with_count = df.join(state_count, on="state", how="left")
  
  # 对行数超过阈值（比如10万）的state添加随机后缀，拆成10个子分区
  df_shuffled = df_with_count.withColumn(
      "state_shuffled",
      when(F.col("count") > 100000, concat(F.col("state"), lit("_"), floor(rand()*10).cast("string")))
      .otherwise(F.col("state"))
  )
  
  # 用新的shuffled键做窗口计算
  window1 = Window.partitionBy("state_shuffled")
      .orderBy(F.col("date").cast('timestamp').cast('long'))
      .rangeBetween(-100,0)
  df = df_shuffled.withColumn('new_variable', F.avg(F.col('binary_variable_01')).over(window1))
  
  # 把拆分的键合并回原state（如果业务需要保留原字段）
  df = df.withColumn("state", F.split(F.col("state_shuffled"), "_")[0]).drop("state_shuffled", "count")
  

• rangeBetween的计算成本：rangeBetween是基于列值的连续范围计算，比rowsBetween（基于行索引）的计算成本高很多。如果你的date字段是离散的（比如按天/小时），业务逻辑允许的话，改成rowsBetween(-100,0)会大幅降低计算压力。

2. 避免重复计算，利用缓存减少开销 #

你在窗口计算后执行了df.count()，这会触发一次全量计算；而后续的write操作又会重新触发一次计算——等于你把窗口计算跑了两遍！这完全是浪费资源，尤其是大表场景下。

解决办法是用persist()把窗口计算后的DataFrame缓存起来，这样count()触发计算后，后续的write直接用缓存的数据：

from pyspark import StorageLevel

df = df.withColumn('new_variable', F.avg(F.col('binary_variable_01')).over(window1))

# 缓存数据，选择MEMORY_AND_DISK级别（内存不够时写到磁盘）
df.persist(StorageLevel.MEMORY_AND_DISK)

df.count()  # 触发缓存加载
# 现在写入就会直接用缓存好的数据，不用重新计算窗口函数
df.write.format('parquet') \
    .option("escape", "\"")
    .option('header', 'true') \
    .mode('overwrite') \
    .save("s3://userName/folderName")

# 用完记得释放缓存，避免占用集群资源
df.unpersist()

3. 优化写入前的分区策略 #

窗口函数会打乱原始的分区分布，可能导致分区数据大小极不均匀（有的分区几MB，有的几GB），写入S3时小文件过多或者大文件处理缓慢。

• 如果不需要按state分区写入：根据你的数据量（900万行，假设每行几十KB，总数据量大概几十GB），设置合适的分区数，比如每个分区控制在1-2GB左右，所以可以设置50-100个分区：

  df = df.repartition(80)  # 可以根据实际数据量调整
  

• 如果需要按state分区写入：提前用repartition("state")把数据按state分区，这样写入时就不会再触发额外的shuffle：

  df = df.repartition("state")
  df.write.format('parquet') \
      .option("escape", "\"")
      .option('header', 'true') \
      .partitionBy("state")
      .mode('overwrite') \
      .save("s3://userName/folderName")
  

4. 优化S3写入的Spark配置 #

Spark写入S3的性能也受配置影响，你可以在初始化SparkSession时调整这些参数：

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("WriteToS3") \
    .config("spark.hadoop.fs.s3a.connection.maximum", "100")  # 增加S3连接数，提升并行写入能力
    .config("spark.hadoop.fs.s3a.fast.upload", "true")  # 开启S3快速上传，减少小文件开销
    .config("spark.sql.shuffle.partitions", "200")  # 调整shuffle分区数，避免分区过少导致倾斜
    .config("spark.sql.parquet.compression.codec", "snappy")  # 开启snappy压缩，减少写入的数据量
    .getOrCreate()

另外，写入时记得开启Parquet压缩，比如在write里加.option("compression", "snappy")，能大幅减少写入S3的数据量，提升速度。

5. 查看执行计划定位瓶颈 #

最后，你可以用explain(True)查看窗口函数和写入的执行计划，看看有没有不必要的shuffle、排序步骤，或者某个stage的任务数异常：

df.explain(True)

重点看窗口操作后的shuffle stage，以及写入前的分区情况，能帮你快速定位到底是计算阶段还是写入阶段卡住。

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