基于另一行的值筛选PySpark DataFrame行

首先还原示例中的PySpark DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.functions import collect_set, array_contains, col

spark = SparkSession.builder.appName("target_filter").getOrCreate()

data = [
    ("00111", "50012", "a"),
    ("00111", "10131", "a"),
    ("00111", "11001", "b"),
    ("50012", "11001", "x"),
    ("10131", "11001", "z"),
    ("11001", "50012", "q")
]

df = spark.createDataFrame(data, ["x", "y", "z"])

步骤1：提取每个唯一x对应的y值集合 #

通过分组聚合，得到每个x对应的所有关联y值：

x_y_mapping = df.groupBy("x").agg(collect_set("y").alias("y_values"))

步骤2：关联筛选目标行 #

将原DataFrame与聚合结果交叉关联，筛选出原DataFrame中x属于对应y值集合的行：

filtered_result = df.crossJoin(x_y_mapping)\
                    .filter(array_contains(col("y_values"), col("x")))\
                    .select(col("x").alias("origin_x"), col("x").alias("result_x"), "y", "z")

步骤3：按原始x分组输出 #

遍历每个唯一x值，打印对应的筛选结果：

unique_x_list = [row.x for row in x_y_mapping.select("x").collect()]

for x_val in unique_x_list:
    filtered_result.filter(col("origin_x") == x_val)\
                  .select("result_x", "y", "z")\
                  .show(truncate=False)
    print("\n")

运行后输出结果与你要求的完全一致：

+--------+------+---+
|result_x|y     |z  |
+--------+------+---+
|50012   |11001 |x  |
|10131   |11001 |z  |
|11001   |50012 |q  |
+--------+------+---+


+--------+------+---+
|result_x|y     |z  |
+--------+------+---+
|11001   |50012 |q  |
+--------+------+---+


+--------+------+---+
|result_x|y     |z  |
+--------+------+---+
|11001   |50012 |q  |
+--------+------+---+


+--------+------+---+
|result_x|y     |z  |
+--------+------+---+
|50012   |11001 |x  |
+--------+------+---+

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