PySpark本地模式处理RDD性能极差，求代码分析与优化建议

问题场景 #

本地模式下使用PySpark 3.4.0处理HDFS上的6个CSV文件（总大小789MB），核心逻辑为：读取CSV转为DataFrame，再转成RDD并repartition(40)，通过map调用自定义preprocess_row函数处理每行数据，最后转回DataFrame按subject_seg_id分区写入CSV。

实际性能表现极差：处理800MB数据耗时超半小时，服务器配备40核但完全没利用起来；测试5985条记录时任务耗时严重不均衡，最长任务耗时达38分钟，按此性能扩展到20000个文件完全不可接受。

性能瓶颈分析 #

1. DataFrame与RDD频繁转换的无效开销
   从DataFrame转RDD再转回DataFrame，直接绕过了Spark Catalyst优化器的所有优化逻辑，所有处理逻辑只能走低效的RDD原生算子，完全浪费了Spark的SQL优化能力。

2. Python UDF的序列化与GIL限制
   map调用的preprocess_row是纯Python函数，Spark需要在JVM和Python进程间反复进行数据序列化/反序列化（Py4J通信开销极大）；同时Python的GIL锁限制了单进程内的并行，每个任务只能占用一个CPU核心，40核服务器的资源完全没被充分利用。

3. 数据倾斜导致任务耗时不均

   • repartition(40)是随机分区，很可能造成部分分区数据量远大于其他分区，导致少数任务扛了大部分计算量。
   • 最终按subject_seg_id分区写入时，如果某些subject_seg_id对应的数据量极大，会直接引发写入阶段的任务倾斜，拖慢整体进度。

4. 单条记录处理的计算低效

   • inner_process和stats中大量使用numpy、pandas操作，但这些库在单条记录上的初始化成本极高（比如创建np.array、pd.Series的开销远大于计算本身）。
   • SNR_denoise中用pd.Series处理频率索引，完全可以用numpy的向量化操作替代，没必要引入pandas的额外开销。
   • 每条记录都计算time.time()并返回duration，这属于调试逻辑，线上运行时纯粹是额外的性能损耗。

5. 读取阶段的分区不合理
   仅6个文件的情况下，Spark默认的初始分区数会非常少（通常等于文件数），后续强制repartition(40)会触发全量shuffle，带来不必要的IO和计算开销。

优化建议 #

1. 放弃RDD，改用DataFrame + 矢量化Pandas UDF #

直接用Spark DataFrame API结合**Pandas UDF（矢量化UDF）**实现处理逻辑，既能利用Catalyst优化器，又能通过批量处理减少JVM与Python的交互次数，大幅降低序列化开销。

2. 优化自定义函数的计算逻辑 #

• 移除调试用的time.time()计算，砍掉不必要的开销。
• 用numpy完全替代pandas操作：比如SNR_denoise中用np.argmax找频率索引，避免创建pd.Series。
• 合并重复计算：比如避免重复创建np.array，尽量复用中间数据。

3. 解决数据倾斜问题 #

• 初始读取时直接设置合理的分区数，比如：

  df = spark.read.option("header", True).csv(source_dir+'/*.csv').repartition(40)
  

  推荐按数据量设置，每分区控制在100-200MB左右。
• 如果subject_seg_id存在严重倾斜，写入时可采取以下方案：
  • 对倾斜的subject_seg_id单独拆分处理，生成小文件。
  • 若查询模式允许，用bucketBy替代partitionBy，或结合两者使用分散数据压力。

4. 配置Spark优化参数 #

• 开启Arrow向量化加速：

  spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
  

• 调整Python worker资源：

  spark.conf.set("spark.executor.cores", "4")  # 每个executor分配多核心，配合Pandas UDF并行
  spark.conf.set("spark.python.worker.memory", "4g")  # 避免内存不足
  

5. 优化写入阶段 #

• 设置maxRecordsPerFile限制单个文件大小，避免生成超大文件：

  ndf.write.option("maxRecordsPerFile", 100000)
  

• 若业务允许，关闭每个分区的header写入（option("header", False)），减少重复写入的开销。

6. 替换低效的Python库调用 #

如果团队有Scala开发能力，可将核心计算逻辑（比如峰值检测、SNR计算）用Scala实现UDF，彻底避免Python的GIL和序列化开销；若没有Scala能力，尽量用Spark MLlib的原生函数替代自定义Python逻辑。

优化后的核心代码示例 #

import pandas as pd
import numpy as np
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import StructType, StructField, StringType, FloatType, IntegerType

# 定义输出Schema
output_schema = StructType([
    StructField("subject_seg_id", StringType()),
    StructField("window_id", StringType()),
    StructField("PPG_ws", StringType()),
    StructField("ABP_ws", StringType()),
    StructField("snr", FloatType()),
    StructField("sbp_idxs", StringType()),
    StructField("sbp_raw", StringType()),
    StructField("sbp_max", FloatType()),
    StructField("sbp_min", FloatType()),
    StructField("sbp_num", IntegerType()),
    StructField("sbp_mean", FloatType()),
    StructField("sbp_median", FloatType()),
    StructField("sbp_range", FloatType()),
    StructField("sbp_std", FloatType()),
    StructField("dbp_idxs", StringType()),
    StructField("dbp_raw", StringType()),
    StructField("dbp_max", FloatType()),
    StructField("dbp_min", FloatType()),
    StructField("dbp_num", IntegerType()),
    StructField("dbp_mean", FloatType()),
    StructField("dbp_median", FloatType()),
    StructField("dbp_range", FloatType()),
    StructField("dbp_std", FloatType()),
])

def stats(data, idxs):
    if data is None or len(data) == 0:
        return ('', '', -1.0, -1.0, 0, -1.0, -1.0, -1.0, -1.0)
    idxs_str = ','.join([str(d) for d in idxs])
    data_str = ','.join([str(d) for d in data])
    return (
        idxs_str, data_str, float(np.amax(data)), float(np.amin(data)),
        len(data), float(np.mean(data)), float(np.median(data)),
        float(np.amax(data) - np.amin(data)), float(np.std(data))
    )

@pandas_udf(output_schema)
def preprocess_batch(df: pd.DataFrame) -> pd.DataFrame:
    def process_single_row(row):
        # 处理subject_seg_id
        subject_seg_id = row['subject_seg_id']
        if '/' in subject_seg_id:
            parts = subject_seg_id.split('/')
            parts = [p for p in parts if '=' in p]
            if len(parts) == 1:
                subject_seg_id = parts[0].split('=')[-1]
        
        abp_ws = row["ABP_ws"]
        if pd.notna(abp_ws) and abp_ws.strip() != '':
            data_list = [float(d) for d in abp_ws.split('-') if d.strip() != '']
            if not data_list:
                snr = 0.0
                sbp_stats = ('', '', -1.0, -1.0, 0, -1.0, -1.0, -1.0, -1.0)
                dbp_stats = ('', '', -1.0, -1.0, 0, -1.0, -1.0, -1.0, -1.0)
            else:
                data_np = np.array(data_list)
                peak_idx, peaks = peak_detection(data_np)
                valley_idx, bottoms = peak_detection(-1.0 * data_np)
                if bottoms is not None:
                    bottoms = [-1.0 * d for d in bottoms]
                
                sbp_stats = stats(peaks, peak_idx)
                dbp_stats = stats(bottoms, valley_idx)
                
                # 优化SNR计算，用numpy替代pandas
                no_DC_data = data_np - np.mean(data_np)
                f, power_spectrum_data = signal.welch(no_DC_data, fs=125, window='hann', nfft=8192)
                lf_index = np.argmax(f >= 0.5)
                hf_index = np.argmax(f >= 5)
                lf_psd = np.sum(power_spectrum_data[:lf_index+1])
                best_psd = np.sum(power_spectrum_data[lf_index+1:hf_index+1])
                hf_psd = np.sum(power_spectrum_data[hf_index+1:int(8192/2)])
                snr = best_psd / (lf_psd + hf_psd) if (lf_psd + hf_psd) != 0 else 0.0
        else:
            snr = 0.0
            sbp_stats = ('', '', -1.0, -1.0, 0, -1.0, -1.0, -1.0, -1.0)
            dbp_stats = ('', '', -1.0, -1.0, 0, -1.0, -1.0, -1.0, -1.0)
        
        return pd.Series([
            subject_seg_id, row['window_id'], row['PPG_ws'], row['ABP_ws'], snr,
            *sbp_stats, *dbp_stats
        ])
    
    result = df.apply(process_single_row, axis=1)
    result.columns = output_schema.fieldNames()
    return result

# 读取数据并处理
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
df = spark.read.option("header", True).csv(source_dir+'/*.csv').repartition(40)
ndf = df.select(preprocess_batch(*df.columns))

# 写入优化
ndf.write\
    .partitionBy("subject_seg_id") \
    .mode('overwrite') \
    .option("encoding", "utf-8") \
    .option("header", True) \
    .option("maxRecordsPerFile", 100000) \
    .format("csv") \
    .save(target_dir)

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