首页 > 解决方案 > 如果不在 pyspark 中,则保持相同 ID 或第一行的两列中的值匹配

问题描述

我有以下示例数据集。

ID      numb       real
213412  2008       2008
213412  2018       2008
393859  2017       2017
213410  2009       NULL
213410  2020       2013
393859  2021       2021

我需要在 numb = real 时保留具有相同 ID 的行,否则如果不是(例如 id = 213410 ),那么我必须只保留第一行(numb = 2009)。预期输出:

ID      numb       real
213412  2008       2008
393859  2017       2017
213410  2009       NULL
393859  2021       2021

这是我的尝试,但是当 numb != real 并且只保留第一行时它不能处理这种情况:

df = df.groupby('ID').filter(col("numb") == col("real"))

标签: pythonapache-sparkpysparkapache-spark-sql

解决方案


它适用于我的排名,所以想法是按原始顺序对具有相同 ID 的行进行排名,然后如果数字不匹配,则保持排名 1。

(df
    .withColumn('mid', F.monotonically_increasing_id())
    .withColumn('rnk', F.rank().over(W.partitionBy('ID').orderBy('mid')))
    .withColumn('keep', F
        .when(F.col('numb') == F.col('real'), True)
        .when(F.col('rnk') == 1, True)
        .otherwise(False)
    )
#     .where(F.col('keep') == True) # uncomment this to filter your desire rows
    .show()
)

# Output
# +------+----+----+-----------+---+-----+
# |    ID|numb|real|        mid|rnk| keep|
# +------+----+----+-----------+---+-----+
# |213412|2008|2008| 8589934592|  1| true|
# |213412|2018|2008|25769803776|  2|false|
# |393859|2017|2017|42949672960|  1| true|
# |393859|2021|2021|94489280512|  2| true|
# |213410|2009|null|60129542144|  1| true|
# |213410|2020|2013|77309411328|  2|false|
# +------+----+----+-----------+---+-----+

推荐阅读