Dataframe -从Pyspark Dataframe的Dict/Map列中排序并选择前4个键值

Question

我有一个pyspark Dataframe，它有两列，ID和计数，计数列是一个字典/映射<str,int>。计数中的值未排序，我试图对计数列内的值进行排序，仅根据值获得前4，并删除其余的其他键值

我

ID                        count
3004000304    {'A' -> 2, 'B' -> 4, 'C -> 5, 'D' -> 1, 'E' -> 9}
3004002756    {'B' -> 3, 'A' -> 8,'D' -> 3, 'C' -> 8, 'E' -> 1}

我想要这样的东西，在计数列中只选择基于值的前4个

ID                        count
3004000304    {'E' -> 9, 'C' -> 5, 'B' -> 4, 'A' -> 2}
3004002756    {'A' -> 8, 'C' -> 8, 'B' ->, 'D' -> 3}

我的方法

from pyspark.sql import functions as F
def sort_dict_f(x):
sorted_x = sorted(x.items(), key=operator.itemgetter(1))
return sorted_x
SorterUDF = F.udf(sort_dict_f)
df = old_df.withColumn('Sorted', SorterUDF("count"))
df.show()

如果在Pyspark中不可能，是否可以转换为pandas df然后按值排序前4 ?如有任何帮助，不胜感激

Answer 1

您可以首先explode映射，获得每个ID的最高计数的4行，然后将其重建为映射。

df = df.select('id', F.explode('count')) 
.withColumn('rn', F.expr('row_number() over (partition by id order by value desc)')) 
.filter('rn <= 4') 
.groupBy('id') 
.agg(F.map_from_entries(F.collect_list(F.struct('key', 'value'))))
df.show(truncate=False)