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
- creating a unique index in
choices, - replacing any
np.nanby a list of all possible values in that particular column, - using
explodesuccessively on all columns to expand the wildcards, - merging
dataandchoiceson the CC - 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.