PySpark數據框：根據條件同時更改兩個列

Question

我想知道是否有一種方法可以同時更改PySpark數據框的兩個（或更多）列。 現在我正在使用withColumn但是我不知道這是否意味着條件將被檢查兩次（這對於大型數據框來說太昂貴了）。 該代碼基本上檢查其他兩列（對於同一行）中的值，並基於此將兩列更改為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

首先，您可以簡單地將UDF添加為新列，將其用於計算並刪除它：

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")

如果您確實要生成兩列，則可以創建一個struct列並將其展平（當然，需要將列添加到select中）：

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