Spark中monotonically_increasing_id生成ID重复变化的问题咨询

问题背景 #

使用以下代码为包含完全重复行的DataFrame添加连续唯一ID，但后续使用该DataFrame时，ID列值会持续变化，无法固定：

def create_one_partition (df: DataFrame = None):
    if df:
        if df.rdd.getNumPartitions() > 1:
            df = df.coalesce(1)
    return df
    

def add_consecutive_unique_id_column(df: DataFrame, unique_id_col_name: str = "unique_id"):
    """
    Add a consecutive unique identifier column to a DataFrame.

    Args:
    - df: DataFrame to which the unique identifier column will be added
    - unique_id_col_name: Name of the column for the unique identifier (default: "unique_id")

    Returns:
    - DataFrame: DataFrame with a consecutive unique identifier column
    """
    # If the temporary unique_id_tmp column doesn't exist, create it
    if "unique_id_tmp" not in df.columns:
        df = create_one_partition(df)
        df = df.withColumn("unique_id_tmp", monotonically_increasing_id())

    # Use row_number to generate consecutive unique IDs based on the temporary unique_id_tmp column
    window_spec = Window.orderBy("unique_id_tmp")
    df = df.withColumn(unique_id_col_name, row_number().over(window_spec))

    # Drop the original unique_id_tmp column as it's no longer needed
    df = df.drop("unique_id_tmp")

    return df

环境：Azure Databricks（ADB）、DBR 14.3、Spark 3.5.0

原因分析 #

Spark采用惰性求值机制：DataFrame本质是执行计划而非实际数据，每次触发动作（如show()、write()、count()）时，都会从头执行整个计划。代码中monotonically_increasing_id()和row_number()都是动态计算逻辑，每次重新执行都会生成新的ID值。

解决方法 #

核心思路是固化生成ID后的DataFrame数据，避免重复执行ID生成逻辑。以下是几种可行方案：

1. 持久化DataFrame（Cache/Persist） #

在生成ID后，将DataFrame持久化到内存或磁盘，后续操作直接读取持久化的数据，不会重新计算ID：

修改原函数，在返回前添加持久化逻辑：

from pyspark.storagelevel import StorageLevel

def add_consecutive_unique_id_column(df: DataFrame, unique_id_col_name: str = "unique_id"):
    # 原有逻辑不变...
    
    # 持久化DataFrame，选择内存+磁盘的存储级别避免内存溢出
    df = df.persist(storageLevel=StorageLevel.MEMORY_AND_DISK)
    # 触发一次动作立即执行计算并完成持久化
    df.count()
    
    return df

注意：不需要该DataFrame时，调用df.unpersist()释放资源；数据量极大时，优先选择MEMORY_AND_DISK级别。

2. 写入临时视图/表 #

将生成ID后的DataFrame写入临时视图，后续从视图读取数据，确保ID固定：

# 生成带ID的DataFrame
df_with_id = add_consecutive_unique_id_column(original_df)
# 写入临时视图
df_with_id.createOrReplaceTempView("df_with_fixed_id")
# 后续使用时从视图读取
fixed_df = spark.sql("SELECT * FROM df_with_fixed_id")

3. 使用Checkpoint截断执行计划 #

Checkpoint会将DataFrame的计算结果写入磁盘，并截断原执行计划，后续操作基于固化的数据：

# 配置checkpoint目录（使用DBFS路径或提前创建的本地路径）
spark.sparkContext.setCheckpointDir("/dbfs/tmp/checkpoint")

# 生成带ID的DataFrame并执行checkpoint
df_with_id = add_consecutive_unique_id_column(original_df)
fixed_df = df_with_id.checkpoint()

Checkpoint会生成物理文件，需定期清理临时文件；生产环境建议使用持久化的存储路径。

优化原代码的小建议 #

原代码中coalesce(1)会将数据压缩到单个分区，若数据量过大会导致性能瓶颈。如果不需要严格连续的ID（仅需全局唯一），可以去掉coalesce(1)，直接用monotonically_increasing_id()生成唯一ID；若必须连续ID，可考虑先对数据进行稳定列排序，再在分区内生成row_number后合并全局ID，平衡性能与连续性需求。

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