Priority based categorization using pandas/python

I have invoice and code data in the below Dataframes

Invoices

df = pd.DataFrame({
'invoice':[1,1,2,2,2,3,3,3,4,4,4,5,5,6,6,6,7],
'code':[101,104,105,101,106,106,104,101,104,105,111,109,111,110,101,114,112],
'qty':[2,1,1,3,2,4,7,1,1,1,1,4,2,1,2,2,1]
})

+---------+------+-----+
| invoice | code | qty |
+---------+------+-----+
|    1    |  101 |  2  |
+---------+------+-----+
|    1    |  104 |  1  |
+---------+------+-----+
|    2    |  105 |  1  |
+---------+------+-----+
|    2    |  101 |  3  |
+---------+------+-----+
|    2    |  106 |  2  |
+---------+------+-----+
|    3    |  106 |  4  |
+---------+------+-----+
|    3    |  104 |  7  |
+---------+------+-----+
|    3    |  101 |  1  |
+---------+------+-----+
|    4    |  104 |  1  |
+---------+------+-----+
|    4    |  105 |  1  |
+---------+------+-----+
|    4    |  111 |  1  |
+---------+------+-----+
|    5    |  109 |  4  |
+---------+------+-----+
|    5    |  111 |  2  |
+---------+------+-----+
|    6    |  110 |  1  |
+---------+------+-----+
|    6    |  101 |  2  |
+---------+------+-----+
|    6    |  114 |  2  |
+---------+------+-----+
|    7    |  112 |  1  |
+---------+------+-----+


Codes

Hot =  [103,109]
Juice =  [104,105]
Milk =  [106,107,108]
Dessert =  [110,111]


My task is to add a now column, category based on the following priorities:

1. If any invoice has more than $$\10\$$ qty it should be categorized as "Mega".
E.g. The total qty of invoice 3 is $$\12\$$ - $$\4 + 7 + 1\$$.

2. If any of the invoice's codes are in the milk list; the category should be "Healthy".
E.g. Invoice 2 contains the code 106 which is in the milk list. So the entire invoice is categorized as Healthy regardless of other items.

3. If any of the invoices's codes are in the juice list;

1. If the total qty of juices is equal to 1; the category should be "OneJuice".
E.g. Invoice 1 has code 104 and qty 1.

2. Otherwise; the category should be "ManyJuice".
E.g. Invoice 4 has codes 104 and 105 with a total qty of 2 - $$\1 + 1\$$.

4. If any of the invoices's codes are in the hot list; the category should be "HotLovers".

5. If any of the invoices's codes are in the dessert list; the category should be "DessertLovers".

6. All other other invoice should be categorized as "Others".

My desired output is as below.

+---------+------+-----+---------------+
| invoice | code | qty |    category   |
+---------+------+-----+---------------+
|    1    |  101 |  2  |    OneJuice   |
+---------+------+-----+---------------+
|    1    |  104 |  1  |    OneJuice   |
+---------+------+-----+---------------+
|    2    |  105 |  1  |    Healthy    |
+---------+------+-----+---------------+
|    2    |  101 |  3  |    Healthy    |
+---------+------+-----+---------------+
|    2    |  106 |  2  |    Healthy    |
+---------+------+-----+---------------+
|    3    |  106 |  4  |      Mega     |
+---------+------+-----+---------------+
|    3    |  104 |  7  |      Mega     |
+---------+------+-----+---------------+
|    3    |  101 |  1  |      Mega     |
+---------+------+-----+---------------+
|    4    |  104 |  1  |   ManyJuice   |
+---------+------+-----+---------------+
|    4    |  105 |  1  |   ManyJuice   |
+---------+------+-----+---------------+
|    4    |  111 |  1  |   ManyJuice   |
+---------+------+-----+---------------+
|    5    |  109 |  4  |   HotLovers   |
+---------+------+-----+---------------+
|    5    |  111 |  2  |   HotLovers   |
+---------+------+-----+---------------+
|    6    |  110 |  1  | DessertLovers |
+---------+------+-----+---------------+
|    6    |  101 |  2  | DessertLovers |
+---------+------+-----+---------------+
|    6    |  114 |  2  | DessertLovers |
+---------+------+-----+---------------+
|    7    |  112 |  1  |     Others    |
+---------+------+-----+---------------+


I have got the following. It works but it seems pretty naive and not at all Pythonic.
When I apply it to the original dataset the code is also very slow.

# Calculating Priority No.1
L = df.groupby(['invoice'])['qty'].transform('sum') >= 10
df_Large = df[L]['invoice'].to_frame()
df_Large['category'] = 'Mega'
df_Large.drop_duplicates(['invoice'], inplace=True)

# Calculating Priority No.2
df_1 = df[~L] # removing Priority No.1 calculated above
M = (df_1['code'].isin(Milk)
.groupby(df_1['invoice'])
.transform('any'))
df_Milk = df_1[M]['invoice'].to_frame()
df_Milk['category'] = 'Healthy'
df_Milk.drop_duplicates(['invoice'], inplace=True)

# Calculating Priority No.3

# 3.a Part -1

df_2 = df[~L & ~M]  # removing Priority No.1 & 2 calculated above
J_1 = (df_2['code'].isin(Juice)
.groupby(df_2['invoice'])
.transform('sum') == 1)
df_SM = df_2[J_1]['invoice'].to_frame()
df_SM['category'] = 'OneJuice'
df_SM.drop_duplicates(['invoice'], inplace=True)

# 3.b Part -2
J_2 = (df_2['code'].isin(Juice)
.groupby(df_2['invoice'])
.transform('sum') > 1)
df_MM = df_2[J_2]['invoice'].to_frame()
df_MM['category'] = 'ManyJuice'
df_MM.drop_duplicates(['invoice'], inplace=True)

# Calculating Priority No.4
df_3 = df[~L & ~M & ~J_1 & ~J_2]  # removing Priority No.1, 2 & 3 (a & b) calculated above
H = (df_3['code'].isin(Hot)
.groupby(df_3['invoice'])
.transform('any'))
df_Hot = df_3[H]['invoice'].to_frame()
df_Hot['category'] = 'HotLovers'
df_Hot.drop_duplicates(['invoice'], inplace=True)

# Calculating Priority No.5
df_4 = df[~L & ~M & ~J_1 & ~J_2 & ~H ] # removing Priority No.1, 2, 3 (a & b) and 4 calculated above
D = (df_4['code'].isin(Dessert)
.groupby(df_4['invoice'])
.transform('any'))
df_Dessert = df_4[D]['invoice'].to_frame()
df_Dessert['category'] = 'DessertLovers'
df_Dessert.drop_duplicates(['invoice'], inplace=True)

# merge all dfs
category = pd.concat([df_Large,df_Milk,df_SM,df_MM,df_Hot,df_Dessert], axis=0,sort=False, ignore_index=True)

# Final merge to the original dataset
df = df.merge(category,on='invoice', how='left').fillna(value='Others')

• Are the inputs guaranteed to be ordered? Will invoice 1 always come before invoice 2. – Peilonrayz Sep 17 '20 at 23:44
• line items of invoice numbers are guaranteed to be grouped together. regards to invoice number order. I can sort them prior to this operation – Tommy Sep 18 '20 at 5:52
• there is a bug in your code. If an invoice has 1 line with a juice code and a quantity more than 1, it classifies is as OneJuice instead of ManyJuice. Or your description is wrong – Maarten Fabré Sep 18 '20 at 7:21
• I guess the error is in (df_2['code'].isin(Juice) .groupby(df_2['invoice']) .transform('sum') > 1). This counts the number of lines of juices in an invoice, without taking the quantities into account – Maarten Fabré Sep 18 '20 at 8:36
• J_1 = (df_2["qty"] * df_2["code"].isin(Juice)).groupby( df_2["invoice"] ).transform("sum") == 1 and the same change for J_2 seems to fix it – Maarten Fabré Sep 18 '20 at 9:24

Your code is pretty impressive. Many python programmers don't know how to use pandas as well as you. Your code might not look very "Pythonic", but you did a great job utilizing vectorized methods with indexing. In this answer, I include one section on Python code conventions and a second attempting to optimizing your code.

Python Code Conventions

Many companies have standardized style guides that make code easier to read. This is invaluable when many people write to the same code base. Without consistency, the repo would degrade to a mess of idiosyncrasies.

1. Follow standard variable naming conventions: Google Python Style Guide On Naming
2. Include a space after commas: Google Python Style Guide On Spaces
# most python programmers use CaseLikeThis (pascal case) for class names
# constants are often written in CASE_LIKE_THIS (snake case)
SODA =  [101, 102]
HOT =  [103, 109]
JUICE =  [104, 105] # remember spaces after commas
MILK =  [106, 107, 108]
DESSERT =  [110, 111]


Attempt to Optimize

To optimize your code, you should time how long each step takes. This can be done by checking the clock before and after a segment of code.

import time

t0 = time.time() # check clock before (milliseconds elapsed since jan 1, 1970)
# segment you want to measure; something like your group by or merge...
t1 = time.time() # check clock after
time_to_run_step = t1 - t0


By measuring how long each step takes to run, you can focus your energy optimizing the slowest steps. For example, optimizing a 0.1 second operation to be 100x faster is less good than optimizing a 10 second operation to be 2x faster.

When thinking how to optimize your code, two questions came to mind:

1. Can we apply the priorities in backward order to avoid filtering already categorized priorities?
2. Can we perform all the group by work at the same time?

Group by and merge are expensive operations since they generally scale quadratically (# of invoices X # of codes). I bet these are the slowest steps in your code, but you should time it to check.

# Act 1: set up everything for the big group by
# priority 1
# will be setup at the end of Act 2

# priority 2
df['milk'] = df['code'].isin(MILK)

# priority 3.a
# priority 3.b
juice = df['code'].isin(JUICE)
df['juice_qty'] = df['qty']
df.loc[~juice, 'juice_qty'] = 0 # I thought df['juice_qty'][~juice] was intuitive, but it gave a warning https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy
# distinguish single from many juice in Act 2

# priority 4
df['hot'] = df['code'].isin(HOT)

# priority 5
df['dessert'] = df['code'].isin(DESSERT)

# Act 2: the big group by and merge
invoices = df.groupby(['invoice']).agg({
'qty': 'sum',
'milk': 'any',
'juice_qty': 'sum',
'hot': 'any',
'dessert': 'any',
}).rename(columns={
'qty': 'total', # this is renamed because joining with duplicate names leads to qty_x and qty_y
'juice_qty': 'juice_total',
})
# priority 1
invoices['mega'] = invoices['total'] >= 10

# priority 3.a
# priority 3.b
invoices['one_juice'] = invoices['juice_total'] == 1
invoices['many_juice'] = invoices['juice_total'] > 1

df = df.merge(invoices, on='invoice', how='left')

# Act 3: apply the categories
# apply the categories in reverse order to overwrite less important with the more important
df['category'] = 'Others'
df.loc[df['dessert_y'], 'category'] = 'DessertLovers'
df.loc[df['hot_y'], 'category'] = 'HotLovers'
df.loc[df['many_juice'], 'category'] = 'ManyJuice'
df.loc[df['one_juice'], 'category'] = 'OneJuice'
df.loc[df['milk_y'], 'category'] = 'Healthy'
df.loc[df['mega'], 'category'] = 'Mega'

df = df[['invoice', 'code', 'qty', 'category']] # get the columns you care about


@Tommy and @MaartenFabré noticed a bug with how single and many juice was categorized. I edited this answer with a correction.

Edit: There are quite a few answers for this question spanning into stack overflow as well. Below a summary as of 09/20/2020.

Performance was plotted using the code from https://stackoverflow.com/a/63947686/14308614

• Thank you. you've mentioned many valid points. really liked your thought process on applying the priorities on reverse order and doing all the grouping at one go . now I'm not a position to get hold of my laptop. once, available, will apply this to my original dataset and will share the results. tc – Tommy Sep 17 '20 at 8:59
• thank you. this code is much faster than than mine. however, there's a small issue. the category for invoice 4 is not correctly calculated. it should be ManyJuice. but your code calculates to OneJuice. I think it doesn't add up the individual qty total of codes codes 104 & 105 of invoice 4 – Tommy Sep 17 '20 at 13:40
• @MaartenFabré thank you. can you let me know what does invoices_one_juice dataframe means in your code. is it same as invoices['one_juice']? and where does all these codes to be placed? – Tommy Sep 18 '20 at 7:02
• You'll need something like invoices_one_juice = df[juice].groupby("invoice")["qty"].sum() == 1, df['one_juice'] = invoices_one_juice.reindex(df["invoice"]).fillna(False).values, df['many_juice'] = (~invoices_one_juice).reindex(df["invoice"]).fillna(False).values  – Maarten Fabré Sep 18 '20 at 7:05
• The speed may already suit your needs, but one last performance idea would be to batch the invoices for parallelization. Because invoices don't depend on each other, the invoices could easily be split up. – Dawson Sep 20 '20 at 7:53

Instead of grouping by the invoice on each category, I would reverse the logic. Group per invoice, and then classify that invoice.

categories = pd.concat(
classify_invoice(data) for invoice, data in df.groupby("invoice")
)

|    | 0             |
|---:|:--------------|
|  0 | OneJuice      |
|  1 | OneJuice      |
|  2 | Healthy       |
|  3 | Healthy       |
|  4 | Healthy       |
|  5 | Mega          |
|  6 | Mega          |
|  7 | Mega          |
|  8 | ManyJuice     |
|  9 | ManyJuice     |
| 10 | ManyJuice     |
| 11 | HotLovers     |
| 12 | HotLovers     |
| 13 | DessertLovers |
| 14 | DessertLovers |
| 15 | DessertLovers |
| 16 | Others        |


Then to add this to the result, you can assign.

result = df.assign(category=categories)


Here I used assign, which returns a new DataFrame. I do this on purpose, so you can keep your original DataFrame intact. Changes inplace to your original DataFrame can be a source of errors.

Classifier

Then we just need to design the classifier. Here we need a function that accepts a DataFrame that covers exactly 1 invoice, and returns a series with the category, with the same index as the invoice.

Priority 1

The priority 1 then is easy:

def classify_invoice(order: pd.DataFrame) -> pd.Series:
if order["qty"].sum() > 10:
return pd.Series("Mega", index=order.index)


Priority 2

Priority 2 is also very easy:

    milk_codes = {106, 107, 108}
if order["code"].isin(milk_codes).any():
return pd.Series("Healthy", index=order.index)


Notice that I renamed the variable Milk to milk_codes, since that better describes what it means, and that I converted it to a set, since that is the datastructure meant for containment checks

further priorities

def classify_invoice(order: pd.DataFrame) -> pd.Series:
if order["qty"].sum() > 10:
return pd.Series("Mega", index=order.index)

milk_codes = {106, 107, 108}
if order["code"].isin(milk_codes).any():
return pd.Series("Healthy", index=order.index)

juice_codes = {104, 105}
juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
if juices_amount == 1:
return pd.Series("OneJuice", index=order.index)
if juices_amount > 1:
return pd.Series("ManyJuice", index=order.index)

hot_codes = {103, 109}
if order["code"].isin(hot_codes).any():
return pd.Series("HotLovers", index=order.index)

dessert_codes = {110, 111}
if order["code"].isin(dessert_codes).any():
return pd.Series("DessertLovers", index=order.index)

return pd.Series("Others", index=order.index)


Testing

Since you offloaded the categorising to another function, you can test this in isolation

Variation

def classify_invoice2(order: pd.DataFrame) -> pd.Series:
if order["qty"].sum() > 10:
return "Mega"

milk_codes = {106, 107, 108}
if order["code"].isin(milk_codes).any():

return "Healthy"

juice_codes = {104, 105}
juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
if juices_amount == 1:
return "OneJuice"
if juices_amount > 1:
return "ManyJuice"

hot_codes = {103, 109}
if order["code"].isin(hot_codes).any():
return "HotLovers"

dessert_codes = {110, 111}
if order["code"].isin(dessert_codes).any():
return "DessertLovers"

return "Others"

df.join(
df.groupby("invoice")
.apply(classify_invoice2)
.rename("category"),
on = "invoice"
)


This is about as fast as my other solution and slightly simpler to follow.

micro optimizations

Now the codes get defined each groupby. I there are a lot of invoices, it might be faster to define them outside the method:

milk_codes = {106, 107, 108}
juice_codes = {104, 105}
hot_codes = {103, 109}
dessert_codes = {110, 111}

def classify_invoice3(order: pd.DataFrame) -> pd.Series:
if order["qty"].sum() > 10:
return "Mega"

if order["code"].isin(milk_codes).any():

return "Healthy"

juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
if juices_amount == 1:
return "OneJuice"
if juices_amount > 1:
return "ManyJuice"

if order["code"].isin(hot_codes).any():
return "HotLovers"

if order["code"].isin(dessert_codes).any():
return "DessertLovers"
return "Others"


categorical

Working with a categorical might be faster than with a column of strings too:

CATEGORIES = {
0: "Mega",
1: "Healthy",
2: "OneJuice",
3: "ManyJuice",
4: "HotLovers",
5: "DessertLovers",
6: "Others",
}

def classify_invoice4(order: pd.DataFrame) -> pd.Series:
if order["qty"].sum() > 10:
return 0

if order["code"].isin(milk_codes).any():

return 1

juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
if juices_amount == 1:
return 2
if juices_amount > 1:
return 3

if order["code"].isin(hot_codes).any():
return 4

if order["code"].isin(dessert_codes).any():
return 5
return 6

df.join(
(
df.groupby("invoice")
.apply(classify_invoice4)
.rename("category")
.astype(pd.Categorical(list(CATEGORIES)))
.cat.rename_categories(CATEGORIES)
),
on="invoice",
)


In the benchmark with the sample data this was slightly slower, but for larger datasets this might be faster

numpy

You can do this in numpy land too:

def classify_invoice_numpy(invoices, quantities, codes):
SODA = np.array([101, 102])
HOT = np.array([103, 109])
JUICE = np.array([104, 105])  # remember spaces after commas
MILK = np.array([106, 107, 108])
DESSERT = np.array([110, 111])

juices = np.isin(codes, JUICE)
milk = np.isin(codes, MILK)
hot = np.isin(codes, HOT)
dessert = np.isin(codes, DESSERT)

result = -np.ones(len(invoices), dtype=int)

for invoice in np.unique(invoices):
index = invoices == invoice

if quantities[index].sum() >= 10:
result[index] = 0
continue

if milk[index].any():
result[index] = 1
continue

juices_index = index & juices
if juices_index.any():
if quantities[juices_index].sum() == 1:
result[index] = 2
continue
else:
result[index] = 3
continue

if hot[index].any():
result[index] = 4
continue

if dessert[index].any():
result[index] = 5
continue

return result

def solution_maarten_numpy(data):
return data.assign(
category=pd.Series(
classify_invoice_numpy(
data["invoice"].values,
data["qty"].values,
data["code"].values,
),
index=data.index,
).map(CATEGORIES)
)


Benchmarking

I did some benchmarking

dummy data:

def dummy_data(
n: int = 100, lines_per_invoice: int = 3, seed: int = 0
) -> pd.DataFrame:
random_generator = np.random.default_rng(seed=seed)
samples = (
random_generator.normal(loc=lines_per_invoice, scale=2, size=n)
.round()
.astype(int)
)
samples = np.where(samples > 0, samples, 1)
invoices = np.repeat(np.arange(n), samples)
quantities = random_generator.integers(1, 10, size=len(invoices))
codes = random_generator.choice(np.arange(101, 112), size=len(invoices))
return pd.DataFrame(
{"invoice": invoices, "qty": quantities, "code": codes}
)


compare when there is something different

def compare_results(left, right):
differences = (left != right).any(axis=1)
return left[differences].merge(
right.loc[differences, "category"], left_index=True, right_index=True
)


benchmark

def benchmark(functions, size=100, lines_per_invoice=3, seed=0):

data_original = dummy_data(
n=size, lines_per_invoice=lines_per_invoice, seed=seed
)
yield data_original
benchmark_result = categorise_dawson(data_original)

for function in functions:
data = data_original.copy()
result = function(data)
try:
pd.testing.assert_frame_equal(result, benchmark_result)
except AssertionError:
print(f"method {function.__name__} differs from the benchmark")
#             print(result)
#             print(benchmark_result)
print(compare_results(benchmark_result, result))
#             pd.testing.assert_frame_equal(result, benchmark_result)
continue
try:
pd.testing.assert_frame_equal(data, data_original)
except AssertionError:
print(f"method {function.__name__} changes the original data")
continue

time = timeit.timeit(
"function(data)",
globals={"function": function, "data": data},
number=1,
)

yield function.__name__, time


calling it

data_originals = {}
sizes = 10, 100, 1000, 10000
functions = [
solution_maarten_1,
solution_maarten_2,
solution_maarten_3,
solution_maarten4,
solution_maarten_numpy,
categorise_dawson,
categorise_OP,
]

result_df = pd.DataFrame(index=[function.__name__ for function in functions])
for size in sizes:
data_original, *results = benchmark(functions=functions, size=size,)
data_originals[size] = data_original
result_df[size] = pd.Series(dict(results))

|                        |        10 |       100 |      1000 |     10000 |
|:-----------------------|----------:|----------:|----------:|----------:|
| solution_maarten_1     | 0.0077566 | 0.089533  | 0.838123  | 9.03633   |
| solution_maarten_2     | 0.0085086 | 0.0564532 | 0.521976  | 5.17024   |
| solution_maarten_3     | 0.0051805 | 0.0461194 | 0.545553  | 6.22027   |
| solution_maarten4      | 0.0091025 | 0.0647327 | 0.545063  | 5.88994   |
| solution_maarten_numpy | 0.0013638 | 0.0038171 | 0.0156193 | 0.977562  |
| categorise_dawson      | 0.0342312 | 0.0253829 | 0.0320662 | 0.0790319 |
| categorise_OP          | 0.0480042 | 0.0463131 | 0.0542139 | 0.150899  |


So my code starts faster for smaller sizes, but changes almost linearly with the size, while your and @dawsons code are almost constant for size

complete code

#!/usr/bin/env python
# coding: utf-8

# In[1]:

import numpy as np
import pandas as pd
import timeit

# In[2]:

def dummy_data(
n: int = 100, lines_per_invoice: int = 3, seed: int = 0
) -> pd.DataFrame:
random_generator = np.random.default_rng(seed=seed)
samples = (
random_generator.normal(loc=lines_per_invoice, scale=2, size=n)
.round()
.astype(int)
)
samples = np.where(samples > 0, samples, 1)
invoices = np.repeat(np.arange(n), samples)
quantities = random_generator.integers(1, 10, size=len(invoices))
codes = random_generator.choice(np.arange(101, 112), size=len(invoices))
return pd.DataFrame(
{"invoice": invoices, "qty": quantities, "code": codes}
)

# In[3]:

def compare_results(left, right):
differences = (left != right).any(axis=1)
return left[differences].merge(
right.loc[differences, "category"], left_index=True, right_index=True
)

# In[63]:

Soda = [101, 102]
Hot = [103, 109]
Juice = [104, 105]
Milk = [106, 107, 108]
Dessert = [110, 111]

def categorise_OP(df):
# Calculating Priority No.1
L = df.groupby(["invoice"])["qty"].transform("sum") >= 10
df_Large = df[L]["invoice"].to_frame()
df_Large["category"] = "Mega"
df_Large.drop_duplicates(["invoice"], inplace=True)

# Calculating Priority No.2
df_1 = df[~L]  # removing Priority No.1 calculated above
M = df_1["code"].isin(Milk).groupby(df_1["invoice"]).transform("any")
df_Milk = df_1[M]["invoice"].to_frame()
df_Milk["category"] = "Healthy"
df_Milk.drop_duplicates(["invoice"], inplace=True)

# Calculating Priority No.3

# 3.a Part -1

df_2 = df[~L & ~M]  # removing Priority No.1 & 2 calculated above
J_1 = (df_2["qty"] * df_2["code"].isin(Juice)).groupby(
df_2["invoice"]
).transform("sum") == 1
df_SM = df_2[J_1]["invoice"].to_frame()
df_SM["category"] = "OneJuice"
df_SM.drop_duplicates(["invoice"], inplace=True)

# 3.b Part -2
J_2 = (df_2["qty"] * df_2["code"].isin(Juice)).groupby(
df_2["invoice"]
).transform("sum") > 1
df_MM = df_2[J_2]["invoice"].to_frame()
df_MM["category"] = "ManyJuice"
df_MM.drop_duplicates(["invoice"], inplace=True)

# Calculating Priority No.4
df_3 = df[
~L & ~M & ~J_1 & ~J_2
]  # removing Priority No.1, 2 & 3 (a & b) calculated above
H = df_3["code"].isin(Hot).groupby(df_3["invoice"]).transform("any")
df_Hot = df_3[H]["invoice"].to_frame()
df_Hot["category"] = "HotLovers"
df_Hot.drop_duplicates(["invoice"], inplace=True)

# Calculating Priority No.5
df_4 = df[
~L & ~M & ~J_1 & ~J_2 & ~H
]  # removing Priority No.1, 2, 3 (a & b) and 4 calculated above
D = df_4["code"].isin(Dessert).groupby(df_4["invoice"]).transform("any")
df_Dessert = df_4[D]["invoice"].to_frame()
df_Dessert["category"] = "DessertLovers"
df_Dessert.drop_duplicates(["invoice"], inplace=True)

# merge all dfs
category = pd.concat(
[df_Large, df_Milk, df_SM, df_MM, df_Hot, df_Dessert],
axis=0,
sort=False,
ignore_index=True,
)

# Final merge to the original dataset
return df.merge(category, on="invoice", how="left").fillna(value="Others")

# In[7]:

SODA = [101, 102]
HOT = [103, 109]
JUICE = [104, 105]  # remember spaces after commas
MILK = [106, 107, 108]
DESSERT = [110, 111]

def categorise_dawson(df):
df = df.copy()
df["milk"] = df["code"].isin(MILK)

# priority 3.a
juice = df["code"].isin(JUICE)
invoices_one_juice = df[juice].groupby("invoice")["qty"].sum() == 1
df["one_juice"] = (
invoices_one_juice.reindex(df["invoice"]).fillna(False).values
)
# priority 3.b
df["many_juice"] = (
(~invoices_one_juice).reindex(df["invoice"]).fillna(False).values
)

# priority 4
df["hot"] = df["code"].isin(HOT)

# priority 5
df["dessert"] = df["code"].isin(DESSERT)

# Act 2: the big group by and merge
invoices = (
df.groupby(["invoice"])
.agg(
{
"qty": "sum",
"milk": "any",
"one_juice": "any",
"many_juice": "any",
"hot": "any",
"dessert": "any",
}
)
.rename(
columns={
"qty": "total",  # this is renamed because joining with duplicate names leads to qty_x and qty_y
}
)
)
# priority 1
invoices["mega"] = invoices["total"] >= 10

df = df.merge(invoices, on="invoice", how="left")

# Act 3: apply the categories
# apply the categories in reverse order to overwrite less important with the more important
df["category"] = "Others"
df.loc[df["dessert_y"], "category"] = "DessertLovers"
df.loc[df["hot_y"], "category"] = "HotLovers"
df.loc[df["many_juice_y"], "category"] = "ManyJuice"
df.loc[df["one_juice_y"], "category"] = "OneJuice"
df.loc[df["milk_y"], "category"] = "Healthy"
df.loc[df["mega"], "category"] = "Mega"

return df[
["invoice", "qty", "code", "category"]
]  # get the columns you care about

# In[72]:

def classify_invoice1(order: pd.DataFrame) -> pd.Series:
if order["qty"].sum() >= 10:
return pd.Series("Mega", index=order.index)

milk_codes = {106, 107, 108}
if order["code"].isin(milk_codes).any():
return pd.Series("Healthy", index=order.index)

juice_codes = {104, 105}
juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()

if juices_amount == 1:
return pd.Series("OneJuice", index=order.index)
if juices_amount > 1:
return pd.Series("ManyJuice", index=order.index)

hot_codes = {103, 109}
if order["code"].isin(hot_codes).any():
return pd.Series("HotLovers", index=order.index)

dessert_codes = {110, 111}
if order["code"].isin(dessert_codes).any():
return pd.Series("DessertLovers", index=order.index)

return pd.Series("Others", index=order.index)

def solution_maarten_1(data):
categories = pd.concat(
classify_invoice1(data) for invoice, data in data.groupby("invoice")
)
return data.assign(category=categories)

# In[14]:

def classify_invoice2(order: pd.DataFrame) -> pd.Series:
if order["qty"].sum() >= 10:
return "Mega"

milk_codes = {106, 107, 108}
if order["code"].isin(milk_codes).any():

return "Healthy"

juice_codes = {104, 105}
juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
if juices_amount == 1:
return "OneJuice"
if juices_amount > 1:
return "ManyJuice"

hot_codes = {103, 109}
if order["code"].isin(hot_codes).any():
return "HotLovers"

dessert_codes = {110, 111}
if order["code"].isin(dessert_codes).any():
return "DessertLovers"

return "Others"

def solution_maarten_2(data):
return data.join(
data.groupby("invoice").apply(classify_invoice2).rename("category"),
on="invoice",
)

# In[17]:

milk_codes = {106, 107, 108}
juice_codes = {104, 105}
hot_codes = {103, 109}
dessert_codes = {110, 111}

def classify_invoice3(order: pd.DataFrame) -> pd.Series:
if order["qty"].sum() >= 10:
return "Mega"

if order["code"].isin(milk_codes).any():
return "Healthy"

juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
if juices_amount == 1:
return "OneJuice"
if juices_amount > 1:
return "ManyJuice"

if order["code"].isin(hot_codes).any():
return "HotLovers"

if order["code"].isin(dessert_codes).any():
return "DessertLovers"
return "Others"

def solution_maarten_3(data):
return data.join(
data.groupby("invoice").apply(classify_invoice3).rename("category"),
on="invoice",
)

# In[20]:

CATEGORIES = {
0: "Mega",
1: "Healthy",
2: "OneJuice",
3: "ManyJuice",
4: "HotLovers",
5: "DessertLovers",
-1: "Others",
}

def classify_invoice4(order: pd.DataFrame) -> pd.Series:
if order["qty"].sum() >= 10:
return 0

if order["code"].isin(milk_codes).any():
return 1

juices_amount = order.loc[order["code"].isin(juice_codes), "qty"].sum()
if juices_amount == 1:
return 2
if juices_amount > 1:
return 3

if order["code"].isin(hot_codes).any():
return 4

if order["code"].isin(dessert_codes).any():
return 5
return -1

def solution_maarten4(data):
return data.join(
(
data.groupby("invoice")
.apply(classify_invoice4)
.map(CATEGORIES)
.rename("category")
),
on="invoice",
)

# In[24]:

def classify_invoice_numpy(invoices, quantities, codes):
SODA = np.array([101, 102])
HOT = np.array([103, 109])
JUICE = np.array([104, 105])  # remember spaces after commas
MILK = np.array([106, 107, 108])
DESSERT = np.array([110, 111])

juices = np.isin(codes, JUICE)
milk = np.isin(codes, MILK)
hot = np.isin(codes, HOT)
dessert = np.isin(codes, DESSERT)

result = -np.ones(len(invoices), dtype=int)

for invoice in np.unique(invoices):
index = invoices == invoice

if quantities[index].sum() >= 10:
result[index] = 0
continue

if milk[index].any():
result[index] = 1
continue

juices_index = index & juices
if juices_index.any():
if quantities[juices_index].sum() == 1:
result[index] = 2
continue
else:
result[index] = 3
continue

if hot[index].any():
result[index] = 4
continue

if dessert[index].any():
result[index] = 5
continue

return result

# In[25]:

def solution_maarten_numpy(data):
return data.assign(
category=pd.Series(
classify_invoice_numpy(
data["invoice"].values,
data["qty"].values,
data["code"].values,
),
index=data.index,
).map(CATEGORIES)
)

# In[28]:

import timeit

# In[52]:

def benchmark(functions, size=100, lines_per_invoice=3, seed=0):

data_original = dummy_data(
n=size, lines_per_invoice=lines_per_invoice, seed=seed
)
yield data_original
benchmark_result = categorise_dawson(data_original)

for function in functions:
data = data_original.copy()
result = function(data)
try:
pd.testing.assert_frame_equal(result, benchmark_result)
except AssertionError:
print(f"method {function.__name__} differs from the benchmark")
#             print(result)
#             print(benchmark_result)
print(compare_results(benchmark_result, result))
#             pd.testing.assert_frame_equal(result, benchmark_result)
continue
try:
pd.testing.assert_frame_equal(data, data_original)
except AssertionError:
print(f"method {function.__name__} changes the original data")
continue

time = timeit.timeit(
"function(data)",
globals={"function": function, "data": data},
number=1,
)

yield function.__name__, time

# In[89]:

data_originals = {}
sizes = 10, 100, 1000, 10000
functions = [
solution_maarten_1,
solution_maarten_2,
solution_maarten_3,
solution_maarten4,
solution_maarten_numpy,
categorise_dawson,
categorise_OP,
]

result_df = pd.DataFrame(index=[function.__name__ for function in functions])
for size in sizes:
data_original, *results = benchmark(functions=functions, size=size,)
data_originals[size] = data_original
result_df[size] = pd.Series(dict(results))

# In[94]:

print(result_df.to_markdown())

# In[99]:

result_df.T.plot(logx=True, logy=True)

• Thanks. Upvoted. your code works perfectly well. and its' way better written than mine. However, when it comes to speed/performance. this is very very slow. even compared to my code. – Tommy Sep 17 '20 at 13:32
• strange. if I time the different solutions with your sample data, it takes about 9ms with my solution, 19 ms for @josh-dawson's solution, and 32ms for your solution – Maarten Fabré Sep 17 '20 at 15:48
• appreciate the different variations of answers. I'll apply these to my original dataset and will get back to you. tks – Tommy Sep 17 '20 at 16:16
• I even moved completely to numpy space. This keep scaling about linearly with the size, while the other solutions remain almost constant (I checked to a size of 10_000) – Maarten Fabré Sep 18 '20 at 9:57
• Nice. I didn't know about np.select – Maarten Fabré Sep 19 '20 at 6:24

Here I provide a different approach to solve this problem more efficiently. Compared with OP's solution, the primary optimization comes in the following aspects:

• Calling isin four times for each item class (Dessert, Hot, Juice, Milk) is inefficient. A better approach is to join the original DataFrame df with a Series that maps each item to a class, and then apply pd.get_dummies to the new class column to perform one-hot encoding. My solution will operate on the class information directly, therefore the second step is not needed.

• Each item class is assigned a priority value that is aligned with its priority in the computation logic of the category value, i.e. Dessert < Hot < Juice < Milk. The computation logic could then be rewritten to the following:

1. Compute the total quantity, total juice quantity, and maximum priority value of each invoice;
2. If the total quantity > 10, the category value is "Mega";
3. If the maximum priority value is "Juice" and total quantity > 1, the category value is "ManyJuice";
4. Otherwise, assign the category value based on the maximum priority value.

In the implementation, the category column is of a categorical type INVOICE_TYPE and each category value has its corresponding numerical code. The priority value of each item class is the numerical code of the class's corresponding category.

• np.select is utilized to implement the if-elif-else logic in a vectorized manner. (Remark: for if-else logic, np.where / pd.DataFrame.where could be utilized instead.)

Solution:

import pandas as pd
import numpy as np

def add_category(df: pd.DataFrame, mega_threshold: int = 10):
# Invoice categories
INVOICE_TYPE = pd.CategoricalDtype([
"Others", "DessertLovers", "HotLovers", "ManyJuice", "OneJuice", "Healthy", "Mega"
], ordered=True)
CODE_OTHERS = 0  # Numerical code of 'Others' category

# Mapping from item classes to invoice category codes
class_values = pd.Series(
pd.Categorical(["DessertLovers", "HotLovers", "OneJuice", "Healthy"], dtype=INVOICE_TYPE).codes,
index=["Dessert", "Hot", "Juice", "Milk"]
)

# Mapping from item codes to class priority values, which are equivalent to corresponding invoice category codes
item_code_values = pd.Series(
class_values[["Hot", "Juice", "Juice", "Milk", "Milk", "Milk", "Hot", "Dessert", "Dessert"]].to_numpy(),
index=pd.RangeIndex(103, 112), name="item_value"
)

df_item_values = df.join(item_code_values, on="code")
df_item_values["juice_qty"] = (df_item_values["item_value"] == class_values["Juice"]) * df_item_values["qty"]

# Compute total quantity, total juice quantity, and maximum item priority value of each invoice by aggregation
df_invoice_info = df_item_values.groupby("invoice").agg({
"qty": "sum",
"juice_qty": "sum",
"item_value": "max"
})
df_invoice_info.columns = ["total_qty", "total_juice_qty", "max_item_value"]

## This version of aggregation has better readability but it turns out to be 2~3 times slower than the above
# df_invoice_info = df_item_values.groupby("invoice").agg(
#     total_qty=("qty", "sum"),
#     total_juice_qty=("juice_qty", "sum"),
#     max_item_value=("item_value", "max")
# )

max_invoice_item_values = df_invoice_info["max_item_value"]
max_invoice_item_values.fillna(CODE_OTHERS, inplace=True, downcast="int8")
is_mega = df_invoice_info["total_qty"] > mega_threshold
is_many_juice = ((max_invoice_item_values == class_values["Juice"]) &
(df_invoice_info["total_juice_qty"] > 1))

# Compute invoice category codes
invoice_type_codes = pd.Series(np.select(
[is_mega, is_many_juice],
pd.Categorical(["Mega", "ManyJuice"], dtype=INVOICE_TYPE).codes,
max_invoice_item_values),
index=df_invoice_info.index
)

# Join category codes with the original DataFrame and transform them to the categorical type INVOICE_TYPE
df["category"] = pd.Categorical.from_codes(invoice_type_codes[df["invoice"]], dtype=INVOICE_TYPE)

# For performance testing, returning a copy of df instead of modifying it in-place
# return df.assign(category=pd.Categorical.from_codes(invoice_type_codes[df["invoice"]], dtype=INVOICE_TYPE))

if __name__ == "__main__":
df = pd.DataFrame({
'invoice': [1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 6, 6, 6, 7],
'code': [101, 104, 105, 101, 106, 106, 104, 101, 104, 105, 111, 109, 111, 110, 101, 114, 112],
'qty': [2, 1, 1, 3, 2, 4, 7, 1, 1, 1, 1, 4, 2, 1, 2, 2, 1]
})
print(df)


Output:

    invoice  code  qty       category
0         1   101    2       OneJuice
1         1   104    1       OneJuice
2         2   105    1        Healthy
3         2   101    3        Healthy
4         2   106    2        Healthy
5         3   106    4           Mega
6         3   104    7           Mega
7         3   101    1           Mega
8         4   104    1      ManyJuice
9         4   105    1      ManyJuice
10        4   111    1      ManyJuice
11        5   109    4      HotLovers
12        5   111    2      HotLovers
13        6   110    1  DessertLovers
14        6   101    2  DessertLovers
15        6   114    2  DessertLovers
16        7   112    1         Others


Performance Testing Code for Jupyter Notebook execution (in the add_category function, a copy of df is returned instead of in-place modification) vs. @JoshDawson's solution and this solution on SO:

df = pd.DataFrame({
'invoice': [1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 6, 6, 6, 7],
'code': [101, 104, 105, 101, 106, 106, 104, 101, 104, 105, 111, 109, 111, 110, 101, 114, 112],
'qty': [2, 1, 1, 3, 2, 4, 7, 1, 1, 1, 1, 4, 2, 1, 2, 2, 1]
})

# Test input DataFrame from OP
test_input = df

# Test input constructed by duplicating the original DataFrame 10**5 times
# and modifying the output to differentiate the invoice ids in each copy
test_input = pd.concat([df] * 10**5, ignore_index=True)
test_input["invoice"] += test_input.index // df.shape[0] * df["invoice"].max()



Performance testing results on original DataFrame from OP:

11.9 ms ± 422 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
17.5 ms ± 357 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
9.52 ms ± 106 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)


Performance testing results on large DataFrame:

411 ms ± 3.65 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1 s ± 5.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1.1 s ± 10.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

• Thank you & upvoted. please do check the answer I received in SO platform which stands out above all. (stackoverflow.com/questions/63908535/…) – Tommy Sep 20 '20 at 5:58
• The code is shorter and easier to comprehend but it is slower than other solutions on large datasets. I've added the benchmarking results in my post. My solution prioritized performance over code length. For example, a Categorical type rather than string is used as the type of the result column, which will lead to better performance for future processing and also save space. I also tried to keep my code as readable as possible. – GZ0 Sep 20 '20 at 7:08
• when I checked in my actual dataset with 11Mn rows and 10 columns that code took only 48seconds. – Tommy Sep 20 '20 at 7:33
• Would you mind measuring the performance of my solution and JoshDawson's updated code? I'm interested to see the results. I'm perfectly fine if the performance of SO code meets your requirement and you prefer that solution over mine because of its conciseness. Meanwhile, I think that solution still has room for further improvements based on Code Review standards. – GZ0 Sep 20 '20 at 10:50