PySpark中rangeBetween滚动窗口计算的数据倾斜优化咨询

一、加盐分区拆分大分区 #

针对倾斜mid，通过添加随机后缀拆分大分区，避免单分区数据量过大拖慢计算：

1. 先识别倾斜mid（可选，也可直接对所有mid加盐，或仅对高频mid加盐）：

# 统计各mid的数据量，筛选出倾斜的mid（阈值可根据实际情况调整）
mid_count_df = filtereddatadf.groupBy("mid").count()
skewed_mids = mid_count_df.filter(col("count") > 100000).select("mid").rdd.flatMap(lambda x: x).collect()
skewed_mids_broadcast = spark.sparkContext.broadcast(set(skewed_mids))

2. 给倾斜mid添加随机盐值，重新定义分组键：

def get_window_spec(partitioncol, ordercol, start, end):
    return Window.partitionBy(partitioncol).orderBy(ordercol).rangeBetween(start, end)

# 添加盐值列，将倾斜mid拆分为10个子分区（数量可调整）
filtereddatadf = filtereddatadf.withColumn(
    "mid_salt",
    when(col("mid").isin(skewed_mids_broadcast.value), 
         concat(col("mid"), lit("_"), floor(rand() * 10)))
    .otherwise(col("mid"))
)

group_cols = ["mid_salt"]

# 预聚合+窗口计算
midstatsdf_salted = filtereddatadf.groupBy(group_cols + ["epochtime"])\
    .agg(count("*").alias("mid_epoch_lvl_cnt"), 
         sum("amount").alias("mid_epoch_lvl_sum"))\
    .withColumn("mid_occur_num_1day", 
                F.sum(col("mid_epoch_lvl_cnt")).over(get_window_spec(group_cols,"epochtime",-24*60*60, 0))) \
    .withColumn("mid_sum_1day", 
                F.sum(col("mid_epoch_lvl_sum")).over(get_window_spec(group_cols,"epochtime",-24*60*60, 0)))

# 合并盐值分区的结果
midstatsdf = midstatsdf_salted.withColumn("mid", split(col("mid_salt"), "_")[0])\
    .groupBy("mid", "epochtime")\
    .agg(sum("mid_occur_num_1day").alias("mid_occur_num_1day"),
         sum("mid_sum_1day").alias("mid_sum_1day"))\
    .withColumn("mid_avg_1day", col("mid_sum_1day")/col("mid_occur_num_1day"))

二、用预排序+区间Join替代rangeBetween窗口 #

rangeBetween在大分区下的shuffle和计算开销极高，换用Join方式实现滚动窗口：

1. 先对每个mid的epochtime预聚合并排序：

# 预聚合得到mid-epochtime级别的统计
pre_agg_df = filtereddatadf.groupBy("mid", "epochtime")\
    .agg(count("*").alias("cnt"), sum("amount").alias("sum_amt"))\
    .orderBy("mid", "epochtime")

2. 自Join并过滤时间窗口：

# 自Join，关联同一mid下时间在[-86400, 0]范围内的记录
windowed_df = pre_agg_df.alias("a")\
    .join(pre_agg_df.alias("b"), 
          (col("a.mid") == col("b.mid")) & 
          (col("b.epochtime") >= col("a.epochtime") - 24*60*60) & 
          (col("b.epochtime") <= col("a.epochtime")),
          "left")\
    .groupBy("a.mid", "a.epochtime")\
    .agg(sum("b.cnt").alias("mid_occur_num_1day"),
         sum("b.sum_amt").alias("mid_sum_1day"))\
    .withColumn("mid_avg_1day", col("mid_sum_1day")/col("mid_occur_num_1day"))

注：该方式适合时间粒度较粗的场景，若epochtime粒度极细需结合加盐进一步优化。

三、单独处理极端倾斜mid #

如果少数mid占了绝大多数数据，可单独提取这些mid做分片计算：

1. 拆分倾斜mid和正常mid的数据：

# 拆分数据集
skewed_df = filtereddatadf.filter(col("mid").isin(skewed_mids_broadcast.value))
normal_df = filtereddatadf.filter(~col("mid").isin(skewed_mids_broadcast.value))

2. 对倾斜mid按时间分片计算滚动窗口：

# 给倾斜数据添加小时分片列，简化跨窗口计算逻辑
skewed_df = skewed_df.withColumn("hour_slot", floor(col("epochtime") / 3600))

# 倾斜数据单独计算窗口
skewed_windowed_df = skewed_df.groupBy("mid", "epochtime")\
    .agg(count("*").alias("cnt"), sum("amount").alias("sum_amt"))\
    .withColumn("mid_occur_num_1day", 
                F.sum("cnt").over(Window.partitionBy("mid").orderBy("epochtime").rangeBetween(-24*60*60, 0)))\
    .withColumn("mid_sum_1day", 
                F.sum("sum_amt").over(Window.partitionBy("mid").orderBy("epochtime").rangeBetween(-24*60*60, 0)))\
    .withColumn("mid_avg_1day", col("mid_sum_1day")/col("mid_occur_num_1day"))

# 正常数据用原逻辑计算
normal_windowed_df = normal_df.groupBy("mid", "epochtime")\
    .agg(count("*").alias("mid_epoch_lvl_cnt"), sum("amount").alias("mid_epoch_lvl_sum"))\
    .withColumn("mid_occur_num_1day", F.sum("mid_epoch_lvl_cnt").over(get_window_spec(["mid"],"epochtime",-24*60*60, 0)))\
    .withColumn("mid_sum_1day", F.sum("mid_epoch_lvl_sum").over(get_window_spec(["mid"],"epochtime",-24*60*60, 0)))\
    .withColumn("mid_avg_1day", col("mid_sum_1day")/col("mid_occur_num_1day"))

# 合并结果
final_df = skewed_windowed_df.unionByName(normal_windowed_df)

四、Spark资源参数调优 #

• 调大shuffle并行度：spark.sql.shuffle.partitions = 2000（建议设为executor核数的2-3倍）
• 调大默认并行度：spark.default.parallelism = 2000
• 增加executor资源：比如--executor-memory 16G --executor-cores 8 --num-executors 30
• 开启自适应执行：spark.sql.adaptive.enabled = true，让Spark自动调整分区数和执行计划

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