db_benchmark.py

import io
import time

import duckdb
import pandas as pd
import psycopg
import psycopg2
from questdb.ingress import Sender
from sqlalchemy import create_engine
from sqlalchemy import text


pg_engine = create_engine('postgresql://shane:password@localhost:5439/yt?client_encoding=utf8')

aapl_filepath = "aapl.parquet"
pd.set_option("display.max_rows", 7)
pd.set_option("display.min_rows", 6)


def main():
    base_df = pd.read_parquet(aapl_filepath)

    # print(f"Base df:\n{base_df}")
    print(humansize(base_df.memory_usage(index=True, deep=True).sum()))

    parquet_benchmark(base_df)

    # quest_benchmark_psyco(base_df)
    quest_benchmark(base_df)

    duck_benchmark(base_df)
    #
    # # timescale_benchmark_pd(base_df)
    # # timescale_benchmark_psy2(base_df)
    timescale_benchmark_psy_csv(base_df)


def quest_benchmark(df):
    """
    Seems to be overall fastest with quest, but can't have foreign key. Can at least have primary key.
    """
    quest_engine = create_engine("questdb://admin:quest@localhost:8812/qdb")
    quest_conn = quest_engine.connect()
    quest_conn.execute(text("DROP TABLE IF EXISTS data_table"))
    print("QuestDB:")
    start = time.time()
    with Sender.from_conf(f"http::addr=localhost:9000;") as sender:
        sender.dataframe(
            df,
            table_name='data_table',
            symbols=['symbol'],
            at='timestamp',
        )
    end = time.time()
    print(f"\nWrite total time: {(end - start):.2f}s")

    start = time.time()
    result = quest_conn.execute(text("select * from data_table"))
    end = time.time()
    read_time = end - start
    start = time.time()
    quest_df = pd.DataFrame(result.fetchall())
    df.columns = result.keys()
    end = time.time()
    conversion_time = end - start
    print(f"Read total time: {(read_time + conversion_time):.2f}s ({read_time:.2f}s + {conversion_time:.2f}s)")

    print(f"\nQuestDB:\n{quest_df}")


def quest_benchmark_psypg(df):
    quest_engine = create_engine("questdb://admin:quest@localhost:8812/qdb")
    quest_conn = quest_engine.connect()
    quest_conn.execute(text("DROP TABLE IF EXISTS data_table"))
    print("QuestDB pscyo:")
    start = time.time()
    with Sender.from_conf(f"http::addr=localhost:9000;") as sender:
        sender.dataframe(
            df,
            table_name='data_table',
            symbols=['symbol'],
            at='timestamp',
        )
    end = time.time()
    print(f"\nWrite total time: {(end - start):.2f}s")

    start = time.time()
    query = "select * from data_table"
    conn = psycopg.connect(
        dbname="questdb",
        host="127.0.0.1",
        user="admin",
        password="quest",
        port=8812,
    )

    quest_df = pd.read_sql_query(query, conn)
    end = time.time()
    print(f"Read total time: {(end - start):.2f}s")

    print(f"\nQuestDB:\n{quest_df}")


def duck_benchmark(df):
    print("\nDuckDB:")
    duck_conn = duckdb.connect("duck.db")
    start = time.time()

    duck_conn.sql("DROP TABLE IF EXISTS data_table")
    start = time.time()
    duck_conn.execute("CREATE TABLE IF NOT EXISTS data_table AS SELECT * FROM df")
    print(f"Write total time: {(time.time() - start):.2f}s")

    start = time.time()
    duck_df = duck_conn.sql("SELECT * FROM data_table").df()
    print(f"Read total time: {(time.time() - start):.2f}s")

    print(f"\nDuckDB:\n{duck_df}")


def timescale_benchmark_psy_csv(df):
    db_params = {
        'dbname': 'yt',
        'user': 'shane',
        'password': 'password',
        'host': 'localhost',
        'port': '5439'
    }

    conn = psycopg2.connect(**db_params)
    cur = conn.cursor()

    print("\nTimescaleDB:")
    start = time.time()
    df.head(0).to_sql('data_table', pg_engine, if_exists='replace', index=False)

    output = io.StringIO()
    df.to_csv(output, sep='\t', header=False, index=False)
    output.seek(0)
    contents = output.getvalue()
    cur.copy_from(output, 'data_table', null="")  # null values become ''
    conn.commit()
    end = time.time()
    print(f"Write total time: {(end - start):.2f}s")

    start = time.time()
    cur.execute("SELECT * FROM data_table;")
    timescale_df = pd.DataFrame(cur.fetchall())
    end = time.time()
    print(f"Read total time: {(end - start):.2f}s")  # 16.12 seconds

    print(f"\nTimescaleDB with csv:\n{timescale_df}")


def timescale_benchmark_pd(df):
    print("\nTimescaleDB:")
    start = time.time()
    df.to_sql(name="data_table", con=pg_engine, if_exists="replace", index=False)  # Takes forever
    end = time.time()
    print(f"Write total time: {(end - start):.2f}s")

    start = time.time()
    timescale_df = pd.read_sql_query("""SELECT * FROM data_table""", pg_engine)
    end = time.time()
    print(f"Read total time: {(end - start):.2f}s")

    print(f"\nTimescaleDB with alchemy to_sql:\n{timescale_df}")


def timescale_benchmark_psypg2(df):
    db_params = {
        'dbname': 'yt',
        'user': 'shane',
        'password': 'password',
        'host': 'localhost',
        'port': '5439'
    }

    conn = psycopg2.connect(**db_params)
    cur = conn.cursor()

    print("\nTimescaleDB:")
    start = time.time()
    df.to_sql('data_table', pg_engine, if_exists='replace', index=False)  # 156 seconds
    end = time.time()
    print(f"Write total time: {(end - start):.2f}s")

    start = time.time()
    cur.execute("SELECT * FROM data_table;")
    timescale_df = pd.DataFrame(cur.fetchall())  # 13.58 seconds
    end = time.time()
    print(f"Read total time: {(end - start):.2f}s")

    print(f"\nTimescaleDB with psycopg2 to_sql:\n{timescale_df}")


def parquet_benchmark(df):
    print("\nParquet file:")
    start = time.time()
    df.to_parquet(aapl_filepath)
    end = time.time()
    print(f"Write total time: {(end - start):.2f}s")  # 64 seconds
    start = time.time()
    flat_df = pd.read_parquet(aapl_filepath)
    end = time.time()
    print(f"Read total time: {(end - start):.2f}s")  # 14.07 seconds

    print(f"\nParquet file:\n{flat_df}")


def humansize(nbytes):
    suffixes = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']
    i = 0
    while nbytes >= 1024 and i < len(suffixes)-1:
        nbytes /= 1024.
        i += 1
    f = ('%.2f' % nbytes).rstrip('0').rstrip('.')
    return '%s %s' % (f, suffixes[i])


if __name__ == '__main__':
    main()