获取嵌套列表中的前两个元素--pyspark

Question

我们说我有一个列表L=[[a,2],[a,3],[a,4],[b,4],[b,8],[b,9]] 使用pyspark我希望能够删除第三个元素，使其看起来像这样：

[a,2]
[a,3]
[b,4]
[b,8]

我是pyspark的新手，不知道我应该在这做什么。

Answer 1

你可以尝试这样的事情。
第一步是groupby key列和列表中的聚合值。然后使用udf获取列表的前两个值，然后展开该列。

df = sc.parallelize([('a',2),('a',3),('a',4),
                       ('b',4),('b',8),('b',9)]).toDF(['key', 'value'])
from pyspark.sql.functions import collect_list, udf, explode
from pyspark.sql.types import *

foo = udf(lambda x:x[0:2], ArrayType(IntegerType()))
df_list = (df.groupby('key').agg(collect_list('value')).
                   withColumn('values',foo('collect_list(value)')).
                   withColumn('value', explode('values')).
                   drop('values', 'collect_list(value)'))
df_list.show()

结果

+---+-----+
|key|value|
+---+-----+
|  b|    4|
|  b|    8|
|  a|    2|
|  a|    3|
+---+-----+