Scala 基于同一行中的其他列动态选择列内容

我使用的是Spark 1.6.1。假设我的数据框看起来像：

+------------+-----+----+
|categoryName|catA |catB|
+------------+-----+----+
|     catA   |0.25 |0.75|
|     catB   |0.5  |0.5 |
+------------+-----+----+

其中，

categoryName

具有字符串类型，

cat*

为双精度。我想添加一列，该列将包含列中的值，该列的名称位于

categoryName

列中：

+------------+-----+----+-------+
|categoryName|catA |catB| score |
+------------+-----+----+-------+
|     catA   |0.25 |0.75| 0.25  | ('score' has value from column name 'catA')
|     catB   |0.5  |0.7 |  0.7  | ('score' value from column name 'catB')
+------------+-----+----+-------+

我需要这样的提取到以后的一些计算中。有什么想法吗

---

重要提示：我不知道类别列的名称。解决方案需要是动态的。

您有几个选项：

如果您使用的是scala，则可以使用Dataset API，在这种情况下，您只需创建一个执行计算的映射即可

您可以从dataframe移动到RDD并使用映射

您可以创建一个UDF，它接收所有相关列作为输入，并在其中进行计算

您可以使用一组when/other子句进行搜索（例如when（col1==CatA，col（CatA））。other（col（CatB）））

---

Spark 2.0： 您可以通过创建一个临时列来实现这一点（对于任意数量的类别列），该临时列包含categroyName->categoryValue的映射，然后从中选择：

// sequence of any number of category columns
val catCols = input.columns.filterNot(_ == "categoryName")

// create a map of category -> value, and then select from that map using categoryName:
input
  .withColumn("asMap", map(catCols.flatMap(c => Seq(lit(c), col(c))): _*))
  .withColumn("score", $"asMap".apply($"categoryName"))
  .drop("asMap")

Spark 1.6：类似的想法，但使用数组和自定义项从中进行选择：

// sequence of any number of category columns
val catCols = input.columns.filterNot(_ == "categoryName")

// UDF to select from array by index of colName in catCols
val getByColName = udf[Double, String, mutable.WrappedArray[Double]] {
  case (colName, colValues) =>
    val index = catCols.zipWithIndex.find(_._1 == colName).map(_._2)
    index.map(colValues.apply).getOrElse(0.0)
}

// create an array of category values and select from it using UDF:
input
  .withColumn("asArray", array(catCols.map(col): _*))
  .withColumn("score", getByColName($"categoryName", $"asArray"))
  .drop("asArray")

---

1.）数据集API来自Spark 2.0.0，对吗？我正在使用1.6.12。）嗯，也许，我会检查一下。3.）但我会在udf中释放列名上下文，对吗？4.）不是DynamicPark 1.6.1在scala中有dataset API（它在2.0中更改）。移动到UDF时，可以通过设置顺序来使用列名。