138 lines
3.5 KiB
Python
138 lines
3.5 KiB
Python
import psycopg2
|
|
import argparse
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import sys
|
|
import subprocess
|
|
import multiprocessing.pool
|
|
|
|
parser = argparse.ArgumentParser(description='Generate TPC-DS reference results from Postgres.')
|
|
parser.add_argument(
|
|
'--sf', dest='sf', action='store', help='The TPC-DS scale factor reference results to generate', default=1
|
|
)
|
|
parser.add_argument(
|
|
'--query-dir',
|
|
dest='query_dir',
|
|
action='store',
|
|
help='The directory with queries to run',
|
|
default='extension/tpcds/dsdgen/queries',
|
|
)
|
|
parser.add_argument(
|
|
'--answer-dir',
|
|
dest='answer_dir',
|
|
action='store',
|
|
help='The directory where to store the answers',
|
|
default='extension/tpcds/dsdgen/answers/sf${SF}',
|
|
)
|
|
parser.add_argument(
|
|
'--duckdb-path',
|
|
dest='duckdb_path',
|
|
action='store',
|
|
help='The path to the DuckDB executable',
|
|
default='build/reldebug/duckdb',
|
|
)
|
|
parser.add_argument(
|
|
'--skip-load',
|
|
dest='skip_load',
|
|
action='store_const',
|
|
const=True,
|
|
help='Whether or not to skip loading',
|
|
default=False,
|
|
)
|
|
parser.add_argument(
|
|
'--query-list', dest='query_list', action='store', help='The list of queries to run (default = all)', default=''
|
|
)
|
|
parser.add_argument('--nthreads', dest='nthreads', action='store', type=int, help='The number of threads', default=0)
|
|
|
|
args = parser.parse_args()
|
|
|
|
con = psycopg2.connect(database='postgres')
|
|
c = con.cursor()
|
|
if not args.skip_load:
|
|
tpcds_dir = f'tpcds_sf{args.sf}'
|
|
|
|
q = f"""
|
|
CALL dsdgen(sf={args.sf});
|
|
EXPORT DATABASE '{tpcds_dir}' (DELIMITER '|');
|
|
"""
|
|
proc = subprocess.Popen([args.duckdb_path, "-c", q])
|
|
proc.wait()
|
|
if proc.returncode != 0:
|
|
exit(1)
|
|
|
|
# drop the previous tables
|
|
tables = [
|
|
'name',
|
|
'web_site',
|
|
'web_sales',
|
|
'web_returns',
|
|
'web_page',
|
|
'warehouse',
|
|
'time_dim',
|
|
'store_sales',
|
|
'store_returns',
|
|
'store',
|
|
'ship_mode',
|
|
'reason',
|
|
'promotion',
|
|
'item',
|
|
'inventory',
|
|
'income_band',
|
|
'household_demographics',
|
|
'date_dim',
|
|
'customer_demographics',
|
|
'customer_address',
|
|
'customer',
|
|
'catalog_sales',
|
|
'catalog_returns',
|
|
'catalog_page',
|
|
'call_center',
|
|
]
|
|
for table in tables:
|
|
c.execute(f'DROP TABLE IF EXISTS {table};')
|
|
|
|
with open(os.path.join(tpcds_dir, 'schema.sql'), 'r') as f:
|
|
schema = f.read()
|
|
|
|
c.execute(schema)
|
|
|
|
with open(os.path.join(tpcds_dir, 'load.sql'), 'r') as f:
|
|
load = f.read()
|
|
|
|
load = load.replace(f'{tpcds_dir}/', f'{os.getcwd()}/{tpcds_dir}/')
|
|
|
|
c.execute(load)
|
|
|
|
con.commit()
|
|
|
|
# get a list of all queries
|
|
queries = os.listdir(args.query_dir)
|
|
queries.sort()
|
|
|
|
answer_dir = args.answer_dir.replace('${SF}', args.sf)
|
|
|
|
if len(args.query_list) > 0:
|
|
passing_queries = [x + '.sql' for x in args.query_list.split(',')]
|
|
queries = [x for x in queries if x in passing_queries]
|
|
queries.sort()
|
|
|
|
|
|
def run_query(q):
|
|
print(q)
|
|
with open(os.path.join(args.query_dir, q), 'r') as f:
|
|
sql_query = f.read()
|
|
answer_path = os.path.join(os.getcwd(), answer_dir, q.replace('.sql', '.csv'))
|
|
c.execute(f'DROP TABLE IF EXISTS "query_result{q}"')
|
|
c.execute(f'CREATE TABLE "query_result{q}" AS ' + sql_query)
|
|
c.execute(f"COPY \"query_result{q}\" TO '{answer_path}' (FORMAT CSV, DELIMITER '|', HEADER, NULL 'NULL')")
|
|
|
|
|
|
if args.nthreads == 0:
|
|
for q in queries:
|
|
run_query(q)
|
|
else:
|
|
pool = multiprocessing.pool.ThreadPool(processes=args.nthreads)
|
|
|
|
pool.map(run_query, queries)
|