python - 从数据框中获取列总和，包括 map 列 - PySpark

Question

我有一个 PySpark 数据框，如下所示，我有一个 map 数据类型列 Map<Str,Int>

Date              Item (Map<Str,int>)                           Total Items      ColA

2021-02-01    Item_A -> 3, Item_B -> 10, Item_C -> 2                 15            10
2021-02-02    Item_A -> 1, Item_D -> 5, Item_E ->  7                 13            20
2021-02-03    Item_A -> 8, Item_E -> 3, Item_C ->  1                 12            30

我想要对包括 map 列在内的所有列进行求和。对于映射列，应根据键计算总和。

我想要这样的东西:

[[Item_A -> 12, Item_B -> 10, Item_C -> 3, Item_D -> 5, Item_E -> 10], 40, 60]

不一定是列表的列表，但我想要列的总和。

我的方法:

df.rdd.map(lambda x: (1,x[1])).reduceByKey(lambda x,y: x + y).collect()[0][1]

Answer 1

您可以分别对 map 列和其他列进行聚合，因为您需要在项目列上进行爆炸，然后您需要求和的其他列将变得难以处理。

示例数据框:

from pyspark.sql import functions as F

df = spark.createDataFrame(
    [('2021-02-01', {'Item_A': 3, 'Item_B': 10, 'Item_C': 2}, 15, 10),
     ('2021-02-02', {'Item_A': 1, 'Item_D': 5, 'Item_E': 7}, 13, 20),
     ('2021-02-03', {'Item_A': 8, 'Item_E': 3, 'Item_C': 1}, 12, 30)],
    ['Date', 'Item', 'Total Items', 'ColA'])
df.show(truncate=0)
# +----------+----------------------------------------+-----------+----+
# |Date      |Item                                    |Total Items|ColA|
# +----------+----------------------------------------+-----------+----+
# |2021-02-01|{Item_C -> 2, Item_B -> 10, Item_A -> 3}|15         |10  |
# |2021-02-02|{Item_E -> 7, Item_D -> 5, Item_A -> 1} |13         |20  |
# |2021-02-03|{Item_E -> 3, Item_C -> 1, Item_A -> 8} |12         |30  |
# +----------+----------------------------------------+-----------+----+

脚本:

aggs = df.agg(F.sum('Total Items'), F.sum('ColA')).head()

df = (df
    .select('*', F.explode('Item'))
    .groupBy('key')
    .agg(F.sum('value').alias('value'))
    .select(
        F.map_from_entries(F.collect_set(F.struct('key', 'value'))).alias('Item'),
        F.lit(aggs[0]).alias('Total Items'),
        F.lit(aggs[1]).alias('ColA'),
    )
)
df.show(truncate=0)
# +--------------------------------------------------------------------+-----------+----+
# |Item                                                                |Total Items|ColA|
# +--------------------------------------------------------------------+-----------+----+
# |{Item_C -> 3, Item_E -> 10, Item_A -> 12, Item_B -> 10, Item_D -> 5}|40         |60  |
# +--------------------------------------------------------------------+-----------+----+