# Add integer keys in rocksdb and get highest and lowest keys using a custom comparator

rocks_int_comparator.py

# Add integer keys in rocksdb and get highest and lowest keys using a custom comparator

import random
import string
from time import perf_counter
import rocksdb
from rocksdb.interfaces import Comparator


class IntegerComparator(Comparator):
    def compare(self, a, b):
        a = int(a)
        b = int(b)
        if (a < b):
            # print(a, b)
            return -1
        if (a > b):
            # print(a, b)
            return 1
        if (a == b):
            # print(a, b)
            return 0

    def name(self):
        return b'IntegerComparator'


num_records = 10000
val_size = 100

data = {i: ''.join(random.choice(string.printable) for _ in range(val_size)).encode() for i in range(1, num_records+1)}



db = rocksdb.DB("test.db", rocksdb.Options(create_if_missing=True, comparator=IntegerComparator()))
for k, v in data.items():
    db.put(str(k).encode(), v)


start = perf_counter()
it = db.iterkeys()
end = perf_counter()
print('Creating iterator in {} records takes {}'.format(num_records, end-start))

start = perf_counter()
it.seek_to_last()
print(it.get())
assert it.get() == str(num_records).encode() 
end = perf_counter()
print('Seek to last in {} records takes {}'.format(num_records, end-start))

start = perf_counter()
it.seek_to_first()
print(it.get())
assert it.get() == b'1' 
end = perf_counter()
print('Seek to first in {} records takes {}'.format(num_records, end-start))


more_records = 3*num_records

for i in range(num_records+1, num_records+more_records+1):
    data[i] = ''.join(random.choice(string.printable) for _ in range(val_size)).encode()
    db.put(str(i).encode(), data[i])


start = perf_counter()
it = db.iterkeys()
end = perf_counter()
print('Creating iterator in {} records takes {}'.format(num_records+more_records, end-start))

start = perf_counter()
it.seek_to_last()
print(it.get())
assert it.get() == str(num_records+more_records).encode()
end = perf_counter()
print('Seek to last in {} records takes {}'.format(num_records+more_records, end-start))

start = perf_counter()
it.seek_to_first()
print(it.get())
assert it.get() == b'1'
end = perf_counter()
print('Seek to first in {} records takes {}'.format(num_records+more_records, end-start))

num_records += more_records