Benchmarking Postgres timeseries queries

import os, random, json, pytz
from datetime import datetime, timedelta, time
from time import time as t
from functools import lru_cache
from itertools import zip_longest

from IPython.display import Javascript
import plotly.express as px
import tqdm

# Install python dependencies
!conda install -y --quiet psycopg2
# Install postgres (if not already installed)
!sudo apt-get install postgresql -yqq
# It needs to run in the background
!sudo service postgresql start
!psql --version

# Optionally drop the user and/or database if starting anew
# !sudo su - postgres -c "psql -c \"DROP DATABASE timeseries_benchmark\""
# !sudo su - postgres -c "psql -c \"DROP USER tse\""
# Create user, db, and give access
!sudo su - postgres -c "psql -c \"CREATE USER tse WITH PASSWORD 'abc'\" && psql -c \"CREATE DATABASE timeseries_benchmark OWNER tse\""
# Needed to be able to COPY FROM a file
!sudo su - postgres -c "psql -c \"ALTER USER tse WITH SUPERUSER\""

import psycopg2
conn = psycopg2.connect(user="tse", password="abc", database="timeseries_benchmark", host="localhost")

# Drop tables if running anew
# !sudo su - postgres -c "psql -d timeseries_benchmark -c \"DROP TABLE IF EXISTS actuals_array_double; DROP TABLE IF EXISTS actuals_array; DROP TABLE IF EXISTS actuals_cols;\""

# FLOAT means DOUBLE PRECISION (8B)
with conn as conn:
    with conn.cursor() as curs:
        curs.execute("""
        CREATE TABLE actuals_array_double (
          id serial PRIMARY KEY,
          asset_id INTEGER NOT NULL,
          day DATE NOT NULL,
          created TIMESTAMP NOT NULL,
          values FLOAT[] NOT NULL,
          CONSTRAINT actuals_array_double_asset_day UNIQUE (asset_id, day, created)
        );
        """)

# REAL is single precision floating point
with conn as conn:
    with conn.cursor() as curs:
        curs.execute("""
        CREATE TABLE actuals_array (
          id serial PRIMARY KEY,
          asset_id INTEGER NOT NULL,
          day DATE NOT NULL,
          created TIMESTAMP NOT NULL,
          values REAL[] NOT NULL,
          CONSTRAINT actuals_array_asset_day UNIQUE (asset_id, day, created)
        );
        """)

with conn as conn:
    with conn.cursor() as curs:
        pte_str = ",\n      ".join(f"pte{i:02d} REAL" for i in range(100))
        curs.execute(f"""
        CREATE TABLE actuals_cols (
          id serial PRIMARY KEY,
          asset_id INTEGER NOT NULL,
          day DATE NOT NULL,
          created TIMESTAMP NOT NULL,
          {pte_str},
          CONSTRAINT actuals_cols_asset_day UNIQUE (asset_id, day, created)
        );
        """)

# Use a cache to speed up the function call, the randomness is not critical
@lru_cache(maxsize=10000)
def randvals(d: datetime, tz = pytz.timezone("Europe/Amsterdam"), values_arr=True):
    d1 = tz.localize(datetime.combine(d, time()))
    d2 = tz.localize(datetime.combine(d + timedelta(days=1), time()))
    periods = int((d2 - d1).total_seconds() / 3600 * 4)
    values = [random.random() for i in range(periods)]
    if values_arr:
        return '"{' + ",".join(f"{v:.3f}" for v in values) + '}"'
    else:
        # Always make 100 columns, but fill with unquoted empty string
        # See parameters -> NULL here: https://www.postgresql.org/docs/9.2/sql-copy.html
        return ",".join(f"{v:.3f}" if v else "" for i, v in zip_longest(range(100), values))

def generate_fake_data(filename, n_assets=5E3, n_years=3, n_versions=3, stop_dt=datetime.now(), values_arr=True):
    with open(filename, "w") as f:
        i = 0
        for days_ago in tqdm.tqdm(range(n_years * 365)):
            for asset_id in range(int(n_assets)):
                day = (stop_dt - timedelta(days=days_ago))
                for i_v in range(n_versions):
                    i += 1
                    created = day - i_v * timedelta(days=3)
                    values = randvals(day, values_arr=values_arr)
                    line = f"{i:d}, {asset_id:d}, '{day:%Y-%m-%d}', '{created.isoformat()}', {values}"
                    f.write(line + "\n")

generate_fake_data("actuals_array.csv")

generate_fake_data("actuals_cols.csv", values_arr=False)

!head -1 actuals_array.csv
!head -1 actuals_cols.csv
!wc -l actuals_array.csv

def copy_from(conn, filename, table):
    path = os.path.abspath(filename)
    print(f"Importing '{filename}' to '{table}'")
    with conn as conn:
        with conn.cursor() as curs:
            curs.execute(f"DELETE FROM {table};")
            curs.execute(f"COPY {table} FROM '{path}' (FORMAT CSV);")
            curs.execute(f"SELECT COUNT(*) FROM {table};")
            n_rows = curs.fetchone()[0]
            curs.execute(f"SELECT pg_size_pretty(pg_relation_size('{table}'));")
            print(f"Table '{table}' rows and size:", n_rows, ",", curs.fetchone()[0])

%time copy_from(conn, "actuals_array.csv", "actuals_array_double")
%time copy_from(conn, "actuals_array.csv", "actuals_array")
%time copy_from(conn, "actuals_cols.csv", "actuals_cols")

%%javascript
    function addScript(url) {
        var script = document.createElement("script");
        script.type = "text/javascript";
        script.src = url;
        //document.head.appendChild(script);
        // For use on website, as anything added to <head> is not available on jupyter reload
        element.append(script);
    }
    function addCSS(url) {
        var link = document.createElement("link");

        link.type = "text/css";
        link.rel = "stylesheet";
        link.href = url;

        //document.head.appendChild(link);
        element.append(link);
    }
    addCSS("https://unpkg.com/bootstrap@4.5.0/dist/css/bootstrap.min.css");
    addCSS("https://unpkg.com/@fortawesome/fontawesome-free@5.13.0/css/all.css");
    addScript("https://cdn.jsdelivr.net/npm/vue@2.6.12/dist/vue.min.js");
    // Only works if you're running jupyter lab with this notebook as root repo and the compiled file in this folder with name pev2.umd.js
    addScript("/files/pev2.umd.js");

def pev2(query, plan):
    """
    Visualize a Postgres query plan with pev2. Try to pass a text query plan, not JSON,
    as text has less issues with escaping double quotes.
    """
    s = """
    var el = document.createElement('div');
    element.append(el);

    var v = new Vue({
        el: el,
        components: {
            'pev2': window.pev2
        },
        data: function() {
            return {
                plan: `""" + plan + """`,
                query: `""" + query + """`,
            }
        },
        template: "<pev2 :plan-source='plan' :plan-query='query'></pev2>"
    });
    """
    return Javascript(s)

def query_array_unnest(table="actuals_array", explain=True):
    with conn.cursor() as curs:
        query = f"""
        {'EXPLAIN (ANALYZE, COSTS, VERBOSE, BUFFERS)' if explain else ''}

        -- Put aggregated load result in array per day. max/min for unique assets
        -- does not matter, we just need to pick one as the same value
        -- has been distributed over the rows
        SELECT
        day,
        MAX(count_asset) as count_asset,
        array_agg(s ORDER BY timeblock) as total_value
        FROM (
        SELECT
            day,
            timeblock,
            SUM(elem) as s,
            COUNT(*) as count_asset

        FROM (
            -- Select only latest rows
            SELECT DISTINCT ON (asset_id, day)
                day,
                values
            FROM {table}
            -- Sort descending by reference_datetime so it picks latest distinct value,
            -- choosing a value over NULL, and the highest ID if there are ties
            ORDER BY asset_id, day, created DESC NULLS LAST
        ) AS t, LATERAL UNNEST(t.values) WITH ORDINALITY x(elem, timeblock)

        GROUP BY day, timeblock
        ) as a
        GROUP BY day;
        """
        curs.execute(query)
        if explain:
            plan = "\n".join([l[0] for l in curs.fetchall()])
        else:
            plan = curs.fetchall()
    
    return query, plan

pev2(*query_array_unnest())

with conn.cursor() as curs:
    def create_idx(table):
        curs.execute(f"CREATE INDEX IF NOT EXISTS {table}_asset_id_day_created "
                     f"ON {table} (asset_id, day, created DESC NULLS LAST); "
                     f"ANALYZE {table};")
    create_idx("actuals_array_double")
    create_idx("actuals_array")
    create_idx("actuals_cols")

pev2(*query_array_unnest())

def query_array_subscript(table='actuals_array'):
    with conn.cursor() as curs:
        # Maximum amount of columns
        n_cols = 100
        array_index_str = ", ".join(f'values[{i+1}] as t{i:02d}' for i in range(n_cols))
        array_sum_str = ", ".join(f"SUM(t{i:02d})" for i in range(n_cols))
        query = f"""
        EXPLAIN (ANALYZE, COSTS, VERBOSE, BUFFERS)

        SELECT
        day,
        COUNT(*) as count_asset,
        {array_sum_str}

        FROM (
        -- Select only latest rows
        SELECT DISTINCT ON (asset_id, day)
            day,
            {array_index_str}
        FROM {table}
        ORDER BY asset_id, day, created DESC NULLS LAST
        ) AS t

        GROUP BY day;
        """
        curs.execute(query)
        plan = "\n".join([l[0] for l in curs.fetchall()])
    
    return query, plan

pev2(*query_array_subscript())

with conn.cursor() as curs:
    table = 'actuals_cols'
    n_cols = 100
    col_str = ", ".join(f'pte{i:02d}' for i in range(n_cols))
    array_sum_str = ", ".join(f"SUM(pte{i:02d})" for i in range(n_cols))
    query = f"""
    EXPLAIN (ANALYZE, COSTS, VERBOSE, BUFFERS)

    SELECT
    day,
    COUNT(*) as count_asset,
    {array_sum_str}
    FROM (
    -- Select only latest rows
    SELECT DISTINCT ON (asset_id, day)
        day,
        {col_str}
    FROM {table}
    ORDER BY asset_id, day, created DESC NULLS LAST
    ) AS t
    GROUP BY day;
    """
    curs.execute(query)
    plan = "\n".join([l[0] for l in curs.fetchall()])

pev2(query, plan)

pev2(*query_array_subscript('actuals_array_double'))

query, result = query_array_unnest(explain=False)

dt = []
val = []

for date, n_assets, values in result[:2]:
    for i, v in enumerate(values):
        dt.append(datetime.combine(date, time()) + timedelta(minutes=i*15))
        val.append(v)

fig = px.line(x=dt, y=val, title="Last two days of query result")
# Need to include it like this because plotly does not include its javascript by default any more
display(HTML(fig.to_html("plot.html", include_plotlyjs='cdn')))