PySpark DataFrame索引优化、性能调优及数据profiling问题求助

• 数据集：约20000个小Parquet文件，总大小1.3GB（转CSV约15GB），包含1466764行、900列，来自AWS Appflow同步的Salesforce数据，无法合并源文件。
• 需求：读取为PySpark DataFrame后添加索引列，再执行数据profiling与质量检查。
• 索引测试结果：
  • monotonically_increasing_id+row_number方法：耗时0.104秒，但生成的mono_df分区数降至1。
  • zipWithIndex方法：耗时超4分钟（250.91秒），但zip_df保留原DataFrame的350个分区。
• Profiling性能差异：
  • mono_df单列平均耗时15秒。
  • zip_df单列平均耗时30分钟。
• 已尝试：调整Spark配置（如maxPartitionBytes），合并Parquet文件不可行。

索引实现代码 #

spark = SparkSession.builder.appName("Example").getOrCreate()
partition_size = spark.conf.get("spark.sql.files.maxPartitionBytes").replace("b","")
print(f"Partition Size: {int(partition_size) / 1024 / 1024} MB")

df_no_schema = spark.read.parquet('dbfs:parquet_folder/')
print(f"Number of Partition: {df_no_schema.rdd.getNumPartitions()}")
print(df_no_schema.count())

columns = df_no_schema.columns
row_with_index = Row(*columns, "index")

def create_new_schema(df_no_schema):
    new_schema = StructType(df_no_schema.schema.fields[:] + [StructField("index", LongType(), False)])
    return new_schema

def zip_rdd(df_no_schema, new_schema):
    zipped_rdd = df_no_schema.rdd.zipWithIndex()
    df = (zipped_rdd.map(lambda ri: row_with_index(*list(ri[0]) + [ri[1]])).toDF(new_schema))
    return df

def mono_id(df_no_schema):
    window_spec = Window().orderBy(F.monotonically_increasing_id())
    df = df_no_schema.withColumn("index", F.row_number().over(window_spec))
    return df

new_schema = create_new_schema(df_no_schema)
mono_df = mono_id(df_no_schema)
print(f"Number of Partition: {mono_df.rdd.getNumPartitions()}")
zip_df = zip_rdd(df_no_schema, new_schema)
print(f"Number of Partition: {zip_df.rdd.getNumPartitions()}")

数据Profiling代码 #

for column in df.columns:
    start_time = time.time()    
    unique_count = df.select(column).distinct().count()
    unique_percentage = (unique_count / total_rows) * 100
    duplicate_count = total_rows - unique_count
    duplicate_percentage = (duplicate_count / total_rows) * 100 if duplicate_count > 0 else 0
    null_count = df.filter(df[column].isNull()).count()
    null_percentage = (null_count / total_rows) * 100
    quality = int(unique_percentage / 10)
    quality_string = f"{quality}/10"

    unique_payload = {"count": unique_count, "percentage": unique_percentage}
    duplicate_payload = {"count": duplicate_count, "percentage": duplicate_percentage}
    quality_payload = {"count": "", "percentage": quality_string}
    null_payload = {"count": null_count, "percentage": null_percentage}

一、问题根源 #

1. 索引方法的性能差异 #

• mono_id分区骤降原因：row_number().over(Window.orderBy(...))会强制全局排序，默认将所有数据shuffle到单个分区（无业务字段可用于分区排序）。虽然生成索引快，但后续操作只能单分区串行执行。
• zipWithIndex耗时过长原因：该方法需要先全量扫描计算总行数，再为每个分区分配连续行号，会产生大量shuffle和数据传输开销。保留分区是因为zipWithIndex不改变RDD分区结构，但后续profiling慢是因为循环触发900次独立Job，每次都要重新扫描全量数据。

2. Profiling性能差异核心 #

• mono_df是单分区，单次单列计算无需跨节点shuffle，虽然无法并行，但避免了任务调度开销，所以单次耗时短。
• zip_df是350个分区，循环900列意味着触发900次全量扫描Job，叠加调度和数据传输开销，耗时呈指数级增长。

二、性能优化方案 #

1. 索引列优化：兼顾速度与分区 #

非连续全局索引（推荐）

不需要连续索引的话，直接使用monotonically_increasing_id()，不套row_number()，既保留原分区结构，又能快速生成唯一递增索引：

mono_df = df_no_schema.withColumn("index", F.monotonically_increasing_id())

连续全局索引

必须要连续索引时，采用分区内递增+分区偏移量的方式，避免全局shuffle：

from pyspark.sql.types import StructField, LongType

# 给每个分区分配ID
df_with_partition_id = df_no_schema.withColumn("partition_id", F.spark_partition_id())
# 计算每个分区的行数，生成偏移量字典
partition_counts = df_with_partition_id.groupBy("partition_id").count().orderBy("partition_id").collect()
offset_dict = {}
current_offset = 0
for row in partition_counts:
    offset_dict[row["partition_id"]] = current_offset
    current_offset += row["count"]
# 广播偏移量
broadcast_offset = spark.sparkContext.broadcast(offset_dict)
# 分区内生成连续索引，加上偏移量得到全局索引
def add_index(pid, iterator):
    offset = broadcast_offset.value[pid]
    for idx, row in enumerate(iterator, start=offset+1):
        yield Row(*row[:-1], idx)
# 转换为DataFrame
df_with_index = df_with_partition_id.rdd.mapPartitionsWithIndex(add_index).toDF(df_no_schema.schema.add(StructField("index", LongType(), False)))

2. Profiling性能优化 #

方案1：批量聚合计算所有列统计量

利用Spark的agg函数一次性计算所有列的所需指标，仅触发2次Job（一次count，一次agg）：

from pyspark.sql import functions as F

# 构建所有列的聚合表达式
agg_exprs = []
for col_name in df.columns:
    agg_exprs.append(F.countDistinct(col_name).alias(f"{col_name}_unique_count"))
    agg_exprs.append(F.count(F.when(F.col(col_name).isNull(), 1)).alias(f"{col_name}_null_count"))

# 一次性执行聚合
stats_df = df.agg(*agg_exprs)
stats_dict = stats_df.collect()[0].asDict()
total_rows = df.count()

# 生成最终profiling结果
profiling_results = {}
for col_name in df.columns:
    unique_count = stats_dict[f"{col_name}_unique_count"]
    null_count = stats_dict[f"{col_name}_null_count"]
    unique_percentage = (unique_count / total_rows) * 100
    duplicate_count = total_rows - unique_count
    duplicate_percentage = (duplicate_count / total_rows) * 100 if duplicate_count > 0 else 0
    null_percentage = (null_count / total_rows) * 100
    quality = int(unique_percentage / 10)
    
    profiling_results[col_name] = {
        "unique": {"count": unique_count, "percentage": unique_percentage},
        "duplicate": {"count": duplicate_count, "percentage": duplicate_percentage},
        "null": {"count": null_count, "percentage": null_percentage},
        "quality": f"{quality}/10"
    }

方案2：使用SQL批量查询

将DataFrame注册为临时视图，通过SQL一次性查询所有列统计量，效果与方案1一致：

df.createOrReplaceTempView("salesforce_data")
total_rows = df.count()

# 构建SQL查询语句
sql_select = []
for col_name in df.columns:
    sql_select.append(f"COUNT(DISTINCT {col_name}) AS {col_name}_unique_count")
    sql_select.append(f"COUNT(CASE WHEN {col_name} IS NULL THEN 1 END) AS {col_name}_null_count")

sql_query = f"SELECT {', '.join(sql_select)} FROM salesforce_data"
stats_df = spark.sql(sql_query)

# 后续处理同方案1

方案3：Databricks内置工具

直接使用Databricks提供的dbutils.data.summarize()，自动计算空值数、唯一值数等统计量，无需手动编码：

display(dbutils.data.summarize(df))

三、额外优化建议 #

• 小文件处理：读取时设置recursiveFileLookup为true，调整spark.sql.files.maxPartitionBytes为64MB或128MB，让Spark自动合并小文件为合理分区，减少任务调度开销。
• 缓存优化：执行Profiling前缓存DataFrame，避免重复扫描源文件：

df.cache()
df.count()  # 触发缓存加载

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