请求将Python脚本转换为PySpark脚本，新增new_col字段

需求说明 #

需要将一段Python循环处理逻辑迁移到PySpark DataFrame上，核心规则：

• 当flag为'N'时，new_col设为id - 1
• 当flag为'Y'时，new_col沿用最近一次flag为'N'时计算出的值

原Python脚本 #

data = [(1, 'N'),
        (2, 'N'),
        (3, 'N'),
        (4, 'Y'),
        (5, 'Y'),
        (6, 'N'),
        (7, 'N'),
        (8, 'Y'),
        (9, 'Y'),
        (10, 'N')]
modified_data = []
new_col = 0  # Initialize new_col
for id_, flag in data:
    if flag == 'N':
        new_col = id_ - 1
    modified_data.append((id_, flag, new_col))
print(modified_data)

预期结果 #

[(1, 'N', 0), (2, 'N', 1), (3, 'N', 2), (4, 'Y', 2), (5, 'Y', 2), (6, 'N', 5), (7, 'N', 6), (8, 'Y', 6), (9, 'Y', 6), (10, 'N', 9)]

PySpark实现方案 #

核心思路 #

PySpark无法像Python循环那样逐行维护变量，需通过窗口函数+分组标记实现：

1. 先为flag='N'的行计算id-1作为候选值，其他行留空
2. 用累加标记将连续的Y行与最近的N行划分为同一组
3. 在分组内取最近的非空候选值填充所有行的new_col

代码实现 #

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

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

# 创建原始DataFrame
data = [(1, 'N'), (2, 'N'), (3, 'N'), (4, 'Y'), (5, 'Y'),
        (6, 'N'), (7, 'N'), (8, 'Y'), (9, 'Y'), (10, 'N')]
df = spark.createDataFrame(data, ["id", "flag"])

# 步骤1：生成new_col候选值，仅flag=N时赋值id-1
df = df.withColumn("temp_new_col", F.when(F.col("flag") == "N", F.col("id") - 1))

# 步骤2：创建分组标记——每遇到flag=N，分组ID加1
window_order = Window.orderBy("id")
df = df.withColumn("group_id", F.sum(F.when(F.col("flag") == "N", 1).otherwise(0)).over(window_order))

# 步骤3：在分组内取最近的非空候选值作为最终new_col
window_group = Window.partitionBy("group_id").orderBy("id")
df = df.withColumn("new_col", F.last("temp_new_col", ignorenulls=True).over(window_group))

# 查看结果
df.select("id", "flag", "new_col").show()

# 转换为列表验证（可选）
result = df.select("id", "flag", "new_col").collect()
print([(row.id, row.flag, row.new_col) for row in result])

代码解释 #

• temp_new_col：仅为需要更新的行生成目标值，避免无效计算
• group_id：通过累加flag='N'的次数，实现“最近一次N行+后续Y行”的分组逻辑
• last(..., ignorenulls=True)：在分组内向前追溯最近的有效候选值，完美匹配原逻辑中“沿用之前值”的需求

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