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

Question

我有两个 pyspark 数据框

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

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

以及如下所示的成本表数据框

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，成本表约为 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函数，正如你所说的成本数据框很小）

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|
+----+----+----+-----+-----+

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

希望回答对你有帮助