如何用PySpark生成仅ID递增的扩展DataFrame记录

方法一：基于ID序列构造新数据 #

先读取原始ORC数据，提取基础字段值后，生成指定范围的ID序列，再将固定字段与ID序列组合成新DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.functions import lit, col

# 初始化SparkSession
spark = SparkSession.builder.appName("GenerateIncrementalRecords").getOrCreate()

# 读取原始ORC文件
df = spark.read.option("inferSchema", True).orc("abc/path/abc.snappy.orc")

# 获取原始记录的非ID字段值
base_fields = df.select([c for c in df.columns if c != "ID"]).first().asDict()
# 获取原始ID值
original_id = df.select("ID").first()[0]
# 定义要生成的记录总数（包含原始记录）
total_records = 6

# 生成递增ID序列
id_sequence = spark.range(original_id, original_id + total_records).withColumnRenamed("id", "ID")

# 组合固定字段与ID序列
new_df = id_sequence
for field_name, field_value in base_fields.items():
    new_df = new_df.withColumn(field_name, lit(field_value))

# 保持列顺序与原始DataFrame一致
new_df = new_df.select(df.columns)

# 查看结果
new_df.show()

方法二：通过交叉连接复制记录再更新ID #

先将原始记录复制指定次数，再通过行号计算递增ID：

from pyspark.sql import SparkSession
from pyspark.sql.functions import lit, row_number
from pyspark.sql.window import Window

spark = SparkSession.builder.appName("GenerateIncrementalRecords").getOrCreate()

# 读取原始ORC文件
df = spark.read.option("inferSchema", True).orc("abc/path/abc.snappy.orc")

# 定义要生成的记录总数
total_records = 6

# 交叉连接生成重复的基础记录
repeat_df = df.withColumn("temp_id", lit(1)) \
              .crossJoin(spark.range(total_records).withColumnRenamed("id", "temp_id"))

# 计算递增ID
window_spec = Window.orderBy("temp_id")
new_df = repeat_df.withColumn("ID", col("ID") + row_number().over(window_spec) - 1) \
                  .drop("temp_id")

# 查看结果
new_df.show()

两种方法都能实现仅ID递增、其余字段保持不变的需求，可根据实际场景选择：

• 方法一适合明确知道ID范围的场景，性能更优；
• 方法二适合需要基于原始记录批量复制的场景，逻辑更直观。

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