eine Funktion für eine gefilterte Datenrahmen in pyspark Rufen

Question

Ich bin Daten wie

Daten

ID   filter 
1    A 
2    A 
3    A 
4    A 
5    B 
6    B 
7    B 
8    B 

Ich möchte mit einer Funktion für den Datenrahmen anzuwenden,

def add(x): 
    y = x+1 
    return(y) 

from pyspark.sql.functions import *     
from pyspark.sql.functions import udf 

ol_val = udf(add, StringType()) 

data = data.withColumn("sum",ol_val(data.ID)) 

das gibt ein Ausgang,

Daten

ID   filter  sum 
1    A   2 
2    A   3 
3    A   4 
4    A   5 
5    B   6 
6    B   7 
7    B   8 
8    B   9 

Ich möchte diese Funktion ist nur anzuwenden, wenn Filter A und für den verbleibenden Teil = ich es NULL sein wollen. Der Ausgang ich hier will, ist,

Daten

ID   filter  sum 
1    A   2 
2    A   3 
3    A   4 
4    A   5 
5    B   NULL 
6    B   NULL 
7    B   NULL 
8    B   NULL 

Hier ist der Wert NULL ist, weil sie den Zustand Filter erfüllen didnt = A. Ich möchte die Funktion nur, wenn filter = A.

angewendet werden

Kann jemand mir helfen, den Code zu ändern, um diese Ausgabe in pyspark zu erhalten?

Answer 1

Was Sie brauchen, ist when und otherwise zu verwenden. Übrigens müssen Sie das nicht erstellen.

df = sc.parallelize([ 
    (1, "a"), 
    (1, "b"), 
    (3, "c") 
    ]).toDF(["id", "filter"]) 

df.select("*", when(col("filter") == lit("a"), col("id") + 1).otherwise(None).alias("result")).show() 

Wenn Sie wirklich diese Funktion aufrufen müssen, können Sie einfach ersetzen col("id") + 1 durch yourUDF(col("id"))