PySpark DataFrame两列数组生成笛卡尔积列（无UDF/拆分连接）

可以利用PySpark 3.0及以上版本支持的数组高阶函数，通过transform和flatten的组合直接生成目标列，完全不需要拆分DataFrame或自定义UDF。

代码实现 #

首先构造示例DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.functions import transform, flatten, struct, col

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

# 构造示例数据
data = [
    ([1, 2], [3, 4, 5]),
    ([1, 2], [7, 8])
]
df = spark.createDataFrame(data, ["a1", "a2"])

# 生成笛卡尔积新列a3
result_df = df.withColumn(
    "a3",
    flatten(
        transform(
            col("a1"), 
            lambda x: transform(col("a2"), lambda y: struct(x, y))
        )
    )
)

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

逻辑说明 #

• 内层transform：遍历a2数组的每个元素y，将其与a1数组的当前遍历元素x组合成结构体struct(x, y)，得到每个x对应的a2元素配对数组。
• 外层transform：遍历a1数组的每个元素x，执行上述内层操作，最终得到一个嵌套数组（每个元素是一组x与a2的配对集合）。
• flatten：将嵌套数组展平为一维数组，即两列数组的笛卡尔积集合。

输出结果 #

+------+---------+----------------------------------------------------+
|a1    |a2       |a3                                                  |
+------+---------+----------------------------------------------------+
|[1, 2]|[3, 4, 5]|[{1, 3}, {1, 4}, {1, 5}, {2, 3}, {2, 4}, {2, 5}]|
|[1, 2]|[7, 8]   |[{1, 7}, {1, 8}, {2, 7}, {2, 8}]                  |
+------+---------+----------------------------------------------------+

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