Dynamische Spaltenauswahl in Funken (basierend auf einer anderen Spalte Wert)

Question

Mit dem gegebenen Funkendatenrahmen:

> 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| 
+---+-----+---+---+---+---+ 

Wie eine Spalte dynamisch wählen, so dass die neue, col Spalte ist das Ergebnis der p{delay} vorhandenen Spalte?

> 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 
+---+-----+---+---+---+---+----+ 

Answer 1

Die einfachste Lösung, die ich ist denken kann array zu verwenden mit delay als index:

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

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

Answer 2

Eine Möglichkeit ist es, eine Karte von Nummer Spaltennamen erstellen und dann foldLeft verwenden, um die col Spalte mit den entsprechenden Werten zu aktualisieren:

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| 
+---+-----+---+---+---+---+---+