如何为1664599行3935列的超大pandas DataFrame计算Jaccard距离k近邻

问题原因 #

你当前使用BallTree的方案内存溢出主要有两个核心原因：

1. sklearn的BallTree不支持稀疏矩阵输入，会将你166万行、3935列的二值数据直接转为密集矩阵存储，再加上树结构本身的额外开销，60G内存无法承载。
2. 百万级样本、三千维度下的精确k近邻计算本身算力和内存开销极高，非定制化的通用实现很难在有限内存下跑通。

可行解决方案 #

方案1：用近似最近邻（ANN）实现（推荐）

如果没有强精确结果要求，优先选该方案，是工业界大规模近邻计算的标准实践，精度损失通常低于1%，内存和速度提升可达数十倍，60G内存完全足够。
推荐使用pynndescent库，原生支持Jaccard度量和稀疏矩阵输入，代码示例如下：

import pandas as pd
import numpy as np
import gc
from scipy.sparse import csr_matrix
import pynndescent

# 第一步：将原始DataFrame转为稀疏矩阵，大幅降低内存占用
sparse_data = csr_matrix(df.values.astype(bool))
# 释放原始df的内存
del df
gc.collect()

# 第二步：构建近似近邻索引，n_neighbors设为你需要的k=10
index = pynndescent.NNDescent(
    sparse_data, 
    metric="jaccard", 
    n_neighbors=10, 
    n_jobs=-1,
    random_state=42
)
index.prepare()

# 第三步：查询得到所有用户的近邻索引和对应Jaccard距离
neighbors, distances = index.query(sparse_data, k=10)

方案2：分块计算精确近邻

如果必须要精确结果，不能用近似方案，可以使用sklearn的分块距离计算接口，逐块计算只保留每个用户的top10近邻，避免一次性加载全量距离矩阵：

import numpy as np
import gc
from sklearn.metrics import pairwise_distances_chunked

k = 10
all_neighbors = []
all_distances = []

# working_memory单位为MB，可根据剩余内存调整，每次处理的样本量会自动适配
for dist_chunk in pairwise_distances_chunked(
    df, 
    metric="jaccard", 
    n_jobs=-1, 
    working_memory=20000
):
    # 取当前块每个用户的前k小距离的索引
    topk_idx = np.argpartition(dist_chunk, k, axis=1)[:, :k]
    topk_dist = np.take_along_axis(dist_chunk, topk_idx, axis=1)
    # 按距离从小到大排序
    sort_idx = np.argsort(topk_dist, axis=1)
    all_neighbors.append(np.take_along_axis(topk_idx, sort_idx, axis=1))
    all_distances.append(np.take_along_axis(topk_dist, sort_idx, axis=1))
    # 释放当前块的内存
    del dist_chunk
    gc.collect()

# 合并所有块的结果
neighbors = np.concatenate(all_neighbors, axis=0)
distances = np.concatenate(all_distances, axis=0)

额外优化建议 #

• 可以提前过滤有效行为过少的用户：比如交互商品数少于3个的用户，这类用户的Jaccard近邻参考价值很低，过滤后可以进一步降低计算规模。
• 如果对速度要求更高，也可以使用Faiss库的二进制索引，将二值特征编码为二进制向量后计算近似Jaccard近邻，速度比pynndescent更快。

内容的提问来源于stack exchange，提问作者rambutan