Databricks PySpark：基于双值共存性修改DataFrame列的方法

核心逻辑 #

仅当指定列中同时存在两个目标值时，将对应值的行的comment列标记为mark；若任一目标值不存在，则保持原DataFrame完全不变。全程使用PySpark原生API，避免转换为Pandas以支持分布式运行。

实现步骤与代码 #

1. 定义参数与创建示例DataFrame #

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

# Databricks环境可省略SparkSession初始化
spark = SparkSession.builder.appName("MarkRowsByDualValues").getOrCreate()

# 示例输入数据
data = [
    ('E1', 'A1',''), 
    ('E2', 'A2',''),
    ('F1', 'A3',''),
    ('F2', 'B1',''),
    ('F3', 'B2',''),
    ('G1', 'B3',''),
    ('G2', 'C1',''),
    ('G3', 'C2',''),
    ('G4', 'C3',''),
    ('H1', 'C4',''),
    ('H2', 'D1',''),
]
columns = ['old_comp_id', 'db_id', 'comment']
df = spark.createDataFrame(data, columns)

# 定义目标列与目标值
target_col = "old_comp_id"
target_vals = ["E1", "C1"]

2. 检查双值是否同时存在 #

通过聚合操作快速判断目标列中是否包含两个指定值，避免冗余的全表扫描：

# 收集目标列唯一值并验证双值存在性
value_check = df.select(
    F.array_contains(F.collect_set(target_col), target_vals[0]).alias("has_val1"),
    F.array_contains(F.collect_set(target_col), target_vals[1]).alias("has_val2")
).select((F.col("has_val1") & F.col("has_val2")).alias("has_both"))

# 获取布尔标记（仅执行一次聚合计算）
has_both_vals = value_check.collect()[0]["has_both"]

# 广播标记，让所有Worker节点高效获取，避免Shuffle开销
broadcast_flag = F.broadcast(F.lit(has_both_vals))

3. 标记符合条件的行 #

仅当双值均存在时，修改对应行的comment列：

result_df = df.withColumn(
    "comment",
    F.when(
        # 仅满足双值存在+当前行值为目标值时标记
        broadcast_flag & F.col(target_col).isin(target_vals),
        F.lit("mark")
    ).otherwise(F.col("comment"))  # 其余情况保留原comment值
)

# 查看处理结果
result_df.show()

预期输出 #

执行上述代码后，E1和C1对应的行comment列会被标记为mark，其余行保持不变：

+-----------+-----+-------+
|old_comp_id|db_id|comment|
+-----------+-----+-------+
|         E1|   A1|   mark|
|         E2|   A2|       |
|         F1|   A3|       |
|         F2|   B1|       |
|         F3|   B2|       |
|         G1|   B3|       |
|         G2|   C1|   mark|
|         G3|   C2|       |
|         G4|   C3|       |
|         H1|   C4|       |
|         H2|   D1|       |
+-----------+-----+-------+

特殊场景处理 #

若目标列中不存在任一目标值（例如无C1），则has_both_vals为False，result_df与原DataFrame完全一致，不会进行任何修改。

内容的提问来源于stack exchange，提问作者Leonardo Kanashiro Felizardo