如何在PySpark中按Flag列规则拆分双1行并保留其他行

需求说明 #

当数据行的Flag1和Flag2均为1时，将该行拆分为两行，分别设置(Flag1=1, Flag2=0)和(Flag1=0, Flag2=1)；其余情况保留原行数据。

源数据 #

ID      Col1    Flag1   Flag2
1234    Value1  1       1
5678    Value2  0       1

期望结果 #

ID      Col1    Flag1   Flag2
1234    Value1  1       0
1234    Value1  0       1
5678    Value2  0       1

实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import when, array, explode, col, struct, lit

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

# 创建测试DataFrame
data = [
    (1234, "Value1", 1, 1),
    (5678, "Value2", 0, 1)
]
columns = ["ID", "Col1", "Flag1", "Flag2"]
df = spark.createDataFrame(data, columns)

# 处理逻辑：根据条件生成待展开的行数组
processed_df = df.withColumn(
    "split_rows",
    when(
        (col("Flag1") == 1) & (col("Flag2") == 1),
        array(
            struct(col("ID"), col("Col1"), lit(1).alias("Flag1"), lit(0).alias("Flag2")),
            struct(col("ID"), col("Col1"), lit(0).alias("Flag1"), lit(1).alias("Flag2"))
        )
    ).otherwise(
        array(struct(col("ID"), col("Col1"), col("Flag1"), col("Flag2")))
    )
)

# 展开数组并提取字段
final_df = processed_df.select(explode(col("split_rows")).alias("row")) \
    .select("row.ID", "row.Col1", "row.Flag1", "row.Flag2")

# 查看结果
final_df.show()

代码说明 #

• 用when函数判断Flag1和Flag2同时为1的触发条件
• 满足条件时生成包含两个目标行的struct数组，不满足条件时生成包含原行的数组
• 通过explode函数将数组展开为多行数据
• 最后提取struct中的字段，得到符合要求的最终结果

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