使用用户定义的函数进行Pyspark数据帧聚合

时间:2016-05-17 18:25:20

标签: group-by pyspark-sql

如何使用'groupby(key).agg('带有用户定义的函数?具体来说,我需要每个键的所有唯一值列表[不计数]。

2 个答案:

答案 0 :(得分:2)

collect_set和collect_list(分别用于无序和有序结果)可用于后处理groupby结果。从简单的火花数据框开始

df = sqlContext.createDataFrame(
[('first-neuron', 1, [0.0, 1.0, 2.0]), 
('first-neuron', 2, [1.0, 2.0, 3.0, 4.0])], 
("neuron_id", "time", "V"))

让我们说目标是为每个神经元返回V列表的最长长度(按名称分组)

from pyspark.sql import functions as F
grouped_df = tile_img_df.groupby('neuron_id').agg(F.collect_list('V'))

我们现在已将V列表分组到列表列表中。由于我们想要最长的长度,我们可以运行

import pyspark.sql.types as sq_types
len_udf = F.udf(lambda v_list: int(np.max([len(v) in v_list])),
                  returnType = sq_types.IntegerType())

max_len_df = grouped_df.withColumn('max_len',len_udf('collect_list(V)'))

使用V列表的最大长度添加max_len列

答案 1 :(得分:1)

我发现pyspark.sql.functions.collect_set(col)能完成我想要的工作。