Fix tidb chbench

experiments/17-chbenchmark-baseline/common.sh

@@ -14,7 +14,6 @@ function run_tpcc_tidb() {
   mkdir -p $results_dir
   RECORD_DETAILED_STATS=1 COND_OUT=$results_dir python3 -m pytpcc.tpcc tidb "${args[@]}" &
   tpcc_pid=$!
-  popd
 }
 
 function log_workload_point() {

load_baseline.py

@@ -6,7 +6,9 @@
     PostgresCompatibleLoader,
     TiDBLoader,
     redshift_stress_test,
+    make_tidb_conn,
 )
+from workloads.IMDB_extended.workload_utils.database import PyodbcDatabase
 import time
 import pickle
 import numpy as np
@@ -129,9 +131,9 @@ def main():
         with open("good_adhoc_queries.sql", "w", encoding="utf-8") as f:
             f.write("\n".join(res))
 
-# if __name__ == "__main__":
-#     main()
-#     sys.exit(0)
+if __name__ == "__main__":
+    main()
+    sys.exit(0)
 
 import yaml
 

workloads/IMDB_extended/workload_utils/database.py

@@ -39,7 +39,7 @@ def __init__(self, connection, engine: Optional[Engine] = None) -> None:
         self._cursor = None
 
     def should_fetch(self, query, cur) -> bool:
-        is_select = query.strip().lower().startswith("select")
+        is_select = query.strip().lower().startswith("select") or query.strip().lower().startswith("with")
         has_rows = cur.rowcount is not None and cur.rowcount > 0
         if self._engine == Engine.Aurora or self._engine == Engine.Redshift:
             return is_select

workloads/chbenchmark/queries_tidb.sql

@@ -10,7 +10,7 @@ select n_name, extract(year from o_entry_d) as l_year, sum(ol_amount) as sum_pro
 select c_id, c_last, sum(ol_amount) as revenue, c_city, c_phone, n_name from customer, orders, order_line, nation where c_id = o_c_id and c_w_id = o_w_id and c_d_id = o_d_id and ol_w_id = o_w_id and ol_d_id = o_d_id and ol_o_id = o_id and o_entry_d <= ol_delivery_d and n_nationkey = ascii(cast(substring(c_state,1,1) as char(1))) group by c_id, c_last, c_city, c_phone, n_name order by revenue desc;
 select s_i_id, sum(s_order_cnt) as ordercount from stock, supplier, nation where mod((s_w_id * s_i_id),10000) = su_suppkey and su_nationkey = n_nationkey and n_name = 'Germany' group by s_i_id having sum(s_order_cnt) > (select sum(s_order_cnt) * .005 from stock, supplier, nation where mod((s_w_id * s_i_id),10000) = su_suppkey and su_nationkey = n_nationkey and n_name = 'Germany') order by ordercount desc;
 select o_ol_cnt,  sum(case when o_carrier_id = 1 or o_carrier_id = 2 then 1 else 0 end) as high_line_count, sum(case when o_carrier_id <> 1 and o_carrier_id <> 2 then 1 else 0 end) as low_line_count from orders, order_line where ol_w_id = o_w_id and ol_d_id = o_d_id and ol_o_id = o_id and o_entry_d <= ol_delivery_d group by o_ol_cnt order by o_ol_cnt;
-select c_count, count(*) as custdist from (select c_id, count(o_id) from customer left outer join orders on ( c_w_id = o_w_id and c_d_id = o_d_id and c_id = o_c_id and o_carrier_id > 8) group by c_id) as c_orders (c_id, c_count) group by c_count order by custdist desc, c_count desc;
+select c_count, count(*) as custdist from (select c_id, count(o_id) AS c_count from customer left outer join orders on ( c_w_id = o_w_id and c_d_id = o_d_id and c_id = o_c_id and o_carrier_id > 8) group by c_id) as c_orders group by c_count order by custdist desc, c_count desc;
 select 100.00 * sum(case when i_data like 'PR%' then ol_amount else 0 end) / (1+sum(ol_amount)) as promo_revenue from order_line, item where ol_i_id = i_id;
 with revenue (supplier_no, total_revenue) as (select mod((s_w_id * s_i_id),10000) as supplier_no, sum(ol_amount) as total_revenue from order_line, stock where ol_i_id = s_i_id and ol_supply_w_id = s_w_id group by mod((s_w_id * s_i_id),10000)) select su_suppkey, su_name, su_address, su_phone, total_revenue from supplier, revenue where su_suppkey = supplier_no and total_revenue = (select max(total_revenue) from revenue) order by su_suppkey;
 select i_name, substring(i_data, 1, 3) as brand, i_price, count(distinct (mod((s_w_id * s_i_id),10000))) as supplier_cnt from stock, item where i_id = s_i_id and i_data not like 'zz%' and (mod((s_w_id * s_i_id),10000) not in (select su_suppkey from supplier where su_comment like '%bad%')) group by i_name, substring(i_data, 1, 3), i_price order by supplier_cnt desc;