pyspark; 检查元素是否在collect_list中

Question

我正在处理一个数据框df，例如以下数据框：

df.show()

输出：

+----+------+
|keys|values|
+----+------+
|  aa| apple|
|  bb|orange|
|  bb|  desk|
|  bb|orange|
|  bb|  desk|
|  aa|   pen|
|  bb|pencil|
|  aa| chair|
+----+------+

我使用collect_set聚合来获得一组消除了重复元素的对象（或collect_list获得对象列表）。

df_new = df.groupby('keys').agg(collect_set(df.values).alias('collectedSet_values'))

结果数据帧如下：

df_new.show()

输出：

+----+----------------------+
|keys|collectedSet_values   |
+----+----------------------+
|bb  |[orange, pencil, desk]|
|aa  |[apple, pen, chair]   |
+----+----------------------+

我正在努力寻找一种方法来查看结果集中的对象（在列中collectedSet_values）中是否存在特定的关键字（例如“ chair” ）。我不想udf解决。

请评论您的解决方案/想法。

亲切的问候。

Answer 1

实际上有一个很好的函数array_contains可以为我们做到这一点。我们将它用于一组对象的方式与此处相同。要知道每组对象中是否存在单词“chair”，我们可以简单地执行以下操作：

df_new.withColumn('contains_chair', array_contains(df_new.collectedSet_values, 'chair')).show()

输出：

+----+----------------------+--------------+
|keys|collectedSet_values   |contains_chair|
+----+----------------------+--------------+
|bb  |[orange, pencil, desk]|false         |
|aa  |[apple, pen, chair]   |true          |
+----+----------------------+--------------+

这同样适用于 的结果collect_list。