比较2个pyspark数据框列,并根据它更改另一列的值

时间:2020-01-14 22:19:14

标签: pyspark

我有一个问题,我从我写的图算法生成了一个数据框。事实是,我希望在每次运行图形代码后,基本组件的值基本上保持不变。

这是生成的示例数据帧:

df = spark.createDataFrame(
    [
        (1, 'A1'), 
        (1, 'A2'),
        (1, 'A3'),
        (2, 'B1'),
        (2, 'B2'),
        (3, 'B3'),
        (4, 'C1'),
        (4, 'C2'),
        (4, 'C3'),
        (4, 'C4'),
        (5, 'D1'),
    ],
    ['old_comp_id', 'db_id'] 
)

下一次运行时,值将完全更改,因此新运行具有类似的值

df2 = spark.createDataFrame(
    [
        (2, 'A1'), 
        (2, 'A2'),
        (2, 'A3'),
        (3, 'B1'),
        (3, 'B2'),
        (3, 'B3'),
        (1, 'C1'),
        (1, 'C2'),
        (1, 'C3'),
        (1, 'C4'),
        (4, 'D1'),
    ],
    ['new_comp_id', 'db_id'] 
)

所以我要做的是比较以上两个数据框之间的值,并根据关联的数据库ID更改组件ID的值。

  1. 如果database_id相同,则将组件ID更新为来自第一个数据帧
  2. 如果它们不同,则分配一个全新的comp_id(new_comp_id = max(old_comp_id)+1)

这是我到目前为止提出的:

old_ids = df.groupBy("old_comp_id").agg(F.collect_set(F.col("db_id")).alias("old_db_id"))
new_ids = df2.groupBy("new_comp_id").agg(F.collect_set(F.col("db_id")).alias("new_db_id"))

joined = new_ids.join(old_ids,old_ids.old_comp_id == new_ids.new_comp_id,"outer")

joined.withColumn("update_comp", F.when( F.col("new_db_id") == F.col("old_db_id"), F.col('old_comp_id')).otherwise(F.max(F.col("old_comp_id")+1))).show()

1 个答案:

答案 0 :(得分:0)

要在非聚合列中使用聚合函数,应使用窗口函数。

首先,将DF与db_id外部连接:

from pyspark.sql.functions import when, col, max
joinedDF = df.join(df2, df["db_id"] == df2["new_db_id"], "outer")

然后,开始构建窗口(您在其中按db_id分组,并按old_comp_id进行排序,以使第一行中的old_comp_id具有最高的值。

from pyspark.sql.window import Window
from pyspark.sql.functions import desc
windowSpec = Window\
.partitionBy("db_id")\
.orderBy(desc("old_comp_id"))\
.rowsBetween(Window.unboundedPreceding, Window.currentRow)

然后,您使用windowSpec构建max列

from pyspark.sql.functions import max
maxCompId = max(col("old_comp_id")).over(windowSpec)

然后,将其应用于选择

joinedDF.select(col("db_id"), when(col("new_db_id").isNotNull(), col("old_comp_id")).otherwise(maxCompId+1).alias("updated_comp")).show()

有关更多信息,请参阅文档(http://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.Window

希望这会有所帮助