1from __future__ import print_function
2import os
3import os.path
4from time import time
5import numpy
6import random
7
8# in order to always generate the same random sequence
9random.seed(19)
10
11
12def fill_arrays(start, stop):
13    col_i = numpy.arange(start, stop, dtype=numpy.int32)
14    if userandom:
15        col_j = numpy.random.uniform(0, nrows, stop - start)
16    else:
17        col_j = numpy.array(col_i, dtype=numpy.float64)
18    return col_i, col_j
19
20# Generator for ensure pytables benchmark compatibility
21
22
23def int_generator(nrows):
24    step = 1000 * 100
25    j = 0
26    for i in range(nrows):
27        if i >= step * j:
28            stop = (j + 1) * step
29            if stop > nrows:  # Seems unnecessary
30                stop = nrows
31            col_i, col_j = fill_arrays(i, stop)
32            j += 1
33            k = 0
34        yield (col_i[k], col_j[k])
35        k += 1
36
37
38def int_generator_slow(nrows):
39    for i in range(nrows):
40        if userandom:
41            yield (i, float(random.randint(0, nrows)))
42        else:
43            yield (i, float(i))
44
45
46def open_db(filename, remove=0):
47    if remove and os.path.exists(filename):
48        os.remove(filename)
49    con = sqlite.connect(filename)
50    cur = con.cursor()
51    return con, cur
52
53
54def create_db(filename, nrows):
55    con, cur = open_db(filename, remove=1)
56    cur.execute("create table ints(i integer, j real)")
57    t1 = time()
58    # This is twice as fast as a plain loop
59    cur.executemany("insert into ints(i,j) values (?,?)", int_generator(nrows))
60    con.commit()
61    ctime = time() - t1
62    if verbose:
63        print("insert time:", round(ctime, 5))
64        print("Krows/s:", round((nrows / 1000.) / ctime, 5))
65    close_db(con, cur)
66
67
68def index_db(filename):
69    con, cur = open_db(filename)
70    t1 = time()
71    cur.execute("create index ij on ints(j)")
72    con.commit()
73    itime = time() - t1
74    if verbose:
75        print("index time:", round(itime, 5))
76        print("Krows/s:", round(nrows / itime, 5))
77    # Close the DB
78    close_db(con, cur)
79
80
81def query_db(filename, rng):
82    con, cur = open_db(filename)
83    t1 = time()
84    ntimes = 10
85    for i in range(ntimes):
86        # between clause does not seem to take advantage of indexes
87        # cur.execute("select j from ints where j between %s and %s" % \
88        cur.execute("select i from ints where j >= %s and j <= %s" %
89                    # cur.execute("select i from ints where i >= %s and i <=
90                    # %s" %
91                    (rng[0] + i, rng[1] + i))
92        results = cur.fetchall()
93    con.commit()
94    qtime = (time() - t1) / ntimes
95    if verbose:
96        print("query time:", round(qtime, 5))
97        print("Mrows/s:", round((nrows / 1000.) / qtime, 5))
98        print(results)
99    close_db(con, cur)
100
101
102def close_db(con, cur):
103    cur.close()
104    con.close()
105
106if __name__ == "__main__":
107    import sys
108    import getopt
109    try:
110        import psyco
111        psyco_imported = 1
112    except:
113        psyco_imported = 0
114
115    usage = """usage: %s [-v] [-p] [-m] [-i] [-q] [-c] [-R range] [-n nrows] file
116            -v verbose
117            -p use "psyco" if available
118            -m use random values to fill the table
119            -q do query
120            -c create the database
121            -i index the table
122            -2 use sqlite2 (default is use sqlite3)
123            -R select a range in a field in the form "start,stop" (def "0,10")
124            -n sets the number of rows (in krows) in each table
125            \n""" % sys.argv[0]
126
127    try:
128        opts, pargs = getopt.getopt(sys.argv[1:], 'vpmiqc2R:n:')
129    except:
130        sys.stderr.write(usage)
131        sys.exit(0)
132
133    # default options
134    verbose = 0
135    usepsyco = 0
136    userandom = 0
137    docreate = 0
138    createindex = 0
139    doquery = 0
140    sqlite_version = "3"
141    rng = [0, 10]
142    nrows = 1
143
144    # Get the options
145    for option in opts:
146        if option[0] == '-v':
147            verbose = 1
148        elif option[0] == '-p':
149            usepsyco = 1
150        elif option[0] == '-m':
151            userandom = 1
152        elif option[0] == '-i':
153            createindex = 1
154        elif option[0] == '-q':
155            doquery = 1
156        elif option[0] == '-c':
157            docreate = 1
158        elif option[0] == "-2":
159            sqlite_version = "2"
160        elif option[0] == '-R':
161            rng = [int(i) for i in option[1].split(",")]
162        elif option[0] == '-n':
163            nrows = int(option[1])
164
165    # Catch the hdf5 file passed as the last argument
166    filename = pargs[0]
167
168    if sqlite_version == "2":
169        import sqlite
170    else:
171        from pysqlite2 import dbapi2 as sqlite
172
173    if verbose:
174        print("pysqlite version:", sqlite.version)
175        if userandom:
176            print("using random values")
177
178    if docreate:
179        if verbose:
180            print("writing %s krows" % nrows)
181        if psyco_imported and usepsyco:
182            psyco.bind(create_db)
183        nrows *= 1000
184        create_db(filename, nrows)
185
186    if createindex:
187        index_db(filename)
188
189    if doquery:
190        query_db(filename, rng)
191