如何使用 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 逐行与内置函数相矛盾逐列 em>。

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

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

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

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

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

Answer 2

您可以通过map 发送select。定义将名称转换为列值的映射：

import org.apache.spark.sql.functions.{col, concat_ws, lit, map}

val dataMap = map(
  df.columns.diff(Seq("A", "B")).flatMap(c => lit(c) :: col(c) :: Nil): _*
)

df.select(dataMap).show(false)

+---------------------------+
|map(A_1, A_1, B_2, B_2)    |
+---------------------------+
|Map(A_1 -> 0.1, B_2 -> 0.3)|
|Map(A_1 -> 0.2, B_2 -> 0.4)|
+---------------------------+

然后用apply从中选择：

df.withColumn("C", dataMap(concat_ws("_", $"A", $"B"))).show

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

您也可以尝试映射，但我怀疑它在非常宽的数据中表现不佳：

import org.apache.spark.sql.catalyst.encoders.RowEncoder
import org.apache.spark.sql.types._
import org.apache.spark.sql.Row

val outputEncoder = RowEncoder(df.schema.add(StructField("C", DoubleType)))

df.map(row => {
   val a = row.getAs[String]("A")
   val b = row.getAs[String]("B")
   val key = s"${a}_${b}"
   Row.fromSeq(row.toSeq :+ row.getAs[Double](key))
})(outputEncoder).show

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

一般来说我不会推荐这种方法。

如果数据来自csv，您可以考虑跳过默认的csv 阅读器并使用自定义逻辑将列选择直接推送到解析过程。用伪代码：

spark.read.text(...).map { line => {
  val a = ???  // parse A
  val b = ???  // parse B
  val c = ???  // find c, based on a and b
  (a, b, c)
}}