Pyspark --- Hinzufügen neuer Spalte mit Werten pro Gruppe von

Question

Angenommen, ich habe folgende Datensatz:

a | b 
1 | 0.4 
1 | 0.8 
1 | 0.5 
2 | 0.4 
2 | 0.1 

Ich möchte eine neue Spalte hinzufügen „label“ genannt, wo die Werte für jede lokal bestimmt werden Gruppe von Werten in a. Der höchste Wert von b in einer Gruppe ein 1 markiert und alle anderen sind 0 markiert

Die Ausgabe würde wie folgt aussehen:

a | b | label 
1 | 0.4 | 0 
1 | 0.8 | 1 
1 | 0.5 | 0 
2 | 0.4 | 1 
2 | 0.1 | 0 

Wie kann ich diese effizient zu nutzen PySpark tun?

Answer 1

Sie können es mit Fensterfunktionen tun. Zunächst werden Sie ein paar Importe benötigen:

from pyspark.sql.functions import desc, row_number, when 
from pyspark.sql.window import Window 

und Fensterdefinition:

w = Window().partitionBy("a").orderBy(desc("b")) 

Schließlich verwenden Sie diese:

df.withColumn("label", when(row_number().over(w) == 1, 1).otherwise(0)) 

Zum Beispiel Daten:

df = sc.parallelize([ 
    (1, 0.4), (1, 0.8), (1, 0.5), (2, 0.4), (2, 0.1) 
]).toDF(["a", "b"]) 

das Ergebnis ist:

+---+---+-----+ 
| a| b|label| 
+---+---+-----+ 
| 1|0.8| 1| 
| 1|0.5| 0| 
| 1|0.4| 0| 
| 2|0.4| 1| 
| 2|0.1| 0| 
+---+---+-----+