Simple benchmark for zdiff/zunion/zinter

bench_zdiff.py

import random
import re
import string
import pickle
import redis

ZADD_BATCH_SIZE = 2000

r = redis.Redis(port=6380)

def randomkey(N=16):
    return ''.join(random.choice(string.ascii_lowercase) for _ in range(N))

def randomscore():
    return random.random()

def fillZset(key, size):
    # print(key, size, ZADD_BATCH_SIZE, size//ZADD_BATCH_SIZE)
    r.delete(key)
    batchsize = min(size, ZADD_BATCH_SIZE)
    fullset = {}
    for i in range(size//batchsize):
        batch = {}
        for j in range(batchsize):
            k = randomkey()
            # s = i*batchsize + j
            s = randomscore()
            batch[k] = s
            fullset[k] = s
        r.zadd(key, batch)
    return fullset

def fillZsetFromSample(key, size, sample):
    r.delete(key)
    batchsize = min(size, ZADD_BATCH_SIZE)
    all_keys = list(sample.keys())
    for i in range(size//batchsize):
        batch = {}
        for j in range(batchsize):
            k = random.choice(all_keys)
            s = sample[k]
            batch[k] = s
        r.zadd(key, batch)

def benchZdiff(store=False):
    if store:
        r.execute_command("zdiffstore benchtarget 2 benchz1 benchzsmall")
    else:
        r.execute_command("zdiff 2 benchz1 benchzsmall")
    # print(r.zcount("benchtarget", "-inf", "+inf"))
    r.delete("benchtarget")

def benchZdiffAlgo2(store=False):
    if store:
        r.execute_command("zdiffstore benchtarget 3 benchz1 benchzsmall benchzsmall2")
    else:
        r.execute_command("zdiff 3 benchz1 benchzsmall benchzsmall2")
    # print(r.zcount("benchtarget", "-inf", "+inf"))
    r.delete("benchtarget")

def benchZunion(store=False):
    if store:
        r.zunionstore("benchtarget", ["benchz1", "benchz2", "benchz3"])
    else:
        r.execute_command("zunion 3 benchz1 benchz2 benchz3")
    # print(r.zcount("benchtarget", "-inf", "+inf"))
    r.delete("benchtarget")

def benchZinter(store=False):
    if store:
        r.zinterstore("benchtarget", ["benchz1", "benchz1sample"])
    else:
        r.execute_command("zinter 2 benchz1 benchz1sample")
    # print(r.zcount("benchtarget", "-inf", "+inf"))
    r.delete("benchtarget")

def buildZsets(base_size):
    print("benchz1")
    fullset = fillZset("benchz1", base_size)
    print("benchz1sample")
    fillZsetFromSample("benchz1sample", base_size//2, fullset)
    print("benchz2")
    fillZset("benchz2", base_size)
    print("benchz3")
    fillZset("benchz3", base_size)
    print("benchzsmall")
    fillZset("benchzsmall", 100)
    print("benchzsmall2")
    fillZset("benchzsmall2", 100)
    return {
        "benchz1": r.dump("benchz1"),
        "benchz1sample": r.dump("benchz1sample"),
        "benchz2": r.dump("benchz2"),
        "benchz3": r.dump("benchz3"),
        "benchzsmall": r.dump("benchzsmall"),
        "benchzsmall2": r.dump("benchzsmall2"),
    }

def loadZsets(values):
    for k, v in values.items():
        r.restore(k, 0, v)

def clearZsets(values):
    for k, v in values.items():
        r.delete(k)

def printTimes():
    for k, v in r.info("commandstats").items():
        if re.search("zdiff", k) or re.search("zinter", k) or re.search("zunion", k):
            print(k, v)

def run():
    r.config_resetstat()

    try:
        with open("dump_values.pickle", "rb") as f:
            values = pickle.load(f)
    except Exception as e:
        print(e)
        values = buildZsets(100000)
        f = open("dump_values.pickle", "wb")
        pickle.dump(values, f)
        f.close()

    clearZsets(values)

    # for i in range(30):
    #     print("zdiffstore it", i)
    #     loadZsets(values)
    #     benchZdiff(store=True)
    #     clearZsets(values)
    # for i in range(30):
    #     print("zunionstore it", i)
    #     loadZsets(values)
    #     benchZunion(store=True)
    #     clearZsets(values)
    # for i in range(30):
    #     print("zinterstore it", i)
    #     loadZsets(values)
    #     benchZinter(store=True)
    #     clearZsets(values)

    # for i in range(30):
    #     print("zdiff it", i)
    #     loadZsets(values)
    #     benchZdiff()
    #     clearZsets(values)
    for i in range(30):
        # print("zunion it", i)
        loadZsets(values)
        benchZunion()
        clearZsets(values)
    # for i in range(30):
    #     print("zinter it", i)
    #     loadZsets(values)
    #     benchZinter()
    #     clearZsets(values)

    printTimes()

run()