如何使用UDF添加多列？

Question

我想将 UDF 的返回值添加到单独列中的现有数据帧中。我如何以足智多谋的方式实现这一目标？

这是我到目前为止所拥有的一个例子。

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StructType, StructField, IntegerType  

df = spark.createDataFrame([("Alive",4)],["Name","Number"])
df.show(1)

+-----+------+
| Name|Number|
+-----+------+
|Alive|     4|
+-----+------+

def example(n):
        return [[n+2], [n-2]]

#  schema = StructType([
#          StructField("Out1", ArrayType(IntegerType()), False),
#          StructField("Out2", ArrayType(IntegerType()), False)])

example_udf = udf(example)

现在我可以向数据框中添加一列，如下所示

newDF = df.withColumn("Output", example_udf(df["Number"]))
newDF.show(1)
+-----+------+----------+
| Name|Number|Output    |
+-----+------+----------+
|Alive|     4|[[6], [2]]|
+-----+------+----------+

但是，我不希望这两个值位于同一列中，而是位于不同的列中。

理想情况下，我想现在分割输出列，以避免调用示例函数两次（每个返回值一次），如下所示here https://stackoverflow.com/questions/34362412/pyspark-split-a-column-to-multiple-columns-without-pandas and here https://stackoverflow.com/questions/39235704/split-spark-dataframe-string-column-into-multiple-columns，但是在我的情况下，我得到了一个数组数组，但我看不到拆分如何在那里工作（请注意，每个数组将包含多个值，用“，”分隔。

结果应该是什么样子

我最终想要的是这个

+-----+------+----+----+
| Name|Number|Out1|Out2|
+-----+------+----+----+
|Alive|     4|   6|   2|
+-----+------+----+----+

请注意，StructType 返回类型的使用是可选的，并且不一定是解决方案的一部分。

编辑：我注释掉了 StructType 的使用（并编辑了 udf 分配），因为示例函数的返回类型不需要它。但是，如果返回值类似于

return [6,3,2],[4,3,1]

返回一个StructType，只需使用Row

from pyspark.sql.types import StructType,StructField,IntegerType,Row
from pyspark.sql import functions as F

df = spark.createDataFrame([("Alive", 4)], ["Name", "Number"])


def example(n):
    return Row('Out1', 'Out2')(n + 2, n - 2)


schema = StructType([
    StructField("Out1", IntegerType(), False),
    StructField("Out2", IntegerType(), False)])

example_udf = F.UserDefinedFunction(example, schema)

newDF = df.withColumn("Output", example_udf(df["Number"]))
newDF = newDF.select("Name", "Number", "Output.*")

newDF.show(truncate=False)