如何在Spark中拆分对齐DataFrame并在Pandas/Numpy分批读取大Parquet

1. 多DataFrame对齐拆分方案（Spark侧） #

你可以完全避免使用randomSplit，采用以下两种方案保证多DataFrame拆分后仍行对齐：

• 方案1：通过控制输出分片大小直接拆分
  Spark原生支持控制单Parquet文件的最大记录数，不会修改数据原有的行顺序，所有行对齐的DataFrame采用同一配置输出即可得到分片一一对应的小文件，完全保证行对齐。
  示例代码：

# 配置单文件最大输出记录数，可根据单行大小调整数值，最终得到的单文件大小约为6GB即可
spark.conf.set("spark.sql.files.maxRecordsPerFile", 1000000)
# 正常写入即可，输出会自动拆分为多个大小均匀的小文件
FULL_MAIN.write.mode("overwrite").parquet(PATH+"/FULL_MAIN.parquet")

• 方案2：新增行号按范围拆分
  如果需要显式控制拆分逻辑，可以给所有行对齐的DataFrame新增全局唯一的连续行号，按行号范围拆分即可：

from pyspark.sql.functions import row_number, monotonically_increasing_id
from pyspark.sql.window import Window

# 按固定规则生成行号，只要所有DataFrame使用相同的排序规则，生成的行号完全对齐
w = Window.orderBy(monotonically_increasing_id())
FULL_MAIN = FULL_MAIN.withColumn("row_id", row_number().over(w) - 1)
# 按行号范围拆分，比如总共有100万行，拆10份的话每份10万行，所有DataFrame采用同一拆分规则即可保持对齐
split_ranges = [(i*100000, (i+1)*100000) for i in range(10)]
for idx, (start, end) in enumerate(split_ranges):
    split_df = FULL_MAIN.filter((FULL_MAIN.row_id >= start) & (FULL_MAIN.row_id < end)).drop("row_id")
    split_df.write.mode("overwrite").parquet(f"{PATH}/FULL_MAIN_RDD{idx+1}.parquet")

2. pandas/numpy分批读取Parquet方案 #

无需全量加载即可实现分批读取，推荐使用pyarrow的内置接口，实现逻辑如下：

import pyarrow.parquet as pq

# 打开Parquet文件，不加载全量数据
parquet_file = pq.ParquetFile("/path/to/FULL_MAIN.parquet")
# 按批次迭代读取，batch_size为每次读取的行数，可根据内存大小调整
for batch in parquet_file.iter_batches(batch_size=100000):
    # 转为pandas DataFrame处理
    df_batch = batch.to_pandas()
    # 此处插入你对当前批次的处理逻辑
    # 如果需要转numpy数组，直接调用df_batch.to_numpy()即可

如果不需要读取所有列，可以在iter_batches中指定columns参数，进一步降低内存占用：

for batch in parquet_file.iter_batches(batch_size=100000, columns=["col1", "col2", "col3"]):
    # 仅读取需要的列处理
    ...

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