PySpark中不移动行统计group=1的连续周期数

核心思路 #

要统计连续的group=1周期，关键是给每一段连续的1分配唯一的组ID，最后统计不同组ID的数量即可。这里用窗口函数实现，步骤如下：

实现代码 #

首先构造示例DataFrame（和你给出的示例一致）：

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

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

# 构造示例数据
data = [("A", 0), ("A", 0), ("A", 1), ("A", 1), ("A", 1),
        ("A", 0), ("A", 0), ("A", 1), ("A", 1), ("A", 0), ("A", 0)]
df = spark.createDataFrame(data, ["letter", "group"])

然后执行统计逻辑：

# 1. 添加行号，保证数据顺序（PySpark默认不保证行顺序，必须显式指定）
df_with_row = df.withColumn("row_num", F.row_number().over(Window.orderBy(F.monotonically_increasing_id())))

# 2. 计算当前行与前一行的group差异，标记新组的开始
window_spec = Window.orderBy("row_num")
df_with_diff = df_with_row.withColumn(
    "is_new_group",
    F.when(
        (F.col("group") == 1) & (F.lag("group").over(window_spec) != 1),
        1
    ).when(
        (F.col("group") == 1) & (F.lag("group").over(window_spec).isNull()),
        1  # 处理第一行就是1的情况
    ).otherwise(0)
)

# 3. 累加新组标记，得到每个连续1的组ID
df_with_group_id = df_with_diff.withColumn(
    "continuous_group_id",
    F.sum("is_new_group").over(window_spec.rangeBetween(Window.unboundedPreceding, 0))
)

# 4. 过滤出group=1的行，统计不同组ID的数量
result = df_with_group_id.filter(F.col("group") == 1).select(F.countDistinct("continuous_group_id").alias("continuous_1_groups_count"))

# 查看结果
result.show()

代码说明 #

• 添加行号：用monotonically_increasing_id()生成唯一排序键，再用row_number()生成连续行号，确保数据按你给出的顺序处理。
• 标记新组：用lag()函数获取前一行的group值，当当前行是1且前一行不是1（或者是第一行）时，标记为新组开始。
• 生成组ID：对新组标记做累加，每遇到一个新组开始，组ID就加1，这样连续的1会拥有同一个组ID。
• 统计数量：过滤出group=1的行，统计不同组ID的数量，就是你要的连续周期数。

运行后结果会显示continuous_1_groups_count为2，符合你的期望。

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