Spark高效读取大型Zstandard压缩文件的方案探讨

我在Databricks中用Spark读取30GB的Zstandard压缩PGN棋谱文件（.pgn.zst），当前用spark.read.text加载时，因为Zstandard文件不可拆分，整个文件被塞进单个分区，导致数据溢出；后续的pivot转换操作也因单分区运行，集群核心大量闲置，仅靠加大内存缓解并非高效方案。

方案1：预拆分Zstandard文件（推荐） #

Zstandard本身不支持分片读取，先通过外部工具拆分压缩文件，让Spark能并行处理：

• 使用zstd命令行工具拆分文件，先解压再按行数拆分后重新压缩：

zstd -d -c input.pgn.zst | split -l 100000 - --filter='zstd -z > $FILE.zst'

拆分后Spark可自动将多个.zst文件分配到不同分区，充分利用集群资源。

方案2：优化现有Spark代码逻辑 #

若无法预拆分文件，调整代码减少单分区压力：

2.1 替换全局窗口的GameID生成逻辑 #

当前用monotonically_increasing_id()全局排序生成GameID会强制单分区，改用分区内局部ID+全局偏移量的方式：

from pyspark.sql import Row
from pyspark.sql.functions import lit, when

def assign_local_game_ids(iterator):
    game_id = 0
    for row in iterator:
        if row.Line.startswith("[Event"):
            game_id +=1
        yield Row(Line=row.Line, LocalGameID=game_id)

# 分区内生成局部GameID
df_local = df.rdd.mapPartitions(assign_local_game_ids).toDF()

# 计算各分区的ID偏移量，生成全局GameID
partition_max_ids = df_local.groupBy().agg(max("LocalGameID")).collect()
offset = 0
for idx, max_id_row in enumerate(partition_max_ids):
    df_local = df_local.withColumn("PartitionID", lit(idx))
    df_local = df_local.withColumn("GameID", when(col("PartitionID") == idx, col("LocalGameID") + offset))
    offset += max_id_row[0]

2.2 替换低效的pivot操作 #

pivot在单分区运行效率极低，改用键值对收集+结构化解析的方式：

from pyspark.sql.functions import collect_list, struct, to_json, from_json, map_from_entries
from pyspark.sql.types import StructType, StructField, StringType

# 定义目标结构化Schema
target_schema = StructType([
    StructField("UTCDate", StringType()),
    StructField("Event", StringType()),
    StructField("TimeControl", StringType()),
    StructField("Opening", StringType()),
    StructField("Moves", StringType()),
    StructField("White", StringType()),
    StructField("Black", StringType())
    # 补充其他需要的字段
])

# 收集每个GameID的键值对，转JSON后解析为结构化数据
df_agg = df.groupBy("GameID").agg(collect_list(struct("Key", "Value")).alias("kv_pairs"))
df_json = df_agg.withColumn("json_data", to_json(map_from_entries(col("kv_pairs"))))
df_final = df_json.withColumn("game_data", from_json(col("json_data"), target_schema)).select("GameID", "game_data.*")

方案3：使用专用PGN预处理工具 #

若Spark不是必须选项，先用专用工具将PGN转换为Spark友好的格式：

• 使用pgn-extract命令行工具批量转换PGN为CSV：

zstd -d -c input.pgn.zst | pgn-extract -C -o output.csv

再用Spark读取预处理后的CSV文件，天然支持并行处理，无需处理单分区瓶颈。

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