Speeding up group apply in python

In my code it often happens that I need to calculate values for a group. For example, suppose there is the following data:

groups = [A, A, A, B, B, B]
values = [1, 2, 3, 0, 1, 1]


And I want to calculate the cumsum by group:

cumsum = [1, 3, 6, 0, 1, 2]


I always code this in the same way, as follows:

from collections import defaultdict
import numpy as np

N = 1000000
values = np.arange(N)
groups = np.random.choice(np.arange(100), N)

def group_apply(values, groups, func):
output = np.repeat(np.nan, len(values))
ixs = get_group_name_and_rows(groups)
for ix in ixs.itervalues():
output[ix] = func(values[ix])
return output

def get_group_name_and_rows(groups):
mapper = defaultdict(list)
for i, group in enumerate(groups):
mapper[group].append(i)
return mapper

group_apply(values, groups, np.cumsum)


This is now a big bottleneck in my code, and I was wondering if you know of any way to speed this up (perhaps cythonise it?)

Thanks!

You don't have to cythonize anything, it would be sufficient to rethink the way you represent data (e.g. groups) and use numpy to do the hard work (i.e. iterating over arrays of length N) efficiently. Here is my take on the problem, definitely not optimal but enough to give you an idea. Your version was limited by the defaultdict construction in the get_group_name_and_rows function, mine is limited by the two argsorts, so if you can have sorted data in the whole program (or at least sort only once for multiple invokations), then you can do much better.

from collections import defaultdict
import numpy as np

N = 1000000
values = np.arange(N)
groups = np.random.choice(np.arange(100), N)

def group_apply(values, groups, func):
output = np.repeat(np.nan, len(values))
ixs = get_group_name_and_rows(groups)
for ix in ixs.values():
output[ix] = func(values[ix])
return output

def get_group_name_and_rows(groups):
mapper = defaultdict(list)
for i, group in enumerate(groups):
mapper[group].append(i)
return mapper

def my_group_apply(values, groups, func):
output = np.repeat(np.nan, len(values))

# Sort values and groups. Note that a *stable* algorithm has to be used.
perm = groups.argsort(kind="mergesort")
sorted_values = values[perm]
sorted_groups = groups[perm]

# Indices of the non-zero elements in the diffs array represent where each group ends.
diffs = np.diff(sorted_groups)
separators = np.append(np.flatnonzero(diffs), len(values))

# Iterate over the groups and compute the func on each group.
left = 0
for right in separators:
output[left:right+1] = func(sorted_values[left:right+1])
left = right + 1

# Permute the output array into the original order.
iperm = perm.argsort()
return output[iperm]

output = group_apply(values, groups, np.cumsum)
my_output = my_group_apply(values, groups, np.cumsum)

assert (output == my_output).all()


Profiler output with N = 1000000:

Timer unit: 1e-06 s

Total time: 0.926312 s
File: test.py
Function: group_apply at line 8

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
8                                           @profile
9                                           def group_apply(values, groups, func):
10         1         5339   5339.0      0.6      output = np.repeat(np.nan, len(values))
11         1       684709 684709.0     73.9      ixs = get_group_name_and_rows(groups)
12       101          399      4.0      0.0      for ix in ixs.values():
13       100       235864   2358.6     25.5          output[ix] = func(values[ix])
14         1            1      1.0      0.0      return output

Total time: 0.207693 s
File: test.py
Function: my_group_apply at line 22

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
22                                           @profile
23                                           def my_group_apply(values, groups, func):
24         1         4770   4770.0      2.3      output = np.repeat(np.nan, len(values))
25
26                                               # Sort values and groups. Note that a *stable* algorithm has to be used.
27         1        94151  94151.0     45.3      perm = groups.argsort(kind="mergesort")
28         1        10533  10533.0      5.1      sorted_values = values[perm]
29         1        10652  10652.0      5.1      sorted_groups = groups[perm]
30
31                                               # Indices of the non-zero elements in the diffs array represent where each group ends.
32         1         2856   2856.0      1.4      diffs = np.diff(sorted_groups)
33         1         7574   7574.0      3.6      separators = np.append(np.flatnonzero(diffs), len(values))
34
35                                               # Iterate over the groups and compute the func on each group.
36         1            1      1.0      0.0      left = 0
37       101          120      1.2      0.1      for right in separators:
38       100         4967     49.7      2.4          output[left:right+1] = func(sorted_values[left:right+1])
39       100          170      1.7      0.1          left = right + 1
40
41                                               # Permute the output array into the original order.
42         1        63462  63462.0     30.6      iperm = perm.argsort()
43         1         8437   8437.0      4.1      return output[iperm]


And with N = 10000000:

Timer unit: 1e-06 s

Total time: 12.6474 s
File: test.py
Function: group_apply at line 8

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
8                                           @profile
9                                           def group_apply(values, groups, func):
10         1        63560  63560.0      0.5      output = np.repeat(np.nan, len(values))
11         1      7986616 7986616.0     63.1      ixs = get_group_name_and_rows(groups)
12       101          602      6.0      0.0      for ix in ixs.values():
13       100      4596623  45966.2     36.3          output[ix] = func(values[ix])
14         1            1      1.0      0.0      return output

Total time: 2.63137 s
File: test.py
Function: my_group_apply at line 22

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
22                                           @profile
23                                           def my_group_apply(values, groups, func):
24         1        64885  64885.0      2.5      output = np.repeat(np.nan, len(values))
25
26                                               # Sort values and groups. Note that a *stable* algorithm has to be used.
27         1      1354211 1354211.0     51.5      perm = groups.argsort(kind="mergesort")
28         1       130488 130488.0      5.0      sorted_values = values[perm]
29         1       144869 144869.0      5.5      sorted_groups = groups[perm]
30
31                                               # Indices of the non-zero elements in the diffs array represent where each group ends.
32         1        27008  27008.0      1.0      diffs = np.diff(sorted_groups)
33         1        74815  74815.0      2.8      separators = np.append(np.flatnonzero(diffs), len(values))
34
35                                               # Iterate over the groups and compute the func on each group.
36         1            2      2.0      0.0      left = 0
37       101          199      2.0      0.0      for right in separators:
38       100        40465    404.6      1.5          output[left:right+1] = func(sorted_values[left:right+1])
39       100          363      3.6      0.0          left = right + 1
40
41                                               # Permute the output array into the original order.
42         1       707766 707766.0     26.9      iperm = perm.argsort()
43         1        86302  86302.0      3.3      return output[iperm]

• Thans a lot, exactly what I was looking for! – mrz May 30 '17 at 8:40