PySpark实现分组批量生成JSON文档写入NoSQL数据库

嘿，我来帮你搞定这个Spark分组分批写入NoSQL的难题——10亿行的时序数据确实得靠分布式框架才能撑住，之前用Pandas卡壳太正常了，毕竟单进程根本扛不住这么大的数据量。下面我给你一步步拆解解决方案：

核心思路 #

我们需要在Spark中实现：按唯一ID分组→组内按批次大小n拆分→每个批次转JSON→批量写入NoSQL。Spark的分布式特性决定了不能像Pandas那样直接在单分组里循环拆分，得用窗口函数+分组聚合的方式实现分布式分批。

具体实现步骤 #

1. 高效加载XML数据 #

首先用Spark的XML连接器加载数据，确保配置正确的行标签（rowTag），这是处理XML的关键：

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import col, row_number, floor, collect_list, to_json, struct

# 初始化SparkSession，记得根据集群资源调整配置
spark = SparkSession.builder \
    .appName("TimedDataBatchWriter") \
    .config("spark.executor.memory", "16g") \
    .config("spark.driver.memory", "8g") \
    .getOrCreate()

# 加载XML文件，替换成你的row标签和文件路径
df = spark.read.format("xml") \
    .option("rowTag", "your_record_tag") \
    .load("hdfs://path/to/large/xml/files")

2. 为每个分组分配批次编号 #

用窗口函数给每个唯一ID组内的行排序（时序数据必须按时间排序，保证批次顺序正确），然后计算每个行所属的批次ID：

batch_size = n  # 替换成你的目标批次大小n

# 定义窗口：按唯一ID分区，按时序字段排序（比如timestamp）
window_spec = Window.partitionBy("unique_id").orderBy("timestamp")

# 生成组内行号，再通过floor计算批次ID
df_with_batch = df.withColumn("row_num", row_number().over(window_spec)) \
                  .withColumn("batch_id", floor((col("row_num") - 1) / batch_size))

3. 按批次聚合并生成JSON文档 #

现在按(unique_id, batch_id)分组，把每个批次的行收集成列表，再转成结构化的JSON：

# 将批次内的所有字段打包成struct，再收集为列表，最后转成JSON
batch_json_df = df_with_batch.groupBy("unique_id", "batch_id") \
    .agg(collect_list(struct(*df.columns)).alias("batch_records")) \
    .withColumn("json_document", to_json(struct("unique_id", "batch_id", "batch_records")))

4. 批量写入NoSQL数据库 #

避免单条写入的性能损耗，用foreachBatch或foreachPartition实现批量写入。下面以MongoDB为例（其他NoSQL如Cassandra、Elasticsearch逻辑类似，只需替换客户端代码）：

def batch_write_to_nosql(batch_df, batch_idx):
    # 在每个批次中复用数据库连接，减少开销
    import pymongo
    from json import loads

    # 初始化MongoDB连接
    client = pymongo.MongoClient("mongodb://your-host:27017/")
    db = client["your_database"]
    collection = db["your_collection"]

    # 将Spark DataFrame的JSON字符串转成字典列表
    docs = batch_df.select("json_document") \
                   .rdd.map(lambda row: loads(row.json_document)) \
                   .collect()

    # 批量插入
    if docs:
        collection.insert_many(docs)

# 执行写入，记得设置checkpoint路径（用于容错）
batch_json_df.write \
    .foreachBatch(batch_write_to_nosql) \
    .option("checkpointLocation", "/hdfs/path/to/checkpoint") \
    .mode("append") \
    .save()

关键优化点 #

• 避免数据倾斜：如果存在热点ID（某个ID的行数特别多），可以给unique_id加盐（比如unique_id + rand()）后再分区，缓解单分区压力。
• 内存优化：调整Spark executor内存和分区数，确保每个分区的数据量在合理范围（建议每个分区1-2GB数据）。
• 连接复用：在foreachPartition或foreachBatch中创建数据库连接，不要在每条数据上创建，减少连接开销。
• 序列化优化：用Spark内置的to_json函数，比自定义UDF序列化更高效，避免性能瓶颈。

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