python - 从数据框 Spark 中过滤大量 ID

Question

我有一个格式类似于

的大型数据框

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

我需要根据包含大约 500 万个 ID 的列表过滤行。直截了当的方法是使用 isin 进行过滤，但性能确实很差。这个过滤器怎么做到的？

Answer 1

如果您致力于使用 Spark SQL 并且 isin 不再可扩展，那么内部等值连接应该是一个不错的选择。

首先将 id 列表转换为单列 DataFrame。如果这是本地集合

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

和加入:

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