PySpark大数据集按主键分片后丢失RETAIL_SITE_ID问题求助

核心问题定位 #

1. 分片文件写入错误：循环中错误地将累积的Final_table写入当前segment文件，而非仅写入当前segment的处理结果segment_final_table。这会导致后续文件包含前置所有segment的数据，不仅引发大量重复，还可能因循环中断导致部分segment数据未正确持久化。
2. 分布式分片逻辑隐患：依赖monotonically_increasing_id()和row_number()进行分片，在Spark分布式环境下可能因shuffle、数据分区问题导致ID分配不均，甚至遗漏部分RETAIL_SITE_ID。
3. Inner Join的潜在过滤：若原数据集Mkt_all存在RETAIL_SITE_ID为null的记录，尽管unique_ids包含该值，但后续处理可能因业务逻辑过滤这类数据（若主键不应为null，需提前清理）。

修复后的完整代码实现 #

1. 优化分片逻辑（更可靠的Driver端分组） #

直接收集唯一ID到Driver，手动分组，避免分布式分片的不确定性：

import pandas as pd
import numpy as np
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StringType, IntegerType

# 初始化Spark Session（保留原有配置）
spark = SparkSession.builder \
    .appName("concatDataFrames") \
    .config("spark.jars.packages", "com.crealytics:spark-excel_2.12:<version>") \
    .config("spark.memory.fraction", "0.6") \
    .config("spark.memory.storageFraction", "0.5") \
    .config("spark.rdd.compress", "true") \
    .config("spark.memory.offHeap.enabled", "true") \
    .config("spark.memory.offHeap.size", "<offHeapSizeInMB>") \
    .config("spark.storage.level", "DISK_ONLY") \
    .config("spark.driver.memory", "16g") \
    .getOrCreate()

# 启用Arrow优化
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
spark.conf.set("spark.sql.execution.arrow.pyspark.fallback.enabled", "true")

# 读取Widget参数
dbutils.widgets.text("Market Code", "")
dbutils.widgets.text("Time Period", "")
dbutils.widgets.text("Time Block", "5")
Market_Code = dbutils.widgets.get("Market Code")
Time_Period = dbutils.widgets.get("Time Period")
Time_Block = int(dbutils.widgets.get("Time Block"))

##### Mkt_all 数据集初始化逻辑保留 ##############

# 提前过滤无效主键（若主键不应为null）
Mkt_all = Mkt_all.filter(F.col("RETAIL_SITE_ID").isNotNull())

# 收集所有唯一RETAIL_SITE_ID到Driver（确保ID数量在Driver内存承受范围内）
unique_ids_list = [row.RETAIL_SITE_ID for row in Mkt_all.select("RETAIL_SITE_ID").distinct().collect()]

# 按每125个ID分组
id_groups = [unique_ids_list[i:i+125] for i in range(0, len(unique_ids_list), 125)]
num_segments = len(id_groups)

2. 修正分片处理与文件写入 #

循环处理每个ID组，仅写入当前组的处理结果：

# 处理每个分片
for idx, id_group in enumerate(id_groups):
    # 过滤当前组的记录
    segment_df = Mkt_all.filter(F.col("RETAIL_SITE_ID").isin(id_group))
    # 转换为Pandas DataFrame
    segment_pd_df = segment_df.toPandas()
    
    ##### 此处保留你的Pandas处理逻辑，生成segment_final_table ######
    
    # 仅写入当前segment的处理结果，避免累积写入
    output_path = f"/dbfs/mnt/RAAP_ADLS2/CRM_Network/Forecourt Reconfiguration/Output 2023/New_Trial/{Market_Code}_Pump_Estimation_{idx+1}{Time_Period}.csv"
    segment_final_table.to_csv(output_path, sep=',', index=False)

3. 合并与去重（优化去重逻辑） #

# 合并所有分片文件
final_dfs = []
for idx in range(num_segments):
    segment_path = f"/dbfs/mnt/RAAP_ADLS2/CRM_Network/Forecourt Reconfiguration/Output 2023/New_Trial/{Market_Code}_Pump_Estimation_{idx+1}{Time_Period}.csv"
    segment_df = pd.read_csv(segment_path)
    final_dfs.append(segment_df)

concatenated_Final_Table = pd.concat(final_dfs, ignore_index=True)
# 按主键去重（而非全列，避免因处理后的细微差异丢失有效数据）
concatenated_Final_Table = concatenated_Final_Table.drop_duplicates(subset=["RETAIL_SITE_ID"], keep='first')
# 写入最终文件
concatenated_Final_Table.to_csv(f"/dbfs/mnt/RAAP_ADLS2/CRM_Network/Forecourt Reconfiguration/Output 2023/New_Trial/COUNTRY/{Market_Code}_Pump_Estimation{Time_Period}.csv", sep=',', index=False)

额外建议 #

• 优先使用Spark原生处理：若业务逻辑可通过Spark API实现，建议避免频繁转换为Pandas，提升处理效率并降低数据丢失风险。
• 验证分片完整性：处理前可对比原数据集与分片后的数据量：

  # 原数据集记录数
  original_count = Mkt_all.count()
  # 分片总记录数
  segment_total = sum([Mkt_all.filter(F.col("RETAIL_SITE_ID").isin(group)).count() for group in id_groups])
  print(f"原记录数：{original_count}，分片总记录数：{segment_total}")
  

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