import array, cPickle
class pynumgrid:
'In-memory number array database'
def __init__( self, schema=None, fname=None ):
'''schema -> 'field1:atype field2:atype ...' where atype is a
python array type code (see eg below)'''
if fname:
self.load( fname )
else:
if type(schema) == str:
schema = [ f.split(':') for f in schema.split() ]
self.schema = schema
self.flist = [ f for f,t in schema ]
self.data = dict( [(f, array.array(t)) for f,t in schema] )
self.rowcount = 0
def insert( self, record, recnum=None ):
'Insert one record (or append if no rec number)'
if recnum == None:
for i, value in enumerate( record ):
self.data[ self.flist[i] ].append( value )
else:
for i, value in enumerate( record ):
self.data[ self.flist[i] ].insert( value, recnum )
self.rowcount += 1
def update( self, record, recnum ):
'Change the record at recnum with updated values'
for i, value in enumerate( record ):
self.data[ self.flist[i] ][ recnum ] = value
def delete( self, recnum ):
'Delete the record ar recnum, and return it'
record = []
for field in self.flist:
record.append( self.data[ field ].pop( recnum ) )
self.rowcount -= 1
return record
def extend( self, colset ):
'Add a set of columns to the database (fast)'
for i, col in enumerate( colset ):
self.data[ self.flist[i] ].extend( col )
self.rowcount += len( colset[0] )
def select( self, query, ns={} ):
'Execute a python expression using data and ns namespace'
ns.update( dict( _en=enumerate, _rc=xrange(self.rowcount),
array=array.array ) )
return eval( query, self.data, ns )
def fetchlist( self, reclist, cols=[] ):
'Return a list of records given a list of record numbers'
cols = cols or self.flist
colset = [ self.data[f] for f in cols ]
return [ [d[i] for d in colset] for i in reclist ]
def save( self, fname ):
'Write the db to disk files (FAST!)'
cfg = (self.schema, self.rowcount)
cPickle.dump( cfg, open(fname+'.cfg','wb'), -1 )
for field in self.flist:
fh = open( '%s__%s.dat' % (fname,field), 'wb' )
self.data[ field ].tofile( fh )
def load( self, fname ):
'Read and/or build db on to this database (FAST!)'
cfg = cPickle.load( open(fname+'.cfg','rb') )
schema, newrowcount = cfg
if not hasattr( self, 'data' ):
self.__init__( schema )
for field in self.flist:
fh = open( '%s__%s.dat' % (fname,field), 'rb' )
try: self.data[ field ].fromfile( fh, newrowcount )
except EOFError: pass
self.rowcount += newrowcount
if __name__ == '__main__':
# Test with "prog -i -s" to build then "prog -l -q" to test a query.
import sys
# Age (years), height (cm), weight (kg) for a million people:
pyg = pynumgrid( 'age:B height:B weight:B' )
blocksize = 10**5
blocks = 10
if '-i' in sys.argv:
agelist = [ (i*1779+517)%80+18 for i in xrange(blocksize) ]
heightlist = [ (i*5111+711)%90+120 for i in xrange(blocksize) ]
weightlist = [ (i*3113+797)%100+50 for i in xrange(blocksize) ]
for j in xrange(blocks):
pyg.extend( (agelist, heightlist, weightlist) )
if '-s' in sys.argv:
pyg.save( '/tmp/pynumgrid' )
if '-l' in sys.argv:
pyg.load( '/tmp/pynumgrid' )
if '-q' in sys.argv:
print 'The maximum height (in cm) of people in their 30s:',
print pyg.select('max(height[i] for i in _rc if 29<age[i]<40)')
print pyg.fetchlist( range(100,110), cols=['weight','height'] )
