benchmark_queries.py

"""
benchmark_queries.py — Project 04: Time Series Databases
Fullstack Data — Kenneth

Responsibility:
  Run the same 5 benchmark queries against all THREE databases:
    - Vanilla PostgreSQL  (port 5432) — the baseline, no extensions
    - TimescaleDB         (port 5433) — PostgreSQL + hypertable + continuous aggregate
    - InfluxDB            (port 8086) — purpose-built TSDB, Flux query language

  Time each query (best of 3 runs). Print a three-column comparison table.
  Identify the winner per query. Explain why each result looks the way it does.

The 5 queries:
  Q1  Last value per device         — most recent reading for all 30 devices
  Q2  5-minute rolling average      — window function over time per device
  Q3  Gap detection                 — find missing readings > 2 minutes
  Q4  Hourly downsample             — aggregate to hourly averages (CA vs raw vs raw)
  Q5  Range scan                    — last 24 hours for a specific device

What the three-way comparison answers:
  Vanilla PG vs TimescaleDB  — is the hypertable extension worth it on the same engine?
  Vanilla PG vs InfluxDB     — how much does a purpose-built TSDB beat a general DB?
  TimescaleDB vs InfluxDB    — which specialist wins per query type?

Q4 is the critical query:
  Vanilla PG     — aggregates 18M raw rows live every time
  TimescaleDB    — reads from sensor_readings_hourly (pre-computed continuous aggregate)
  InfluxDB       — aggregates live via aggregateWindow()
  Expected: TimescaleDB wins Q4 by the largest margin

Each query is run 3 times. Best time recorded (eliminates cold cache noise).

Output:
  - Per-query section with timings, row counts, and winner across all 3 DBs
  - Final three-column comparison table
  - Notes explaining each result
  - data/benchmark_results.json

Run after: python ingest_postgres.py AND ingest_timescale.py AND ingest_influx.py
"""

import json
import os
import time

import psycopg2
import pandas as pd
from influxdb_client import InfluxDBClient
from tabulate import tabulate

# ── Config ────────────────────────────────────────────────────────────────────

POSTGRES_VANILLA = {
    "host":            "localhost",
    "port":            5432,
    "dbname":          "sensors",
    "user":            "engineer",
    "password":        "engineer",
    "connect_timeout": 10,
}

TIMESCALE = {
    "host":            "localhost",
    "port":            5433,
    "dbname":          "postgres",
    "user":            "postgres",
    "password":        "engineer",
    "connect_timeout": 10,
}

INFLUX_URL    = "http://localhost:8086"
INFLUX_TOKEN  = os.environ.get("INFLUX_TOKEN", "")
INFLUX_ORG    = "fullstackdata"
INFLUX_BUCKET = "sensors"
MEASUREMENT   = "sensor_readings"

RUNS_PER_QUERY = 3
OUTPUT_PATH    = "data/benchmark_results.json"

# ── Helpers ───────────────────────────────────────────────────────────────────

def section(title: str):
    print(f"\n{'─' * 60}")
    print(f"  {title}")
    print(f"{'─' * 60}\n")


def time_query(fn, runs: int = RUNS_PER_QUERY):
    """Run fn() N times. Return (best_time_s, last_result)."""
    best   = float("inf")
    result = None
    for _ in range(runs):
        t0     = time.perf_counter()
        result = fn()
        elapsed = time.perf_counter() - t0
        best   = min(best, elapsed)
    return round(best, 4), result


def pg_query(conn, sql: str) -> pd.DataFrame:
    with conn.cursor() as cur:
        cur.execute(sql)
        cols = [d[0] for d in cur.description]
        rows = cur.fetchall()
    return pd.DataFrame(rows, columns=cols)


def flux_query(query_api, flux: str) -> pd.DataFrame:
    tables = query_api.query_data_frame(flux)
    if isinstance(tables, list):
        return pd.concat(tables, ignore_index=True) if tables else pd.DataFrame()
    return tables


def three_way_winner(t_vanilla: float, t_tsdb: float, t_influx: float) -> str:
    """Return label of the fastest database. Ties within 10% treated as equal."""
    times = {
        "Vanilla PG":  t_vanilla,
        "TimescaleDB": t_tsdb,
        "InfluxDB":    t_influx,
    }
    # Remove errored entries
    valid = {k: v for k, v in times.items() if v != float("inf")}
    if not valid:
        return "all errors"
    best_name  = min(valid, key=valid.get)
    best_time  = valid[best_name]
    # Check if any other is within 10% — call it a tie
    close = [k for k, v in valid.items() if v <= best_time * 1.10 and k != best_name]
    if close:
        return f"Tie: {best_name} / {close[0]}"
    return f"{best_name} ✓"


# ── SQL queries — Vanilla PostgreSQL & TimescaleDB share the same SQL ─────────
# TimescaleDB is PostgreSQL — same syntax. Q4 differs: TimescaleDB reads the
# continuous aggregate view, vanilla PG aggregates raw rows.

# Q1 — Last value per device
Q1_SQL = """
    SELECT DISTINCT ON (device_id)
        device_id,
        ts,
        temperature,
        battery_pct
    FROM sensor_readings
    ORDER BY device_id, ts DESC;
"""

Q1_FLUX = f"""
    from(bucket: "{INFLUX_BUCKET}")
      |> range(start: -8d)
      |> filter(fn: (r) => r._measurement == "{MEASUREMENT}")
      |> filter(fn: (r) => r._field == "temperature" or r._field == "battery_pct")
      |> last()
      |> pivot(rowKey: ["_time", "device_id"], columnKey: ["_field"], valueColumn: "_value")
"""

# Q2 — 5-minute rolling average, sensor_001, last 1 hour of actual data
# Anchored to data window (2026-05-13 → 2026-05-19) — not NOW() which falls
# outside the static dataset generated by generate_sensor.py
Q2_SQL = """
    SELECT
        device_id,
        ts,
        AVG(temperature) OVER (
            PARTITION BY device_id
            ORDER BY ts
            RANGE BETWEEN INTERVAL '5 minutes' PRECEDING AND CURRENT ROW
        ) AS rolling_avg_temp
    FROM sensor_readings
    WHERE device_id = 'sensor_001'
      AND ts >= '2026-05-19 23:00:00+00'
      AND ts <= '2026-05-19 23:59:59+00'
    ORDER BY ts;
"""

Q2_FLUX = f"""
    from(bucket: "{INFLUX_BUCKET}")
      |> range(start: 2026-05-19T23:00:00Z, stop: 2026-05-20T00:00:00Z)
      |> filter(fn: (r) => r._measurement == "{MEASUREMENT}")
      |> filter(fn: (r) => r._field == "temperature")
      |> filter(fn: (r) => r.device_id == "sensor_001")
      |> timedMovingAverage(every: 1m, period: 5m)
"""

# Q3 — Gap detection: readings missing > 2 minutes, last 7 days
Q3_SQL = """
    WITH gaps AS (
        SELECT
            device_id,
            ts AS gap_start,
            LEAD(ts) OVER (PARTITION BY device_id ORDER BY ts) AS gap_end,
            LEAD(ts) OVER (PARTITION BY device_id ORDER BY ts) - ts AS gap_size
        FROM sensor_readings
        WHERE ts >= '2026-05-13 00:00:00+00'
          AND ts <= '2026-05-19 23:59:59+00'
    )
    SELECT
        device_id,
        gap_start,
        gap_end,
        gap_size
    FROM gaps
    WHERE gap_size > INTERVAL '2 minutes'
    ORDER BY gap_size DESC
    LIMIT 20;
"""

Q3_FLUX = f"""
    from(bucket: "{INFLUX_BUCKET}")
      |> range(start: 2026-05-13T00:00:00Z, stop: 2026-05-20T00:00:00Z)
      |> filter(fn: (r) => r._measurement == "{MEASUREMENT}")
      |> filter(fn: (r) => r._field == "temperature")
      |> elapsed(unit: 1s)
      |> filter(fn: (r) => r.elapsed > 120)
      |> keep(columns: ["_time", "device_id", "elapsed"])
      |> sort(columns: ["elapsed"], desc: true)
      |> limit(n: 20)
"""

# Q4 — Hourly downsample, last 7 days
# KEY DIFFERENCE:
#   Vanilla PG     → aggregates 18M raw rows live (no pre-computation)
#   TimescaleDB    → reads sensor_readings_hourly continuous aggregate (pre-computed)
#   InfluxDB       → aggregates live via aggregateWindow()
Q4_SQL_VANILLA = """
    SELECT
        device_id,
        date_trunc('hour', ts)   AS bucket,
        AVG(temperature)         AS avg_temperature,
        AVG(humidity)            AS avg_humidity,
        COUNT(*)                 AS reading_count
    FROM sensor_readings
    WHERE ts >= '2026-05-13 00:00:00+00'
      AND ts <= '2026-05-19 23:59:59+00'
    GROUP BY device_id, bucket
    ORDER BY device_id, bucket
    LIMIT 50;
"""

Q4_SQL_TIMESCALE = """
    SELECT
        device_id,
        bucket,
        avg_temperature,
        avg_humidity,
        reading_count
    FROM sensor_readings_hourly
    WHERE bucket >= '2026-05-13 00:00:00+00'
      AND bucket <= '2026-05-19 23:59:59+00'
    ORDER BY device_id, bucket
    LIMIT 50;
"""

Q4_FLUX = f"""
    from(bucket: "{INFLUX_BUCKET}")
      |> range(start: 2026-05-13T00:00:00Z, stop: 2026-05-20T00:00:00Z)
      |> filter(fn: (r) => r._measurement == "{MEASUREMENT}")
      |> filter(fn: (r) => r._field == "temperature" or r._field == "humidity")
      |> aggregateWindow(every: 1h, fn: mean, createEmpty: false)
      |> pivot(rowKey: ["_time", "device_id"], columnKey: ["_field"], valueColumn: "_value")
      |> limit(n: 50)
"""

# Q5 — Range scan: last full day of actual data, sensor_001
# Anchored to 2026-05-19 — the final day in the static dataset
Q5_SQL = """
    SELECT
        device_id,
        ts,
        temperature,
        humidity,
        battery_pct
    FROM sensor_readings
    WHERE device_id = 'sensor_001'
      AND ts >= '2026-05-19 00:00:00+00'
      AND ts <= '2026-05-19 23:59:59+00'
    ORDER BY ts;
"""

Q5_FLUX = f"""
    from(bucket: "{INFLUX_BUCKET}")
      |> range(start: 2026-05-19T00:00:00Z, stop: 2026-05-20T00:00:00Z)
      |> filter(fn: (r) => r._measurement == "{MEASUREMENT}")
      |> filter(fn: (r) => r._field == "temperature" or
                           r._field == "humidity"    or
                           r._field == "battery_pct")
      |> filter(fn: (r) => r.device_id == "sensor_001")
      |> pivot(rowKey: ["_time", "device_id"], columnKey: ["_field"], valueColumn: "_value")
      |> sort(columns: ["_time"])
      |> pivot(rowKey: ["_time", "device_id"], columnKey: ["_field"], valueColumn: "_value")
      |> sort(columns: ["_time"])
"""

# ── Query registry ────────────────────────────────────────────────────────────

QUERIES = [
    {
        "id":             "Q1",
        "name":           "Last value per device",
        "sql_vanilla":    Q1_SQL,
        "sql_timescale":  Q1_SQL,   # same SQL — both are PostgreSQL
        "flux":           Q1_FLUX,
        "note":           "All three use index lookups. DISTINCT ON (PG/TS) vs last() (Influx). Gap should be small.",
    },
    {
        "id":             "Q2",
        "name":           "5-min rolling average (sensor_001, 23:00-23:59 May 19)",
        "sql_vanilla":    Q2_SQL,
        "sql_timescale":  Q2_SQL,
        "flux":           Q2_FLUX,
        "note":           "Small window (~3,600 rows). All three should be fast. Window function vs timedMovingAverage.",
    },
    {
        "id":             "Q3",
        "name":           "Gap detection (> 2 min, May 13-19)",
        "sql_vanilla":    Q3_SQL,
        "sql_timescale":  Q3_SQL,
        "flux":           Q3_FLUX,
        "note":           "Full 7-day scan across all devices. TimescaleDB prunes chunks. Vanilla PG scans the full ts index.",
    },
    {
        "id":             "Q4",
        "name":           "Hourly downsample (May 13-19)",
        "sql_vanilla":    Q4_SQL_VANILLA,
        "sql_timescale":  Q4_SQL_TIMESCALE,  # reads continuous aggregate
        "flux":           Q4_FLUX,
        "note":           "THE KEY QUERY. Vanilla PG and InfluxDB aggregate 18M rows live. TimescaleDB reads pre-computed view.",
    },
    {
        "id":             "Q5",
        "name":           "Range scan — sensor_001, May 19 (full day)",
        "sql_vanilla":    Q5_SQL,
        "sql_timescale":  Q5_SQL,
        "flux":           Q5_FLUX,
        "note":           "24h window = ~86,400 rows for one device. Hypertable touches 24 chunks max. Vanilla scans full index.",
    },
]


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    print("\n" + "═" * 60)
    print("  PROJECT 04 — TIME SERIES  |  benchmark_queries.py")
    print("  Three-way: Vanilla PG vs TimescaleDB vs InfluxDB")
    print("  Fullstack Data — Kenneth")
    print("═" * 60)

    if not INFLUX_TOKEN:
        print("\n  ✗  INFLUX_TOKEN not set. Export it before running:\n"
              "     export INFLUX_TOKEN='your-token'\n")
        raise SystemExit(1)

    # ── Connect ───────────────────────────────────────────────────────────────
    section("CONNECTING TO ALL THREE DATABASES")

    conn_vanilla = psycopg2.connect(**POSTGRES_VANILLA)
    conn_vanilla.autocommit = True
    print("  ✓  Vanilla PostgreSQL connected  (port 5432)")

    conn_ts = psycopg2.connect(**TIMESCALE)
    conn_ts.autocommit = True
    print("  ✓  TimescaleDB connected          (port 5433)")

    influx_client = InfluxDBClient(url=INFLUX_URL, token=INFLUX_TOKEN, org=INFLUX_ORG)
    query_api     = influx_client.query_api()
    print("  ✓  InfluxDB connected             (port 8086)")

    # ── Run benchmarks ────────────────────────────────────────────────────────
    results = []

    for q in QUERIES:
        section(f"{q['id']} — {q['name']}")
        print(f"  Running {RUNS_PER_QUERY}× per database, taking best time …\n")

        # Vanilla PostgreSQL
        try:
            t_vanilla, df_vanilla = time_query(
                lambda sql=q["sql_vanilla"]: pg_query(conn_vanilla, sql))
            vanilla_rows   = len(df_vanilla)
            vanilla_status = f"{t_vanilla:.4f}s"
        except Exception as e:
            t_vanilla, df_vanilla, vanilla_rows = float("inf"), pd.DataFrame(), 0
            vanilla_status = f"ERROR: {e}"

        # TimescaleDB
        try:
            t_ts, df_ts = time_query(
                lambda sql=q["sql_timescale"]: pg_query(conn_ts, sql))
            ts_rows   = len(df_ts)
            ts_status = f"{t_ts:.4f}s"
        except Exception as e:
            t_ts, df_ts, ts_rows = float("inf"), pd.DataFrame(), 0
            ts_status = f"ERROR: {e}"

        # InfluxDB
        try:
            t_influx, df_influx = time_query(
                lambda flux=q["flux"]: flux_query(query_api, flux))
            influx_rows   = len(df_influx)
            influx_status = f"{t_influx:.4f}s"
        except Exception as e:
            t_influx, df_influx, influx_rows = float("inf"), pd.DataFrame(), 0
            influx_status = f"ERROR: {e}"

        w = three_way_winner(t_vanilla, t_ts, t_influx)

        # Per-query print
        rows_table = [
            {"database": "Vanilla PostgreSQL", "time":  vanilla_status, "rows": vanilla_rows},
            {"database": "TimescaleDB",        "time":  ts_status,      "rows": ts_rows},
            {"database": "InfluxDB",           "time":  influx_status,  "rows": influx_rows},
        ]
        print(tabulate(rows_table, headers="keys", tablefmt="rounded_outline"))
        print(f"\n  Winner: {w}")
        print(f"  Note:   {q['note']}")

        # Sample rows from vanilla PG (all three return equivalent data)
        if not df_vanilla.empty:
            print(f"\n  Sample result (Vanilla PG — first 3 rows):")
            print(tabulate(df_vanilla.head(3), headers="keys",
                           tablefmt="rounded_outline", showindex=False))

        results.append({
            "query_id":      q["id"],
            "query_name":    q["name"],
            "vanilla_s":     t_vanilla  if t_vanilla  != float("inf") else None,
            "timescale_s":   t_ts       if t_ts       != float("inf") else None,
            "influx_s":      t_influx   if t_influx   != float("inf") else None,
            "vanilla_rows":  vanilla_rows,
            "timescale_rows": ts_rows,
            "influx_rows":   influx_rows,
            "winner":        w,
            "note":          q["note"],
        })

    # ── Final comparison table ────────────────────────────────────────────────
    section("BENCHMARK RESULTS — THREE-WAY COMPARISON")
    display = [
        {
            "query":       r["query_id"],
            "name":        r["query_name"],
            "vanilla_pg":  f"{r['vanilla_s']:.4f}s"  if r["vanilla_s"]  else "ERROR",
            "timescale":   f"{r['timescale_s']:.4f}s" if r["timescale_s"] else "ERROR",
            "influxdb":    f"{r['influx_s']:.4f}s"    if r["influx_s"]   else "ERROR",
            "winner":      r["winner"],
        }
        for r in results
    ]
    print(tabulate(display, headers="keys", tablefmt="rounded_outline"))

    # ── Speedup table vs vanilla baseline ─────────────────────────────────────
    section("SPEEDUP OVER VANILLA POSTGRESQL BASELINE")
    speedup = []
    for r in results:
        if r["vanilla_s"] and r["vanilla_s"] > 0:
            ts_x     = round(r["vanilla_s"] / r["timescale_s"], 1) if r["timescale_s"] else "—"
            influx_x = round(r["vanilla_s"] / r["influx_s"],    1) if r["influx_s"]    else "—"
        else:
            ts_x = influx_x = "—"
        speedup.append({
            "query":          r["query_id"],
            "timescale_speedup": f"{ts_x}×" if isinstance(ts_x, float) else ts_x,
            "influx_speedup":    f"{influx_x}×" if isinstance(influx_x, float) else influx_x,
        })
    print(tabulate(speedup, headers="keys", tablefmt="rounded_outline"))
    print(f"""
  HOW TO READ THIS:
    2.0× means the TSDB is twice as fast as vanilla PostgreSQL for that query.
    <1.0× means vanilla PostgreSQL was actually faster (possible for simple queries).
    Q4 is where TimescaleDB's continuous aggregate advantage should be most visible.
    """)

    # ── Notes ─────────────────────────────────────────────────────────────────
    section("QUERY NOTES")
    notes = [{"query": r["query_id"], "note": r["note"]} for r in results]
    print(tabulate(notes, headers="keys", tablefmt="rounded_outline"))

    # ── Teardown prompts ──────────────────────────────────────────────────────
    section("WHAT TO DOCUMENT IN TEARDOWN")
    print("""
  Q1  Did vanilla PG and TimescaleDB return the same time? They should —
      DISTINCT ON uses the same index strategy on both. Any gap shows
      the overhead of the TimescaleDB extension layer itself.

  Q2  All three should be close. Small window = small scan regardless of DB.
      If InfluxDB is slower here, Flux query overhead is the reason.

  Q3  Watch the vanilla PG vs TimescaleDB gap. This is pure chunk pruning —
      TimescaleDB skips irrelevant time chunks, vanilla PG walks the full index.

  Q4  THE DEFINING RESULT. Document the exact numbers.
      Vanilla PG time / TimescaleDB time = the continuous aggregate multiplier.
      This number is your concrete argument for why continuous aggregates exist.

  Q5  86,400 rows for one device over 24h. Hypertable touches 24 chunks max.
      Does the speedup materialise? At 18M total rows it should be visible.
    """)

    # ── Write results ─────────────────────────────────────────────────────────
    os.makedirs("data", exist_ok=True)
    with open(OUTPUT_PATH, "w") as f:
        json.dump(results, f, indent=2, default=str)
    print(f"  Benchmark results written → {OUTPUT_PATH}")
    print(f"  Project 04 complete. Write teardown.md next.\n")

    conn_vanilla.close()
    conn_ts.close()
    influx_client.close()


if __name__ == "__main__":
    main()