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

Question

假设我有一个清单L=[[a,2],[a,3],[a,4],[b,4],[b,8],[b,9]]，我希望能够删除第三个元素，使它看起来如下所示：

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

我对火种很陌生，不知道我应该在这里做什么。

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|
+---+-----+