如何编写PySpark或SQL查询实现条件计数生成ID列

PySpark 实现 #

核心逻辑 #

通过窗口函数对非空Data值进行分组，再在每个分组内计算行号：

1. 标记非空行并累加得到分组标识；
2. 在分组内用行号生成递增ID。

代码示例 #

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

# 初始化会话
spark = SparkSession.builder.appName("ResetIDGenerator").getOrCreate()

# 构造示例数据
sample_data = [("this",), (None,), ("that",), ("those",), (None,), (None,)]
df = spark.createDataFrame(sample_data, ["Data"])

# 生成分组ID：每遇到非空Data就开启新分组
window_group = Window.orderBy(F.monotonically_increasing_id())
df = df.withColumn("group_id", F.sum(F.when(F.col("Data").isNotNull(), 1).otherwise(0)).over(window_group))

# 分组内生成递增ID
window_row = Window.partitionBy("group_id").orderBy(F.monotonically_increasing_id())
df = df.withColumn("ID", F.row_number().over(window_row))

# 查看结果
df.select("Data", "ID").show()

SQL 实现 #

核心逻辑 #

用窗口函数的累加和进行分组，再在分组内生成行号：

1. 用CASE标记非空行，累加得到分组标识；
2. 分组内调用ROW_NUMBER()生成ID。

代码示例 #

假设目标表名为data_table，执行以下查询：

WITH grouped_rows AS (
    SELECT 
        Data,
        SUM(CASE WHEN Data IS NOT NULL THEN 1 ELSE 0 END) OVER (ORDER BY (SELECT NULL)) AS group_id
    FROM data_table
)
SELECT 
    Data,
    ROW_NUMBER() OVER (PARTITION BY group_id ORDER BY (SELECT NULL)) AS ID
FROM grouped_rows;

注意事项 #

• ORDER BY (SELECT NULL)用于保留原数据顺序，若表中有明确的排序字段（如自增ID、时间戳），建议替换为该字段以确保顺序稳定；
• 每次Data出现非空值时，group_id会递增，后续空行归为同一分组，ID从1开始依次递增。

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