Spark DataFrame按记录数拆分文件时如何保证同一ID不跨文件？

要保证同一ID仅出现在单个文件中，同时满足单文件记录数上限，你需要先按ID分组并分配"文件组"，确保每个组的总记录数不超过设定上限，再按组分区写入文件。具体步骤如下：

实现步骤 #

1. 统计每个ID的记录数 #

先计算每个ID对应的条目数量，后续用来分配文件组：

from pyspark.sql.functions import count, lit, col

max_records = 4  # 单文件记录数上限

# 统计每个id的记录数
id_counts = df.groupBy("id").agg(count("*").alias("id_row_count"))

2. 手动分配文件组 #

遍历所有ID，根据每个ID的记录数，将ID分配到不同的文件组中，确保每个组的总记录数不超过max_records：

# 收集ID及其记录数到本地列表
id_list = id_counts.collect()

current_group = 0
current_total = 0
id_group_map = {}

for id_row in id_list:
    id_val = id_row["id"]
    row_count = id_row["id_row_count"]
    
    # 如果当前组加入该ID后超过上限，新建组
    if current_total + row_count > max_records:
        current_group += 1
        current_total = row_count
    else:
        current_total += row_count
    
    id_group_map[id_val] = current_group

3. 将文件组映射回原DataFrame #

把每个ID对应的文件组号添加到原DataFrame中：

df_with_group = df.withColumn("file_group", lit(id_group_map[col("id")]))

4. 按文件组分区写入 #

通过repartition按文件组分区，每个分区对应一个文件，确保同一ID只在一个文件中，且单文件记录数不超过上限：

# 按file_group分区写入，每个分区生成一个文件
df_with_group.repartition(col("file_group")).write.csv(path, header=True)

说明 #

• 原方法maxRecordsPerFile仅按写入时的记录流切割，不考虑ID的关联性，因此会出现同一ID跨文件的情况。
• 上述方法通过预分组确保同一ID被分配到同一个文件组，再按组写入，既满足单文件记录数上限，又保证ID不跨文件。

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