根据列表中的值过滤 pyspark 数据框答案

【问题标题】：Filter the pyspark dataframe based on values in list根据列表中的值过滤 pyspark 数据框
【发布时间】：2020-09-22 15:45:48
【问题描述】：

我对 pyspark 还很陌生。我有 pyspark 数据框，其中包含有关特定人从品牌获得消息的次数的信息。它包含id、brand 和count 三列，如下所示。

|  id |  brand  | Count |
|:---:|:-------:|:-----:|
| 143 |  AD-ABC |   3   |
| 314 | AX-DEFG |   8   |
| 381 |  AD-ABC |   6   |
| 425 | AD-XYZP |   7   |
| 432 |  AD-GAF |   8   |
| 102 |  AD-GAF |   1   |
| 331 |  AX-ABC |   10  |
| 191 |  AD-GAF |   9   |
| 224 |  AD-GAF |   6   |

品牌列有点复杂，我想从品牌列中派生新列brand2，如下所示（在 - 后保留字符）

+-----+---------+-------+--------+
| id  |  brand  | Count | brand2 |
+-----+---------+-------+--------+
| 143 | AD-ABC  |     3 | ABC    |
| 314 | AX-DEFG |     8 | DEFG   |
| 381 | AD-ABC  |     6 | ABC    |
| 425 | AD-XYZP |     7 | XYZP   |
| 432 | AD-GAF  |     8 | GAF    |
| 102 | AD-GAF  |     1 | GAF    |
| 331 | AX-ABC  |    10 | ABC    |
| 191 | AD-GAF  |     9 | GAF    |
| 224 | AD-GAF  |     6 | GAF    |
+-----+---------+-------+--------+

我有一个非常大的列表，其中包含我想从数据框中过滤掉的品牌，如下所示

brand_subset = ['ABC', 'DEF', 'XYZP'] #The list is very large !!

我想要的数据框如下

+-----+---------+-------+--------+
| id  |  brand  | Count | brand2 |
+-----+---------+-------+--------+
| 143 | AD-ABC  |     3 | ABC    |
| 381 | AD-ABC  |     6 | ABC    |
| 425 | AD-XYZP |     7 | XYZP   |
| 331 | AX-ABC  |    10 | ABC    |
+-----+---------+-------+--------+

以上只是一个示例场景，实际上列表和表格都非常大。

任何帮助将不胜感激。（如果考虑到数据库的大小来优化解决方案会很好）

【问题讨论】：

标签： python pyspark

【解决方案1】：

拆分品牌列，得到第二个元素，然后用isin检查brand2是否在列表中：

import pyspark.sql.functions as F
brand_subset = ['ABC', 'DEF', 'XYZP']

(df.withColumn("brand2",F.split("brand","-")[1]).where(F.col("brand2")
                                          .isin(brand_subset))).show()

或：

(df.withColumn("brand2",F.split("brand","-")[1]).filter(F.col("brand2")
                                            .isin(brand_subset)).show()

+---+-------+-----+------+
| id|  brand|Count|brand2|
+---+-------+-----+------+
|143| AD-ABC|    3|   ABC|
|381| AD-ABC|    6|   ABC|
|425|AD-XYZP|    7|  XYZP|
|331| AX-ABC|   10|   ABC|
+---+-------+-----+------+

【讨论】：