#!/usr/bin/env python3
"""Column store vs row store: the analytics scan, DDIA 2e Ch. 4.

A fact table is often ~100 columns wide, but an analytics query touches only a
handful (Example 4-1 needs date_key, product_sk, quantity). A row store keeps
each row's columns contiguous, so scanning the table reads every column of
every row -- even the ones the query never mentions. A column store keeps each
column separately, so it reads only the columns the query asks for.

We build one 30-column fact table (10M rows), store it as a row store (SQLite)
and a column store (Parquet), and measure how much data each must READ to answer
    SELECT date_key, sum(quantity) FROM facts GROUP BY date_key
which references just 2 of the 30 columns.

Run it:
    uv run --with duckdb python3 column_scan.py
(Needs DuckDB; sqlite3 is in the standard library. Deterministic data -> the
sizes and bytes-read reproduce exactly; wall-clock timings will vary.)
"""
import os
import sqlite3
import tempfile
import time

import duckdb

N = 10_000_000
FILLER = 26          # 4 real columns + 26 filler = 30 columns
QUERY = "SELECT date_key, sum(quantity) AS q FROM {t} GROUP BY date_key"


def mb(path):
    return os.path.getsize(path) / 1e6


def best_ms(fn, reps=3):
    best = float("inf")
    for _ in range(reps):
        t = time.perf_counter()
        fn()
        best = min(best, time.perf_counter() - t)
    return best * 1000


def main():
    work = tempfile.mkdtemp(prefix="ddia_colscan_")
    def p(name):
        return os.path.join(work, name)

    con = duckdb.connect()
    con.execute("PRAGMA threads=4")

    # 30-column fact table. The 26 filler columns are high-cardinality so they
    # take real space -- the row store cannot avoid reading them.
    filler = ", ".join(f"(hash(id*{k + 7}) % 1000000)::bigint AS c{k}" for k in range(FILLER))
    con.execute(f"""
        CREATE TABLE wide AS
        SELECT id,
               (id % 731)::int              AS date_key,
               (hash(id) % 5000)::int        AS product_sk,
               ((hash(id * 3) % 10) + 1)::int AS quantity,
               {filler}
        FROM range(1, {N} + 1) t(id)
    """)
    con.execute("CREATE TABLE narrow AS SELECT id, date_key, product_sk, quantity FROM wide")

    # Column-store artifacts.
    con.execute(f"COPY wide   TO '{p('wide.parquet')}'   (FORMAT parquet, COMPRESSION snappy)")
    con.execute(f"COPY narrow TO '{p('narrow.parquet')}' (FORMAT parquet, COMPRESSION snappy)")

    # Row-store artifacts (separate SQLite files so we can size each table).
    con.execute("INSTALL sqlite; LOAD sqlite;")
    for tbl in ("wide", "narrow"):
        con.execute(f"ATTACH '{p(tbl + '.db')}' AS s (TYPE sqlite);"
                    f"CREATE TABLE s.{tbl} AS SELECT * FROM {tbl}; DETACH s;")

    # Bytes each store must READ to answer the query (references date_key, quantity).
    def parquet_needed(path, cols):
        rows = con.execute(f"""SELECT path_in_schema, sum(total_compressed_size)
                               FROM parquet_metadata('{path}') GROUP BY path_in_schema""").fetchall()
        total = sum(b for _, b in rows)
        needed = sum(b for name, b in rows if name in cols)
        return total, needed

    wide_total, wide_needed = parquet_needed(p("wide.parquet"), {"date_key", "quantity"})
    row_read = mb(p("wide.db"))       # a full table scan reads the whole file

    print(f"fact table: {N:,} rows x {4 + FILLER} columns")
    print(f"query: SELECT date_key, sum(quantity) GROUP BY date_key   (touches 2 of {4 + FILLER} columns)\n")

    print("on-disk size of the wide table:")
    print(f"    row store    (SQLite) = {mb(p('wide.db')):8.1f} MB")
    print(f"    column store (Parquet)= {wide_total / 1e6:8.1f} MB\n")

    print("bytes the query must READ:")
    print(f"    row store    scans every column of every row = {row_read:8.1f} MB  (100%)")
    print(f"    column store reads only date_key + quantity  = {wide_needed / 1e6:8.2f} MB  "
          f"({100 * wide_needed / (wide_total):.1f}% of the file)")
    print(f"    -> the column store reads {row_read / (wide_needed / 1e6):.0f}x less data\n")

    # Wall-clock (row store via SQLite's own engine; column store via DuckDB).
    sq = sqlite3.connect(p("wide.db"))
    row_ms = best_ms(lambda: sq.execute(QUERY.format(t="wide")).fetchall())
    sq.close()
    col_ms = best_ms(lambda: con.execute(QUERY.format(t=f"read_parquet('{p('wide.parquet')}')")).fetchall())
    print("wall-clock, best of 3 (yours will differ; timings are not the point, bytes are):")
    print(f"    row store    (SQLite) = {row_ms:8.1f} ms")
    print(f"    column store (DuckDB) = {col_ms:8.1f} ms   -> {row_ms / col_ms:.0f}x faster\n")

    # The unused columns are free in a column store.
    _, narrow_needed = parquet_needed(p("narrow.parquet"), {"date_key", "quantity"})
    narrow_ms = best_ms(lambda: con.execute(QUERY.format(t=f"read_parquet('{p('narrow.parquet')}')")).fetchall())
    print("the 26 unused columns are free in the column store:")
    print(f"    narrow (4-col) Parquet: query reads {narrow_needed / 1e6:.2f} MB, {narrow_ms:.1f} ms")
    print(f"    wide  (30-col) Parquet: query reads {wide_needed / 1e6:.2f} MB, {col_ms:.1f} ms   "
          f"(same bytes read; the extra columns cost nothing)")

    for f in os.listdir(work):
        os.remove(p(f))
    os.rmdir(work)


if __name__ == "__main__":
    main()