按组查找金额

Question

请帮助我理解 PySpark

我有一个数据框，如何通过 PySpark 找到每个参数的总和？

+----+-----+-----+-----+-----+-----+-----+
|o_id|col_1|col_3|col_9|ser_1|ser_2|ser_9|
+----+-----+-----+-----+-----+-----+-----+
|   0|  103|   17|   73|   c1|   c2|   c2|
|   1|  323|  245|   66|   c2|   c1|   c3|
|   2|  112|   96| 1452|   c3|   c2|   c1|
|   3|   46|  746|    4|   c3|   c2|   c3|
|   4|   82|  379|   78|   c1|   c1|   c2|
+----+-----+-----+-----+-----+-----+-----+

如何求所有出现的 c1 等的总和。

如何获得这种类型的数据透视表？

+----+-----+
| ser|  sum|
+----+-----+
|  c1| 2261| 
|  c2| 1333| 
|  c3|  228| 
+----+-----+

如果列更少，那么我会很容易地做.groupby()，但是有很多列（更多），我不明白如何在不使用 python 循环的情况下对所有内容进行分组，这将显着降低代码速度。

df_1 = (df.groupby('ser_1').sum('col_1'))
...
df_x = (df.groupby('ser_x').sum('col_x'))

还有更优雅的解决方案吗？

Answer 1

可以使用以下策略优化计算。

1. 创建ser_* 和col_* 列之间的映射以创建结构数组。
2. 将结构数组分解为行。
3. 按key 分组，并对分解后的行的value 求和。

工作示例

数据设置

from pyspark.sql import functions as F

data = [(0, 103, 17, 73, "c1", "c2", "c2",),
        (1, 323, 245, 66, "c2", "c1", "c3",),
        (2, 112, 96, 1452, "c3", "c2", "c1",),
        (3, 46, 746, 4, "c3", "c2", "c3",),
        (4, 82, 379, 78, "c1", "c1", "c2",), ]

df = spark.createDataFrame(data, ("o_id", "col_1", "col_3", "col_9", "ser_1", "ser_2", "ser_9",))

在ser_* 和col_* 列之间创建映射

这里的映射是硬编码的，因为col 和ser 之间没有基于它们的索引的一对一映射，因为col_3 从问题映射到ser_2。如果存在一对一映射，则可以通过代码生成组合。

ser_arrays = F.array(F.col("ser_1"), F.col("ser_2"), F.col("ser_9")).alias("ser")
col_arrays = F.array(F.col("col_1"), F.col("col_3"), F.col("col_9")).alias("col_val")

df_col_ser_map = df.select(F.arrays_zip(ser_arrays, col_arrays).alias("mapped_col_ser"))

df_col_ser_map.show(truncate=False)

#+---------------------------------+
#|mapped_col_ser                   |
#+---------------------------------+
#|[{c1, 103}, {c2, 17}, {c2, 73}]  |
#|[{c2, 323}, {c1, 245}, {c3, 66}] |
#|[{c3, 112}, {c2, 96}, {c1, 1452}]|
#|[{c3, 46}, {c2, 746}, {c3, 4}]   |
#|[{c1, 82}, {c1, 379}, {c2, 78}]  |
#+---------------------------------+

分解映射列

df_col_ser_exploded = df_col_ser_map.select(F.explode("mapped_col_ser").alias("exploded_col_ser"))

df_col_ser_exploded.show()

#+----------------+
# |exploded_col_ser|
# +----------------+
# |       {c1, 103}|
# |        {c2, 17}|
# |        {c2, 73}|
# |       {c2, 323}|
# |       {c1, 245}|
# |        {c3, 66}|
# |       {c3, 112}|
# |        {c2, 96}|
# |      {c1, 1452}|
# |        {c3, 46}|
# |       {c2, 746}|
# |         {c3, 4}|
# |        {c1, 82}|
# |       {c1, 379}|
# |        {c2, 78}|
# +----------------+

按键和聚合值分组

df_col_ser_exploded.groupBy((df_col_ser_exploded["exploded_col_ser"]["ser"]).alias("ser")) \
    .agg(F.sum(df_col_ser_exploded["exploded_col_ser"]["col_val"]).alias("sum")).show()


# +---+----+
# |ser| sum|
# +---+----+
# | c1|2261|
# | c3| 228|
# | c2|1333|
# +---+----+