PySpark DataFrame异常：重复ID与唯一ID计数波动问题

问题描述 #

现有包含ID和QUARTER字段的PySpark DataFrame，数据构造代码如下：

import pandas as pd
from pyspark.sql import SparkSession

# 构造原始数据
pandas_df = pd.DataFrame({
    "ID":[1, 2, 3,4, 5, 3,5,6,3,7,2,6,8,9,1,7,5,1,10],
    "QUARTER":[1, 1, 1, 1, 1,2,2,2,3,3,3,3,3,4,4,5,5,5,5]
})
spark = SparkSession.builder.appName('UniqueIDSelection').getOrCreate()
spark_df = spark.createDataFrame(pandas_df)
spark_df.createOrReplaceTempView('spark_df')

另有列表numbers=[2,1,3,1,2]，指定从季度1到季度5中分别选取的行数，要求最终结果中ID唯一，即某季度选中的ID不能在其他季度出现。

自行编写的代码存在两个问题：

1. 结果中出现跨季度的重复ID，不符合唯一性要求；
2. 执行print(df0.select('ID').distinct().count())时，每次得到的唯一ID计数结果不同，大数据集下差异更明显。

问题分析 #

1. 重复ID问题：未在每一步选取季度数据时排除已选中的ID，导致同一ID可能被多个季度选中。
2. 计数不稳定问题：
   • Spark是惰性求值框架，若未持久化中间结果，每次查询都会重新执行整个数据处理流程；
   • 若使用随机选取逻辑（如sample）未固定随机种子，会导致每次选取的结果不同；
   • union操作后的DataFrame未被持久化，多次查询会重复计算。

解决方案 #

核心思路：按季度顺序依次处理，每处理一个季度时，先过滤掉已选中的ID，再从当前季度剩余的ID中选取指定数量，同时固定随机种子保证结果可复现，最后持久化结果避免重复计算。

完整实现代码：

import pandas as pd
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, IntegerType
from pyspark.sql.functions import col

# 初始化SparkSession
spark = SparkSession.builder.appName('UniqueIDSelection').getOrCreate()

# 构造原始数据并统一字段类型
pandas_df = pd.DataFrame({
    "ID":[1, 2, 3,4, 5, 3,5,6,3,7,2,6,8,9,1,7,5,1,10],
    "QUARTER":[1, 1, 1, 1, 1,2,2,2,3,3,3,3,3,4,4,5,5,5,5]
})
spark_df = spark.createDataFrame(pandas_df)
spark_df = spark_df.withColumn("ID", col("ID").cast(IntegerType()))\
                   .withColumn("QUARTER", col("QUARTER").cast(IntegerType()))

# 定义每个季度要选取的行数（对应季度1-5）
numbers = [2,1,3,1,2]
# 初始化已选中的ID集合
selected_ids = spark.createDataFrame([], StructType([StructField("ID", IntegerType())]))
# 初始化结果DataFrame
result_df = spark.createDataFrame([], spark_df.schema)

# 按季度顺序处理
for quarter, num in enumerate(numbers, start=1):
    # 过滤当前季度未被选中的ID，先去重当前季度的ID避免重复选取
    available_data = spark_df.filter(
        (col("QUARTER") == quarter) & ~col("ID").isin([row.ID for row in selected_ids.collect()])
    ).distinct()
    
    # 选取指定数量的ID，用orderBy保证结果稳定；若需随机选取，替换为sample并固定seed
    if available_data.count() >= num:
        selected = available_data.orderBy(col("ID")).limit(num)
        # 更新已选中ID集合
        selected_ids = selected_ids.union(selected.select("ID"))
        # 更新结果集
        result_df = result_df.union(selected)

# 持久化结果集，避免重复计算
result_df.persist()

# 验证结果
print("最终结果：")
result_df.show()
print("唯一ID数量：", result_df.select("ID").distinct().count())
print("各季度选取行数：")
result_df.groupBy("QUARTER").count().show()

关键优化点 #

• 排除已选ID：每处理一个季度前，用~col("ID").isin(...)过滤掉已选中的ID，确保唯一性；
• 固定结果稳定性：用orderBy保证选取结果固定，若需随机选取，使用sample(withReplacement=False, fraction=..., seed=42)固定随机种子；
• 持久化结果：用persist()将最终结果存入内存/磁盘，避免多次查询重复计算；
• 类型一致性：统一ID和QUARTER字段类型，避免隐式转换导致的异常。

验证结果 #

执行上述代码后，会得到：

1. 无跨季度重复ID的结果集；
2. 每次执行result_df.select("ID").distinct().count()得到的数值一致。

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