cython function returns all values in a single cell after groupby apply

Question

I was looking to speed up processing of a groupby operation and while it now processes much faster, the resulting dataframe is not what I want.

Make MultiIndexed dataframe with some data:

import pandas as pd
import numpy as np
import cython

data = np.round(np.random.randn(4, 3), 1)

df = pd.DataFrame(data, columns=pd.MultiIndex.from_product([['TOP'], ['A', 'B','C']]))
df

My cython function looks like this:

%load_ext cython
%%cython
import numpy as np

def func4(x):
    result = np.zeros([len(x)])
    for i in range(len(x)):
        result[i] = x[i][2] + x[i][1] - x[i][0]
    return result

And when I execute this function:

g = df.groupby(level=0, axis=1).apply(lambda x: func4(x.to_numpy()))

I get a series with all values in one cell:

What I would like to get instead is an indexed series like this:

I am able to get the result I want if I convert result to a dataframe at the end of the function, but then the groupby + apply becomes slow again:

def func4(x):
    result = np.zeros([len(x)])
    for i in range(len(x)):
        result[i] = x[i][2] + x[i][1] - x[i][0]
    return pd.DataFrame(result)

EDIT:

You can use this to produce a large dataframe for testing:

import pandas as pd
import itertools
import numpy as np
import string
import cython

col = [f'{a}{b}{c}' for a,b,c in list(itertools.product(string.ascii_uppercase[:20], repeat = 3))]
col1 = [str(i) for i in range(1, 10)] # change the 10 to a higher number to produce even more data
col2 = list(i for i in string.ascii_uppercase[:3])

all_keys = ['.'.join(i) for i in itertools.product(col, col1, col2)]
rng = pd.date_range(end=pd.Timestamp.today().date(), periods=6, freq='M')
df = pd.DataFrame(np.random.randint(0, 1000, size=(len(rng), len(all_keys))), columns=all_keys, index=rng)

def top_key(x):
    split = x.split('.', 2)
    return f'{split[0]}.{split[1]}'

df.columns = pd.MultiIndex.from_tuples([(top_key(c), c) for c in df.columns])

print(len(df.columns.get_level_values(0).unique()))
df.head(5)

UPDATE 2, performance of solutions:

Fastest, but resulting series has the issue described above.

Fastest with the correct result daatframe.

user_4745529 · Answer

Since Henry already answered the pandas part of your question, let me address the performance aspect. I don't really see the need for Cython here. As a rule of thumb, try to avoid loops over np.ndarrays whenever possible, i.e. use vectorized operations/functions instead of loops:

def func4_py_no_loop(x):
    result = np.zeros(x.shape[0])
    result = x[:, 2] + x[:, 1] - x[:, 0]
    return result

However, in case you really want to use Cython, here a few key points for further performance improvements of your initial function:

Type your loop variable i to reduce python overhead.
Don't use len multiple times. It's a python function and thus, has python overhead. Instead, use the .shape attribute of the np.ndarray.
Use typed memoryviews for fast and efficient access to np.ndarray's underlying memory.
Disable index checks by means of the boundscheck directive.
Similarly, you can disable checks for negative indices by the wraparound directive

%%cython

cimport numpy as np
import numpy as np
from cython cimport boundscheck, wraparound

@wraparound(False)
@boundscheck(False)
def func4_cy_loop(double[:, ::1] x):
    cdef int i, N = x.shape[0]
    cdef double[::1] result = np.zeros(N)
    for i in range(N):
        result[i] = x[i, 2] + x[i, 1] - x[i, 0]
    return result

Timing all the functions for

data = np.round(np.random.randn(20000, 3), 1)

with Henry's proposed solution yields:

In [10]: %timeit df.groupby(level=0, axis=1).aggregate(lambda x: func4_cy_your_version(x.to_numpy()))
8.13 ms ± 160 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [11]: %timeit df.groupby(level=0, axis=1).aggregate(lambda x: func4_py_no_loop(x.to_numpy()))
1.18 ms ± 24.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [12]: %timeit df.groupby(level=0, axis=1).aggregate(lambda x: func4_cy_loop(x.to_numpy()))
1.3 ms ± 116 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

user_15497888 · Answer

groupby.apply is a tricky function because it can produce both aggregated and non-aggregated values (and it doesn't always choose correctly).

This is specifically an aggregate so groupby.aggregate should be used:

def func4(x):
    result = np.zeros([len(x)])
    for i in range(len(x)):
        result[i] = x[i][2] + x[i][1] - x[i][0]
    return result


g = df.groupby(level=0, axis=1).aggregate(lambda x: func4(x.to_numpy()))

g:

   TOP
0  1.7
1  2.0
2  0.5
3 -1.1

(reproduceable with np.random.seed(5))

import numpy as np
import pandas as pd

np.random.seed(5)
data = np.round(np.random.randn(4, 3), 1)

df = pd.DataFrame(
    data,
    columns=pd.MultiIndex.from_product([['TOP'], ['A', 'B', 'C']])
)

The slower way (seriously don't do this) would be to pass the index from x and reconstruct a Series:

def func4(x, idx):
    result = np.zeros([len(x)])
    for i in range(len(x)):
        result[i] = x[i][2] + x[i][1] - x[i][0]
    return pd.Series(result, index=idx)


g = df.groupby(level=0, axis=1).apply(lambda x: func4(x.to_numpy(), x.index))

g:

   TOP
0  1.7
1  2.0
2  0.5
3 -1.1

cython function returns all values in a single cell after groupby apply

2 Answers