如何在Spark DataFrame中添加频率列且不影响索引列与行顺序

实现前提 #

假设你的原始DataFrame命名为raw_df，包含唯一有序索引列id，以及其他用于判断记录重复的特征列。

完整实现代码 #

from pyspark.sql import Window
import pyspark.sql.functions as F

# ---------------------- 步骤1：统计重复记录频率 ----------------------
# 提取所有非id列作为重复判断的分组依据
feature_cols = [col for col in raw_df.columns if col != "id"]
# 定义窗口：仅按特征列分区，不指定排序避免额外开销、保留原始行序
freq_window = Window.partitionBy(*feature_cols)
# 新增Freq列：统计每个特征重复组的总出现次数
df_with_freq = raw_df.withColumn("Freq", F.count("*").over(freq_window))

# ---------------------- 步骤2：Freq列0-1归一化 ----------------------
# 分布式统计全量Freq的最大最小值，仅返回两个标量，无OOM风险
freq_stats = df_with_freq.agg(
    F.max("Freq").alias("max_freq"),
    F.min("Freq").alias("min_freq")
).collect()[0]
freq_max, freq_min = freq_stats["max_freq"], freq_stats["min_freq"]

# 处理边界情况：所有频率一致时分母为0，避免生成空值
if freq_max == freq_min:
    # 所有记录频率相同，归一化值可按需设为0或1
    df_final = df_with_freq.withColumn("Freq_norm", F.lit(0.0))
else:
    df_final = df_with_freq.withColumn(
        "Freq_norm",
        (F.col("Freq") - freq_min) / (freq_max - freq_min)
    )

方案说明 #

• 解决了窗口函数分组的矛盾：分组时排除唯一id列，避免因id唯一导致Freq全为1的问题；未引入全局排序/重分区操作，完全保留原始行顺序和id索引。
• 归一化阶段加入边界判断，不会出现空值。
• 整个流程无Driver端全量数据拉取操作，完全基于Spark分布式计算，支持TB级大数据量处理，不会触发OOM。

顺序一致性验证（可选） #

# 对比原始DF和最终DF的前N条id，确认行序未错乱
sample_size = 10
raw_ids = [row["id"] for row in raw_df.select("id").limit(sample_size).collect()]
final_ids = [row["id"] for row in df_final.select("id").limit(sample_size).collect()]
assert raw_ids == final_ids, "行顺序发生异常"

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