在python（native、pandas、numpy）中提取基于索引的值/行的最有效方法是什么？

Question

我试着寻找这个，但我的问题比简单的“这种方式”稍微复杂一些，我想。我正在寻找优化以下问题的方法：

我有一个文本文件，其中包含N行（数亿）和几列数据。问题是，出于某种原因，列1有索引，而其他列有值，如下所示：

1  2.3  4.7
2  2.8  2.4
1  1.9  3.1
2  6.7  3.1
... # and so forth (first column = index, thousands of unique indexes)

所以我想要的是读取这些文件并连接它们，然后选择所有具有唯一索引的行，并将它们放入每列的单独向量中。以上是：

# Vector 1
1  2.3  4.7
1  1.9  3.1
... # and so on
# Vector 2
2  2.8  2.4
2  6.7  3.1
... # and so on

我有一个可行的解决方案，但它需要很多时间，所以我正在寻找改进它的方法，因此我的标题（这是一个索引问题）。我正在寻找解决方案使用任何包，但我想熊猫是一个很好的候选人。下面是我目前的代码（它的相关部分）。你知道吗

# Load data
data = pd.concat([pd.read_csv(path,sep=r'\t',header=None,engine='python') for f in files])
# Sort data
for col in columns:
    d_dict[name][col] = [data[col][data[0] == i] for i in range(min,max+1)] # range min/max is the min/max of possible index values in column 1

数据的加载和排序都需要花费大量的时间，但是它的格式和我想要的一样，而且我认为它还保持了原始数据中行的原始顺序（如果这个假设是错误的，请告诉我：p）。你知道吗

我希望你有什么好主意，如何加快这个过程，因为现在只需要大约40分钟，这只是一个样本数量的数据，我必须处理。最终的数据集大约是原来的10倍大。然而，我在这方面只使用了20%的系统内存，所以我有空间在那里工作（但是如果需要的话我可以转储一些数据）。我也可以考虑与之平行。你知道吗

干杯！你知道吗

Answer 1

您可以对第一列进行argsort，并使用结果对其他列进行索引。你知道吗

由于您的索引不是太大的整数，我们可以使用一个技巧来获得argsort在我相信O（n）中。你知道吗

>>> from scipy import sparse
>>> import numpy as np
>>> 
# mock first column
>>> idx = np.random.randint(5_000, 15_000, (50_000_000,))
>>> 
# construct sparse one-hot matrix and convert from csr to csc
# for this conversion scipy must stably argsort the column indices
# but because it can exploit certain properties of the index set
# this is faster than using argsort directly
>>> imn, imx = idx.min(), idx.max()+1
>>> rng = np.arange(idx.size + 1)
>>> spM = sparse.csr_matrix((rng[:-1], idx-imn, rng), (idx.size, imx-imn)).tocsc()
>>> 
# extract the sorting index and the group boundaries
>>> sidx, bnds = spM.indices, spM.indptr
>>>
# use them to extract the groups, here we are using the first column
# itself as an example, the result will - sanity check - be groups
# consisting of copies of the group id 
# in practice, you would use another column in place of `idx` below
>>> groups = np.split(idx[sidx], bnds[1:-1])
>>> groups
# [array([5000, 5000, 5000, ..., 5000, 5000, 5000]), array([5001, 5001, 5001, ..., 5001, 5001, 5001]), array([5002, 5002, 5002, ..., 
#
# ... VERY long list