在 python 中从大型数据框中快速采样大量行

Question

我有一个非常大的数据框（大约 110 万行），我正在尝试对其进行采样。

我有一个索引列表（大约 70,000 个索引），我想从整个数据框中 select。

这是我到目前为止尝试过的方法，但是所有这些方法都花费了太多时间：

方法 1 - 使用 pandas :

sample = pandas.read_csv("data.csv", index_col = 0).reset_index()
sample = sample[sample['Id'].isin(sample_index_array)]

方法二：

我试图将所有采样行写入另一个 csv。

f = open("data.csv",'r')

out  = open("sampled_date.csv", 'w')
out.write(f.readline())

while 1:
    total += 1
    line = f.readline().strip()

    if line =='':
        break
    arr = line.split(",")

    if (int(arr[0]) in sample_index_array):
        out.write(",".join(e for e in (line)))

谁能推荐一个更好的方法？或者我如何修改它以使其更快？

谢谢

Answer 1

我们没有你的数据，所以这里有一个有两个选项的例子：

阅读后：使用pandas Index 对象通过.iloc selection method select 一个子集]
while reading: 带有skiprows parameter

给定

一组索引和一个（大）样本 DataFrame 写入 test.csv：

import pandas as pd
import numpy as np


indices = [1, 2, 3, 10, 20, 30, 67, 78, 900, 2176, 78776]

df = pd.DataFrame(np.random.randint(0, 100, size=(1000000, 4)), columns=list("ABCD"))
df.to_csv("test.csv", header=False)
df.info()

输出

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 1000000 entries, 0 to 999999
Data columns (total 4 columns):
A    1000000 non-null int32
B    1000000 non-null int32
C    1000000 non-null int32
D    1000000 non-null int32
dtypes: int32(4)
memory usage: 15.3 MB

代码

选项 1 - 阅读后

将示例索引列表转换为 Index 对象并切片加载的 DataFrame:

idxs = pd.Index(indices)   
subset = df.iloc[idxs, :]
print(subset)

甚至更快，但需要标量索引。

选项 2 - 阅读时（推荐）

我们可以编写一个谓词，在读取文件时保留 selected 索引（效率更高）：

pred = lambda x: x not in indices
data = pd.read_csv("test.csv", skiprows=pred, index_col=0, names="ABCD")
print(data)

另见 issue that led to extending skiprows。

结果

后面的选项产生相同的输出：

        A   B   C   D
1      74  95  28   4
2      87   3  49  94
3      53  54  34  97
10     58  41  48  15
20     86  20  92  11
30     36  59  22   5
67     49  23  86  63
78     98  63  60  75
900    26  11  71  85
2176   12  73  58  91
78776  42  30  97  96

在 python 中从大型数据框中快速采样大量行

Quickly sampling large number of rows from large dataframes in python

python

bigdata

sampling

dataframe

pandas