!pwd

import time
import pandas as pd
import polars as pl
import duckdb
from contextlib import contextmanager
 
# Configure your data path here
DATA_PATH = "/Users/bswr/smithy/portfolio/data/artifacts/iowa_liquor_sales_monthly/"
 
# Timing results stored for final comparison
BENCHMARKS = {}
 
@contextmanager
def benchmark(name: str, library: str):
    """Context manager to time operations consistently."""
    start = time.perf_counter()
    yield
    elapsed = time.perf_counter() - start
    if name not in BENCHMARKS:
        BENCHMARKS[name] = {}
    BENCHMARKS[name][library] = elapsed
    print(f"{library}: {elapsed:.4f}s")

%%time
# Pandas: Notice the 'object' dtype for strings
print("=== Pandas Schema ===")
df_temp = pd.read_parquet(DATA_PATH)
print(df_temp.dtypes)
print(f"\nShape: {df_temp.shape}")
del df_temp

%%time
# Polars: Explicit types - Utf8 for strings, proper numeric types
print("=== Polars Schema ===")
print(pl.scan_parquet(DATA_PATH).collect_schema())

%%time
# DuckDB: SQL-style types with DESCRIBE
print("=== DuckDB Schema ===")
print(duckdb.sql(f"DESCRIBE SELECT * FROM '{DATA_PATH}*'"))

print("=== Benchmark #1: Full Data Read ===\n")
 
with benchmark("full_read", "pandas"):
    df_pandas = pd.read_parquet(DATA_PATH)
 
with benchmark("full_read", "polars"):
    df_polars = pl.read_parquet(DATA_PATH)
 
with benchmark("full_read", "duckdb"):
    df_duckdb = duckdb.read_parquet(DATA_PATH + "*").df()
 
    print(f"\nRows loaded: {len(df_pandas):,}")

# Polars lazy scan - builds query plan only (instant!)
start = time.perf_counter()
lf_polars = pl.scan_parquet(DATA_PATH)
print(f"Polars scan_parquet: {time.perf_counter() - start:.6f}s (no data loaded!)")
print(f"Type: {type(lf_polars).__name__}")
 
# DuckDB relation - also lazy
start = time.perf_counter()
rel_duckdb = duckdb.read_parquet(DATA_PATH+"*")
print(f"DuckDB read_parquet: {time.perf_counter() - start:.6f}s (lazy relation)")
print(f"Type: {type(rel_duckdb).__name__}")

# Polars shows you the query plan before execution
query = (
    pl.scan_parquet(DATA_PATH)
    .filter(pl.col("category_name") == "VODKA")
    .select(["store_name", "sale_dollars"])
    .group_by("store_name")
    .agg(pl.col("sale_dollars").sum())
)
 
print("=== Polars Query Plan ===")
print(query.explain())

print("=== Benchmark #2: Filtered Read (Predicate Pushdown) ===\n")
 
# Pandas: Must load ALL data, then filter
with benchmark("filter_category", "pandas"):
    df = pd.read_parquet(DATA_PATH)
    result_pandas = df[df["category_name"] == "VODKA"]
    del df
 
# Polars: Pushes filter to parquet reader
with benchmark("filter_category", "polars"):
    result_polars = (
        pl.scan_parquet(DATA_PATH)
        .filter(pl.col("category_name") == "VODKA")
        .collect()
    )
 
# DuckDB: Same pushdown via SQL
with benchmark("filter_category", "duckdb"):
    result_duckdb = duckdb.sql(
        f"SELECT * FROM '{DATA_PATH}' WHERE category_name = 'VODKA'"
    ).df()
 
print(f"\nFiltered rows: {len(result_pandas):,} (VODKA only)")

print("=== Benchmark #3: Column Selection (Projection Pushdown) ===\n")
 
# Pandas: Reads all columns, then selects (some optimization with 'columns' param)
with benchmark("select_columns", "pandas"):
    result_pandas = pd.read_parquet(DATA_PATH, columns=["store_name", "sale_dollars"])
 
# Polars: Only reads 2 columns from disk
with benchmark("select_columns", "polars"):
    result_polars = (
        pl.scan_parquet(DATA_PATH)
        .select(["store_name", "sale_dollars"])
        .collect()
    )
 
# DuckDB: Same projection pushdown
with benchmark("select_columns", "duckdb"):
    result_duckdb = duckdb.sql(
        f"SELECT store_name, sale_dollars FROM '{DATA_PATH}'"
    ).df()
 
print(f"\nColumns selected: {list(result_pandas.columns)}")

print("=== Benchmark #4: Row Calculations ===\n")
 
# First, let's see how bad apply() really is (on a sample to not wait forever)
sample_size = 100_000
df_sample = df_pandas.head(sample_size).copy()
 
print(f"Testing on {sample_size:,} rows first (apply is SLOW)...\n")
 
with benchmark("row_calc_apply", "pandas"):
    df_sample["margin"] = df_sample.apply(
        lambda row: row["sale_dollars"] - row["state_bottle_cost"], axis=1
    )

print("Now vectorized operations on FULL dataset:\n")
 
# Pandas vectorized (the right way)
with benchmark("row_calc_vectorized", "pandas"):
    df_pandas["margin"] = df_pandas["sale_dollars"] - df_pandas["state_bottle_cost"]
 
# Polars expression
with benchmark("row_calc_vectorized", "polars"):
    df_polars = df_polars.with_columns(
        (pl.col("sale_dollars") - pl.col("state_bottle_cost")).alias("margin")
    )
 
# DuckDB SQL expression
with benchmark("row_calc_vectorized", "duckdb"):
    result = duckdb.sql("""
        SELECT *, sale_dollars - state_bottle_cost AS margin 
        FROM df_polars
    """).df()

print("=== Benchmark #5: Window Functions ===\n")
 
# Pandas: Sequential window operations
with benchmark("window_ops", "pandas"):
    df_temp = df_pandas.copy()
    df_temp["store_total"] = df_temp.groupby("store_name")["sale_dollars"].transform("sum")
    df_temp["category_avg"] = df_temp.groupby("category_name")["bottles_sold"].transform("mean")
    del df_temp
 
# Polars: Parallel window operations with .over()
with benchmark("window_ops", "polars"):
    result_polars = df_polars.with_columns([
        pl.col("sale_dollars").sum().over("store_name").alias("store_total"),
        pl.col("bottles_sold").mean().over("category_name").alias("category_avg"),
    ])
 
# DuckDB: Window functions via SQL
with benchmark("window_ops", "duckdb"):
    result_duckdb = duckdb.sql("""
        SELECT *,
            SUM(sale_dollars) OVER (PARTITION BY store_name) as store_total,
            AVG(bottles_sold) OVER (PARTITION BY category_name) as category_avg
        FROM df_polars
    """).df()

print("=== Benchmark #6: Complex Pipeline ===\n")
 
# Pandas: Method chaining gets awkward
with benchmark("complex_pipeline", "pandas"):
    result_pandas = (
        df_pandas
        .assign(year=df_pandas["date"].dt.year)
        .groupby(["year", "category_name"])
        .agg({"sale_dollars": "sum"})
        .reset_index()
        .sort_values("sale_dollars", ascending=False)
    )
 
# Polars: Natural chaining with lazy evaluation
with benchmark("complex_pipeline", "polars"):
    result_polars = (
        df_polars.lazy()
        .with_columns(pl.col("date").dt.year().alias("year"))
        .group_by(["year", "category_name"])
        .agg(pl.col("sale_dollars").sum())
        .sort("sale_dollars", descending=True)
        .collect()
    )
 
# DuckDB: SQL is naturally declarative
with benchmark("complex_pipeline", "duckdb"):
    result_duckdb = duckdb.sql("""
        SELECT 
            YEAR(date) as year, 
            category_name, 
            SUM(sale_dollars) as sale_dollars
        FROM df_polars 
        GROUP BY 1, 2 
        ORDER BY sale_dollars DESC
    """).df()
 
print(f"\nTop 5 results:")
print(result_polars.head())

print("=== Benchmark #7: Group-By Aggregation ===\n")
 
# Pandas
with benchmark("groupby_agg", "pandas"):
    result_pandas = (
        df_pandas
        .groupby(["store_name", "category_name"])
        .agg(
            total_sales=("sale_dollars", "sum"),
            avg_bottles=("bottles_sold", "mean"),
            txn_count=("invoice_line_no", "count")
        )
        .reset_index()
        .sort_values("total_sales", ascending=False)
    )
 
# Polars
with benchmark("groupby_agg", "polars"):
    result_polars = (
        df_polars.lazy()
        .group_by(["store_name", "category_name"])
        .agg([
            pl.col("sale_dollars").sum().alias("total_sales"),
            pl.col("bottles_sold").mean().alias("avg_bottles"),
            pl.col("invoice_line_no").count().alias("txn_count"),
        ])
        .sort("total_sales", descending=True)
        .collect()
    )
 
# DuckDB
with benchmark("groupby_agg", "duckdb"):
    result_duckdb = duckdb.sql("""
        SELECT 
            store_name, 
            category_name,
            SUM(sale_dollars) as total_sales,
            AVG(bottles_sold) as avg_bottles,
            COUNT(invoice_line_no) as txn_count
        FROM df_polars 
        GROUP BY 1, 2 
        ORDER BY total_sales DESC
    """).df()
 
print(f"\nUnique store-category combinations: {len(result_pandas):,}")

print("=== BENCHMARK RESULTS ===\n")
 
results = []
for op, times in BENCHMARKS.items():
    row = {"operation": op}
    row["pandas (s)"] = f"{times.get('pandas', 0):.4f}"
    row["polars (s)"] = f"{times.get('polars', 0):.4f}"
    row["duckdb (s)"] = f"{times.get('duckdb', 0):.4f}"
    
    if "pandas" in times:
        pandas_time = times["pandas"]
        polars_time = times.get("polars", pandas_time)
        duckdb_time = times.get("duckdb", pandas_time)
        row["polars speedup"] = f"{pandas_time / polars_time:.1f}x"
        row["duckdb speedup"] = f"{pandas_time / duckdb_time:.1f}x"
    results.append(row)
 
summary_df = pd.DataFrame(results)
print(summary_df.to_markdown(index=False))