ing-bank/sparse_dot_topn

What it solves

比较非常大的特征向量以寻找最佳匹配项时，通常需要执行稀疏矩阵乘法，随后进行前 n 个结果的选择。标准方法在处理海量数据集时，可能会非常耗费内存且速度缓慢。本项目提供了一个高性能、并行化的实现，将前 n 个结果的选择直接集成到乘法过程中，以减少内存占用并提高速度。

How it works

该库使用带有 OpenMP 并行化的 C++ 扩展。它支持 CSR, CSC, 和 COO 稀疏矩阵格式（将后两者转换为 CSR 以进行处理）。在内部，它采用 max-heap（最大堆积）来收集前 n 个值，将选择的复杂度从 $O(n ext{ columns})$ 降低至 $O(top ext{-}n)$。用户也可以应用一个阈值来过滤掉低值结果，或指定一个预期的密度来优化内存预分配。

Who it’s for

处理大规模稀疏数据集（例如 TF-IDF 矩阵）的数据科学家和工程师，他们需要高效地检索 Top-k 最相似的项，或是在大规模尺度下进行高维向量特征匹配。

Highlights

• Integrated Top-N Selection: 集成了前 n 个结果的选择功能，将乘法与选择一步完成以节省内存。
• Parallelized Execution: 通过 OpenMP 实现并行化执行。
• Flexible Formats: Supports 32-bit 和 64-bit 整数与浮点数，支持 CSR, CSC, 和 COO 格式。
• Cluster Distribution: 包含如 zip_sp_matmul_topn 等工具，可将大型矩阵（10M+ 行）拆分到集群中并合并结果。
• High Performance: 在特定硬件（Apple M2 Pro）上处理大型 TF-IDF 矩阵时，比标准方法快上 6 倍。