如何在PySpark中基于其他列元素两两配对生成新列

针对大数据量级的PySpark DataFrame场景，提供两种实现方案，分别适配性能优先和写法简洁的需求：

方案1：内置高阶函数实现（无UDF，性能最优） #

纯原生Spark算子实现，没有跨进程序列化开销，适合超大规模数据集，输出结果和itertools.combinations完全一致：

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

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

# 构造示例测试数据
data = [
    (1, ["summer", "book", "hot"]),
    (2, ["g", "o", "p"])
]
df = spark.createDataFrame(data, schema=["id", "col"])

# 核心逻辑
# 1. 炸开数组同时保留每个元素的索引
df_explode = df.select("id", F.posexplode("col").alias("pos", "item"))

# 2. 自join筛选出索引递增的元素对，避免重复配对/自身配对
df_combine = df_explode.alias("a") \
    .join(df_explode.alias("b"), (F.col("a.id") == F.col("b.id")) & (F.col("a.pos") < F.col("b.pos"))) \
    .select("a.id", F.array("a.item", "b.item").alias("pair"))

# 3. 按主键分组，收集所有配对为新列，关联回原表
df_result = df_combine.groupBy("id").agg(F.collect_list("pair").alias("new_col"))
df_final = df.join(df_result, on="id", how="left")

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

输出效果：

+---+---------------------+------------------------------------------------+
|id |col                  |new_col                                         |
+---+---------------------+------------------------------------------------+
|1  |[summer, book, hot]  |[[summer, book], [summer, hot], [book, hot]]    |
|2  |[g, o, p]            |[[g, o], [g, p], [o, p]]                        |
+---+---------------------+------------------------------------------------+

方案2：UDF实现（写法简洁，适合中小规模数据集） #

和pandas写法逻辑一致，代码更短，适合数据量不大的场景快速开发：

import itertools
from pyspark.sql.types import ArrayType, StringType, StructType, StructField

# 定义UDF，封装itertools.combinations逻辑
combination_udf = F.udf(
    lambda x: list(itertools.combinations(x, 2)),
    ArrayType(StructType([
        StructField("_1", StringType(), nullable=False),
        StructField("_2", StringType(), nullable=False)
    ]))
)

# 直接生成新列
df_final = df.withColumn("new_col", combination_udf(F.col("col")))

注意：UDF存在Python和JVM进程的序列化开销，TB级以上超大数据集优先选择方案1。

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