如何使用UDF返回多列？

Question

是否可以创建一个返回列集的UDF？

即具有如下数据框:

| Feature1 | Feature2 | Feature 3 |
| 1.3      | 3.4      | 4.5       |

现在我想提取一个新特征,可以将其描述为两个元素的向量(例如,在线性回归中看到 - 斜率和偏移).所需数据集应如下所示:

| Feature1 | Feature2 | Feature 3 | Slope | Offset |
| 1.3      | 3.4      | 4.5       | 0.5   | 3      |

是否可以使用单个UDF创建多个列,还是需要遵循以下规则:"每个UDF单列"？

Answer 1

结构方法

您可以将udf函数定义为

def myFunc: (String => (String, String)) = { s => (s.toLowerCase, s.toUpperCase)}

import org.apache.spark.sql.functions.udf
val myUDF = udf(myFunc)

并.*用作

val newDF = df.withColumn("newCol", myUDF(df("Feature2"))).select("Feature1", "Feature2", "Feature 3", "newCol.*")

我已经返回Tuple2测试目的(根据需要多少列,可以使用更高阶的元组),udf它将被视为struct列.然后,您可以使用.*选择单独列中的所有元素,最后重命名它们.

你应该输出为

+--------+--------+---------+---+---+
|Feature1|Feature2|Feature 3|_1 |_2 |
+--------+--------+---------+---+---+
|1.3     |3.4     |4.5      |3.4|3.4|
+--------+--------+---------+---+---+

你可以重命名_1和_2

数组方法

udf 功能应该返回一个 array

def myFunc: (String => Array[String]) = { s => Array("s".toLowerCase, s.toUpperCase)}

import org.apache.spark.sql.functions.udf
val myUDF = udf(myFunc)

您可以选择其中的元素array并使用alias它们重命名

val newDF = df.withColumn("newCol", myUDF(df("Feature2"))).select($"Feature1", $"Feature2", $"Feature 3", $"newCol"(0).as("Slope"), $"newCol"(1).as("Offset"))

你应该有

+--------+--------+---------+-----+------+
|Feature1|Feature2|Feature 3|Slope|Offset|
+--------+--------+---------+-----+------+
|1.3     |3.4     |4.5      |s    |3.4   |
+--------+--------+---------+-----+------+

Answer 2

此外，您可以返回案例类：

case class NewFeatures(slope: Double, offset: Int)

val getNewFeatures = udf { s: String =>
      NewFeatures(???, ???)
    }

df
  .withColumn("newF", getNewFeatures($"Feature1"))
  .select($"Feature1", $"Feature2", $"Feature3", $"newF.slope", $"newF.offset")