TPCX Blog Bodo Code

TPCX_Blog_bodo_code.py

import bodo
import pandas as pd
import numpy as np

@bodo.jit
def tpcx_bb_q26():
    store_sales = pd.read_parquet('s3://...')
    item = pd.read_parquet('s3://...')
    item2 = item[item['i_category'] == 'Books']
    sale_items = pd.merge(
        store_sales, item2, left_on='ss_item_sk', right_on='i_item_sk'
    )
    count1 = sale_items.groupby('ss_customer_sk')['ss_item_sk'].count()
    gp1 = sale_items.groupby('ss_customer_sk')['i_class_id']
    def id1(x): return (x == 1).sum()
    def id2(x): return (x == 2).sum()
    ...
    def id15(x): return (x == 15).sum()
    customer_i_class = gp1.agg((id1, id2, ... id15))
    customer_i_class['ss_item_count'] = count1
    customer_i_class = customer_i_class[customer_i_class.ss_item_count > 5]
    customer_i_class = customer_i_class.drop(['ss_item_count'], axis=1)
    customer_i_class = customer_i_class.sort_values('ss_customer_sk')
    return customer_i_class