非Databricks环境下Spark SQL中Delta Identity Column功能失效求助

问题根源 #

开源Delta Lake社区版与Spark 3.4.2本身不支持GENERATED ALWAYS AS IDENTITY语法——该特性是Databricks Delta的专有扩展，并非Spark SQL标准或开源Delta的内置功能，因此直接编写会触发ParseException。

替代实现方案（贴近原生Identity体验） #

以下方案无需依赖Databricks，可实现自动递增、唯一主键的效果，模拟原生Identity Column的行为：

方案1：Delta预写入钩子（Pre-Write Hook）自动生成ID

通过注册Delta表的预写入钩子，在数据写入前自动计算并填充递增ID，无需手动干预：

1. 创建基础Delta表
   先定义不含自增逻辑的表结构：

spark.sql(f'''
    CREATE TABLE {database}.{table_name} (
        regionkey BIGINT, 
        regionname STRING
    )
    USING DELTA
    LOCATION '{target_csv_delta_path}'
''')

2. 注册预写入钩子
   编写Python钩子函数，自动为新数据分配递增ID：

from delta.tables import DeltaTable
from pyspark.sql.functions import max, lit, row_number
from pyspark.sql.window import Window

def auto_assign_identity(df, delta_table):
    # 获取当前表的最大ID，表为空时默认0
    current_max = delta_table.toDF().select(max("regionkey")).first()[0] or 0
    # 为新数据生成连续递增ID
    window = Window.orderBy(lit(1))
    return df.withColumn("regionkey", row_number().over(window) + current_max)

# 绑定钩子到目标Delta表
delta_table = DeltaTable.forPath(spark, target_csv_delta_path)
delta_table.registerPreWriteHook(auto_assign_identity)

3. 写入数据
   写入时无需指定regionkey，钩子会自动填充：

# 示例：写入新数据
new_regions = spark.createDataFrame([("North America",), ("Africa",)], ["regionname"])
new_regions.write.mode("append").save(target_csv_delta_path)

方案2：序列表+原子更新实现全局唯一ID

通过维护一个独立的序列表记录当前ID值，利用原子更新确保ID唯一且连续：

1. 创建序列表
   用于存储下一个待分配的ID：

CREATE TABLE {database}.region_id_seq (next_id BIGINT) 
USING DELTA 
LOCATION '{sequence_table_path}'

2. 初始化序列值
   设置初始ID为1：

INSERT INTO {database}.region_id_seq VALUES (1)

3. 写入数据时分配ID
   通过原子操作获取并更新序列值，为新数据分配连续ID：

def get_next_id_batch(record_count):
    # 原子操作获取ID范围，避免并发冲突
    with spark._jvm.scala.concurrent.Lock():
        current_id = spark.sql(f"SELECT next_id FROM {database}.region_id_seq").first()[0]
        new_next_id = current_id + record_count
        spark.sql(f"UPDATE {database}.region_id_seq SET next_id = {new_next_id}")
        return range(current_id, new_next_id)

# 示例写入流程
new_regions = spark.createDataFrame([("Europe",), ("Asia",)], ["regionname"])
record_count = new_regions.count()
id_range = get_next_id_batch(record_count)

# 为数据绑定ID
regions_with_id = new_regions.rdd.zipWithIndex()\
    .map(lambda x: (list(id_range)[x[1]], x[0][0]))\
    .toDF(["regionkey", "regionname"])

regions_with_id.write.mode("append").save(target_csv_delta_path)

说明 #

• 两种方案均实现了类似原生Identity Column的自动递增、唯一特性，且完全基于开源Spark与DeltaLake，无需依赖Databricks。
• 方案1更贴近原生体验，无需额外维护序列表；方案2适合高并发写入场景，原子更新可避免ID冲突。

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