Skip to main content
 首页 » 编程设计

python之pySpark/Python 遍历数据框列,检查条件并填充另一列

2024年08月20日35Leo_wl

我在 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 = "None",所有日期都在日期格式,真正的数据框也有 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 

提前致谢

请您参考如下方法:

我会使用合并:

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 快得多。