2017-09-22 29 views
1

動的な列選択

> df.show() 

+---+-----+---+---+---+---+ 
| id|delay| p1| p2| p3| p4| 
+---+-----+---+---+---+---+ 
| 1| 3| a| b| c| d| 
| 2| 1| m| n| o| p| 
| 3| 2| q| r| s| t| 
+---+-----+---+---+---+---+ 

新しい、col列がp{delay}既存の列の結果になるように動的に列を選択するには?

> df.withColumn("col", /* ??? */).show() 

+---+-----+---+---+---+---+----+ 
| id|delay| p1| p2| p3| p4| col| 
+---+-----+---+---+---+---+----+ 
| 1| 3| a| b| c| d| c| // col = p3 
| 2| 1| m| n| o| p| m| // col = p1 
| 3| 2| q| r| s| t| r| // col = p2 
+---+-----+---+---+---+---+----+ 

答えて

6

私は考えることができる最も簡単な解決策はarrayを使用することですdelayを指標として:

import org.apache.spark.sql.functions.array 

df.withColumn("col", array($"p1", $"p2", $"p3", $"p4")($"delay" - 1)) 
+1

これは非常に巧妙です。 – Jivan

+1

優秀です。私はそれを2回upvoteすることができたらいいなあ。 – philantrovert

1

1つのオプションは、列名に番号からマップを作成し、対応する値でcol列を更新するためにfoldLeftを使用している:

val cols = (1 to 4).map(i => i -> s"p$i") 

(cols.foldLeft(df.withColumn("col", lit(null))){ 
    case (df, (k, v)) => df.withColumn("col", when(df("delay") === k, df(v)).otherwise(df("col"))) 
}).show 
+---+-----+---+---+---+---+---+  
| id|delay| p1| p2| p3| p4|col| 
+---+-----+---+---+---+---+---+ 
| 1| 3| a| b| c| d| c| 
| 2| 1| m| n| o| p| m| 
| 3| 2| q| r| s| t| r| 
+---+-----+---+---+---+---+---+