Testing storage of millions of keys in Redis (and sets)

redis-size-test.py

#! /usr/bin/env python26
# original from https://gist.github.com/1329319

import redis
import random
import sys

r = redis.Redis(host = 'localhost', port = 6380)

REDIS_SETGET = False
REDIS_HSET = False
REDIS_RPUSH_HSET = False
REDIS_SSET = False
REDIS_SSET_HSET = False

NUM_ENTRIES =  1200000
MAX_VAL     = 12000000
AVG_VAL     =       13
BIG_SET_VAL =    20000

if len(sys.argv) != 2 or sys.argv[1] not in ('redis-normal', 'redis-hashes', 'redis-listnhash', 'redis-setnhash', 'redis-sset'):
   print 'Specify a test: redis-normal, redis-hashes, redis-listnhash redis-setnhash redis-sset'
   print 'NOTE: clear out memcached (restart) or Redis (FLUSHALL) before running'
   sys.exit(2)
if sys.argv[1] == 'redis-normal':
   REDIS_SETGET = True
elif sys.argv[1] == 'redis-hashes':
   REDIS_HSET = True
elif sys.argv[1] == 'redis-listnhash':
   REDIS_RPUSH_HSET = True
elif sys.argv[1] == 'redis-setnhash':
   REDIS_SSET_HSET = True
elif sys.argv[1] == 'redis-sset':
   REDIS_SSET = True

p = r.pipeline()
total = 0
for i in range(0, NUM_ENTRIES):
   value = random.randint(0, MAX_VAL)
   if REDIS_SETGET:
       r.set(str(i), value)
   elif REDIS_HSET:
       bucket = int(i / 500)
       p.hset(bucket, i, value)
   elif REDIS_RPUSH_HSET:
               p.hset(i, 'score', value)
               p.rpush('list-to-score', i)
   elif REDIS_SSET_HSET:
               #name = 'object_%d' % i
               #p.hset(name, 'score', float(value))
               r.zadd('list-to-score', i, float(value))
   elif REDIS_SSET:
               if (random.randint(1, 10000) == 10000):
                       print "creating big set"
                       for x in range(0, BIG_SET_VAL):
                               total = total + 1
                               value = random.randint(0, MAX_VAL)
                               p.zadd('sset-%d' % i, float(value), x)
               else:
                       for x in range(0, random.randint(1, 2 * AVG_VAL)):
                               total = total + 1
                               value = random.randint(0, MAX_VAL)
                               p.zadd('sset-%d' % i, float(value), x)

   if i % (NUM_ENTRIES/50) == 0:
               p.execute()
               p = r.pipeline()
               print i

# one final clear out
p.execute()

# get size
size = int(r.info()['used_memory'])

if REDIS_SSET:
       print 'avg sset size %2.2f' % total / NUM_ENTRIES

print '%s bytes, %s MB' % (size, size / 1024 / 1024)