python - 从 groupby 返回列表的最有效方法

Question

我有一个 130M 行的数据框，这是一个示例:

    id              id2     date         value
0   33208381500016  1927637 2014-07-31   120.0
1   77874276700016  3418498 2014-11-22   10.5
2   77874276700016  1174018 2014-11-22   8.4
3   77874276700016  1174018 2014-11-20   1.4
4   77874276700016  1643839 2014-06-27   4.2
5   77874276700016  1972929 2014-06-27   6.7
6   77874276700016  1972929 2014-06-27   12.7
7   77874276700016  1588191 2014-02-20   123.4
8   77874276700016  1966627 2014-02-20   973.1
9   77874276700016  1830252 2014-02-20   0.5

我需要对该数据帧(称为数据)执行groupby。对于像 sum 这样的简单 groupby 没问题:

data[['id','value']].groupby('id',as_index=False).sum()
time: 11.19s

但现在我需要检索另一列(或其长度)中的值列表。以下代码有效，但需要很长时间，是否有更有效的方法？

temp = data[['id','date','id2']].drop_duplicates()
temp.groupby('id',as_index = False).agg({'date': lambda x: set(x.tolist()),'id2':lambda x: len(set(x.tolist()))})
time: 159s

第一个问题:

是否有更有效的方法来计算每个 id 的唯一 id2 的数量，但仍然使用此 groupby？我的意思是我不想拆分两个 groupby，因为它可能需要更长的时间(执行一个 groupby 和 2 个聚合大约需要一个单独的 grouby 的 1.5 倍)。

第二个问题:

是否有更有效的方法来检索唯一日期列表？我知道它已在 this question 中得到解决但我不能简单地使用 .apply(list)。

Answer 1

要获取唯一日期，请使用 SeriesGroupBy.unique() .要计算每个组中唯一 id2 的数量，请使用 SeriesGroupBy.nunique() .

temp = data[['id', 'date', 'id2']].drop_duplicates()
temp.groupby('id', as_index=False).agg({'date': 'unique', 'id2': 'nunique'})

事先不删除重复项可能会更快——pandas 只需对所有数据迭代一次而不是两次。

data.groupby('id', as_index=False).agg({'date': 'unique', 'id2': 'nunique'})

编辑:

这里有一些基准。有趣的是，SeriesGroupBy.unique() 和 SeriesGroupBy.nunique() 似乎并不比使用集合快。但之前不要删除重复项。

import io

import pandas as pd

raw = io.StringIO("""\
id              id2     date         value
0   33208381500016  1927637 2014-07-31   120.0
1   77874276700016  3418498 2014-11-22   10.5
2   77874276700016  1174018 2014-11-22   8.4
3   77874276700016  1174018 2014-11-20   1.4
4   77874276700016  1643839 2014-06-27   4.2
5   77874276700016  1972929 2014-06-27   6.7
6   77874276700016  1972929 2014-06-27   12.7
7   77874276700016  1588191 2014-02-20   123.4
8   77874276700016  1966627 2014-02-20   973.1
9   77874276700016  1830252 2014-02-20   0.5
""")

data = pd.read_csv(raw, delim_whitespace=True)

def using_sets_drop_then_group():
    temp = data[['id', 'date', 'id2']].drop_duplicates()
    temp.groupby('id', as_index=False).agg({'date': lambda x: set(x),
                                            'id2': lambda x: len(set(x))})

def using_sets_drop_just_group():
    data.groupby('id', as_index=False).agg({'date': lambda x: set(x),
                                            'id2': lambda x: len(set(x))})

def using_unique_drop_then_group():
    temp = data[['id', 'date', 'id2']].drop_duplicates()
    temp.groupby('id', as_index=False).agg({'date': 'unique', 'id2': 'nunique'})

def using_unique_just_group():
    data.groupby('id', as_index=False).agg({'date': 'unique', 'id2': 'nunique'})

%timeit using_sets_drop_then_group()   # => 100 loops, best of 3: 4.82 ms per loop
%timeit using_sets_drop_just_group()   # => 100 loops, best of 3: 2.91 ms per loop
%timeit using_unique_drop_then_group() # => 100 loops, best of 3: 5.14 ms per loop
%timeit using_unique_just_group()      # => 100 loops, best of 3: 3.26 ms per loop

编辑:

在评论中，@ptrj 建议 SeriesGroupBy.unique() 和 SeriesGroupBy.nunique() 如果将日期转换为 datetime64 可能会更快>。唉，情况似乎并非如此，至少对于这个小数据样本而言。

data['parsed_date'] = pd.to_datetime(data['date'])

def using_sets_and_datetime64():
    data.groupby('id', as_index=False).agg({'parsed_date': lambda x: set(x),
                                            'id2': lambda x: len(set(x))})

def using_unique_and_datetime64():
    data.groupby('id', as_index=False).agg({'parsed_date': 'unique',
                                            'id2': 'nunique'})

%timeit using_sets_and_datetime64()    # => 100 loops, best of 3: 3.2 ms per loop
%timeit using_unique_and_datetime64()  # => 100 loops, best of 3: 3.53 ms per loop

编辑:

@MaxU 关于连接 100,000 个样本数据副本的建议确实导致 SeriesGroupBy.unique() 和 SeriesGroupBy.nunique() 优于 set.

large_data = pd.concat([data] * 10**5, ignore_index=True)

def using_sets():
    large_data.groupby('id', as_index=False).agg({'date': lambda x: set(x),
                                                  'id2': lambda x: len(set(x))})

def using_unique():
    large_data.groupby('id', as_index=False).agg({'date': 'unique',
                                                  'id2': 'nunique'})

def using_sets_and_datetime64():
    large_data.groupby('id', as_index=False).agg({'parsed_date': lambda x: set(x),
                                                  'id2': lambda x: len(set(x))})

def using_unique_and_datetime64():
    large_data.groupby('id', as_index=False).agg({'parsed_date': 'unique',
                                                  'id2': 'nunique'})

%timeit using_sets()                   # => 1 loops, best of 3: 295 ms per loop
%timeit using_unique()                 # => 1 loops, best of 3: 327 ms per loop
%timeit using_sets_and_datetime64()    # => 1 loops, best of 3: 5.02 s per loop
%timeit using_unique_and_datetime64()  # => 1 loops, best of 3: 248 ms per loop