如何高效将Python生成的2.3亿条GRIB气象数据写入PySpark表？

一、避免Driver端全量加载，将计算推至Executor #

你之前循环创建DataFrame、重塑超大数组的操作，本质是把所有数据先加载到Driver内存，必然引发内存溢出。正确思路是让每个Executor并行处理单个时间步的网格数据：

• 将80个时间步拆分为RDD分区（比如sc.parallelize(range(80), numSlices=80)），每个分区对应一个时间步。
• 在每个分区的map操作中，直接在Executor端读取对应时间步的GRIB网格（1201x2400），生成结构化条目，而非将全量数据传回Driver。

示例代码：

def process_single_timestep(timestep_idx):
    # Executor端加载单时间步GRIB数据
    grib_data = load_grib_timestep(timestep_idx)  # 自定义加载函数，返回1201x2400网格
    structured_rows = []
    # 遍历网格生成结构化数据（后续可改用矢量化操作提速）
    for lat_idx in range(1201):
        lat_min = get_lat_min(lat_idx)
        lat_max = get_lat_max(lat_idx)
        for lon_idx in range(2400):
            lon_min = get_lon_min(lon_idx)
            lon_max = get_lon_max(lon_idx)
            structured_rows.append((grib_data[lat_idx][lon_idx], lat_min, lat_max, lon_min, lon_max))
    return structured_rows

# 并行处理所有时间步，每个时间步对应一个Executor任务
rdd = sc.parallelize(range(80), numSlices=80).flatMap(process_single_timestep)
df = rdd.toDF(["value", "lat_min", "lat_max", "lon_min", "lon_max"])

二、用矢量化操作替代Python循环，降低计算开销 #

Python循环处理网格效率极低，改用numpy矢量化批量生成结构化数据，减少解释器开销：

import numpy as np

def process_single_timestep(timestep_idx):
    grib_data = load_grib_timestep(timestep_idx)  # 返回numpy数组(1201,2400)
    # 批量生成lat/lon的min/max数组
    lat_mins = np.array([get_lat_min(i) for i in range(1201)])[:, np.newaxis]
    lat_maxs = np.array([get_lat_max(i) for i in range(1201)])[:, np.newaxis]
    lon_mins = np.array([get_lon_min(i) for i in range(2400)])[np.newaxis, :]
    lon_maxs = np.array([get_lon_max(i) for i in range(2400)])[np.newaxis, :]
    
    # 广播数组并扁平化生成全量条目
    values = grib_data.flatten()
    lat_min_flat = lat_mins.repeat(2400).flatten()
    lat_max_flat = lat_maxs.repeat(2400).flatten()
    lon_min_flat = lon_mins.repeat(1201, axis=0).flatten()
    lon_max_flat = lon_maxs.repeat(1201, axis=0).flatten()
    
    return list(zip(values, lat_min_flat, lat_max_flat, lon_min_flat, lon_max_flat))

三、优化Spark内存与分区配置 #

调整集群参数适配大数据量：

• 分区数：设置为Executor核心数的2-3倍（如100核心集群设200-300分区），让每个分区数据量维持在100万-1000万条区间。可通过rdd.repartition(240)调整。
• 内存配置：
  • 给Executor分配足够内存：--executor-memory 16G --driver-memory 8G
  • 配置堆外内存避免OOM：--conf spark.executor.memoryOverhead=4G
  • 启用Kryo序列化：--conf spark.serializer=org.apache.spark.serializer.KryoSerializer，大幅减少内存占用与序列化开销。

四、写入阶段优化：列式存储+合理分区/分桶 #

• 文件格式：优先选择Parquet或ORC列式存储，支持压缩与谓词下推。启用Snappy压缩：

  df.write.mode("overwrite").option("compression", "snappy").parquet("/path/to/output")
  

• 避免小文件：通过coalesce（无shuffle）或repartition（有shuffle）合并分区，比如合并至100个分区：df.coalesce(100).write.parquet(...)
• 按时间分区写入：若后续需按时间查询，可在结构化数据中加入timestep字段，再按分区写入：df.write.partitionBy("timestep").parquet(...)

五、减少不必要的数据转换 #

直接生成Row对象而非元组，减少DataFrame转换时的类型推断开销：

from pyspark.sql import Row

def process_single_timestep(timestep_idx):
    # ... 矢量化处理逻辑 ...
    return [Row(value=v, lat_min=lat_min, lat_max=lat_max, lon_min=lon_min, lon_max=lon_max) 
            for v, lat_min, lat_max, lon_min, lon_max in zip(values, lat_min_flat, lat_max_flat, lon_min_flat, lon_max_flat)]

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