b-trees-and-traverse.py

# -*- coding: utf-8 -*-

import matplotlib.pyplot as plt
from statistics import mean
import numpy as np
import psycopg2
import random
import time

ONE_TIME_INSERT = 1000000

def connect():

    connection = psycopg2.connect(
        user = "postgres",
        password = "password",
        host = "localhost",
        port = 5432,
        database = "dataset"
    )
    return connection


def insert(cursor, x, y):

    sql = f'''
    INSERT INTO t_numbers
    SELECT
    i, i * 10
    FROM generate_series({x}, {y}) AS i;
    '''
    cursor.execute(sql)


def query(cursor, value):

    val_times = []
    for _ in range(3):
        t1 = time.time()
        sql = f'''SELECT * FROM t_numbers WHERE val = {value} '''
        cursor.execute(sql)
        t2 = time.time()
        val_times.append((t2 - t1) * 100)

    val_indexed_times = []
    for _ in range(3):
        t3 = time.time()
        sql = f'''SELECT * FROM t_numbers WHERE val_indexed = {value * 10} '''
        cursor.execute(sql)
        t4 = time.time()
        val_indexed_times.append((t4 - t3) * 100)

    return round(mean(val_times), 2), round(mean(val_indexed_times), 2)


def draw(x, y1, y2):

    plt.plot(x, y1, color="gray", label="Traverse")
    plt.plot(x, y2, color="#2CBDFE", label="B-Trees")

    plt.title('B-Trees vs Traverse')
    plt.xlabel('Records')
    plt.ylabel('Time(ms)')

    plt.legend()
    plt.savefig('./save.jpg')


def main():

    conn = connect()
    cursor = conn.cursor()
    cursor.execute("SELECT VERSION();")
    record = cursor.fetchone()
    print(record)

    x, y1, y2 = []
    for i in range(5):
        insert(cursor, i * ONE_TIME_INSERT + 1, (i + 1) * ONE_TIME_INSERT)
        conn.commit()

        traverse_time, indexed_time = query(cursor, random.randint(i, (i + 1) * ONE_TIME_INSERT))

        x.append((i + 1) * ONE_TIME_INSERT)
        y1.append(traverse_time)
        y2.append(indexed_time)

    draw(x, y_val, y_val_indexed)

    cursor.close()
    conn.close()


if __name__ == "__main__":

    main()