Priority based categorization using pandas/python

Question

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')

Answer 1

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.

You should consider adopting the following code conventions to make your code easier to read:

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.

• original Priority based categorization using pandas/python
• one_group_by https://codereview.stackexchange.com/a/249481/230673
• np_select https://stackoverflow.com/a/63947686/14308614
• np_select_where https://codereview.stackexchange.com/a/249586/230673
• https://codereview.stackexchange.com/a/249486/230673 was not plotted because the time complexity was different

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

Answer 2

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)

Answer 3

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]
    })
    add_category(df)
    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

%timeit add_category(test_input)
%timeit add_category_dawson(test_input)
%timeit add_category_SO(test_input)

# 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()

%timeit add_category(test_input)
%timeit add_category_dawson(test_input)
%timeit add_category_SO(test_input)

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)