MongoDB index benchmark

mongodb_perf.py

from pymongo import MongoClient, UpdateOne
import random
import time

mongo_client = MongoClient()

# drop old stuff
mongo_client['test']['without_index'].drop()
mongo_client['test']['with_index'].drop()
mongo_client['test']['index'].drop()

mongo_client['test']['with_index'].create_index("ints")

insert1 = 0
insert2 = 0
insert3 = 0

j = 0
while True:
    j += 1
    num_ints = random.randint(10, 2000)
    ints = [random.randint(0, 1000000) for _ in range(num_ints)]

    time1_before = time.time()
    mongo_client['test']['without_index'].insert_one({"ints": ints})
    time1_after = time.time()
    insert1 += time1_after - time1_before

    time2_before = time.time()
    _id = mongo_client['test']['with_index'].insert_one({"ints": ints}).inserted_id
    time2_after = time.time()
    insert2 += time2_after - time2_before

    time3_before = time.time()
    ops = []
    for i in ints:
        ops.append(UpdateOne({'_id': i}, {'$addToSet': {'docs': {'$each': [_id]}}}, upsert=True))
    mongo_client['test']['index'].bulk_write(ops) 
    time3_after = time.time()
    insert3 += time3_after - time3_before

    if j % 100 == 0:
        print(f"Time to insert in collection without index: {insert1 / j}")
        print(f"Time to insert in collection with index {insert2 / j}")
        print(f"Time to insert in auxiliary collection: {insert3 / j}")