Spark DataFrame：移除lst_second非lst_first元素并同步删对应位置

需求概述 #

给定分组聚合后的Spark DataFrame，需过滤lst_second中不存在于lst_first的元素，同时删除lst_first中对应位置的元素，保留位置匹配且符合条件的元素对。

解决方案 #

利用Spark内置的数组操作函数，通过「配对-过滤-拆分」的流程实现需求，具体步骤如下：

1. 导入依赖函数

from pyspark.sql import functions as F

2. 核心逻辑实现

# 将两个列表按位置配对成结构体数组
sdf3 = sdf2.withColumn(
    "pairs",
    F.array_zip(F.col("lst_first"), F.col("lst_second"))
)

# 过滤出lst_second元素存在于当前组lst_first中的元素对
sdf4 = sdf3.withColumn(
    "filtered_pairs",
    F.filter(
        F.col("pairs"),
        lambda x: F.array_contains(F.col("lst_first"), x["id_second"])
    )
)

# 拆分过滤后的元素对，重新生成目标列表
result_df = sdf4.select(
    "group",
    F.transform(F.col("filtered_pairs"), lambda x: x["id_first"]).alias("lst_first"),
    F.transform(F.col("filtered_pairs"), lambda x: x["id_second"]).alias("lst_second")
).sort("group")

# 查看结果
result_df.show(truncate=False)

最终结果 #

执行后输出符合需求的DataFrame：

+-----+------------------+------------------+
|group|lst_first         |lst_second        |
+-----+------------------+------------------+
|0    |[p4, p1, p2, p3]  |[p2, p1, p1, p1]  |
|1    |[p6, p9, p10, p8] |[p6, p6, p9, p9]  |
|2    |[p9]              |[p9]              |
+-----+------------------+------------------+

逻辑说明 #

• array_zip：按索引位置将两个列表配对，生成包含id_first和id_second的结构体数组；
• filter + array_contains：判断每个配对中的id_second是否存在于当前组的lst_first中，保留符合条件的元素对；
• transform：从过滤后的元素对数组中分别提取id_first和id_second，重新生成目标列表。

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