时间戳

首页 > 解决方案 > 从列表中获取用于计算分布的 StructType 输出

python - 从列表中获取用于计算分布的 StructType 输出

问题描述

问题：

我有一列包含冗余值列表，我需要在 PySpark 数据框的新列中将其转换为类似字典的格式。

场景：

这是我的 PySpark 数据框：

一种 C all_classes
10 RDK [1, 1, 1, 2, 2]
10 USW [1, 2, 2, 2, 2, 2, 2]
8 RDK [21、21、21、21、20]
8 RDJ [20, 20, 21]
10 RDK [45、45、45、45、2、2、2]
7 SSW [6、6、6、19、19]

我使用以下行创建：

my_df.groupBy(['A', 'C']).agg(collect_list("Class").alias("all_classes"))

现在我想要一个新列来整理数据，如下所示：
考虑上列第一行的数据：[1, 1, 1, 2, 2]

那需要转换成{1: 3, 2: 2}基本上是数字本身而不是数字。它被重复的次数。

我的尝试：

因为，我对 python 很熟悉，所以我选择编写这样的 UDF：

custom_collect_function = udf(lambda li: {k: li.count(k) for k in set(li)}, StructType(li, ArrayType(elementType=IntegerType()), False))

display(my_df.withColumn("Distribution", custom_collect_function(my_df.all_classes)))

我显然没有在这里设置 StructType 并且我欢迎任何其他/优化的方法。我也欢迎任何 Spark 方式的方法。

预期输出：

一种 C 分配 all_classes
10 RDK {1:3, 2:2} [1, 1, 1, 2, 2]
10 USW {1:1, 2:6} [1, 2, 2, 2, 2, 2, 2]
8 RDK {21:4, 20:1} [21、21、21、21、20]
8 RDJ {20:2, 21:1} [20, 20, 21]
10 RDK {45：4、2：3} [45、45、45、45、2、2、2]
7 SSW {6:3, 19:2} [6、6、6、19、19]

标签： pythonapache-sparkpysparkapache-spark-sql

解决方案

你可以试试这个：

from pyspark.sql import functions as F, Window

result = my_df.withColumn(
    'all_classes',
    F.collect_list('Class').over(Window.partitionBy('A', 'C'))
).groupBy(
    'A', 'C', 'all_classes', 'Class'
).agg(
    F.count('Class').alias('cnt_class')
).groupBy(
    'A', 'C', 'all_classes'
).agg(
    F.map_from_entries(
        F.collect_list(F.struct('Class', 'cnt_class'))
    ).alias('distribution')
)

result.show(truncate=False)
+---+---+----------------------------------------+-------------------------+
|A  |C  |all_classes                             |distribution             |
+---+---+----------------------------------------+-------------------------+
|7  |SSW|[6, 6, 6, 19, 19]                       |[6 -> 3, 19 -> 2]        |
|8  |RDK|[21, 21, 21, 21, 20]                    |[21 -> 4, 20 -> 1]       |
|10 |RDK|[1, 1, 1, 2, 2, 45, 45, 45, 45, 2, 2, 2]|[1 -> 3, 2 -> 5, 45 -> 4]|
|8  |RDJ|[20, 20, 21]                            |[20 -> 2, 21 -> 1]       |
|10 |USW|[1, 2, 2, 2, 2, 2, 2]                   |[1 -> 1, 2 -> 6]         |
+---+---+----------------------------------------+-------------------------+

结果有点不同，因为您的分组数据框以某种方式由带有(A, C) = (10, RDK). 另请注意，我使用了 MapType。获取 StructType 是不可能的，因为结构必须在给定列中具有相同的字段，在这种情况下并非如此。

推荐阅读