import time
from sqlglot.optimizer import optimize
TPCH_SCHEMA = {
"lineitem": {
"l_orderkey": "bigint",
"l_partkey": "bigint",
"l_suppkey": "bigint",
"l_linenumber": "bigint",
"l_quantity": "double",
"l_extendedprice": "double",
"l_discount": "double",
"l_tax": "double",
"l_returnflag": "string",
"l_linestatus": "string",
"l_shipdate": "string",
"l_commitdate": "string",
"l_receiptdate": "string",
"l_shipinstruct": "string",
"l_shipmode": "string",
"l_comment": "string",
},
"orders": {
"o_orderkey": "bigint",
"o_custkey": "bigint",
"o_orderstatus": "string",
"o_totalprice": "double",
"o_orderdate": "string",
"o_orderpriority": "string",
"o_clerk": "string",
"o_shippriority": "int",
"o_comment": "string",
},
"customer": {
"c_custkey": "bigint",
"c_name": "string",
"c_address": "string",
"c_nationkey": "bigint",
"c_phone": "string",
"c_acctbal": "double",
"c_mktsegment": "string",
"c_comment": "string",
},
"part": {
"p_partkey": "bigint",
"p_name": "string",
"p_mfgr": "string",
"p_brand": "string",
"p_type": "string",
"p_size": "int",
"p_container": "string",
"p_retailprice": "double",
"p_comment": "string",
},
"supplier": {
"s_suppkey": "bigint",
"s_name": "string",
"s_address": "string",
"s_nationkey": "bigint",
"s_phone": "string",
"s_acctbal": "double",
"s_comment": "string",
},
"partsupp": {
"ps_partkey": "bigint",
"ps_suppkey": "bigint",
"ps_availqty": "int",
"ps_supplycost": "double",
"ps_comment": "string",
},
"nation": {
"n_nationkey": "bigint",
"n_name": "string",
"n_regionkey": "bigint",
"n_comment": "string",
},
"region": {
"r_regionkey": "bigint",
"r_name": "string",
"r_comment": "string",
},
}
TPCDS_SCHEMA = {
"catalog_sales": {
"cs_sold_date_sk": "bigint",
"cs_sold_time_sk": "bigint",
"cs_ship_date_sk": "bigint",
"cs_bill_customer_sk": "bigint",
"cs_bill_cdemo_sk": "bigint",
"cs_bill_hdemo_sk": "bigint",
"cs_bill_addr_sk": "bigint",
"cs_ship_customer_sk": "bigint",
"cs_ship_cdemo_sk": "bigint",
"cs_ship_hdemo_sk": "bigint",
"cs_ship_addr_sk": "bigint",
"cs_call_center_sk": "bigint",
"cs_catalog_page_sk": "bigint",
"cs_ship_mode_sk": "bigint",
"cs_warehouse_sk": "bigint",
"cs_item_sk": "bigint",
"cs_promo_sk": "bigint",
"cs_order_number": "bigint",
"cs_quantity": "bigint",
"cs_wholesale_cost": "double",
"cs_list_price": "double",
"cs_sales_price": "double",
"cs_ext_discount_amt": "double",
"cs_ext_sales_price": "double",
"cs_ext_wholesale_cost": "double",
"cs_ext_list_price": "double",
"cs_ext_tax": "double",
"cs_coupon_amt": "double",
"cs_ext_ship_cost": "double",
"cs_net_paid": "double",
"cs_net_paid_inc_tax": "double",
"cs_net_paid_inc_ship": "double",
"cs_net_paid_inc_ship_tax": "double",
"cs_net_profit": "double",
},
"catalog_returns": {
"cr_returned_date_sk": "bigint",
"cr_returned_time_sk": "bigint",
"cr_item_sk": "bigint",
"cr_refunded_customer_sk": "bigint",
"cr_refunded_cdemo_sk": "bigint",
"cr_refunded_hdemo_sk": "bigint",
"cr_refunded_addr_sk": "bigint",
"cr_returning_customer_sk": "bigint",
"cr_returning_cdemo_sk": "bigint",
"cr_returning_hdemo_sk": "bigint",
"cr_returning_addr_sk": "bigint",
"cr_call_center_sk": "bigint",
"cr_catalog_page_sk": "bigint",
"cr_ship_mode_sk": "bigint",
"cr_warehouse_sk": "bigint",
"cr_reason_sk": "bigint",
"cr_order_number": "bigint",
"cr_return_quantity": "bigint",
"cr_return_amount": "double",
"cr_return_tax": "double",
"cr_return_amt_inc_tax": "double",
"cr_fee": "double",
"cr_return_ship_cost": "double",
"cr_refunded_cash": "double",
"cr_reversed_charge": "double",
"cr_store_credit": "double",
"cr_net_loss": "double",
},
"inventory": {
"inv_date_sk": "bigint",
"inv_item_sk": "bigint",
"inv_warehouse_sk": "bigint",
"inv_quantity_on_hand": "bigint",
},
"store_sales": {
"ss_sold_date_sk": "bigint",
"ss_sold_time_sk": "bigint",
"ss_item_sk": "bigint",
"ss_customer_sk": "bigint",
"ss_cdemo_sk": "bigint",
"ss_hdemo_sk": "bigint",
"ss_addr_sk": "bigint",
"ss_store_sk": "bigint",
"ss_promo_sk": "bigint",
"ss_ticket_number": "bigint",
"ss_quantity": "bigint",
"ss_wholesale_cost": "double",
"ss_list_price": "double",
"ss_sales_price": "double",
"ss_ext_discount_amt": "double",
"ss_ext_sales_price": "double",
"ss_ext_wholesale_cost": "double",
"ss_ext_list_price": "double",
"ss_ext_tax": "double",
"ss_coupon_amt": "double",
"ss_net_paid": "double",
"ss_net_paid_inc_tax": "double",
"ss_net_profit": "double",
},
"store_returns": {
"sr_returned_date_sk": "bigint",
"sr_return_time_sk": "bigint",
"sr_item_sk": "bigint",
"sr_customer_sk": "bigint",
"sr_cdemo_sk": "bigint",
"sr_hdemo_sk": "bigint",
"sr_addr_sk": "bigint",
"sr_store_sk": "bigint",
"sr_reason_sk": "bigint",
"sr_ticket_number": "bigint",
"sr_return_quantity": "bigint",
"sr_return_amt": "double",
"sr_return_tax": "double",
"sr_return_amt_inc_tax": "double",
"sr_fee": "double",
"sr_return_ship_cost": "double",
"sr_refunded_cash": "double",
"sr_reversed_charge": "double",
"sr_store_credit": "double",
"sr_net_loss": "double",
},
"web_sales": {
"ws_sold_date_sk": "bigint",
"ws_sold_time_sk": "bigint",
"ws_ship_date_sk": "bigint",
"ws_item_sk": "bigint",
"ws_bill_customer_sk": "bigint",
"ws_bill_cdemo_sk": "bigint",
"ws_bill_hdemo_sk": "bigint",
"ws_bill_addr_sk": "bigint",
"ws_ship_customer_sk": "bigint",
"ws_ship_cdemo_sk": "bigint",
"ws_ship_hdemo_sk": "bigint",
"ws_ship_addr_sk": "bigint",
"ws_web_page_sk": "bigint",
"ws_web_site_sk": "bigint",
"ws_ship_mode_sk": "bigint",
"ws_warehouse_sk": "bigint",
"ws_promo_sk": "bigint",
"ws_order_number": "bigint",
"ws_quantity": "bigint",
"ws_wholesale_cost": "double",
"ws_list_price": "double",
"ws_sales_price": "double",
"ws_ext_discount_amt": "double",
"ws_ext_sales_price": "double",
"ws_ext_wholesale_cost": "double",
"ws_ext_list_price": "double",
"ws_ext_tax": "double",
"ws_coupon_amt": "double",
"ws_ext_ship_cost": "double",
"ws_net_paid": "double",
"ws_net_paid_inc_tax": "double",
"ws_net_paid_inc_ship": "double",
"ws_net_paid_inc_ship_tax": "double",
"ws_net_profit": "double",
},
"web_returns": {
"wr_returned_date_sk": "bigint",
"wr_returned_time_sk": "bigint",
"wr_item_sk": "bigint",
"wr_refunded_customer_sk": "bigint",
"wr_refunded_cdemo_sk": "bigint",
"wr_refunded_hdemo_sk": "bigint",
"wr_refunded_addr_sk": "bigint",
"wr_returning_customer_sk": "bigint",
"wr_returning_cdemo_sk": "bigint",
"wr_returning_hdemo_sk": "bigint",
"wr_returning_addr_sk": "bigint",
"wr_web_page_sk": "bigint",
"wr_reason_sk": "bigint",
"wr_order_number": "bigint",
"wr_return_quantity": "bigint",
"wr_return_amt": "double",
"wr_return_tax": "double",
"wr_return_amt_inc_tax": "double",
"wr_fee": "double",
"wr_return_ship_cost": "double",
"wr_refunded_cash": "double",
"wr_reversed_charge": "double",
"wr_account_credit": "double",
"wr_net_loss": "double",
},
"call_center": {
"cc_call_center_sk": "bigint",
"cc_call_center_id": "string",
"cc_rec_start_date": "string",
"cc_rec_end_date": "string",
"cc_closed_date_sk": "bigint",
"cc_open_date_sk": "bigint",
"cc_name": "string",
"cc_class": "string",
"cc_employees": "bigint",
"cc_sq_ft": "bigint",
"cc_hours": "string",
"cc_manager": "string",
"cc_mkt_id": "bigint",
"cc_mkt_class": "string",
"cc_mkt_desc": "string",
"cc_market_manager": "string",
"cc_division": "bigint",
"cc_division_name": "string",
"cc_company": "bigint",
"cc_company_name": "string",
"cc_street_number": "string",
"cc_street_name": "string",
"cc_street_type": "string",
"cc_suite_number": "string",
"cc_city": "string",
"cc_county": "string",
"cc_state": "string",
"cc_zip": "string",
"cc_country": "string",
"cc_gmt_offset": "double",
"cc_tax_percentage": "double",
},
"catalog_page": {
"cp_catalog_page_sk": "bigint",
"cp_catalog_page_id": "string",
"cp_start_date_sk": "bigint",
"cp_end_date_sk": "bigint",
"cp_department": "string",
"cp_catalog_number": "bigint",
"cp_catalog_page_number": "bigint",
"cp_description": "string",
"cp_type": "string",
},
"customer": {
"c_customer_sk": "bigint",
"c_customer_id": "string",
"c_current_cdemo_sk": "bigint",
"c_current_hdemo_sk": "bigint",
"c_current_addr_sk": "bigint",
"c_first_shipto_date_sk": "bigint",
"c_first_sales_date_sk": "bigint",
"c_salutation": "string",
"c_first_name": "string",
"c_last_name": "string",
"c_preferred_cust_flag": "string",
"c_birth_day": "bigint",
"c_birth_month": "bigint",
"c_birth_year": "bigint",
"c_birth_country": "string",
"c_login": "string",
"c_email_address": "string",
"c_last_review_date": "string",
},
"customer_address": {
"ca_address_sk": "bigint",
"ca_address_id": "string",
"ca_street_number": "string",
"ca_street_name": "string",
"ca_street_type": "string",
"ca_suite_number": "string",
"ca_city": "string",
"ca_county": "string",
"ca_state": "string",
"ca_zip": "string",
"ca_country": "string",
"ca_gmt_offset": "double",
"ca_location_type": "string",
},
"customer_demographics": {
"cd_demo_sk": "bigint",
"cd_gender": "string",
"cd_marital_status": "string",
"cd_education_status": "string",
"cd_purchase_estimate": "bigint",
"cd_credit_rating": "string",
"cd_dep_count": "bigint",
"cd_dep_employed_count": "bigint",
"cd_dep_college_count": "bigint",
},
"date_dim": {
"d_date_sk": "bigint",
"d_date_id": "string",
"d_date": "string",
"d_month_seq": "bigint",
"d_week_seq": "bigint",
"d_quarter_seq": "bigint",
"d_year": "bigint",
"d_dow": "bigint",
"d_moy": "bigint",
"d_dom": "bigint",
"d_qoy": "bigint",
"d_fy_year": "bigint",
"d_fy_quarter_seq": "bigint",
"d_fy_week_seq": "bigint",
"d_day_name": "string",
"d_quarter_name": "string",
"d_holiday": "string",
"d_weekend": "string",
"d_following_holiday": "string",
"d_first_dom": "bigint",
"d_last_dom": "bigint",
"d_same_day_ly": "bigint",
"d_same_day_lq": "bigint",
"d_current_day": "string",
"d_current_week": "string",
"d_current_month": "string",
"d_current_quarter": "string",
"d_current_year": "string",
},
"household_demographics": {
"hd_demo_sk": "bigint",
"hd_income_band_sk": "bigint",
"hd_buy_potential": "string",
"hd_dep_count": "bigint",
"hd_vehicle_count": "bigint",
},
"income_band": {
"ib_income_band_sk": "bigint",
"ib_lower_bound": "bigint",
"ib_upper_bound": "bigint",
},
"item": {
"i_item_sk": "bigint",
"i_item_id": "string",
"i_rec_start_date": "string",
"i_rec_end_date": "string",
"i_item_desc": "string",
"i_current_price": "double",
"i_wholesale_cost": "double",
"i_brand_id": "bigint",
"i_brand": "string",
"i_class_id": "bigint",
"i_class": "string",
"i_category_id": "bigint",
"i_category": "string",
"i_manufact_id": "bigint",
"i_manufact": "string",
"i_size": "string",
"i_formulation": "string",
"i_color": "string",
"i_units": "string",
"i_container": "string",
"i_manager_id": "bigint",
"i_product_name": "string",
},
"promotion": {
"p_promo_sk": "bigint",
"p_promo_id": "string",
"p_start_date_sk": "bigint",
"p_end_date_sk": "bigint",
"p_item_sk": "bigint",
"p_cost": "double",
"p_response_target": "bigint",
"p_promo_name": "string",
"p_channel_dmail": "string",
"p_channel_email": "string",
"p_channel_catalog": "string",
"p_channel_tv": "string",
"p_channel_radio": "string",
"p_channel_press": "string",
"p_channel_event": "string",
"p_channel_demo": "string",
"p_channel_details": "string",
"p_purpose": "string",
"p_discount_active": "string",
},
"reason": {"r_reason_sk": "bigint", "r_reason_id": "string", "r_reason_desc": "string"},
"ship_mode": {
"sm_ship_mode_sk": "bigint",
"sm_ship_mode_id": "string",
"sm_type": "string",
"sm_code": "string",
"sm_carrier": "string",
"sm_contract": "string",
},
"store": {
"s_store_sk": "bigint",
"s_store_id": "string",
"s_rec_start_date": "string",
"s_rec_end_date": "string",
"s_closed_date_sk": "bigint",
"s_store_name": "string",
"s_number_employees": "bigint",
"s_floor_space": "bigint",
"s_hours": "string",
"s_manager": "string",
"s_market_id": "bigint",
"s_geography_class": "string",
"s_market_desc": "string",
"s_market_manager": "string",
"s_division_id": "bigint",
"s_division_name": "string",
"s_company_id": "bigint",
"s_company_name": "string",
"s_street_number": "string",
"s_street_name": "string",
"s_street_type": "string",
"s_suite_number": "string",
"s_city": "string",
"s_county": "string",
"s_state": "string",
"s_zip": "string",
"s_country": "string",
"s_gmt_offset": "double",
"s_tax_precentage": "double",
},
"time_dim": {
"t_time_sk": "bigint",
"t_time_id": "string",
"t_time": "bigint",
"t_hour": "bigint",
"t_minute": "bigint",
"t_second": "bigint",
"t_am_pm": "string",
"t_shift": "string",
"t_sub_shift": "string",
"t_meal_time": "string",
},
"warehouse": {
"w_warehouse_sk": "bigint",
"w_warehouse_id": "string",
"w_warehouse_name": "string",
"w_warehouse_sq_ft": "bigint",
"w_street_number": "string",
"w_street_name": "string",
"w_street_type": "string",
"w_suite_number": "string",
"w_city": "string",
"w_county": "string",
"w_state": "string",
"w_zip": "string",
"w_country": "string",
"w_gmt_offset": "double",
},
"web_page": {
"wp_web_page_sk": "bigint",
"wp_web_page_id": "string",
"wp_rec_start_date": "string",
"wp_rec_end_date": "string",
"wp_creation_date_sk": "bigint",
"wp_access_date_sk": "bigint",
"wp_autogen_flag": "string",
"wp_customer_sk": "bigint",
"wp_url": "string",
"wp_type": "string",
"wp_char_count": "bigint",
"wp_link_count": "bigint",
"wp_image_count": "bigint",
"wp_max_ad_count": "bigint",
},
"web_site": {
"web_site_sk": "bigint",
"web_site_id": "string",
"web_rec_start_date": "string",
"web_rec_end_date": "string",
"web_name": "string",
"web_open_date_sk": "bigint",
"web_close_date_sk": "bigint",
"web_class": "string",
"web_manager": "string",
"web_mkt_id": "bigint",
"web_mkt_class": "string",
"web_mkt_desc": "string",
"web_market_manager": "string",
"web_company_id": "bigint",
"web_company_name": "string",
"web_street_number": "string",
"web_street_name": "string",
"web_street_type": "string",
"web_suite_number": "string",
"web_city": "string",
"web_county": "string",
"web_state": "string",
"web_zip": "string",
"web_country": "string",
"web_gmt_offset": "string",
"web_tax_percentage": "double",
},
}
def rewrite_fixtures(in_path, out_path, schema, num, kind):
with open(out_path, "w", encoding="UTF-8") as fixture:
for i in range(num):
i = i + 1
with open(in_path.format(i=i), encoding="UTF-8") as file:
original = "\n".join(
line.rstrip()
for line in file.read().split(";")[0].split("\n")
if not line.startswith("--")
)
original = original.replace("`", '"').strip()
now = time.time()
try:
optimized = optimize(original, schema=schema)
except Exception as e:
print("****", i, e, "****")
continue
fixture.write(
f"""--------------------------------------
-- TPC-{kind} {i}
--------------------------------------
{original};
{optimized.sql(pretty=True)};
"""
)
print(i, time.time() - now)
rewrite_fixtures(
"/home/toby/dev/tpch/{i}.sql",
"/home/toby/dev/sqlglot/tests/fixtures/optimizer/tpc-h/tpc-h.sql",
TPCH_SCHEMA,
22,
"H",
)
rewrite_fixtures(
"/home/toby/dev/tpcds/query{i}.sql",
"/home/toby/dev/sqlglot/tests/fixtures/optimizer/tpc-ds/tpc-ds.sql",
TPCDS_SCHEMA,
99,
"DS",
)