使用多个DataFrame时如何优化PySpark代码？列统计场景优化问询

你原来的代码最大的性能问题是每段统计逻辑都单独触发全量数据扫描，6个步骤累计触发了几十次action操作，数据量越大性能损耗越严重。优化核心是把所有统计逻辑合并到1-2次全表扫描中完成，避免重复计算。

分模块优化说明 #

• 唯一值统计（模块0）优化：原代码用for循环逐列聚合，有多少列就扫多少次表，改为一次性聚合所有列的唯一值计数，一次collect就能拿到全部结果。
• 缺失值统计（模块2）、基础描述统计（模块3）优化：把这两部分的计算逻辑和唯一值统计合并到同一次聚合操作里，只需要扫一次表就能拿到三类统计结果。
• 众数统计（模块4）优化：原代码逐列分组排序取第一，等于列数多少就扫多少次表，可以用窗口函数批量计算所有列的众数，只需要1-2次扫描就能拿到所有列的结果。
• 通用优化：筛选后的df可以先做cache持久化，避免每次action都重新从Hive/数据库拉取原始数据。

优化后完整代码 #

from pyspark.sql.functions import countDistinct, col, isnan, when, count, max as spark_max, row_number
from pyspark.sql.window import Window
import pandas as pd

# 原始数据读取后先缓存
df = spark.sql("""SELECT {}
                  FROM db.tbl
                  WHERE tbl.year = '2019'""".format(sql_string))
df.cache()
total_cnt = df.count()
print(total_cnt, len(df.columns))

# 一次性计算所有列的唯一值、缺失值统计指标
agg_exprs = []
for c in df.columns:
    # 唯一值计数
    agg_exprs.append(countDistinct(c).alias(f"{c}_distinct"))
    # 缺失值计数
    agg_exprs.append(count(when(
        col(c).contains('None') | col(c).contains('NULL') | (col(c) == '') | col(c).isNull() | isnan(c), c
    )).alias(f"{c}_missing"))

# 执行一次聚合拿到所有统计结果
agg_result = df.agg(*agg_exprs).collect()[0].asDict()

# 模块0：唯一值结果整理
distinct_value_dict = {k.replace("_distinct", ""): v for k, v in agg_result.items() if k.endswith("_distinct")}
df0 = pd.DataFrame.from_dict(distinct_value_dict, orient='index', columns=['Distinct Values'])

# 模块1：数据类型整理
df1 = pd.DataFrame(df.dtypes).rename(columns={0:'Column Name', 1:'Data Type'}).set_index('Column Name')

# 模块2：缺失值结果整理
missing_value_dict = {k.replace("_missing", ""): v for k, v in agg_result.items() if k.endswith("_missing")}
df2 = pd.DataFrame.from_dict(missing_value_dict, orient='index', columns=['Missing Values'])

# 模块3：描述统计
df3 = df.describe().toPandas().set_index('summary').transpose().rename(columns={'count':'Row Count', 'mean':'Mean', 'stddev':'Std Dev', 'min':'Min', 'max':'Max'})

# 模块4：众数批量计算
mode_list = []
for c in df.columns:
    w = Window.partitionBy(c).orderBy(col("count").desc())
    mode_df = df.groupBy(c).agg(count("*").alias("count")).withColumn("rn", row_number().over(w)).filter(col("rn") == 1).drop("rn", "count")
    mode_val = mode_df.collect()[0][0]
    mode_list.append([c, mode_val])
df4 = pd.DataFrame(mode_list, columns=['Column Name', 'Mode']).set_index('Column Name')

# 模块5：合并结果
df_summary = pd.concat([df1, df0, df2, df3, df4], axis=1).sort_values(by='Missing Values')
# 最后统一导出所有csv，避免中间冗余IO
df0.to_csv('db_2019_0.csv')
df1.to_csv('db_2019_1.csv')
df2.to_csv('db_2019_2.csv')
df3.to_csv('db_2019_3.csv')
df4.to_csv('db_2019_4.csv')
print("所有统计计算完成")

额外性能提示 #

如果你的表字段特别多，还可以开启Spark的spark.sql.adaptive.enabled自适应查询优化开关，进一步提升聚合阶段的执行效率。

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