Spark过滤低频次值后DataFrame性能骤降的原因排查

你的核心问题是多次重复的Shuffle操作+冗余的模型拟合步骤导致Spark执行计划爆炸，后续操作性能急剧下降。以下是具体问题点和优化代码：

主要性能瓶颈 #

1. 循环Join导致重复Shuffle：对4列分别执行groupBy+join+filter，每一次Join都会触发一次Shuffle，4次操作后Shuffle开销指数级增长，且left join后再过滤的逻辑可以大幅简化。
2. 重复拟合模型：PCA和StringIndexer都循环调用fit(df)，导致Spark重复扫描全量数据多次，完全违背Pipeline的设计初衷。
3. 执行计划膨胀：多次变换后DataFrame的 lineage 过长，Spark优化器难以有效优化，计算效率被大幅拉低。

优化后的代码 #

第一步：优化低频次过滤逻辑（一次性完成，避免多次Join） #

把循环Join改成用窗口函数一次性计算所有目标列的频次，避免重复Shuffle：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

columns_to_check = ['is_zero', 'enc_loans', 'enc_limit', 'ratio_overdue']

# 用窗口函数批量计算列频次，替代多次groupBy+join
for col_name in columns_to_check:
    window_spec = Window.partitionBy(col_name)
    df = df.withColumn(f"{col_name}_count", F.count("*").over(window_spec))
    df = df.filter(F.col(f"{col_name}_count") >= 20).drop(f"{col_name}_count")

# 若目标列基数小，也可以用广播Join（效率更高）
# counts = {}
# for col_name in columns_to_check:
#     count_df = df.groupBy(col_name).agg(F.count("*").alias("count")).filter(F.col("count") >=20)
#     counts[col_name] = spark.sparkContext.broadcast(count_df.select(col_name).collect())
# for col_name in columns_to_check:
#     allowed_values = [row[col_name] for row in counts[col_name].value]
#     df = df.filter(F.col(col_name).isin(allowed_values))

第二步：统一用Pipeline处理所有特征变换（避免重复拟合） #

把VectorAssembler、PCA、StringIndexer全部整合到一个Pipeline中，只需一次拟合，大幅减少数据扫描次数：

from pyspark.ml.feature import VectorAssembler, PCA, StringIndexer
from pyspark.ml import Pipeline
from pyspark.ml.feature import SQLTransformer

# 定义特征分组与输出列映射（匹配原逻辑）
feature_mapping = [
    (['enc_paym'], 'enc_paym'),
    (['is_zero_loans'], 'is_zero'),
    (['pre_over', 'pre_maxover'], 'enc_limit'),
    (['enc_loans'], 'enc_loans'),
    (['pre_loans'], 'pre_loans'),
    (['is_zero'], 'ratio_overdue'),
    (['pre_since_opened', 'pre_since_confirmed'], 'confirm_days'),
    (['pre_pterm', 'pre_fterm'], 'open_days'),
    (['pre_till_pclose', 'pre_till_fclose'], 'collection_days')
]

# 构建Pipeline阶段列表
stages = []

# 1. VectorAssembler阶段
for input_patterns, output_col in feature_mapping:
    actual_inputs = [col for col in df.columns if any(pattern in col for pattern in input_patterns)]
    if actual_inputs:
        stages.append(VectorAssembler(inputCols=actual_inputs, outputCol=output_col))

# 2. PCA阶段
pca_cols = ['enc_paym', 'is_zero', 'enc_limit', 'enc_loans', 'pre_loans', 'ratio_overdue', 'confirm_days', 'open_days','collection_days']
for col in pca_cols:
    stages.append(PCA(k=1, inputCol=col, outputCol=f"{col}_encoded"))

# 3. StringIndexer阶段
lbe_cols = ['is_zero_encoded', 'ratio_overdue_encoded','enc_limit_encoded', 'enc_loans_encoded']
for col in lbe_cols:
    output_col = f"{col.split('_')[0]}_{col.split('_')[1]}"
    stages.append(StringIndexer(inputCol=col, outputCol=output_col))

# 4. 删除中间冗余列
drop_cols = pca_cols + lbe_cols
stages.append(SQLTransformer(statement=f"SELECT * EXCEPT ({','.join(drop_cols)}) FROM __THIS__"))

# 执行Pipeline
pipeline = Pipeline(stages=stages)
df = pipeline.fit(df).transform(df)

第三步：平衡数据分区 #

过滤后可能出现数据倾斜，执行以下命令重新平衡分区：

# 根据集群资源调整分区数（建议200-500之间）
df = df.repartition(300)
# 若存在明显倾斜列，可按该列分区
# df = df.repartition('your_skew_column', 300)

额外建议 #

• 执行df.explain()对比优化前后的执行计划，查看Shuffle次数和数据扫描量的变化
• 低频次过滤时，列基数大优先用窗口函数，基数小用广播Join更高效
• 所有特征变换尽量整合到一个Pipeline中，减少Spark重复计算开销

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