tuplesort_btree_test.py

#!/usr/bin/env python
# tuplesort_test.py : Test the effectiveness of fast-path tuplesorting

import psycopg2
import subprocess
import time
import csv
from optparse import OptionParser

queries = [
			# Queries from dellstore database
			("create index test on orderlines (prod_id);",				"drop index test; checkpoint;"), 
			("create index test on orderlines (prod_id, quantity);",	"drop index test; checkpoint;"), 
			#("create unique index test on orderlines (prod_id);", "drop index test"), 
			#("create unique index test on orderlines (prod_id, quantity);", "drop index test"), 
			]

# store results of a given run in a dedicated csv file
def serialize_to_file(vals_dic, filename):
	wrt = csv.writer(open(filename, 'wb'), delimiter=',')

	# mark median run for runs of this 
	# query (or if there is an even number of elements, near enough)
	median_i = (len(vals_dic) + 1) / 2 - 1 

	for i, k in enumerate(vals_dic):
		wrt.writerow([k[0][0], time.ctime(k[0][1]), str(k[1]), '*' if i == median_i else 'n'])

def test_main():
	parser = OptionParser(description="")
	parser.add_option('-c', '--conninfo', type=str, help="libpq-style connection info string of database to connect to. "
						"Can be omitted, in which case we get details from our environment. "
						"You'll probably want to put this in double-quotes, like this: --conninfo \"hostaddr=127.0.0.1 port=5432 dbname=postgres user=postgres\". ", default="")
	parser.add_option('-r', '--runs', type=int, help="Number of times to run each query", default=5)
	parser.add_option('-w', '--warmcache', type=int, help="Number of times to execute each statement without recording, to warm the cache.", default=3)
	parser.add_option('-d', '--description', type=str, help="Mandatory description for this run of the tool, such as 'inlining', 'non-inlining' or 'HEAD'", default=None)

	args = parser.parse_args()[0]
	conn_str = args.conninfo

	conn = psycopg2.connect(conn_str)
	cur = conn.cursor()

	warm_cache_n = args.warmcache
	run_n = args.runs

	description = args.description
	
	if description is None:
		raise SystemExit("You must specify a description for this run of"
		"tuplesort_bree_test.py!") 


	for qry_pair in queries:
		vals = {}
		qry = qry_pair[0]
		qry_clean = qry_pair[1]
		print qry
		for i, j in enumerate(range(run_n + warm_cache_n)):
				before = time.time() 
				cur.execute(qry)
				duration = time.time() - before
				cur.execute(qry_clean)
				if i < warm_cache_n:
					print "Skip recording query that just warms cache"
					continue
				vals[(qry, before)] = duration
				print 'run %(run)d duration: %(dur)fms' % {"run":(i - warm_cache_n) + 1, "dur":duration}
				# Sort values for reference, and to locate the median value	
				sort_vals = sorted(vals.iteritems(), key=lambda tot_time: tot_time[1])
		
		serialize_to_file(sort_vals, description + " " + qry + ".csv")

if __name__== "__main__":
	test_main()