Script to Benchmark TimescaleDB Multivalue Inserts

benchmark.py

from numpy import random
import logging
import datetime
import os
from functools import reduce
import os
import multiprocessing

import tqdm


startDate = datetime.datetime(2018, 9, 11, 13, 00)

def random_date():
    return startDate + datetime.timedelta(days=random.randint(0, 30*12*10), hours=random.randint(0, 6),
                                          minutes=random.randint(0, 60), seconds=random.randint(0, 60))

def gen_tick():
    return [random_date(), random.randint(1, 10000), random.uniform(1000.0, 8000.0)]


def batch_inserts(file):
    size = per_batch
    print(query_prefix, end='', file=open("batch"+str(file)+".sql", "a"))
    for i in range(0, size):
        print("('{}',{},{})".format(
            *gen_tick()), end='', file=open("batch"+str(file)+".sql", "a"))
        if i != size - 1:
            print(",", end='', file=open("batch"+str(file)+".sql", "a"))
    print(query_suffix, end='', file=open("batch"+str(file)+".sql", "a"))


def save_to_db(file):
    # print("saving {}".format(file))
    os.system("bash -c 'time sudo -u postgres psql postgres < batch" +
              str(file)+".sql' > /dev/null 2>&1")


if __name__ == '__main__':
    query_suffix = " ;"
    query_prefix = "INSERT INTO bench_sim(timestamp,alpha,beta) VALUES "

    # Config: edit these
    total_rows = 100000000
    per_batch = 500
    
    num_files = int(total_rows/per_batch)
    tasks = range(0, num_files)
    
    pool = multiprocessing.Pool(32)
    for _ in tqdm.tqdm(pool.imap_unordered(batch_inserts, tasks), total=len(tasks)):
        pass
    pool.close()
    pool.join()

    pool = multiprocessing.Pool(16)
    for _ in tqdm.tqdm(pool.imap_unordered(save_to_db, tasks), total=len(tasks)):
        pass
    pool.close()
    pool.join()