scala - 如何使用 Scala Spark 中 withColumn 的另一列值组成列名

Question

我正在尝试向 DataFrame 添加一个新列.此列的值是另一个列的值，该列的名称依赖于同一 DataFrame 中的其他列.

例如，鉴于此:

+---+---+----+----+
|  A|  B| A_1| B_2|
+---+---+----+----+
|  A|  1| 0.1| 0.3|
|  B|  2| 0.2| 0.4|
+---+---+----+----+

我想得到这个:

+---+---+----+----+----+
|  A|  B| A_1| B_2|   C|
+---+---+----+----+----+
|  A|  1| 0.1| 0.3| 0.1|
|  B|  2| 0.2| 0.4| 0.4|
+---+---+----+----+----+

也就是说，我添加了 C 列，其值来自 A_1 列或 B_2 列。源列 A_1 的名称来自连接列 A 和 B 的值。

我知道我可以添加一个基于另一个和常量的新列，如下所示:

df.withColumn("C", $"B" + 1)

我也知道列的名称可以来自这样的变量:

val name = "A_1"
df.withColumn("C", col(name) + 1)

但是，我想做的是这样的:

df.withColumn("C", col(s"${col("A")}_${col("B")}"))

这是行不通的。

注意:我正在使用 Scala 2.11 和 Spark 2.2 进行编码。

Answer 1

您可以通过编写 udf 来满足您的要求功能。 我建议udf ，因为您的要求是处理 dataframe逐行与逐列运行的内置函数相矛盾 .

但在此之前，您需要列名数组

val columns = df.columns

然后写一个 udf作为

import org.apache.spark.sql.functions._
def getValue = udf((A: String, B: String, array: mutable.WrappedArray[String]) => array(columns.indexOf(A+"_"+B)))

在哪里

A is the first column value
B is the second column value
array is the Array of all the columns values

现在只需调用 udf函数使用 withColumn api

df.withColumn("C", getValue($"A", $"B", array(columns.map(col): _*))).show(false)

你应该得到你想要的输出 dataframe .