Priority based categorization using pandas/python

Question

I've invoice related data in the below Dataframe and lists of codes

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    |  104 |  2  |
+---------+------+-----+

code lists are,

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

My task is to add a new category column based on the below specified Order of Priority.

1. Priority No.1 : if any invoice has more than 10 qty should be categorized as Mega. eg : sum of qty of invoice 3 is 12

2. Priority No.2 : from the rest of the invoice. if any code of the invoice is in the Milk list, then the category should be Healthy. eg : in invoice 2 code 106 is in Milk. hence, the Full invoice is categorized as Healthy. Irrespective of other items (code 101 & 105) are present in the invoice. As priorities are applied to the full invoice.

3. Priority No.3 : from the rest of the invoice, if any code of the invoice is in Juice list, then this has 2 parts

(3.1) if the sum of that juices qty is equal to 1, then category should be OneJuice. eg : invoice 1 has code 104 and qty 1.this invoice 1 will get OneJuice irrespective of other items (code 101) are present in the invoice. As priorities are applied to the full invoice.

(3.2) if the sum of that juices qty is greater than 1, category should be ManyJuice. eg : invoice 4 has code 104 & 105 and qty 1 + 1 = 2.

4. Priority No.4 : from the rest of the invoice, if any code of the invoice is in Hot list, then it should be categorized as HotLovers. Irrespective of other items are present in the invoice.

5. Priority No.5 : from the rest of the invoice, if any code of the invoice is in Dessert list, then it should be categorized as DessertLovers.

6. Finally, rest of all the 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    |  104 |  2  |     ManyJuice |
+---------+------+-----+---------------+

so far I have tried below. it works. but pretty naive and not pythonic at all. also when i applied this to the original datatset, the code is very 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')

So need help to cleanup this code for speed/efficiency and pythonic way.

Answer 1

You can try to use np.select

df['category'] = np.select([
    df.groupby('invoice')['qty'].transform('sum') >= 10,
    df['code'].isin(Milk).groupby(df.invoice).transform('any'),
    (df['qty']*df['code'].isin(Juice)).groupby(df.invoice).transform('sum') == 1,
    (df['qty']*df['code'].isin(Juice)).groupby(df.invoice).transform('sum') > 1,
    df['code'].isin(Hot).groupby(df.invoice).transform('any'),
    df['code'].isin(Dessert).groupby(df.invoice).transform('any')
],
    ['Mega','Healthy','OneJuice','ManyJuice','HotLovers','DessertLovers'],
    'Other'
)
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   104    2      ManyJuice

---

Micro-Benchmark

pd.show_versions()

commit           : None
python           : 3.7.5.final.0
python-bits      : 64
OS               : Linux
OS-release       : 4.4.0-18362-Microsoft
machine          : x86_64
processor        : x86_64
byteorder        : little
LC_ALL           : None
LANG             : C.UTF-8
LOCALE           : en_US.UTF-8

pandas           : 0.25.3
numpy            : 1.17.4

Data was created with

def make_data(n):
     return pd.DataFrame({
    'invoice':np.arange(n)//3,
    'code':np.random.choice(np.arange(101,112),n),
    'qty':np.random.choice(np.arange(1,8), n, p=[10/25,10/25,1/25,1/25,1/25,1/25,1/25])
})

Results

perfplot.show(
    setup=make_data,
    kernels=[get_category, get_with_np_select],
    n_range=[2**k for k in range(8, 20)],
    logx=True,
    logy=True,
    equality_check=False,
    xlabel='len(df)')