PySpark Dataframe: Ändern von zwei Spalten zur gleichen Zeit basierend auf der Bedingung

Question

Ich frage mich, ob es eine Möglichkeit gibt, zwei (oder mehr) Spalten eines PySpark Dataframe zur gleichen Zeit zu ändern. Im Moment verwende ich withColumn, aber ich weiß nicht, ob das bedeutet, dass die Bedingung zweimal überprüft wird (was für einen großen Datenrahmen zu teuer sein kann). Dieser Code sucht im Grunde nach Werten in anderen zwei Spalten (für die gleiche Zeile) und ändert dann zwei Spalten in None/null.

condition = is_special_id_udf(col("id"))) & should_hide_response_udf(col("response_created")) 


    new_df = df.withColumn(
      "response_text", 
      when(condition, None) 
      .otherwise(col("response_text")) 
     ) 

    new_df = df.withColumn(
      "response_created", 
      when(condition, None) 
      .otherwise(col("response_created")) 
     ) 

Answer 1

Das erste, was Sie können einfach die UDF als eine neue Spalte hinzuzufügen, es für Berechnungen verwenden und legen Sie es:

condition = is_special_id_udf(col("id"))) & should_hide_response_udf(col("response_created")) 

new_df = df.withColumn("tmp", condition).withColumn(
     "response_text", 
     when(col("tmp"), None) 
     .otherwise(col("response_text")) 
    ).withColumn(
     "response_created", 
     when(col("tmp"), None) 
     .otherwise(col("response_created")) 
    ).drop("tmp") 

Wenn Sie wirklich zwei Spalten generieren möchten, dann können Sie eine Struktur tun erstellen Spalte und flattern Sie es (natürlich addieren Sie Spalten, die Sie zur Auswahl brauchen):

new_df = df.withColumn(
     "myStruct", 
     when(condition, None) 
     .otherwise(struct(col("response_text"), col("response_created"))) 
    ).select("myStruct.*")