使用pandas,如何有效地按组对大型DataFrame进行抽样?

12 投票
1 回答
4065 浏览
提问于 2025-04-17 03:14

我正在尝试根据分组从一个数据框中抽取一些行。这里有个例子。假设我定义了以下数据:

from pandas import *
df = DataFrame({'group1' : ["a","b","a","a","b","c","c","c","c",
                            "c","a","a","a","b","b","b","b"],
                'group2' : [1,2,3,4,1,3,5,6,5,4,1,2,3,4,3,2,1],
                'value'  : ["apple","pear","orange","apple",
                            "banana","durian","lemon","lime",
                            "raspberry","durian","peach","nectarine",
                            "banana","lemon","guava","blackberry","grape"]})

如果我按 group1group2 分组,那么每个组中的行数如下:

In [190]: df.groupby(['group1','group2'])['value'].agg({'count':len})
Out[190]: 
      count
a  1  2    
   2  1    
   3  2    
   4  1    
b  1  2    
   2  2    
   3  1    
   4  1    
c  3  1    
   4  1    
   5  2    
   6  1    

(如果有更简洁的方法来计算这个,请告诉我。)

现在我想构建一个数据框,从每个组中随机选择一行。我的想法是这样做:

In [215]: from random import choice
In [216]: grouped = df.groupby(['group1','group2'])
In [217]: subsampled = grouped.apply(lambda x: df.reindex(index=[choice(range(len(x)))]))
In [218]: subsampled.index = range(len(subsampled))
In [219]: subsampled
Out[219]: 
    group1  group2  value
0   b       2       pear 
1   a       1       apple
2   b       2       pear 
3   a       1       apple
4   a       1       apple
5   a       1       apple
6   a       1       apple
7   a       1       apple
8   a       1       apple
9   a       1       apple
10  a       1       apple
11  a       1       apple

这样是可行的。不过,我的真实数据大约有250万行和12列。如果我用比较简单的方法自己构建数据结构,这个操作可以在几秒钟内完成。然而,我上面的实现却在30分钟内都没有完成(而且看起来并不是因为内存不足)。顺便提一下,当我尝试在R中实现这个时,最开始用的是 plyr,结果也没有在合理的时间内完成;不过,使用 data.table 的解决方案却很快就完成了。

我该如何让这个在 pandas 中快速运行呢?我想要喜欢这个包,所以请帮帮我!

1 个回答

9

我测试了一下用apply这个方法,发现当有很多子组的时候,速度会很慢。grouped的groups属性是一个字典,你可以直接从中选择索引:

subsampled = df.ix[(choice(x) for x in grouped.groups.itervalues())]

补充说明:从pandas版本0.18.1开始,itervalues在groupby对象上不再有效了,你可以直接使用.values

subsampled = df.ix[(choice(x) for x in grouped.groups.values())]

撰写回答