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中）。我不想采用udf解决方案。

请在评论中提出你的解决方案/想法。

此致敬礼。

- Ala Tarighati

---

分享信息总是一个好主意，所以谢谢你的分享 :) 但是，你能否重新格式化成问答格式？确保你的问题只包含实际的问题，并将你的解决方案作为答案添加。 - Florian

@Florian 没问题，我会修复的 :) - Ala Tarighati

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的结果也是适用的。

- Ala Tarighati