Question

如何将数据框 A 中的值添加到包含给定数据框 sum 对的数据框 B 中的新列 (A)？最好使用 UDF？

输出应如下所示：

数据框 A：

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

数据框 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|
#+---+---+-----+

将数据帧中的值添加到另一个数据帧 pyspark 中的列

2 个答案: