import time,psycopg2,random

NROWS = 10000000

def query(s,*args,**kwargs):
    print(s)
    c.execute( s,*args,**kwargs )
    print( "%d rows", c.rowcount )

def bench( n, sql, args, name=None ):
    t = time.time()
    try:
        for i in range(n):
            c.execute( sql, args )
    finally:
        dt = time.time()-t
        print("%8.03fms %6d rows %10s %s" % (1e3*dt/n, c.rowcount, name, sql[:40]))

    # c.execute( "EXPLAIN "+sql, args )
    # for row in c:
    #     print( row[0][:60] )
def drop():
    try:
        query( "DROP TABLE test_array" );
        query("COMMIT")
    except Exception as e:
        print(e)
        query("ROLLBACK")

def build():
    try:
        query( "CREATE UNLOGGED TABLE test_array( id INTEGER NOT NULL, s TEXT )" )
        query("COMMIT")
    except Exception as e:
        print(e)
        query("ROLLBACK")

def fill():
    try:
        query( "INSERT INTO test_array SELECT n,n FROM generate_series(1,%s) n", (NROWS,) )
        query("COMMIT")
    except Exception as e:
        print(e)
        query("ROLLBACK")

db = psycopg2.connect("dbname=test")
c = db.cursor()
if 0:
    drop()
    build()
    fill()

    query("ALTER TABLE test_array ADD PRIMARY KEY(id)" )
    query("COMMIT")
    query("VACUUM ANALYZE test_array" )

for length in 1, 10, 100, 1000, 10000, 100000:
    if length <= 1000:
        ntest = 100
    else:
        ntest = 1
    ids = ",".join( str(n) for n in range(1, 1+length))
    bench( ntest, "SELECT * FROM test_array WHERE id =ANY(ARRAY[%s])"%ids, (), "Correlated" )
    bench( ntest, "SELECT * FROM test_array WHERE id IN (%s)"%ids, (), "Correlated" )
    bench( ntest, "SELECT * FROM unnest(ARRAY[%s]) id JOIN test_array USING (id)"%ids, (), "Correlated" )

    ids = ",".join( str(random.randint(1,NROWS)) for n in range(1, 1+length))
    bench( ntest, "SELECT * FROM test_array WHERE id =ANY(ARRAY[%s])"%ids, (), "Random" )
    bench( ntest, "SELECT * FROM test_array WHERE id IN (%s)"%ids, (), "Random" )
    bench( ntest, "SELECT * FROM unnest(ARRAY[%s]) id JOIN test_array USING (id)"%ids, (), "Random" )