pySpark/Python 迭代数据帧列，检查条件并填充另一列

Question

我正在 Jupyter Notebook 中使用 python/pySpark，我试图弄清楚以下内容：

我有一个像这样的数据框

MainDate      Date1        Date2        Date3         Date4
2015-10-25    2015-09-25   2015-10-25   2015-11-25    2015-12-25
2012-07-16    2012-04-16   2012-05-16   2012-06-16    2012-07-16
2005-03-14    2005-07-14   2005-08-14   2005-09-14    2005-10-14

我需要将 MainDate 与 Date1-Date4 中的每个进行比较，如果 MainDate == Date# 则创建一个新列 REAL = Date#，如果没有匹配则 REAL =“无”，所有日期均为日期格式，另外，真实的数据帧有 Date1 到 Date72，并且可能只有一个匹配项（如果有的话）

最后结果：

MainDate      Date1        Date2        Date3         Date4        REAL
2015-10-25    2015-09-25   2015-10-25   2015-11-25    2015-12-25   Date2
2012-07-16    2012-04-16   2012-05-16   2012-06-16    2012-07-16   Date4
2005-03-14    2005-07-14   2005-08-14   2005-09-14    2005-10-14   None

提前致谢

Answer 1

我会用coalesce：

from pyspark.sql.functions import col, when, coalesce, lit

df = spark.createDataFrame([
    ("2015-10-25", "2015-09-25", "2015-10-25", "2015-11-25", "2015-12-25"),
    ("2012-07-16", "2012-04-16", "2012-05-16", "2012-06-16", "2012-07-16"),
    ("2005-03-14", "2005-07-14", "2005-08-14", "2005-09-14", "2005-10-14"),],
    ("MainDate", "Date1", "Date2", "Date3", "Date4")
)


df.withColumn("REAL", 
    coalesce(*[when(col(c) == col("MainDate"), lit(c)) for c in df.columns[1:]])
).show()


+----------+----------+----------+----------+----------+-----+
|  MainDate|     Date1|     Date2|     Date3|     Date4| REAL|
+----------+----------+----------+----------+----------+-----+
|2015-10-25|2015-09-25|2015-10-25|2015-11-25|2015-12-25|Date2|
|2012-07-16|2012-04-16|2012-05-16|2012-06-16|2012-07-16|Date4|
|2005-03-14|2005-07-14|2005-08-14|2005-09-14|2005-10-14| null|
+----------+----------+----------+----------+----------+-----+

在哪里

when(col(c) == col("MainDate"), lit(c))

lit(c)如果存在匹配则返回列名 ( )，NULL否则返回。

这应该比udf转换为RDD.