如何根据pyspark数据帧中多列的笛卡尔乘积创建新列

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, explode
from pyspark.sql.types import ArrayType, StructType, StructField, StringType

spark = SparkSession.builder.appName("Cartesian Product").getOrCreate()

data = [("A", [1, 2, 3]), ("B", [4, 5]), ("C", [6])]
df = spark.createDataFrame(data, ["col1", "col2"])
df.show()

+----+---------+
|col1|     col2|
+----+---------+
|   A|[1, 2, 3]|
|   B|   [4, 5]|
|   C|      [6]|
+----+---------+

def cartesian_product(col1, col2):
    return [(c1, c2) for c1 in col1 for c2 in col2]
cartesian_product_udf = spark.udf.register("cartesian_product", cartesian_product, ArrayType(StructType([
    StructField("col1", StringType(), True),
    StructField("col2", StringType(), True)
])))

df.withColumn("cartesian_product", explode(cartesian_product_udf(col("col1"), col("col2")))).show(truncate=False)

+----+---------+----------------+
|col1|col2     |cartesian_product|
+----+---------+----------------+
|A   |[1, 2, 3]|[A, 1]          |
|A   |[1, 2, 3]|[A, 2]          |
|A   |[1, 2, 3]|[A, 3]          |
|B   |[4, 5]   |[B, 4]          |
|B   |[4, 5]   |[B, 5]          |
|C   |[6]      |[C, 6]          |
+----+---------+----------------+