How can I do an aggregation with a lot of different conditions in pandas efficiently?

Viewed 68

Problem

I have a dataframe data with a variable amount of categorial columns (CC) and one numerical column (NC) containing numbers. The combination of all CC is unique, I could also use them as an index. Now I am given a second dataframe choices with the same CC, but some cells are set to np.nan which is understood as a wildcard. I wish to aggregate the NC in data based on CC in choices.

Example

Here is an example in case of a summation:

data Fruit Color Size NC
0 Apple Green Small 1
1 Apple Yellow Medium 2
2 Banana Yellow Small 3
choices Fruit Color Size
0 NaN Yellow NaN
1 Apple Green Small
2 Apple NaN Large
3 Banana Red Small

What I want to obtain is:

Fruit Color Size Sum
0 NaN Yellow NaN 5
1 Apple Green Small 1
2 Apple NaN Large 0
3 Banana Red Small 0

In reality there are ~10 categories and up to several thousand rows in data and choices.

Using DataFrame.apply

For a single row in choices I can obtain the result doing something like (won't run, just the idea)

def aggregate_single_row(row, data):
    mask = (data == data)
    for c in row.columns:
        if not pd.isna(row[c]):
            mask &= (data[c] == row[c])
    return data[mask].sum()

and DataFrame.apply that function to choices. However, that requires Python-looping over choices which is rather inefficient. In fact, if len(choices) is considerably larger than len(data), it is much faster (up to 100x) to just convert everything to Python lists and do nested for-loops.

Simulating a "Wildcard Merge"

If there was a "wildcard-merge", it would be rather easy to obtain the sum using

CC = ['Fruit', 'Color', 'Size']
data.merge(choices, on=CC, wildcard=np.nan).groupby(CC).sum()

Inspired by that idea, I followed this SO answer to skip Python loops entirly by

  1. creating a unique index in choices,
  2. replacing any np.nan by a list of all possible values in that particular column,
  3. using explode successively on all columns to expand the wildcards,
  4. merging data and choices on the CC
  5. and grouping the merge result by the unique index of step 1.

Turns out, that this method is even slower than the apply-method, probably because of the huge amount of wildcards in choices.

Summary

So my observation is, that the code runs faster the more Python for-loops I use. Since the tale is that it should be the other way around, I wonder if you can come up with a better solution, that outperforms the naive "nested for-loop" approach.

0 Answers
Related