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和here所述，但在我的情况下我和＃39;我得到了一个数组数组，我无法看到分割如何在那里工作（请注意，每个数组将包含多个值，用＆＃34;，＃34;分隔。

结果应该如何

我最终想要的是这个

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

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

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

，则必须使用它

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

Answer 1

要返回StructType，只需使用Row

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)

Answer 2

解决上述问题的更好方法是将输出转换为数组，然后将其分解

import pyspark.sql.functions as f
import pyspark.sql.types as t

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


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


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

example_udf = f.udf(example, schema)

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

newDF.show(truncate=False)
newDF.explain()

注意explain的输出，您会发现示例方法实际上只被调用了一次！

Answer 3

在Scala中

import spark.implicits
val df = Seq(("Alive", 4)).toDF("Name", "Number")

没有UDF

df.
  withColumn("OutPlus",  $"Number" + 2).
  withColumn("OutMinus", $"Number" - 2).
  show
+-----+------+-------+--------+
| Name|Number|OutPlus|OutMinus|
+-----+------+-------+--------+
|Alive|     4|      6|       2|
+-----+------+-------+--------+

使用UDF爆炸

import org.apache.spark.sql.functions.udf
def twoItems(_i: Int) = Seq((_i + 2, _i - 2))
val twoItemsUdf = udf(twoItems(_: Int))

val exploded = df.
  withColumn("Out", explode(twoItemsUdf($"Number"))).
  withColumn("OutPlus", $"Out._1").
  withColumn("OutMinus", $"Out._2")

exploded.printSchema

root
 |-- Name: string (nullable = true)
 |-- Number: integer (nullable = false)
 |-- Out: struct (nullable = true)
 |    |-- _1: integer (nullable = false)
 |    |-- _2: integer (nullable = false)
 |-- OutPlus: integer (nullable = true)
 |-- OutMinus: integer (nullable = true)

  exploded.drop("Out").show

+-----+------+-------+--------+
| Name|Number|OutPlus|OutMinus|
+-----+------+-------+--------+
|Alive|     4|      6|       2|
+-----+------+-------+--------+

如何使用UDF添加多个列？

3 个答案: