从具有目标列值的列更新 pyspark 数据框

Question

我有一个数据框，它有一个列（在本例中为“target_column”），我需要使用“val”列值更新这些目标列。

我尝试过使用 udfs 和 .withcolumn，但它们都希望使用固定的列值。就我而言，它可以是可变的。由于 rdd 是不可变的，因此使用 rdd 映射转换也不起作用。

def test():

    data = [("jose_1", 'mase', "firstname", "jane"), ("li_1", "ken", 'lastname', 'keno'), ("liz_1", 'durn', 'firstname', 'liz')]
    source_df = spark.createDataFrame(data, ["firstname", "lastname", "target_column", "val"])
    source_df.show()


if __name__ == "__main__":
    spark = SparkSession.builder.appName('Name Group').getOrCreate()
    test()
    spark.stop()

输入：

+---------+--------+-------------+----+
|firstname|lastname|target_column| val|
+---------+--------+-------------+----+
|   jose_1|    mase|    firstname|jane|
|     li_1|     ken|     lastname|keno|
|    liz_1|    durn|    firstname| liz|
+---------+--------+-------------+----+

预期输出：

+---------+--------+-------------+----+
|firstname|lastname|target_column| val|
+---------+--------+-------------+----+
|     jane|    mase|    firstname|jane|
|     li_1|    keno|     lastname|keno|
|      liz|    durn|    firstname| liz|
+---------+--------+-------------+----+

例如在输入的第一行中，target_column 是“firstname”，val 是“jane”。所以我需要在该行中用“jane”更新名字。

谢谢

Answer 1

您可以对所有列进行循环：

from pyspark.sql import functions as F

for col in df.columns:
    df = df.withColumn(
        col,
        F.when(
            F.col("target_column")==F.lit(col), 
            F.col("val")
        ).otherwise(F.col(col))
    )