如何按行数拆分Spark Dataset并导出为XML文件？

你原代码存在几个核心问题：

• 反复调用count()和except()会触发多次Spark作业，性能极低，且except()会自动去重，若数据存在重复会导致丢失。
• 每次保存都用mode("overwrite")，后一次保存会覆盖前一次结果，最终只会剩下最后一个文件。

推荐用行号分组的方式实现拆分，既能保证效率，又能确保数据完整：

步骤说明 #

1. 给原数据集添加连续行号，用于分组拆分依据。
2. 根据行号计算分组ID（每500行划分为一组）。
3. 按分组ID遍历，将每组数据保存到独立路径，避免文件覆盖。

完整代码示例 #

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.functions;
import org.apache.spark.sql.Window;
import org.apache.spark.sql.WindowSpec;

// 读取DynamoDB数据
Dataset<Row> citations = sparkSession.read()
    .option("tableName", "Covid19Citation")
    .option("region", "eu-west-1")
    .format("dynamodb")
    .load();

// 添加全局连续行号（用窗口函数保证序号连续）
WindowSpec window = Window.orderBy(functions.monotonically_increasing_id());
Dataset<Row> numberedCitations = citations.withColumn("row_num", functions.row_number().over(window));

// 计算分组ID：每500行一组，row_num从1开始，所以减1后整除500
Dataset<Row> groupedCitations = numberedCitations.withColumn("group_id", functions.floor(functions.col("row_num").minus(1).divide(500)));

// 获取所有分组ID，遍历保存每个分组
long totalGroups = groupedCitations.select("group_id").distinct().count();
for (long i = 0; i < totalGroups; i++) {
    Dataset<Row> groupData = groupedCitations.filter(functions.col("group_id").equalTo(i))
        .drop("row_num", "group_id"); // 移除辅助列
    
    // 保存到独立子路径，避免覆盖
    groupData.coalesce(1)
        .write()
        .format("com.databricks.spark.xml")
        .option("rootTag", "citations")
        .option("rowTag", "citation")
        .mode("overwrite")
        .save("s3a://someoutputfolder/group_" + i + "/");
}

关键细节说明 #

• 用row_number()窗口函数生成连续行号，避免分布式环境下monotonically_increasing_id()生成的ID不连续导致分组不均。
• 每个分组保存到不同子路径（如group_0/、group_1/），彻底解决文件覆盖问题。
• coalesce(1)将每组数据合并到单个分区，确保每个分组只生成一个XML文件。

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