根据列表中的值筛选pyspark数据帧

Question

我对pyspark是相当陌生的。我有pyspark dataframe，它有关于某个人从某个品牌获得信息的次数的信息。它有三列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   |

brand列有点复杂，我想从brand列派生新的列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    |
+-----+---------+-------+--------+

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

任何帮助都将不胜感激。（如果在考虑数据库大小的情况下对解决方案进行优化，这将是一件好事）

Answer 1

拆分brand列并获取第二个元素，然后使用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|
+ -+   -+  -+   +