Динамически выбирайте содержимое столбца на основе другого столбца из той же строки

Question

Я использую Spark 1.6.1. Допустим, мой кадр данных выглядит следующим образом:

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

Где categoryName имеет тип String, и 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') 
+------------+-----+----+-------+ 

мне нужна такая экстракцию некоторых дальнейших расчетов. Есть идеи?

Важно: я не знаю названия столбцов категории. Решение должно быть динамичным.

Answer 1

Спарк 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") 

Спарк 1,6: Похожая идея, но с использованием массива и UDF, чтобы выбрать из него:

// 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") 

Answer 2

У вас есть несколько вариантов:

1. Если вы используете Scala вы можете использовать Dataset API в этом случае вы бы просто создать карту, которая делает расчет.
2. Вы можете перейти к РДУ из dataframe и использовать карту
3. Вы можете создать UDF, который принимает все соответствующие столбцы в качестве входных данных и сделать расчет внутри
4. вы можете использовать кучу, когда/в противном случае положение, чтобы сделать поиск (например, когда (col1 == CATA, цв (CATA)) в противном случае (цв (CatB)).)