#!/usr/bin/env python3
"""Column compression and sort order, DDIA 2e Ch. 4.

A column store keeps each column's values together, and analytic columns tend to
have few distinct values ("billions of transactions, only 100,000 products").
Like values sitting adjacent compress extremely well. And because a column store
is free to store rows in any order, sorting the table turns the sort columns into
long runs -- run-length encoding then shrinks them to almost nothing.

We take one fact table (10M rows, 5 low-cardinality columns) and write it three
ways: a row store (SQLite, uncompressed), a column store (Parquet), and the same
column store with the rows sorted. Then we measure the on-disk sizes, and the
per-column sizes before vs. after sorting.

Run it:
    uv run --with duckdb python3 column_compression.py
(Needs DuckDB. Deterministic data -> sizes reproduce exactly.)
"""
import os
import tempfile

import duckdb

N = 10_000_000
SORT_KEYS = "date_key, store_sk, product_sk"


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


def per_column(con, path):
    rows = con.execute(f"""SELECT path_in_schema, sum(total_compressed_size)
                           FROM parquet_metadata('{path}') GROUP BY path_in_schema""").fetchall()
    return dict(rows)


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

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

    # A star-schema fact table with realistic low-cardinality columns.
    con.execute(f"""
        CREATE TABLE facts AS
        SELECT (id % 731)::int                    AS date_key,      -- 731 distinct
               (hash(id) % 5000)::int              AS product_sk,    -- 5000 distinct, scattered
               (hash(id * 7) % 200)::int           AS store_sk,      -- 200 distinct
               (CASE WHEN hash(id * 11) % 100 < 8
                     THEN (hash(id * 13) % 40)::int ELSE 0 END)  AS promotion_sk,  -- mostly 0
               ((hash(id * 3) % 10) + 1)::int      AS quantity       -- 1..10
        FROM range(1, {N} + 1) t(id)
    """)

    # Row store (SQLite, no compression) as the baseline.
    con.execute("INSTALL sqlite; LOAD sqlite;")
    con.execute(f"ATTACH '{p('row.db')}' AS s (TYPE sqlite);"
                f"CREATE TABLE s.facts AS SELECT * FROM facts; DETACH s;")

    # Column store, unsorted and sorted, same compression codec.
    con.execute(f"COPY facts TO '{p('unsorted.parquet')}' (FORMAT parquet, COMPRESSION zstd)")
    con.execute(f"COPY (SELECT * FROM facts ORDER BY {SORT_KEYS}) "
                f"TO '{p('sorted.parquet')}' (FORMAT parquet, COMPRESSION zstd)")

    row, unsorted, sorted_ = mb(p("row.db")), mb(p("unsorted.parquet")), mb(p("sorted.parquet"))

    print(f"fact table: {N:,} rows x 5 low-cardinality columns")
    print("  distinct values: date_key=731, product_sk=5000, store_sk=200, "
          "promotion_sk=~40 (mostly 0), quantity=10\n")

    print("same data, three ways on disk:")
    print(f"    row store    SQLite (uncompressed) = {row:7.1f} MB")
    print(f"    column store Parquet (unsorted)    = {unsorted:7.1f} MB   "
          f"({row / unsorted:.1f}x smaller than the row store)")
    print(f"    column store Parquet (sorted)      = {sorted_:7.1f} MB   "
          f"({unsorted / sorted_:.1f}x smaller again -- same data, just reordered)\n")

    u = per_column(con, p("unsorted.parquet"))
    s = per_column(con, p("sorted.parquet"))
    print(f"per-column compressed size, unsorted -> sorted by ({SORT_KEYS}):")
    print(f"    {'column':13s}{'unsorted':>12s}{'sorted':>12s}   effect of sorting")
    for col in ("date_key", "store_sk", "product_sk", "promotion_sk", "quantity"):
        uc, sc = u[col] / 1e6, s[col] / 1e6
        factor = f"{uc / sc:.0f}x smaller" if sc > 0 and uc / sc >= 1.5 else "~unchanged"
        print(f"    {col:13s}{uc:9.2f} MB{sc:9.2f} MB   {factor}")

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


if __name__ == "__main__":
    main()