获取嵌套列表中的前两个元素 - 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|
+---+-----+