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 defaultdictimport numpy as npN = 1000000values = 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 outputdef get_group_name_and_rows(groups):    mapper = defaultdict(list)    for i, group in enumerate(groups):        mapper[group].append(i)    return mappergroup_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!
Speeding up group apply in python
python;performance
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 defaultdictimport numpy as npN = 1000000values = 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 outputdef get_group_name_and_rows(groups):    mapper = defaultdict(list)    for i, group in enumerate(groups):        mapper[group].append(i)    return mapperdef 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 sTotal time: 0.926312 sFile: test.pyFunction: group_apply at line 8Line #      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 outputTotal time: 0.207693 sFile: test.pyFunction: my_group_apply at line 22Line #      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 sTotal time: 12.6474 sFile: test.pyFunction: group_apply at line 8Line #      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 outputTotal time: 2.63137 sFile: test.pyFunction: my_group_apply at line 22Line #      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]