从数据帧Spark中过滤大量ID

Question

我有一个大型数据框，其格式类似于

+-----+------+------+
|ID   |Cat   |date  |
+-----+------+------+
|12   | A    |201602|
|14   | B    |201601|
|19   | A    |201608|
|12   | F    |201605|
|11   | G    |201603|
+-----+------+------+

我需要根据一个包含大约5000万个ID的列表来过滤行。直截了当的方式是使用isin进行过滤，但性能确实很差。如何完成此过滤器？

Answer 1

如果您已经承诺使用Spark SQL并且isin不再进行扩展，则内部等连接应该是合适的。

首先将ID列表转换为单列DataFrame。如果这是本地收藏

ids_df = sc.parallelize(id_list).map(lambda x: (x, )).toDF(["id"])

和join：

df.join(ids_df, ["ID"], "inner")