Spark DataFrame替换skills-specialization后续skills值的问题求助

需求说明 #

将skills-specialization之后的所有连续skills替换为skills-specialization，直到遇到非skills的内容为止。

原始数据 #

现有代码 #

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

# Define the data
data = [("skills",), ("skills",), ("skills",), ("skills",), ("skills",), 
        ("skills",), ("skills",), ("skills",), ("skills-specialization",), 
        ("skills",), ("skills",), ("skills",), ("job profile",), ("job description",)]

# Create a DataFrame named 'df'
df = spark.createDataFrame(data, ["group"])

# Add a new column 'partition' indicating presence of "skills-specialization"
df = df.withColumn("partition", F.when(F.col("group") == "skills-specialization", 1).otherwise(0))

# Assign row numbers within each partition
windowSpec = Window.orderBy(F.monotonically_increasing_id())
df = df.withColumn("row_num", F.row_number().over(windowSpec))

# Create a new column 'replace' to mark rows where "skills" follows "skills-specialization"
df = df.withColumn("replace", F.when((F.col("group") == "skills") & (F.lag(F.col("group")).over(windowSpec) == "skills-specialization"), "skills-specialization").otherwise(F.col("group")))

# Update 'replace' column for all rows after the first "skills"
df = df.withColumn("replace", F.when(F.col("replace") == "skills-specialization", "skills-specialization").otherwise(F.col("replace")))

# Select the updated 'replace' column and drop temporary columns
df = df.select("replace").drop("partition", "row_num")

# Display the DataFrame
df.display()

当前运行结果 #

预期目标结果 #

解决方案 #

核心思路是通过累积标记识别需要替换的区间：从skills-specialization出现后，直到遇到非skills内容前，所有skills都替换为skills-specialization。

代码实现：

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

data = [("skills",), ("skills",), ("skills",), ("skills",), ("skills",), 
        ("skills",), ("skills",), ("skills",), ("skills-specialization",), 
        ("skills",), ("skills",), ("skills",), ("job profile",), ("job description",)]

df = spark.createDataFrame(data, ["group"])

# 定义窗口：按行顺序保留原始数据顺序
window_spec = Window.orderBy(F.monotonically_increasing_id())

# 1. 创建标记列：遇到skills-specialization则激活标记，遇到非skills且标记已激活时关闭标记，否则继承前一行标记
df = df.withColumn(
    "should_replace",
    F.coalesce(
        F.when(F.col("group") == "skills-specialization", 1),
        F.when((F.col("group") != "skills") & (F.lag("should_replace").over(window_spec) == 1), 0),
        F.lag("should_replace").over(window_spec)
    ).cast("int")
)

# 2. 根据标记替换数据：当标记激活且当前值是skills时，替换为skills-specialization，否则保留原值
df = df.withColumn(
    "group",
    F.when((F.col("should_replace") == 1) & (F.col("group") == "skills"), "skills-specialization")
    .otherwise(F.col("group"))
)

# 3. 清理临时列并展示结果
df = df.drop("should_replace")
df.display()

代码说明 #

• should_replace列：用lag函数实现累积标记，激活条件是遇到skills-specialization，关闭条件是遇到非skills的内容。
• 替换逻辑：仅当标记激活且当前值是skills时执行替换操作，其他情况保留原始值。

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