2016-03-08 11 views
6

Avere due colonne:matrice pyspark con variabili dummy

ID Text 
1 a 
2 b 
3 c 

Come posso in grado di creare matrici con variabili dummy come questo:

ID a b c 
1 1 0 0 
2 0 1 0 
3 0 0 1 

utilizzando la libreria pyspark e le sue caratteristiche?

risposta

5
from pyspark.sql import functions as F 

df = sqlContext.createDataFrame([ 
    (1, "a"), 
    (2, "b"), 
    (3, "c"), 
], ["ID", "Text"]) 

categories = df.select("Text").distinct().rdd.flatMap(lambda x: x).collect() 

exprs = [F.when(F.col("Text") == category, 1).otherwise(0).alias(category) 
     for category in categories] 

df.select("ID", *exprs).show() 

uscita

+---+---+---+---+ 
| ID| a| b| c| 
+---+---+---+---+ 
| 1| 1| 0| 0| 
| 2| 0| 1| 0| 
| 3| 0| 0| 1| 
+---+---+---+---+ 
+0

Traceback (chiamata più recente scorso): file "", linea 2, in NameError: nome di 'F' non è definito –

+1

Qual è F qui? –

+1

dall'importazione di pyspark.sql funziona come F –