在 pyspark 数据框之间平均分配值

Question

我有两个pyspark数据框

例如，我有一个用户 table 如下所示

users    col1  col2  col3
1      A      AA   AAA
1      A      AA   AAA
2      A      AB   AAB
3      A      AA   AAA
4      A      AA   AAA
6      B      BB   BBB
7      B      BB   BBB
8      A      AA   AAA

和如下所示的成本 table 数据框

col1 col2 col3 cost
A     AA   AAA   1000
A     AB   AAB   200
B     BB   BBB   420

我需要在用户之间平均分配这笔费用。如下所示，我需要如下所示的最终输出列

结果栏

users    col1  col2  col3   cost 
1      A      AA   AAA      200
1      A      AA   AAA      200
2      A      AB   AAB      200
3      A      AA   AAA      250
4      A      AA   AAA      200
6      B      BB   BBB      210
7      B      BB   BBB      210
8      A      AA   AAA      200

如何使用 pysaprk 数据框完成此操作：这是示例数据集。我的用户列大小约为 1GB，成本 table 约为 10 Mb。我是 pyspark 的新手，解决这个问题的最佳方法是什么？

更新：需要更新示例数据框的行数

Answer 1

这是适合您的解决方案

from pyspark.sql.functions import count, broadcast, col
user.join(broadcast(user.groupBy("col1", "col2", "col3")\
    .agg(count("users").alias("count"))\
    .join(broadcast(cost), ["col1", "col2", "col3"])\
    .withColumn('cost', col('cost')/col('count'))), ["col1", "col2", "col3"])\
    .drop('count')\
    .show(truncate=False)

哪个应该给你

+----+----+----+-----+-----+
|col1|col2|col3|users|cost |
+----+----+----+-----+-----+
|A   |AA  |AAA |1    |250.0|
|A   |AB  |AAB |2    |200.0|
|A   |AA  |AAA |3    |250.0|
|A   |AA  |AAA |4    |250.0|
|B   |BB  |BBB |6    |210.0|
|B   |BB  |BBB |7    |210.0|
|A   |AA  |AAA |8    |250.0|
+----+----+----+-----+-----+

说明：

上述解决方案分为三部分

第一部分是聚合

user.groupBy("col1", "col2", "col3")\
    .agg(count("users").alias("count"))

这给了你

+----+----+----+-----+
|col1|col2|col3|count|
+----+----+----+-----+
|A   |AB  |AAB |1    |
|B   |BB  |BBB |2    |
|A   |AA  |AAA |4    |
+----+----+----+-----+

第二个是第一个连接（我使用了broadcast函数，正如你所说的cost dataframe很小）

user.groupBy("col1", "col2", "col3") \
    .agg(count("users").alias("count")) \
    .join(broadcast(cost), ["col1", "col2", "col3"]) \
    .withColumn('cost', col('cost')/col('count'))

哪个应该给你

+----+----+----+-----+-----+
|col1|col2|col3|count|cost |
+----+----+----+-----+-----+
|A   |AB  |AAB |1    |200.0|
|B   |BB  |BBB |2    |210.0|
|A   |AA  |AAA |4    |250.0|
+----+----+----+-----+-----+

而最后一个是上述输出与用户数据帧的连接。 删除额外的计数列

希望回答对你有帮助

在 pyspark 数据框之间平均分配值

Equally distributing values among a pyspark data frame

python

apache-spark

pyspark

pyspark-sql