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 api
调用withColumn函数即可
df.withColumn("C", getValue($"A", $"B", array(columns.map(col): _*))).show(false)

您应该获得所需的输出dataframe。

Answer 2

您可以select map。定义将名称转换为列值的地图：

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)
}}

如何使用Scala Spark中withColumn的另一列值来组合列名

2 个答案: