将 dataframe 中的值添加到另一个 dataframe pyspark 中的列

Question

如何将 dataframe A B的新列（ sum ），其中包含给定的 dataframe A对？ 最好使用UDF？

output 应如下所示：

dataframe

|id|value|
|--|-----|
|1 |   10|
|2 |  0.3|
|3 |  100|

dataframe B：（添加列sum ）

|src|dst|sum  |
|---|---|-----|
|1  |2  |10.3 |
|2  |3  |100.3|
|3  |1  |110  |

我试过这个

dfB = dfB.withColumn('sum', sum(dfB.source,dfB.dst,dfA))

def sum(src,dst,dfA):
    return dfA.filter(dfA.id == src).collect()[0][1][0] + dfA.filter(dfA.id == dst).collect()[0][1][0]

Answer 1

如果dfA足够小以进行广播连接，那么这应该可以工作：

dfB.join(dfA, how="left", on=F.col("src") == F.col("id")).select(
    "src", "dst", F.coalesce(F.col("value"), F.lit(0)).alias("v1")
).join(dfA, how="left", on=F.col("src") == F.col("id")).select(
    "src", "dst", (F.col("v1") + F.coalesce(F.col("value"), F.lit(0))).alias("sum")
)

如果 id 列包含每个 src 和 dst 值，则可以删除.coalesce() 。 有几种方法可以实现这一点，但最好的选择可能是使用.transform() 。

def join_sum(join_df):
    def _(df):
        return (
            df.join(join_df, how="left", on=F.col("src") == F.col("id"))
            .select("src", "dst", F.coalesce(F.col("value"), F.lit(0)).alias("v1"))
            .join(join_df, how="left", on=F.col("src") == F.col("id"))
            .select(
                "src",
                "dst",
                (F.col("v1") + F.coalesce(F.col("value"), F.lit(0))).alias("sum"),
            )
        )

    return _


dfB.transform(join_sum(dfA))

Answer 2

基本上，您需要在条件(id = src OR id = dst)上加入 2 个数据帧，然后按分组对列value求和：

from pyspark.sql import functions as F

output = df_a.join(
    df_b, 
    (F.col("id") == F.col("src")) | (F.col("id") == F.col("dst"))
).groupBy("src", "dst").agg(F.sum("value").alias("sum"))

output.show()
#+---+---+-----+
#|src|dst|  sum|
#+---+---+-----+
#|  2|  3|100.3|
#|  1|  2| 10.3|
#|  3|  1|110.0|
#+---+---+-----+