具有 1500 万条记录的 Spark 数据帧上的 KNN

Question

我有一个这样的 Pyspark 数据框：

0   [0.010904288850724697, -0.010935504920780659, ...   
1   [0.34882408380508423, -0.19240069389343262, -0...
2   [0.13833148777484894, -0.23080679774284363, -0...   
3   [0.12398581206798553, -0.4803846478462219, -0....
4   [0.16033919155597687, -0.06204992160201073, -0.

现在我想为所有这些数组找到 100 个最近邻。

这是我的尝试：

df_collect = df.toPandas()
features = np.array(df_collect.features.to_list())

knnobj = NearestNeighbors(n_neighbors=100).fit(features)
distance_mat, neighbours_mat = knnobj.kneighbors(features)

但由于 df 太大，它需要的时间太长。我知道我可以广播和并行化最后一步，但我无法找到如何将 spark df 拟合到scikit-learn knn 模型。有没有其他方法可以做到？

我还阅读了一些他们提到 ANN(Approximate Nearest Neighbor) Sparkit-Learn spark_sklearn 的文章，但我无法找到他们对最近邻的实现。谁能指导我下一步该怎么做

Answer 1

1. 仅使用 datatable、cuDF 或 dask 等库加载数据。它们总是比 Pandas 快。

2.通过将每列强制转换为可能的最小子类型，最多可减少 90% 的内存消耗。

3. 选择您熟悉或基于您需要的数据操作库。

4. 抽取 10-20% 的数据样本进行快速分析和实验。

5.思考向量并使用向量化函数。

6.选择像 CatBoost 这样的快速 ML 库来构建基线和进行特征工程。