如何从PySpark DataFrame中高效选取唯一X-Y对？

需求说明 #

给定如下PySpark DataFrame：

X Y
1 a
1 b
1 c
2 b
2 a
2 c
3 a
3 c
3 b
4 p

需要生成一组X-Y配对，满足：

1. 每个X仅出现一次，每个Y仅出现一次
2. 配对必须是原始DataFrame中存在的组合
3. 结果可以是任意符合条件的配对（如示例1、示例2所示）

高效实现方案 #

针对大数据量场景，推荐使用基于随机排序的索引匹配法，所有操作均为分布式执行，避免低效的笛卡尔积或单机计算。

代码实现 #

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

# 构建示例DataFrame（实际使用时替换为你的数据源）
df = spark.createDataFrame(
    [(1, "a"), (1, "b"), (1, "c"),
     (2, "b"), (2, "a"), (2, "c"),
     (3, "a"), (3, "c"), (3, "b"),
     (4, "p")],
    ["X", "Y"]
)

# 步骤1：为每个唯一X分配随机排序后的索引
x_rank_df = df.select("X").distinct() \
    .orderBy(F.rand()) \
    .withColumn("x_idx", F.row_number().over(Window.orderBy(F.rand())))

# 步骤2：为每个唯一Y分配随机排序后的索引
y_rank_df = df.select("Y").distinct() \
    .orderBy(F.rand()) \
    .withColumn("y_idx", F.row_number().over(Window.orderBy(F.rand())))

# 步骤3：关联原始数据与索引表，筛选索引匹配的X-Y组合
result_df = df.join(x_rank_df, on="X", how="inner") \
    .join(y_rank_df, on="Y", how="inner") \
    .filter(F.col("x_idx") == F.col("y_idx")) \
    .select("X", "Y") \
    .distinct()

# 查看结果
result_df.show()

方案优势 #

• 分布式执行：所有操作基于Spark分布式引擎，支持TB级数据处理
• 低复杂度：无笛卡尔积操作，时间复杂度为O(n log n)（主要来自排序）
• 随机性保证：通过F.rand()实现随机排序，每次运行可生成不同的合法配对
• 结果合法性：仅保留原始DataFrame中存在的X-Y组合，同时保证X、Y的唯一性

备选方案（适用于重复率极低场景） #

如果X与Y的重复配对极少，可先为每个X随机选取一个Y，再过滤重复Y：

# 为每个X随机选一个Y
temp_df = df.withColumn("rn", F.row_number().over(Window.partitionBy("X").orderBy(F.rand()))) \
    .filter(F.col("rn") == 1) \
    .select("X", "Y")

# 过滤重复Y，确保每个Y仅出现一次
result_df = temp_df.withColumn("rn_y", F.row_number().over(Window.partitionBy("Y").orderBy(F.rand()))) \
    .filter(F.col("rn_y") == 1)

注：此方案可能会丢失部分X（当Y重复时），需额外补充缺失X的配对，适合重复率极低的场景。

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